From 03f6992de38af37179ae1e604040ecbf54c6a22e Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Tue, 20 Oct 2026 20:06:33 +0000 Subject: [PATCH 01/80] Record a snapshot manifest next to every graph index A snapshot's id says which graph and data version it is, but nothing on disk records which engine build wrote the index, which DiskGraph on-disk formats it uses, when, from which per-table source versions, or how many rows and edges it holds. DiskGraph gains format_versions(), exposing its existing on-disk format constants (shard layout, shard commit, temp meta, ingest report, _MANIFEST.json schema). Graph.snapshot() writes snapshot-manifest.json into the index directory under the snapshot lock: ids, source table coordinates and version stamps, row/edge counts from ingest_stats.json, the DiskGraph formats, engine version and build commit, and the scheduler job identity (explicit KUMO_* env first, then SageMaker, Snowflake and Databricks detection). Reusing an older build backfills a reconstructed manifest. Writing is best-effort: any failure warns and snapshot() returns as before; reuse and ingest behaviour are unchanged. GraphSnapshot.manifest reads it; kumo.engine.diff_manifests compares two manifests by dotted key. The atomic metadata writer moves from the trainer spec module to the shared manifest module. Verified with new snapshot-manifest tests (fresh write, reuse of a pre-manifest index without re-ingest, write failure, diff after appending rows, scheduler detection, unreadable/newer manifests), a format_versions pybind test, and the existing snapshot, file-graph, custom-SQL, bundle and registry suites. Co-Authored-By: Claude Opus 5.5 --- .../diskgraph/diskgraph-index/src/graph.rs | 4 +- .../python/diskgraph/__init__.py | 2 + .../python/tests/test_format_versions.py | 15 ++ .../diskgraph/diskgraph-pybind/src/lib.rs | 13 + packages/kumo-engine/kumo/engine/__init__.py | 2 + packages/kumo-engine/kumo/engine/_manifest.py | 103 ++++++++ .../kumo-engine/kumo/engine/_provenance.py | 99 ++++++++ packages/kumo-engine/kumo/engine/graph.py | 67 ++++- .../kumo-engine/kumo/engine/trainer/_spec.py | 32 +-- .../kumo-engine/test/test_snapshot_cleanup.py | 1 + .../test/test_snapshot_manifest.py | 239 ++++++++++++++++++ .../test/test_snapshot_staging_cleanup.py | 5 +- .../kumo-engine/test/test_trainer_registry.py | 6 +- 13 files changed, 551 insertions(+), 37 deletions(-) create mode 100644 packages/diskgraph/diskgraph-pybind/python/tests/test_format_versions.py create mode 100644 packages/kumo-engine/kumo/engine/_manifest.py create mode 100644 packages/kumo-engine/kumo/engine/_provenance.py create mode 100644 packages/kumo-engine/test/test_snapshot_manifest.py diff --git a/packages/diskgraph/diskgraph-index/src/graph.rs b/packages/diskgraph/diskgraph-index/src/graph.rs index 0fcb6c2cc..b75d64a4f 100644 --- a/packages/diskgraph/diskgraph-index/src/graph.rs +++ b/packages/diskgraph/diskgraph-index/src/graph.rs @@ -27,9 +27,9 @@ use std::path::{Path, PathBuf}; use std::sync::Arc; const TEMP_META_MAGIC: [u8; 4] = *b"DGTP"; -const TEMP_META_VERSION: u8 = 2; +pub const TEMP_META_VERSION: u8 = 2; pub const INDEX_SHARD_COMMIT_DIR: &str = "_tmp/index_commits"; -const INDEX_SHARD_COMMIT_VERSION: u32 = 1; +pub const INDEX_SHARD_COMMIT_VERSION: u32 = 1; #[derive(Debug, Deserialize, Serialize)] struct IndexShardCommit { diff --git a/packages/diskgraph/diskgraph-pybind/python/diskgraph/__init__.py b/packages/diskgraph/diskgraph-pybind/python/diskgraph/__init__.py index d157a5b0c..5e2fbc51e 100644 --- a/packages/diskgraph/diskgraph-pybind/python/diskgraph/__init__.py +++ b/packages/diskgraph/diskgraph-pybind/python/diskgraph/__init__.py @@ -40,6 +40,7 @@ verify_snapshot_integrity, has_resumable_ingest, validated_index_resume_progress, + format_versions, compute_col_stats_for_parquet, resolve_temporal, detect_multicat_delimiter, @@ -141,6 +142,7 @@ def __getattr__(name): "verify_snapshot_integrity", "has_resumable_ingest", "validated_index_resume_progress", + "format_versions", "compute_col_stats_for_parquet", "DiskGraphPlanCompiler", "SplitSpec", diff --git a/packages/diskgraph/diskgraph-pybind/python/tests/test_format_versions.py b/packages/diskgraph/diskgraph-pybind/python/tests/test_format_versions.py new file mode 100644 index 000000000..5863b1889 --- /dev/null +++ b/packages/diskgraph/diskgraph-pybind/python/tests/test_format_versions.py @@ -0,0 +1,15 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import diskgraph + + +def test_format_versions_names_every_on_disk_format(): + assert diskgraph.format_versions() == { + "layout": 4, + "shard_commit": 1, + "temp_meta": 2, + "ingest_report": 2, + "manifest_schema": 1, + } diff --git a/packages/diskgraph/diskgraph-pybind/src/lib.rs b/packages/diskgraph/diskgraph-pybind/src/lib.rs index c907055d1..86d304dca 100644 --- a/packages/diskgraph/diskgraph-pybind/src/lib.rs +++ b/packages/diskgraph/diskgraph-pybind/src/lib.rs @@ -1567,6 +1567,18 @@ fn validated_index_resume_progress(output_dir: &str) -> PyResult std::collections::BTreeMap<&'static str, u32> { + std::collections::BTreeMap::from([ + ("layout", diskgraph_index::layout::LAYOUT_VERSION), + ("shard_commit", diskgraph_index::graph::INDEX_SHARD_COMMIT_VERSION), + ("temp_meta", u32::from(diskgraph_index::graph::TEMP_META_VERSION)), + ("ingest_report", diskgraph_ingest::stats::INGEST_REPORT_VERSION), + ("manifest_schema", diskgraph_ingest::snapshot_mirror::MANIFEST_SCHEMA_VERSION), + ]) +} + #[pymodule] #[pyo3(name = "_c")] fn diskgraph_pybind(m: &Bound<'_, PyModule>) -> PyResult<()> { @@ -1608,6 +1620,7 @@ fn diskgraph_pybind(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_function(wrap_pyfunction!(verify_snapshot_integrity, 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!(format_versions, m)?)?; m.add_function(wrap_pyfunction!(compute_col_stats_for_parquet, m)?)?; m.add_function(wrap_pyfunction!(col_stats_export::read_col_stats, m)?)?; m.add_function(wrap_pyfunction!(col_stats_export::read_edge_stats, m)?)?; diff --git a/packages/kumo-engine/kumo/engine/__init__.py b/packages/kumo-engine/kumo/engine/__init__.py index 4d89af512..ca8520525 100644 --- a/packages/kumo-engine/kumo/engine/__init__.py +++ b/packages/kumo-engine/kumo/engine/__init__.py @@ -25,6 +25,7 @@ from kumoapi.plugin import ModelRegistryLocation from kumoapi.table import Column +from kumo.engine._manifest import diff_manifests from kumo.engine._version import __version__ from kumo.engine.config import configure from kumo.engine.connector import ( @@ -59,6 +60,7 @@ __all__ = [ "__version__", "configure", + "diff_manifests", "EmbeddingModelStore", "GloVeVectors", "TokenEmbeddingModel", diff --git a/packages/kumo-engine/kumo/engine/_manifest.py b/packages/kumo-engine/kumo/engine/_manifest.py new file mode 100644 index 000000000..3123d667b --- /dev/null +++ b/packages/kumo-engine/kumo/engine/_manifest.py @@ -0,0 +1,103 @@ +# 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 os +import tempfile +import warnings +from datetime import datetime, timezone +from pathlib import Path +from typing import Any, Dict, Mapping, Optional, Tuple + +SCHEMA_VERSION = 1 + + +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) + + +def new_manifest(kind: str, id: str, parents: Mapping[str, Any], + **fields: Any) -> Dict[str, Any]: + from kumo.engine._provenance import environment, scheduler_context + + return { + 'kind': kind, + 'schema_version': SCHEMA_VERSION, + 'id': id, + 'parents': dict(parents), + 'created_at': datetime.now(timezone.utc).isoformat(timespec='seconds'), + 'environment': environment(), + 'scheduler': scheduler_context(), + **fields, + } + + +def write_manifest(path: Path, manifest: Mapping[str, Any]) -> None: + atomic_write(path, (json.dumps(manifest, indent=2, sort_keys=True) + + '\n').encode('utf-8')) + + +def read_manifest(path: Path) -> Optional[Dict[str, Any]]: + 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 not isinstance(version, int) or 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 manifest + + +def _flatten(value: Mapping[str, Any], prefix: str = '') -> Dict[str, Any]: + flat: Dict[str, Any] = {} + for key, item in value.items(): + name = f'{prefix}{key}' + if isinstance(item, Mapping) and item: + flat.update(_flatten(item, f'{name}.')) + else: + flat[name] = item + return flat + + +def diff_manifests(a: Mapping[str, Any], + b: Mapping[str, Any]) -> Dict[str, Tuple[Any, Any]]: + """Return ``{dotted_key: (a_value, b_value)}`` for every differing key.""" + left, right = _flatten(a), _flatten(b) + return { + key: (left.get(key), right.get(key)) + for key in sorted(left.keys() | right.keys()) + if key not in left or key not in right or left[key] != right[key] + } diff --git a/packages/kumo-engine/kumo/engine/_provenance.py b/packages/kumo-engine/kumo/engine/_provenance.py new file mode 100644 index 000000000..b6c78d701 --- /dev/null +++ b/packages/kumo-engine/kumo/engine/_provenance.py @@ -0,0 +1,99 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Runtime versions and scheduler job identity recorded in manifests.""" + +import json +import os +import platform +import re +import subprocess +from functools import cache +from importlib import metadata +from pathlib import Path +from typing import Any, Dict, Optional + +PROCESSING_JOB_CONFIG = Path('/opt/ml/config/processingjobconfig.json') + + +def _version(distribution: str) -> Optional[str]: + try: + return metadata.version(distribution) + except metadata.PackageNotFoundError: + return None + + +@cache +def _build_commit(version: str) -> Optional[str]: + if '+' in version: + return version.split('+', 1)[1] + root = Path(__file__).resolve().parents[2] + if root.name in ('site-packages', 'dist-packages'): + 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 environment() -> Dict[str, Optional[str]]: + from kumo.engine._version import __version__ + + return { + 'kumo_engine': __version__, + 'build_commit': _build_commit(__version__), + 'kumo_core': _version('kumo-core'), + 'torch': _version('torch'), + 'python': platform.python_version(), + } + + +def _context(platform_name: str, job_id: Optional[str], + job_run_id: Optional[str], source: str) -> Dict[str, Any]: + return { + 'platform': platform_name, + 'job_id': job_id, + 'job_run_id': job_run_id, + 'source': source, + } + + +def _processing_job() -> Optional[Dict[str, Any]]: + try: + config = json.loads(PROCESSING_JOB_CONFIG.read_text()) + except (OSError, ValueError): + return None + return config if isinstance(config, dict) else None + + +def _detect() -> Optional[Dict[str, Any]]: + env = os.environ + if env.get('TRAINING_JOB_NAME'): + return _context('sagemaker', env['TRAINING_JOB_NAME'], + env.get('TRAINING_JOB_ARN'), 'auto') + processing = _processing_job() + if processing is not None and processing.get('ProcessingJobName'): + return _context('sagemaker', processing['ProcessingJobName'], + processing.get('ProcessingJobArn'), 'auto') + service = env.get('SNOWFLAKE_SERVICE_NAME') + if service: + return _context('snowflake', re.sub(r'_[A-Za-z0-9]{8}$', '', service), + service, 'auto') + if env.get('DATABRICKS_RUNTIME_VERSION'): + return _context('databricks', None, None, 'auto') + return None + + +def scheduler_context() -> Optional[Dict[str, Any]]: + env = os.environ + scheduler = env.get('KUMO_SCHEDULER') + job_id = env.get('KUMO_JOB_ID') + job_run_id = env.get('KUMO_JOB_RUN_ID') + detected = _detect() + if scheduler or job_id or job_run_id: + if not scheduler: + scheduler = detected['platform'] if detected else 'other' + return _context(scheduler, job_id, job_run_id, 'env') + return detected diff --git a/packages/kumo-engine/kumo/engine/graph.py b/packages/kumo-engine/kumo/engine/graph.py index b5ecf3299..0751b3e34 100644 --- a/packages/kumo-engine/kumo/engine/graph.py +++ b/packages/kumo-engine/kumo/engine/graph.py @@ -59,6 +59,7 @@ remote/durable cache, e.g. a Snowflake stage) -- out of scope for now. """ +import hashlib import json import warnings from dataclasses import dataclass, field @@ -69,6 +70,7 @@ from kumo_core.pquery import NativeGraph, normalize_graph from kumoapi.graph import Edge +from kumo.engine._manifest import new_manifest, read_manifest, write_manifest from kumo.engine._scratch import ( fingerprint, get_scratch_dir, @@ -77,6 +79,8 @@ from kumo.engine._storage import network_storage from kumo.engine.table import Table +SNAPSHOT_MANIFEST_FILE = 'snapshot-manifest.json' + _warned_scratch: set[Path] = set() @@ -215,6 +219,51 @@ class GraphSnapshot: index_path: Path graph_config: Dict[str, Any] + @property + def manifest(self) -> Dict[str, Any] | None: + """The snapshot manifest stored next to the index, if any.""" + return read_manifest(Path(self.index_path) / SNAPSHOT_MANIFEST_FILE) + + +def _write_snapshot_manifest(index_path: Path, snapshot_id: str, + graph_key: str, data_version: str, + source_tables: Dict[str, Dict[str, str]], *, + fresh: bool) -> None: + path = index_path / SNAPSHOT_MANIFEST_FILE + if not fresh and path.exists(): + return + try: + from diskgraph import format_versions + + stats = json.loads( + (index_path / 'ingest_stats.json').read_text()).get('stats', {}) + rows = stats.get('source_rows_by_table', {}) + fields = { + 'graph_id': graph_key, + 'data_version': data_version, + 'formats': format_versions(), + 'tables': { + alias: { + 'rows': rows.get(alias) + } + for alias in source_tables + }, + 'edges': dict(stats.get('edge_counts', {})), + } + if not fresh: + fields.update(created_at=None, environment=None, scheduler=None, + formats=None, reconstructed=True) + write_manifest( + path, + new_manifest('graph_snapshot', snapshot_id, + {'source_tables': source_tables}, **fields)) + except Exception as error: + warnings.warn( + f"Could not write the snapshot manifest at {path}: {error}", + UserWarning, + stacklevel=4, + ) + @dataclass(frozen=True) class Graph: @@ -418,6 +467,7 @@ def _snapshot_no_lock(self, graph_key: str, persist_dir: Path) \ table_data_versions: Dict[str, str] = {} source_paths: Dict[str, str] = {} + source_tables: Dict[str, Dict[str, str]] = {} try: for alias, table in self._named_tables.items(): @@ -429,10 +479,13 @@ def _snapshot_no_lock(self, graph_key: str, persist_dir: Path) \ dest_dir.mkdir(parents=True, exist_ok=True) name, sql = table.table_name, None + coordinate = name source = table.definition.source_table if isinstance(source, CustomSQLSourceTable): checked = checked_sql(connector, source.sql) name, sql = checked.table, checked.sql + coordinate = ('sql:' + + hashlib.sha256(sql.encode()).hexdigest()[:16]) prepared = source_backend.prepare_ingest( connector, name, @@ -451,6 +504,10 @@ 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': f'{connector.source_type}:{coordinate}', + 'version': prepared.version, + } paths = prepared.source.get('paths') or [] if paths: @@ -499,10 +556,16 @@ def _snapshot_no_lock(self, graph_key: str, persist_dir: Path) \ ), ) + snapshot_id = f'{graph_key}@{data_version[:16]}' + index_path = Path(result.index_path) + _write_snapshot_manifest(index_path, snapshot_id, graph_key, + data_version, source_tables, + fresh=existing is None) + _warn_ingest_data_loss(result, graph_config) return GraphSnapshot( - id=f'{graph_key}@{data_version[:16]}', - index_path=Path(result.index_path), + id=snapshot_id, + index_path=index_path, graph_config=graph_config, ) finally: diff --git a/packages/kumo-engine/kumo/engine/trainer/_spec.py b/packages/kumo-engine/kumo/engine/trainer/_spec.py index 0115532b7..21cefcce4 100644 --- a/packages/kumo-engine/kumo/engine/trainer/_spec.py +++ b/packages/kumo-engine/kumo/engine/trainer/_spec.py @@ -11,14 +11,14 @@ import hashlib import json -import os -import tempfile from dataclasses import asdict, dataclass from pathlib import Path from typing import Any, Dict from kumoapi.model_plan import ModelPlan, TrainingTableGenerationPlan +from kumo.engine._manifest import atomic_write + #: Bundle metadata filename and schema version. SPEC_FILE = 'training-spec.json' SPEC_VERSION = 1 @@ -52,30 +52,6 @@ def _sha256(path: Path) -> str: 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) - - def write_training_spec( bundle_dir: Path, *, @@ -95,7 +71,7 @@ def write_training_spec( training_table_generation=to_json_dict(training_table_generation), ) spec_path = bundle_dir / SPEC_FILE - _atomic_write( + atomic_write( spec_path, (json.dumps(asdict(spec), indent=2, sort_keys=True) + '\n').encode('utf-8'), @@ -113,7 +89,7 @@ def write_training_spec( 'size_bytes': spec_path.stat().st_size, 'sha256': _sha256(spec_path), }] - _atomic_write( + atomic_write( manifest_path, (json.dumps(manifest, indent=2, sort_keys=True) + '\n').encode('utf-8'), ) diff --git a/packages/kumo-engine/test/test_snapshot_cleanup.py b/packages/kumo-engine/test/test_snapshot_cleanup.py index 3e2924e93..06f9c297f 100644 --- a/packages/kumo-engine/test/test_snapshot_cleanup.py +++ b/packages/kumo-engine/test/test_snapshot_cleanup.py @@ -94,6 +94,7 @@ def prepare_ingest(self, connector, table_name, destination, **kwargs): SimpleNamespace(table_name=name, definition=SimpleNamespace(cols=[], source_table=None), connector=SimpleNamespace( + source_type='LOCAL_FILE', source_backend=source, validate_scratch_dir=lambda path: None, )) 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..72c49df70 --- /dev/null +++ b/packages/kumo-engine/test/test_snapshot_manifest.py @@ -0,0 +1,239 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import json +import os +from datetime import datetime + +import numpy as np +import pandas as pd +import pytest +from diskgraph import format_versions +from kumoapi.graph import Edge +from kumoapi.typing import Stype + +import kumo.engine._provenance as provenance +from kumo.engine import ( + FileConnector, + Graph, + Table, + __version__, + diff_manifests, +) +from kumo.engine.graph import SNAPSHOT_MANIFEST_FILE, GraphSnapshot + +_SCHEDULER_ENV = ('KUMO_SCHEDULER', 'KUMO_JOB_ID', 'KUMO_JOB_RUN_ID', + 'TRAINING_JOB_NAME', 'TRAINING_JOB_ARN', + 'SNOWFLAKE_SERVICE_NAME', 'DATABRICKS_RUNTIME_VERSION') + + +@pytest.fixture +def no_scheduler(tmp_path, monkeypatch): + for name in _SCHEDULER_ENV: + monkeypatch.delenv(name, raising=False) + monkeypatch.setattr(provenance, 'PROCESSING_JOB_CONFIG', + tmp_path / 'processingjobconfig.json') + return monkeypatch + + +@pytest.fixture +def make_snapshot(tmp_path, no_scheduler): + no_scheduler.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) -> 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.parquet'], + 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() + + return make + + +def test_fresh_snapshot_writes_manifest(make_snapshot, tmp_path): + snapshot = make_snapshot(30) + manifest = snapshot.manifest + graph_id, short_version = snapshot.id.split('@') + + assert manifest['kind'] == 'graph_snapshot' + assert manifest['schema_version'] == 1 + assert manifest['id'] == snapshot.id + assert manifest['graph_id'] == graph_id + assert manifest['data_version'][:16] == short_version + assert len(manifest['data_version']) == 64 + assert snapshot.index_path.name == manifest['data_version'] + assert manifest['formats'] == format_versions() + sources = manifest['parents']['source_tables'] + assert { + alias: entry['source'] + for alias, entry in sources.items() + } == { + 'ents': 'LOCAL_FILE:ents.parquet', + 'events': 'LOCAL_FILE:events.parquet', + } + assert all(entry['version'] for entry in sources.values()) + assert manifest['tables'] == {'ents': {'rows': 5}, 'events': {'rows': 30}} + assert manifest['edges']['events_ENT_ID_to_ents_ID'] == 30 + assert set(manifest['environment']) == { + 'kumo_engine', 'build_commit', 'kumo_core', 'torch', 'python' + } + assert manifest['environment']['kumo_engine'] == __version__ + assert manifest['scheduler'] is None + assert datetime.fromisoformat(manifest['created_at']).tzinfo is not None + assert str(tmp_path) not in json.dumps(manifest) + + +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() + rebuilt = reused.manifest + + assert reused == snapshot + assert diff_manifests(original, rebuilt) == { + 'created_at': (original['created_at'], None), + 'reconstructed': (None, True), + **{ + f'{group}.{key}': (value, None) + for group in ('environment', 'formats') + for key, value in original[group].items() + }, + 'environment': (None, None), + 'formats': (None, None), + } + + +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() + assert snapshot.index_path.name.startswith(snapshot.id.split('@')[1]) + + +def test_appended_rows_show_in_manifest_diff(make_snapshot): + before = make_snapshot(30).manifest + after = make_snapshot(40).manifest + + diff = diff_manifests(before, after) + + assert diff['data_version'] == (before['data_version'], + after['data_version']) + version = 'parents.source_tables.events.version' + assert version in diff and diff[version][0] != diff[version][1] + assert diff['tables.events.rows'] == (30, 40) + assert 'parents.source_tables.ents.version' not in diff + assert 'tables.ents.rows' not in diff + + +def test_diff_manifests_reports_added_removed_and_whole_lists(): + a = {'kind': 'k', 'x': {'y': 1, 'gone': 2}, 'items': [1, 2]} + b = {'kind': 'k', 'x': {'y': 1, 'new': 3}, 'items': [1, 3]} + + assert diff_manifests(a, b) == { + 'items': ([1, 2], [1, 3]), + 'x.gone': (2, None), + 'x.new': (None, 3), + } + assert diff_manifests({'kind': 'graph_snapshot'}, + {'kind': 'model_bundle'}) == { + 'kind': ('graph_snapshot', 'model_bundle') + } + + +def test_build_commit_comes_from_the_local_version_segment(): + assert provenance._build_commit('1.2.0+abc1234') == 'abc1234' + + +@pytest.mark.parametrize('env, processing, expected', [ + ({}, None, None), + ({ + 'KUMO_SCHEDULER': 'databricks', + 'KUMO_JOB_ID': '11', + 'KUMO_JOB_RUN_ID': '22', + 'TRAINING_JOB_NAME': 'ignored', + }, None, ('databricks', '11', '22', 'env')), + ({ + 'KUMO_JOB_RUN_ID': 'run-1', + 'DATABRICKS_RUNTIME_VERSION': '16.4', + }, None, ('databricks', None, 'run-1', 'env')), + ({ + 'KUMO_JOB_ID': 'job-1' + }, None, ('other', 'job-1', None, 'env')), + ({ + 'TRAINING_JOB_NAME': 'train-1', + 'TRAINING_JOB_ARN': 'arn:train-1', + }, None, ('sagemaker', 'train-1', 'arn:train-1', 'auto')), + ({}, { + 'ProcessingJobName': 'proc-1', + 'ProcessingJobArn': 'arn:proc-1', + }, ('sagemaker', 'proc-1', 'arn:proc-1', 'auto')), + ({ + 'SNOWFLAKE_SERVICE_NAME': 'SCORE_JOB_A1B2C3D4' + }, None, ('snowflake', 'SCORE_JOB', 'SCORE_JOB_A1B2C3D4', 'auto')), + ({ + 'DATABRICKS_RUNTIME_VERSION': '16.4' + }, None, ('databricks', None, None, 'auto')), +]) +def test_scheduler_context(no_scheduler, env, processing, expected): + for name, value in env.items(): + no_scheduler.setenv(name, value) + if processing is not None: + provenance.PROCESSING_JOB_CONFIG.write_text(json.dumps(processing)) + + context = provenance.scheduler_context() + + if expected is None: + assert context is None + else: + assert context == dict( + zip(('platform', 'job_id', 'job_run_id', 'source'), expected)) + + +def test_manifest_from_a_newer_schema_is_ignored(tmp_path): + snapshot = GraphSnapshot('snapshot', tmp_path, {}) + assert snapshot.manifest is None + + (tmp_path / SNAPSHOT_MANIFEST_FILE).write_text( + json.dumps({ + 'kind': 'graph_snapshot', + 'schema_version': 2 + })) + with pytest.warns(UserWarning, match='schema_version 2'): + assert snapshot.manifest is None + + (tmp_path / SNAPSHOT_MANIFEST_FILE).write_text('{"kind": ') + with pytest.warns(UserWarning, match='unreadable'): + assert snapshot.manifest is None diff --git a/packages/kumo-engine/test/test_snapshot_staging_cleanup.py b/packages/kumo-engine/test/test_snapshot_staging_cleanup.py index ea8abe9c0..806b443fd 100644 --- a/packages/kumo-engine/test/test_snapshot_staging_cleanup.py +++ b/packages/kumo-engine/test/test_snapshot_staging_cleanup.py @@ -20,8 +20,9 @@ def prepare_ingest(connector, table_name, dest_dir, *, dtypes, sql=None): return SimpleNamespace( definition=SimpleNamespace(cols=[], source_table=None), - connector=SimpleNamespace(source_backend=SimpleNamespace( - prepare_ingest=prepare_ingest)), + connector=SimpleNamespace( + source_type='LOCAL_FILE', + source_backend=SimpleNamespace(prepare_ingest=prepare_ingest)), table_name='t', ) diff --git a/packages/kumo-engine/test/test_trainer_registry.py b/packages/kumo-engine/test/test_trainer_registry.py index 931200303..c0cb2b52b 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.engine._manifest as manifest_module bundle = _fake_published_bundle(tmp_path) original_manifest = (bundle / MANIFEST_FILE).read_bytes() - real_replace = spec_module.os.replace + real_replace = manifest_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(manifest_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', From a778565db27ec4137e170279f3e5709882b5e6b4 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Tue, 20 Oct 2026 20:26:16 +0000 Subject: [PATCH 02/80] Address snapshot manifest review feedback Snowflake detection no longer strips a suffix from the service name, so names ending in eight alphanumerics are recorded intact. Explicit KUMO_* ids fill missing fields from the detected platform when they agree. diff_manifests treats null like a missing key. The snapshot manifest still records ids, environment and scheduler when ingest_stats.json is unreadable, and reads it from the ingest result's path. build_commit only runs git from a source checkout. Verified with the snapshot-manifest, snapshot cleanup, file-graph, custom-SQL, bundle and registry suites and the pybind format test. Co-Authored-By: Claude Opus 5.5 --- .../python/tests/test_format_versions.py | 14 +++++----- .../diskgraph/diskgraph-pybind/src/lib.rs | 6 ++--- packages/kumo-engine/kumo/engine/_manifest.py | 2 +- .../kumo-engine/kumo/engine/_provenance.py | 15 ++++++----- packages/kumo-engine/kumo/engine/graph.py | 20 +++++++++----- .../test/test_snapshot_manifest.py | 27 ++++++++++++++++--- 6 files changed, 58 insertions(+), 26 deletions(-) diff --git a/packages/diskgraph/diskgraph-pybind/python/tests/test_format_versions.py b/packages/diskgraph/diskgraph-pybind/python/tests/test_format_versions.py index 5863b1889..c8f22422f 100644 --- a/packages/diskgraph/diskgraph-pybind/python/tests/test_format_versions.py +++ b/packages/diskgraph/diskgraph-pybind/python/tests/test_format_versions.py @@ -6,10 +6,12 @@ def test_format_versions_names_every_on_disk_format(): - assert diskgraph.format_versions() == { - "layout": 4, - "shard_commit": 1, - "temp_meta": 2, - "ingest_report": 2, - "manifest_schema": 1, + versions = diskgraph.format_versions() + assert set(versions) == { + "layout", + "shard_commit", + "temp_meta", + "ingest_report", + "manifest_schema", } + assert all(type(value) is int and value > 0 for value in versions.values()) diff --git a/packages/diskgraph/diskgraph-pybind/src/lib.rs b/packages/diskgraph/diskgraph-pybind/src/lib.rs index 86d304dca..14387e8fa 100644 --- a/packages/diskgraph/diskgraph-pybind/src/lib.rs +++ b/packages/diskgraph/diskgraph-pybind/src/lib.rs @@ -45,7 +45,7 @@ create_exception!(diskgraph, DiskgraphSystemError, PyException); // ``_MANIFEST.json``. create_exception!(diskgraph, DiskgraphCorruptionError, PyException); -use std::collections::HashMap; +use std::collections::{BTreeMap, HashMap}; use std::sync::Arc; use std::thread; use tracing::{error, info}; @@ -1569,8 +1569,8 @@ fn validated_index_resume_progress(output_dir: &str) -> PyResult std::collections::BTreeMap<&'static str, u32> { - std::collections::BTreeMap::from([ +fn format_versions() -> BTreeMap<&'static str, u32> { + BTreeMap::from([ ("layout", diskgraph_index::layout::LAYOUT_VERSION), ("shard_commit", diskgraph_index::graph::INDEX_SHARD_COMMIT_VERSION), ("temp_meta", u32::from(diskgraph_index::graph::TEMP_META_VERSION)), diff --git a/packages/kumo-engine/kumo/engine/_manifest.py b/packages/kumo-engine/kumo/engine/_manifest.py index 3123d667b..5db22ab39 100644 --- a/packages/kumo-engine/kumo/engine/_manifest.py +++ b/packages/kumo-engine/kumo/engine/_manifest.py @@ -87,7 +87,7 @@ def _flatten(value: Mapping[str, Any], prefix: str = '') -> Dict[str, Any]: name = f'{prefix}{key}' if isinstance(item, Mapping) and item: flat.update(_flatten(item, f'{name}.')) - else: + elif item is not None: flat[name] = item return flat diff --git a/packages/kumo-engine/kumo/engine/_provenance.py b/packages/kumo-engine/kumo/engine/_provenance.py index b6c78d701..481647269 100644 --- a/packages/kumo-engine/kumo/engine/_provenance.py +++ b/packages/kumo-engine/kumo/engine/_provenance.py @@ -6,7 +6,6 @@ import json import os import platform -import re import subprocess from functools import cache from importlib import metadata @@ -28,7 +27,7 @@ def _build_commit(version: str) -> Optional[str]: if '+' in version: return version.split('+', 1)[1] root = Path(__file__).resolve().parents[2] - if root.name in ('site-packages', 'dist-packages'): + if not (root / 'pyproject.toml').is_file(): return None command = ['git', '-C', str(root), 'rev-parse', '--short', 'HEAD'] try: @@ -79,8 +78,7 @@ def _detect() -> Optional[Dict[str, Any]]: processing.get('ProcessingJobArn'), 'auto') service = env.get('SNOWFLAKE_SERVICE_NAME') if service: - return _context('snowflake', re.sub(r'_[A-Za-z0-9]{8}$', '', service), - service, 'auto') + return _context('snowflake', None, service, 'auto') if env.get('DATABRICKS_RUNTIME_VERSION'): return _context('databricks', None, None, 'auto') return None @@ -93,7 +91,10 @@ def scheduler_context() -> Optional[Dict[str, Any]]: job_run_id = env.get('KUMO_JOB_RUN_ID') detected = _detect() if scheduler or job_id or job_run_id: - if not scheduler: - scheduler = detected['platform'] if detected else 'other' - return _context(scheduler, job_id, job_run_id, 'env') + if detected is not None and scheduler in (None, '', + detected['platform']): + scheduler = detected['platform'] + job_id = job_id or detected['job_id'] + job_run_id = job_run_id or detected['job_run_id'] + return _context(scheduler or 'other', job_id, job_run_id, 'env') return detected diff --git a/packages/kumo-engine/kumo/engine/graph.py b/packages/kumo-engine/kumo/engine/graph.py index 0751b3e34..c58e65ceb 100644 --- a/packages/kumo-engine/kumo/engine/graph.py +++ b/packages/kumo-engine/kumo/engine/graph.py @@ -225,18 +225,26 @@ def manifest(self) -> Dict[str, Any] | None: return read_manifest(Path(self.index_path) / SNAPSHOT_MANIFEST_FILE) -def _write_snapshot_manifest(index_path: Path, snapshot_id: str, - graph_key: str, data_version: str, +def _ingest_stats(result: Any) -> Dict[str, Any]: + try: + stats = json.loads(Path( + result.ingest_stats_path).read_text()).get('stats') + except (AttributeError, OSError, TypeError, ValueError): + return {} + return stats if isinstance(stats, dict) else {} + + +def _write_snapshot_manifest(result: Any, snapshot_id: str, graph_key: str, + data_version: str, source_tables: Dict[str, Dict[str, str]], *, fresh: bool) -> None: - path = index_path / SNAPSHOT_MANIFEST_FILE + path = Path(result.index_path) / SNAPSHOT_MANIFEST_FILE if not fresh and path.exists(): return try: from diskgraph import format_versions - stats = json.loads( - (index_path / 'ingest_stats.json').read_text()).get('stats', {}) + stats = _ingest_stats(result) rows = stats.get('source_rows_by_table', {}) fields = { 'graph_id': graph_key, @@ -558,7 +566,7 @@ def _snapshot_no_lock(self, graph_key: str, persist_dir: Path) \ snapshot_id = f'{graph_key}@{data_version[:16]}' index_path = Path(result.index_path) - _write_snapshot_manifest(index_path, snapshot_id, graph_key, + _write_snapshot_manifest(result, snapshot_id, graph_key, data_version, source_tables, fresh=existing is None) diff --git a/packages/kumo-engine/test/test_snapshot_manifest.py b/packages/kumo-engine/test/test_snapshot_manifest.py index 72c49df70..d08b02609 100644 --- a/packages/kumo-engine/test/test_snapshot_manifest.py +++ b/packages/kumo-engine/test/test_snapshot_manifest.py @@ -124,11 +124,27 @@ def test_index_without_manifest_is_reused_and_backfilled( for group in ('environment', 'formats') for key, value in original[group].items() }, - 'environment': (None, None), - 'formats': (None, None), } +def test_manifest_without_ingest_stats_keeps_identity(tmp_path, no_scheduler): + from types import SimpleNamespace + + from kumo.engine.graph import _write_snapshot_manifest + + result = SimpleNamespace(index_path=str(tmp_path), + ingest_stats_path=str(tmp_path / 'missing.json')) + sources = {'users': {'source': 'LOCAL_FILE:users.parquet', 'version': '1'}} + _write_snapshot_manifest(result, 'g@d', 'g', 'd' * 64, sources, fresh=True) + + manifest = GraphSnapshot('g@d', tmp_path, {}).manifest + assert manifest['id'] == 'g@d' + assert manifest['parents']['source_tables'] == sources + assert manifest['tables'] == {'users': {'rows': None}} + assert manifest['edges'] == {} + assert manifest['environment'] is not None + + def test_manifest_write_failure_leaves_snapshot_unchanged( make_snapshot, monkeypatch): def fail(path, manifest): @@ -192,6 +208,11 @@ def test_build_commit_comes_from_the_local_version_segment(): ({ 'KUMO_JOB_ID': 'job-1' }, None, ('other', 'job-1', None, 'env')), + ({ + 'KUMO_JOB_RUN_ID': 'run-9', + 'TRAINING_JOB_NAME': 'train-1', + 'TRAINING_JOB_ARN': 'arn:train-1', + }, None, ('sagemaker', 'train-1', 'run-9', 'env')), ({ 'TRAINING_JOB_NAME': 'train-1', 'TRAINING_JOB_ARN': 'arn:train-1', @@ -202,7 +223,7 @@ def test_build_commit_comes_from_the_local_version_segment(): }, ('sagemaker', 'proc-1', 'arn:proc-1', 'auto')), ({ 'SNOWFLAKE_SERVICE_NAME': 'SCORE_JOB_A1B2C3D4' - }, None, ('snowflake', 'SCORE_JOB', 'SCORE_JOB_A1B2C3D4', 'auto')), + }, None, ('snowflake', None, 'SCORE_JOB_A1B2C3D4', 'auto')), ({ 'DATABRICKS_RUNTIME_VERSION': '16.4' }, None, ('databricks', None, None, 'auto')), From 55d1ee5c9d9a4b2f05f055e01611d5b87e41330c Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Tue, 20 Oct 2026 23:03:17 +0000 Subject: [PATCH 03/80] Match the backfill diff when no build commit is recorded Installed wheels without a +commit version record build_commit as null, which diff_manifests skips; the expectation now does the same. Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/test/test_snapshot_manifest.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/kumo-engine/test/test_snapshot_manifest.py b/packages/kumo-engine/test/test_snapshot_manifest.py index d08b02609..6c91ae2ca 100644 --- a/packages/kumo-engine/test/test_snapshot_manifest.py +++ b/packages/kumo-engine/test/test_snapshot_manifest.py @@ -122,7 +122,7 @@ def test_index_without_manifest_is_reused_and_backfilled( **{ f'{group}.{key}': (value, None) for group in ('environment', 'formats') - for key, value in original[group].items() + for key, value in original[group].items() if value is not None }, } From c8faa43e3d524454fe0a7bac024d753fe7a906cb Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Wed, 21 Oct 2026 07:41:34 +0000 Subject: [PATCH 04/80] Record each source table's full location in the snapshot manifest parents.source_tables recorded ":", with the bucket or root directory dropped and Custom SQL shortened to 16 hex, so a snapshot's lineage could not point at the data it was built from, and the training-table manifest could not record a custom table in the same shape. Each source table is now {"connector", "location", "version"}: location is the absolute path for a FileConnector, the full s3:// or gs:// URI for object stores (new kumo-core ObjectStoreSessions.table_uri), the fully qualified name for warehouse tables, and "sql:" plus the full sha256 for Custom SQL. Credentials and SQL text are never recorded. The shared kumo.engine._provenance.source_table builds the entry; a location that cannot be determined is null with a warning. Verified with test_snapshot_manifest.py, including S3 rooted and any-URI, warehouse and SQL locations and a failing lookup. Co-Authored-By: Claude Opus 5.5 --- .../sources/object_store_sessions.py | 5 ++ .../kumo-engine/kumo/engine/_provenance.py | 35 +++++++++++++ packages/kumo-engine/kumo/engine/graph.py | 15 ++---- .../test/test_snapshot_manifest.py | 49 ++++++++++++++----- 4 files changed, 83 insertions(+), 21 deletions(-) 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..712df3f36 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 @@ -104,6 +104,11 @@ def _over(self, connector: SourceConnector, bucket: str, root: str, fetch=access.fetch, object_identity=access.object_identity) + def table_uri(self, connector: SourceConnector, table_name: str) -> str: + """Full ``scheme://bucket/key`` of ``table_name``, without credentials.""" + _, root, name, _ = self._placement(connector, table_name) + return f'{self.scheme}://{posixpath.join(root, name).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-engine/kumo/engine/_provenance.py b/packages/kumo-engine/kumo/engine/_provenance.py index 481647269..8c958fa89 100644 --- a/packages/kumo-engine/kumo/engine/_provenance.py +++ b/packages/kumo-engine/kumo/engine/_provenance.py @@ -3,10 +3,12 @@ # SPDX-License-Identifier: Apache-2.0 """Runtime versions and scheduler job identity recorded in manifests.""" +import hashlib import json import os import platform import subprocess +import warnings from functools import cache from importlib import metadata from pathlib import Path @@ -98,3 +100,36 @@ def scheduler_context() -> Optional[Dict[str, Any]]: job_run_id = job_run_id or detected['job_run_id'] return _context(scheduler or 'other', job_id, job_run_id, 'env') return detected + + +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() + root = getattr(connector.config, 'root_dir', None) + if root is not None: + return str(Path(root, table_name)) + table_uri = getattr(connector.source_backend, 'table_uri', None) + if table_uri is not None: + return table_uri(connector, table_name) + namespace = (getattr(connector, field, None) + for field in ('database', 'catalog', 'project_id', + 'schema_name', 'dataset_id')) + return '.'.join([*filter(None, namespace), table_name]) + + +def source_table(connector: Any, table_name: str, sql: Optional[str], + version: Any) -> Dict[str, Optional[str]]: + entry: Dict[str, Optional[str]] = { + 'connector': None, + 'location': None, + 'version': None if version is None else str(version), + } + try: + kind = connector.source_type + entry['connector'] = str(getattr(kind, 'value', kind)) + entry['location'] = source_location(connector, table_name, sql) + except Exception as error: + warnings.warn(f"Could not record a source table's location: {error}", + UserWarning, stacklevel=3) + return entry diff --git a/packages/kumo-engine/kumo/engine/graph.py b/packages/kumo-engine/kumo/engine/graph.py index c58e65ceb..cfb40abc6 100644 --- a/packages/kumo-engine/kumo/engine/graph.py +++ b/packages/kumo-engine/kumo/engine/graph.py @@ -59,7 +59,6 @@ remote/durable cache, e.g. a Snowflake stage) -- out of scope for now. """ -import hashlib import json import warnings from dataclasses import dataclass, field @@ -71,6 +70,7 @@ from kumoapi.graph import Edge from kumo.engine._manifest import new_manifest, read_manifest, write_manifest +from kumo.engine._provenance import source_table from kumo.engine._scratch import ( fingerprint, get_scratch_dir, @@ -236,7 +236,7 @@ def _ingest_stats(result: Any) -> Dict[str, Any]: def _write_snapshot_manifest(result: Any, snapshot_id: str, graph_key: str, data_version: str, - source_tables: Dict[str, Dict[str, str]], *, + source_tables: Dict[str, Dict[str, Any]], *, fresh: bool) -> None: path = Path(result.index_path) / SNAPSHOT_MANIFEST_FILE if not fresh and path.exists(): @@ -475,7 +475,7 @@ def _snapshot_no_lock(self, graph_key: str, persist_dir: Path) \ table_data_versions: Dict[str, str] = {} source_paths: Dict[str, str] = {} - source_tables: Dict[str, Dict[str, str]] = {} + source_tables: Dict[str, Dict[str, Any]] = {} try: for alias, table in self._named_tables.items(): @@ -487,13 +487,10 @@ def _snapshot_no_lock(self, graph_key: str, persist_dir: Path) \ dest_dir.mkdir(parents=True, exist_ok=True) name, sql = table.table_name, None - coordinate = name source = table.definition.source_table if isinstance(source, CustomSQLSourceTable): checked = checked_sql(connector, source.sql) name, sql = checked.table, checked.sql - coordinate = ('sql:' + - hashlib.sha256(sql.encode()).hexdigest()[:16]) prepared = source_backend.prepare_ingest( connector, name, @@ -512,10 +509,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 - source_tables[alias] = { - 'source': f'{connector.source_type}:{coordinate}', - 'version': prepared.version, - } + source_tables[alias] = source_table(connector, name, sql, + prepared.version) paths = prepared.source.get('paths') or [] if paths: diff --git a/packages/kumo-engine/test/test_snapshot_manifest.py b/packages/kumo-engine/test/test_snapshot_manifest.py index 6c91ae2ca..217a79ac7 100644 --- a/packages/kumo-engine/test/test_snapshot_manifest.py +++ b/packages/kumo-engine/test/test_snapshot_manifest.py @@ -2,9 +2,11 @@ # All rights reserved. # SPDX-License-Identifier: Apache-2.0 +import hashlib import json import os from datetime import datetime +from types import SimpleNamespace import numpy as np import pandas as pd @@ -17,6 +19,7 @@ from kumo.engine import ( FileConnector, Graph, + S3Connector, Table, __version__, diff_manifests, @@ -83,13 +86,13 @@ def test_fresh_snapshot_writes_manifest(make_snapshot, tmp_path): assert snapshot.index_path.name == manifest['data_version'] assert manifest['formats'] == format_versions() sources = manifest['parents']['source_tables'] - assert { - alias: entry['source'] - for alias, entry in sources.items() - } == { - 'ents': 'LOCAL_FILE:ents.parquet', - 'events': 'LOCAL_FILE:events.parquet', - } + data = (tmp_path / 'data').resolve() + assert {alias: entry['location'] + for alias, entry in sources.items()} == { + 'ents': str(data / 'ents.parquet'), + 'events': str(data / 'events.parquet'), + } + assert {entry['connector'] for entry in sources.values()} == {'LOCAL_FILE'} assert all(entry['version'] for entry in sources.values()) assert manifest['tables'] == {'ents': {'rows': 5}, 'events': {'rows': 30}} assert manifest['edges']['events_ENT_ID_to_ents_ID'] == 30 @@ -99,7 +102,7 @@ def test_fresh_snapshot_writes_manifest(make_snapshot, tmp_path): assert manifest['environment']['kumo_engine'] == __version__ assert manifest['scheduler'] is None assert datetime.fromisoformat(manifest['created_at']).tzinfo is not None - assert str(tmp_path) not in json.dumps(manifest) + assert str(tmp_path / 'scratch') not in json.dumps(manifest) def test_index_without_manifest_is_reused_and_backfilled( @@ -128,13 +131,17 @@ def test_index_without_manifest_is_reused_and_backfilled( def test_manifest_without_ingest_stats_keeps_identity(tmp_path, no_scheduler): - from types import SimpleNamespace - from kumo.engine.graph import _write_snapshot_manifest result = SimpleNamespace(index_path=str(tmp_path), ingest_stats_path=str(tmp_path / 'missing.json')) - sources = {'users': {'source': 'LOCAL_FILE:users.parquet', 'version': '1'}} + sources = { + 'users': { + 'connector': 'LOCAL_FILE', + 'location': '/data/users.parquet', + 'version': '1' + } + } _write_snapshot_manifest(result, 'g@d', 'g', 'd' * 64, sources, fresh=True) manifest = GraphSnapshot('g@d', tmp_path, {}).manifest @@ -258,3 +265,23 @@ def test_manifest_from_a_newer_schema_is_ignored(tmp_path): (tmp_path / SNAPSHOT_MANIFEST_FILE).write_text('{"kind": ') with pytest.warns(UserWarning, match='unreadable'): assert snapshot.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(), 's3://bucket-a/exports/tt', None, + 's3://bucket-a/exports/tt'), + (SimpleNamespace(config=None, source_backend=None, database='DB', + schema_name='SCHEMA'), 'TT', None, 'DB.SCHEMA.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 == {'connector': None, 'location': None, 'version': '7'} From 8fc3b37d98a8a9ed8147ba13c63e4d398a6feab0 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Wed, 21 Oct 2026 07:44:09 +0000 Subject: [PATCH 05/80] Drop graph_id and data_version from the snapshot manifest; keep the diff private Both fields restate the manifest's id (graph id @ data version prefix) and the index directory name, and diff_manifests was exported from kumo.engine although only internal callers use it. The manifest no longer writes graph_id or data_version, and diff_manifests stays in kumo.engine._manifest without a public export. Verified with test_snapshot_manifest.py, the snapshot cleanup suites and test_trainer_registry.py. Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/__init__.py | 2 -- packages/kumo-engine/kumo/engine/graph.py | 8 ++----- .../test/test_snapshot_manifest.py | 21 +++++-------------- 3 files changed, 7 insertions(+), 24 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/__init__.py b/packages/kumo-engine/kumo/engine/__init__.py index ca8520525..4d89af512 100644 --- a/packages/kumo-engine/kumo/engine/__init__.py +++ b/packages/kumo-engine/kumo/engine/__init__.py @@ -25,7 +25,6 @@ from kumoapi.plugin import ModelRegistryLocation from kumoapi.table import Column -from kumo.engine._manifest import diff_manifests from kumo.engine._version import __version__ from kumo.engine.config import configure from kumo.engine.connector import ( @@ -60,7 +59,6 @@ __all__ = [ "__version__", "configure", - "diff_manifests", "EmbeddingModelStore", "GloVeVectors", "TokenEmbeddingModel", diff --git a/packages/kumo-engine/kumo/engine/graph.py b/packages/kumo-engine/kumo/engine/graph.py index cfb40abc6..d7a5b6fc6 100644 --- a/packages/kumo-engine/kumo/engine/graph.py +++ b/packages/kumo-engine/kumo/engine/graph.py @@ -234,8 +234,7 @@ def _ingest_stats(result: Any) -> Dict[str, Any]: return stats if isinstance(stats, dict) else {} -def _write_snapshot_manifest(result: Any, snapshot_id: str, graph_key: str, - data_version: str, +def _write_snapshot_manifest(result: Any, snapshot_id: str, source_tables: Dict[str, Dict[str, Any]], *, fresh: bool) -> None: path = Path(result.index_path) / SNAPSHOT_MANIFEST_FILE @@ -247,8 +246,6 @@ def _write_snapshot_manifest(result: Any, snapshot_id: str, graph_key: str, stats = _ingest_stats(result) rows = stats.get('source_rows_by_table', {}) fields = { - 'graph_id': graph_key, - 'data_version': data_version, 'formats': format_versions(), 'tables': { alias: { @@ -561,8 +558,7 @@ def _snapshot_no_lock(self, graph_key: str, persist_dir: Path) \ snapshot_id = f'{graph_key}@{data_version[:16]}' index_path = Path(result.index_path) - _write_snapshot_manifest(result, snapshot_id, graph_key, - data_version, source_tables, + _write_snapshot_manifest(result, snapshot_id, source_tables, fresh=existing is None) _warn_ingest_data_loss(result, graph_config) diff --git a/packages/kumo-engine/test/test_snapshot_manifest.py b/packages/kumo-engine/test/test_snapshot_manifest.py index 217a79ac7..d7c1ebf59 100644 --- a/packages/kumo-engine/test/test_snapshot_manifest.py +++ b/packages/kumo-engine/test/test_snapshot_manifest.py @@ -16,14 +16,8 @@ from kumoapi.typing import Stype import kumo.engine._provenance as provenance -from kumo.engine import ( - FileConnector, - Graph, - S3Connector, - Table, - __version__, - diff_manifests, -) +from kumo.engine import FileConnector, Graph, S3Connector, Table, __version__ +from kumo.engine._manifest import diff_manifests from kumo.engine.graph import SNAPSHOT_MANIFEST_FILE, GraphSnapshot _SCHEDULER_ENV = ('KUMO_SCHEDULER', 'KUMO_JOB_ID', 'KUMO_JOB_RUN_ID', @@ -75,15 +69,11 @@ def make(events: int | None = None) -> GraphSnapshot: def test_fresh_snapshot_writes_manifest(make_snapshot, tmp_path): snapshot = make_snapshot(30) manifest = snapshot.manifest - graph_id, short_version = snapshot.id.split('@') assert manifest['kind'] == 'graph_snapshot' assert manifest['schema_version'] == 1 assert manifest['id'] == snapshot.id - assert manifest['graph_id'] == graph_id - assert manifest['data_version'][:16] == short_version - assert len(manifest['data_version']) == 64 - assert snapshot.index_path.name == manifest['data_version'] + assert 'graph_id' not in manifest and 'data_version' not in manifest assert manifest['formats'] == format_versions() sources = manifest['parents']['source_tables'] data = (tmp_path / 'data').resolve() @@ -142,7 +132,7 @@ def test_manifest_without_ingest_stats_keeps_identity(tmp_path, no_scheduler): 'version': '1' } } - _write_snapshot_manifest(result, 'g@d', 'g', 'd' * 64, sources, fresh=True) + _write_snapshot_manifest(result, 'g@d', sources, fresh=True) manifest = GraphSnapshot('g@d', tmp_path, {}).manifest assert manifest['id'] == 'g@d' @@ -172,8 +162,7 @@ def test_appended_rows_show_in_manifest_diff(make_snapshot): diff = diff_manifests(before, after) - assert diff['data_version'] == (before['data_version'], - after['data_version']) + assert diff['id'] == (before['id'], after['id']) version = 'parents.source_tables.events.version' assert version in diff and diff[version][0] != diff[version][1] assert diff['tables.events.rows'] == (30, 40) From e514912b8b58bde024a7f456ca0095f8530480ff Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Wed, 21 Oct 2026 08:07:31 +0000 Subject: [PATCH 06/80] Tighten manifest schema checks and scheduler context read_manifest accepted True, 0 and negative schema versions; it now requires an int from 1 to the supported version. The scheduler context drops its "source" field (auto-detected or from KUMO_* variables), which nothing reads. A reused index whose manifest cannot be backfilled, such as a read-only shared index, warns once per path instead of on every snapshot() call. Verified with test_snapshot_manifest.py. Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/_manifest.py | 2 +- .../kumo-engine/kumo/engine/_provenance.py | 13 +++++----- packages/kumo-engine/kumo/engine/graph.py | 13 ++++++---- .../test/test_snapshot_manifest.py | 24 ++++++++++++------- 4 files changed, 30 insertions(+), 22 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/_manifest.py b/packages/kumo-engine/kumo/engine/_manifest.py index 5db22ab39..731284bd7 100644 --- a/packages/kumo-engine/kumo/engine/_manifest.py +++ b/packages/kumo-engine/kumo/engine/_manifest.py @@ -70,7 +70,7 @@ def read_manifest(path: Path) -> Optional[Dict[str, Any]]: return None version = (manifest.get('schema_version') if isinstance(manifest, dict) else None) - if not isinstance(version, int) or version > SCHEMA_VERSION: + 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}).", diff --git a/packages/kumo-engine/kumo/engine/_provenance.py b/packages/kumo-engine/kumo/engine/_provenance.py index 8c958fa89..096a13d89 100644 --- a/packages/kumo-engine/kumo/engine/_provenance.py +++ b/packages/kumo-engine/kumo/engine/_provenance.py @@ -52,12 +52,11 @@ def environment() -> Dict[str, Optional[str]]: def _context(platform_name: str, job_id: Optional[str], - job_run_id: Optional[str], source: str) -> Dict[str, Any]: + job_run_id: Optional[str]) -> Dict[str, Any]: return { 'platform': platform_name, 'job_id': job_id, 'job_run_id': job_run_id, - 'source': source, } @@ -73,16 +72,16 @@ def _detect() -> Optional[Dict[str, Any]]: env = os.environ if env.get('TRAINING_JOB_NAME'): return _context('sagemaker', env['TRAINING_JOB_NAME'], - env.get('TRAINING_JOB_ARN'), 'auto') + env.get('TRAINING_JOB_ARN')) processing = _processing_job() if processing is not None and processing.get('ProcessingJobName'): return _context('sagemaker', processing['ProcessingJobName'], - processing.get('ProcessingJobArn'), 'auto') + processing.get('ProcessingJobArn')) service = env.get('SNOWFLAKE_SERVICE_NAME') if service: - return _context('snowflake', None, service, 'auto') + return _context('snowflake', None, service) if env.get('DATABRICKS_RUNTIME_VERSION'): - return _context('databricks', None, None, 'auto') + return _context('databricks', None, None) return None @@ -98,7 +97,7 @@ def scheduler_context() -> Optional[Dict[str, Any]]: scheduler = detected['platform'] job_id = job_id or detected['job_id'] job_run_id = job_run_id or detected['job_run_id'] - return _context(scheduler or 'other', job_id, job_run_id, 'env') + return _context(scheduler or 'other', job_id, job_run_id) return detected diff --git a/packages/kumo-engine/kumo/engine/graph.py b/packages/kumo-engine/kumo/engine/graph.py index d7a5b6fc6..e0aff36bf 100644 --- a/packages/kumo-engine/kumo/engine/graph.py +++ b/packages/kumo-engine/kumo/engine/graph.py @@ -82,6 +82,7 @@ SNAPSHOT_MANIFEST_FILE = 'snapshot-manifest.json' _warned_scratch: set[Path] = set() +_warned_backfill: set[Path] = set() def _warn_if_network_scratch(root: Path) -> None: @@ -263,11 +264,13 @@ def _write_snapshot_manifest(result: Any, snapshot_id: str, new_manifest('graph_snapshot', snapshot_id, {'source_tables': source_tables}, **fields)) except Exception as error: - warnings.warn( - f"Could not write the snapshot manifest at {path}: {error}", - UserWarning, - stacklevel=4, - ) + if fresh or path not in _warned_backfill: + _warned_backfill.add(path) + warnings.warn( + f"Could not write the snapshot manifest at {path}: {error}", + UserWarning, + stacklevel=4, + ) @dataclass(frozen=True) diff --git a/packages/kumo-engine/test/test_snapshot_manifest.py b/packages/kumo-engine/test/test_snapshot_manifest.py index d7c1ebf59..a5030b46a 100644 --- a/packages/kumo-engine/test/test_snapshot_manifest.py +++ b/packages/kumo-engine/test/test_snapshot_manifest.py @@ -196,33 +196,33 @@ def test_build_commit_comes_from_the_local_version_segment(): 'KUMO_JOB_ID': '11', 'KUMO_JOB_RUN_ID': '22', 'TRAINING_JOB_NAME': 'ignored', - }, None, ('databricks', '11', '22', 'env')), + }, None, ('databricks', '11', '22')), ({ 'KUMO_JOB_RUN_ID': 'run-1', 'DATABRICKS_RUNTIME_VERSION': '16.4', - }, None, ('databricks', None, 'run-1', 'env')), + }, None, ('databricks', None, 'run-1')), ({ 'KUMO_JOB_ID': 'job-1' - }, None, ('other', 'job-1', None, 'env')), + }, None, ('other', 'job-1', None)), ({ 'KUMO_JOB_RUN_ID': 'run-9', 'TRAINING_JOB_NAME': 'train-1', 'TRAINING_JOB_ARN': 'arn:train-1', - }, None, ('sagemaker', 'train-1', 'run-9', 'env')), + }, None, ('sagemaker', 'train-1', 'run-9')), ({ 'TRAINING_JOB_NAME': 'train-1', 'TRAINING_JOB_ARN': 'arn:train-1', - }, None, ('sagemaker', 'train-1', 'arn:train-1', 'auto')), + }, None, ('sagemaker', 'train-1', 'arn:train-1')), ({}, { 'ProcessingJobName': 'proc-1', 'ProcessingJobArn': 'arn:proc-1', - }, ('sagemaker', 'proc-1', 'arn:proc-1', 'auto')), + }, ('sagemaker', 'proc-1', 'arn:proc-1')), ({ 'SNOWFLAKE_SERVICE_NAME': 'SCORE_JOB_A1B2C3D4' - }, None, ('snowflake', None, 'SCORE_JOB_A1B2C3D4', 'auto')), + }, None, ('snowflake', None, 'SCORE_JOB_A1B2C3D4')), ({ 'DATABRICKS_RUNTIME_VERSION': '16.4' - }, None, ('databricks', None, None, 'auto')), + }, None, ('databricks', None, None)), ]) def test_scheduler_context(no_scheduler, env, processing, expected): for name, value in env.items(): @@ -236,7 +236,7 @@ def test_scheduler_context(no_scheduler, env, processing, expected): assert context is None else: assert context == dict( - zip(('platform', 'job_id', 'job_run_id', 'source'), expected)) + zip(('platform', 'job_id', 'job_run_id'), expected)) def test_manifest_from_a_newer_schema_is_ignored(tmp_path): @@ -251,6 +251,12 @@ def test_manifest_from_a_newer_schema_is_ignored(tmp_path): with pytest.warns(UserWarning, match='schema_version 2'): assert snapshot.manifest is None + for version in (True, 0): + (tmp_path / SNAPSHOT_MANIFEST_FILE).write_text( + json.dumps({'schema_version': version})) + with pytest.warns(UserWarning, match='schema_version'): + assert snapshot.manifest is None + (tmp_path / SNAPSHOT_MANIFEST_FILE).write_text('{"kind": ') with pytest.warns(UserWarning, match='unreadable'): assert snapshot.manifest is None From 83c475606f8023409d86ec7f56c1e6bde8a844c9 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Wed, 21 Oct 2026 08:15:26 +0000 Subject: [PATCH 07/80] Drop snapshot-manifest tests of private helpers The build-commit parser and the synthetic diff test exercised private helpers that the snapshot and appended-rows tests already cover through real manifests; three scheduler cases repeated others. Co-Authored-By: Claude Opus 5.5 --- .../test/test_snapshot_manifest.py | 30 ------------------- 1 file changed, 30 deletions(-) diff --git a/packages/kumo-engine/test/test_snapshot_manifest.py b/packages/kumo-engine/test/test_snapshot_manifest.py index a5030b46a..1228cd99f 100644 --- a/packages/kumo-engine/test/test_snapshot_manifest.py +++ b/packages/kumo-engine/test/test_snapshot_manifest.py @@ -170,25 +170,6 @@ def test_appended_rows_show_in_manifest_diff(make_snapshot): assert 'tables.ents.rows' not in diff -def test_diff_manifests_reports_added_removed_and_whole_lists(): - a = {'kind': 'k', 'x': {'y': 1, 'gone': 2}, 'items': [1, 2]} - b = {'kind': 'k', 'x': {'y': 1, 'new': 3}, 'items': [1, 3]} - - assert diff_manifests(a, b) == { - 'items': ([1, 2], [1, 3]), - 'x.gone': (2, None), - 'x.new': (None, 3), - } - assert diff_manifests({'kind': 'graph_snapshot'}, - {'kind': 'model_bundle'}) == { - 'kind': ('graph_snapshot', 'model_bundle') - } - - -def test_build_commit_comes_from_the_local_version_segment(): - assert provenance._build_commit('1.2.0+abc1234') == 'abc1234' - - @pytest.mark.parametrize('env, processing, expected', [ ({}, None, None), ({ @@ -197,22 +178,11 @@ def test_build_commit_comes_from_the_local_version_segment(): 'KUMO_JOB_RUN_ID': '22', 'TRAINING_JOB_NAME': 'ignored', }, None, ('databricks', '11', '22')), - ({ - 'KUMO_JOB_RUN_ID': 'run-1', - 'DATABRICKS_RUNTIME_VERSION': '16.4', - }, None, ('databricks', None, 'run-1')), - ({ - 'KUMO_JOB_ID': 'job-1' - }, None, ('other', 'job-1', None)), ({ 'KUMO_JOB_RUN_ID': 'run-9', 'TRAINING_JOB_NAME': 'train-1', 'TRAINING_JOB_ARN': 'arn:train-1', }, None, ('sagemaker', 'train-1', 'run-9')), - ({ - 'TRAINING_JOB_NAME': 'train-1', - 'TRAINING_JOB_ARN': 'arn:train-1', - }, None, ('sagemaker', 'train-1', 'arn:train-1')), ({}, { 'ProcessingJobName': 'proc-1', 'ProcessingJobArn': 'arn:proc-1', From 92801ea50c34b26e1372368e4eb7f4181dd1da8c Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Wed, 21 Oct 2026 09:32:50 +0000 Subject: [PATCH 08/80] Trim snapshot manifest PR to essentials The manifest PR carried docstrings on private helpers, a single-use processing-job reader, and tests that restated fields or exercised private helpers. Inline the processing-job read into scheduler detection, drop private docstrings, and fold the manifest tests into behaviour tests: one fresh snapshot plus append diff, backfill, best-effort write failure, unsupported manifests, scheduler detection and source locations. No behaviour change. Verified: test_snapshot_manifest, test_trainer_registry, test_snapshot_cleanup, test_snapshot_staging_cleanup pass; flake8 and isort clean. Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/_manifest.py | 2 - .../kumo-engine/kumo/engine/_provenance.py | 22 +-- packages/kumo-engine/kumo/engine/graph.py | 10 +- .../kumo-engine/test/test_snapshot_cleanup.py | 1 - .../test/test_snapshot_manifest.py | 147 +++++------------- .../test/test_snapshot_staging_cleanup.py | 5 +- 6 files changed, 48 insertions(+), 139 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/_manifest.py b/packages/kumo-engine/kumo/engine/_manifest.py index 731284bd7..57e4efb52 100644 --- a/packages/kumo-engine/kumo/engine/_manifest.py +++ b/packages/kumo-engine/kumo/engine/_manifest.py @@ -15,7 +15,6 @@ 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, @@ -94,7 +93,6 @@ def _flatten(value: Mapping[str, Any], prefix: str = '') -> Dict[str, Any]: def diff_manifests(a: Mapping[str, Any], b: Mapping[str, Any]) -> Dict[str, Tuple[Any, Any]]: - """Return ``{dotted_key: (a_value, b_value)}`` for every differing key.""" left, right = _flatten(a), _flatten(b) return { key: (left.get(key), right.get(key)) diff --git a/packages/kumo-engine/kumo/engine/_provenance.py b/packages/kumo-engine/kumo/engine/_provenance.py index 096a13d89..5cbf5dc6f 100644 --- a/packages/kumo-engine/kumo/engine/_provenance.py +++ b/packages/kumo-engine/kumo/engine/_provenance.py @@ -53,19 +53,8 @@ def environment() -> Dict[str, Optional[str]]: def _context(platform_name: str, job_id: Optional[str], job_run_id: Optional[str]) -> Dict[str, Any]: - return { - 'platform': platform_name, - 'job_id': job_id, - 'job_run_id': job_run_id, - } - - -def _processing_job() -> Optional[Dict[str, Any]]: - try: - config = json.loads(PROCESSING_JOB_CONFIG.read_text()) - except (OSError, ValueError): - return None - return config if isinstance(config, dict) else None + return {'platform': platform_name, 'job_id': job_id, + 'job_run_id': job_run_id} def _detect() -> Optional[Dict[str, Any]]: @@ -73,8 +62,11 @@ def _detect() -> Optional[Dict[str, Any]]: if env.get('TRAINING_JOB_NAME'): return _context('sagemaker', env['TRAINING_JOB_NAME'], env.get('TRAINING_JOB_ARN')) - processing = _processing_job() - if processing is not None and processing.get('ProcessingJobName'): + try: + processing = json.loads(PROCESSING_JOB_CONFIG.read_text()) + except (OSError, ValueError): + processing = None + if isinstance(processing, dict) and processing.get('ProcessingJobName'): return _context('sagemaker', processing['ProcessingJobName'], processing.get('ProcessingJobArn')) service = env.get('SNOWFLAKE_SERVICE_NAME') diff --git a/packages/kumo-engine/kumo/engine/graph.py b/packages/kumo-engine/kumo/engine/graph.py index e0aff36bf..358d8d3dc 100644 --- a/packages/kumo-engine/kumo/engine/graph.py +++ b/packages/kumo-engine/kumo/engine/graph.py @@ -248,12 +248,7 @@ def _write_snapshot_manifest(result: Any, snapshot_id: str, rows = stats.get('source_rows_by_table', {}) fields = { 'formats': format_versions(), - 'tables': { - alias: { - 'rows': rows.get(alias) - } - for alias in source_tables - }, + 'tables': {a: {'rows': rows.get(a)} for a in source_tables}, 'edges': dict(stats.get('edge_counts', {})), } if not fresh: @@ -560,14 +555,13 @@ def _snapshot_no_lock(self, graph_key: str, persist_dir: Path) \ ) snapshot_id = f'{graph_key}@{data_version[:16]}' - index_path = Path(result.index_path) _write_snapshot_manifest(result, snapshot_id, source_tables, fresh=existing is None) _warn_ingest_data_loss(result, graph_config) return GraphSnapshot( id=snapshot_id, - index_path=index_path, + index_path=Path(result.index_path), graph_config=graph_config, ) finally: diff --git a/packages/kumo-engine/test/test_snapshot_cleanup.py b/packages/kumo-engine/test/test_snapshot_cleanup.py index 06f9c297f..3e2924e93 100644 --- a/packages/kumo-engine/test/test_snapshot_cleanup.py +++ b/packages/kumo-engine/test/test_snapshot_cleanup.py @@ -94,7 +94,6 @@ def prepare_ingest(self, connector, table_name, destination, **kwargs): SimpleNamespace(table_name=name, definition=SimpleNamespace(cols=[], source_table=None), connector=SimpleNamespace( - source_type='LOCAL_FILE', source_backend=source, validate_scratch_dir=lambda path: None, )) diff --git a/packages/kumo-engine/test/test_snapshot_manifest.py b/packages/kumo-engine/test/test_snapshot_manifest.py index 1228cd99f..34525caf5 100644 --- a/packages/kumo-engine/test/test_snapshot_manifest.py +++ b/packages/kumo-engine/test/test_snapshot_manifest.py @@ -66,34 +66,32 @@ def make(events: int | None = None) -> GraphSnapshot: return make -def test_fresh_snapshot_writes_manifest(make_snapshot, tmp_path): - snapshot = make_snapshot(30) - manifest = snapshot.manifest +def test_snapshot_manifest_records_sources_and_diffs(make_snapshot, + tmp_path): + before = make_snapshot(30) + manifest = before.manifest assert manifest['kind'] == 'graph_snapshot' - assert manifest['schema_version'] == 1 - assert manifest['id'] == snapshot.id - assert 'graph_id' not in manifest and 'data_version' not in manifest + assert manifest['id'] == before.id assert manifest['formats'] == format_versions() - sources = manifest['parents']['source_tables'] data = (tmp_path / 'data').resolve() - assert {alias: entry['location'] - for alias, entry in sources.items()} == { + assert {a: e['location'] + for a, e in manifest['parents']['source_tables'].items()} == { 'ents': str(data / 'ents.parquet'), 'events': str(data / 'events.parquet'), } - assert {entry['connector'] for entry in sources.values()} == {'LOCAL_FILE'} - assert all(entry['version'] for entry in sources.values()) assert manifest['tables'] == {'ents': {'rows': 5}, 'events': {'rows': 30}} assert manifest['edges']['events_ENT_ID_to_ents_ID'] == 30 - assert set(manifest['environment']) == { - 'kumo_engine', 'build_commit', 'kumo_core', 'torch', 'python' - } assert manifest['environment']['kumo_engine'] == __version__ assert manifest['scheduler'] is None assert datetime.fromisoformat(manifest['created_at']).tzinfo is not None assert str(tmp_path / 'scratch') not in json.dumps(manifest) + diff = diff_manifests(manifest, make_snapshot(40).manifest) + assert diff['tables.events.rows'] == (30, 40) + assert 'parents.source_tables.events.version' in diff + assert 'parents.source_tables.ents.version' not in diff + def test_index_without_manifest_is_reused_and_backfilled( make_snapshot, monkeypatch): @@ -106,40 +104,13 @@ def test_index_without_manifest_is_reused_and_backfilled( lambda *args: pytest.fail('cached index re-ingested')) reused = make_snapshot() - rebuilt = reused.manifest assert reused == snapshot - assert diff_manifests(original, rebuilt) == { - 'created_at': (original['created_at'], None), - 'reconstructed': (None, True), - **{ - f'{group}.{key}': (value, None) - for group in ('environment', 'formats') - for key, value in original[group].items() if value is not None - }, - } - - -def test_manifest_without_ingest_stats_keeps_identity(tmp_path, no_scheduler): - from kumo.engine.graph import _write_snapshot_manifest - - result = SimpleNamespace(index_path=str(tmp_path), - ingest_stats_path=str(tmp_path / 'missing.json')) - sources = { - 'users': { - 'connector': 'LOCAL_FILE', - 'location': '/data/users.parquet', - 'version': '1' - } - } - _write_snapshot_manifest(result, 'g@d', sources, fresh=True) - - manifest = GraphSnapshot('g@d', tmp_path, {}).manifest - assert manifest['id'] == 'g@d' - assert manifest['parents']['source_tables'] == sources - assert manifest['tables'] == {'users': {'rows': None}} - assert manifest['edges'] == {} - assert manifest['environment'] is not None + changed = diff_manifests(original, reused.manifest) + assert changed['reconstructed'] == (None, True) + assert all(key.split('.')[0] in ('created_at', 'reconstructed', + 'environment', 'formats') + for key in changed) def test_manifest_write_failure_leaves_snapshot_unchanged( @@ -153,46 +124,18 @@ def fail(path, manifest): assert snapshot.manifest is None assert (snapshot.index_path / '_SUCCESS').is_file() - assert snapshot.index_path.name.startswith(snapshot.id.split('@')[1]) - - -def test_appended_rows_show_in_manifest_diff(make_snapshot): - before = make_snapshot(30).manifest - after = make_snapshot(40).manifest - - diff = diff_manifests(before, after) - - assert diff['id'] == (before['id'], after['id']) - version = 'parents.source_tables.events.version' - assert version in diff and diff[version][0] != diff[version][1] - assert diff['tables.events.rows'] == (30, 40) - assert 'parents.source_tables.ents.version' not in diff - assert 'tables.ents.rows' not in diff @pytest.mark.parametrize('env, processing, expected', [ ({}, None, None), - ({ - 'KUMO_SCHEDULER': 'databricks', - 'KUMO_JOB_ID': '11', - 'KUMO_JOB_RUN_ID': '22', - 'TRAINING_JOB_NAME': 'ignored', - }, None, ('databricks', '11', '22')), - ({ - 'KUMO_JOB_RUN_ID': 'run-9', - 'TRAINING_JOB_NAME': 'train-1', - 'TRAINING_JOB_ARN': 'arn:train-1', - }, None, ('sagemaker', 'train-1', 'run-9')), - ({}, { - 'ProcessingJobName': 'proc-1', - 'ProcessingJobArn': 'arn:proc-1', - }, ('sagemaker', 'proc-1', 'arn:proc-1')), - ({ - 'SNOWFLAKE_SERVICE_NAME': 'SCORE_JOB_A1B2C3D4' - }, None, ('snowflake', None, 'SCORE_JOB_A1B2C3D4')), - ({ - 'DATABRICKS_RUNTIME_VERSION': '16.4' - }, None, ('databricks', None, None)), + ({'KUMO_SCHEDULER': 'databricks', 'KUMO_JOB_ID': '11', + 'KUMO_JOB_RUN_ID': '22', 'TRAINING_JOB_NAME': 'ignored'}, None, + ('databricks', '11', '22')), + ({'KUMO_JOB_RUN_ID': 'run-9', 'TRAINING_JOB_NAME': 'train-1'}, None, + ('sagemaker', 'train-1', 'run-9')), + ({}, {'ProcessingJobName': 'proc-1', 'ProcessingJobArn': 'arn:proc-1'}, + ('sagemaker', 'proc-1', 'arn:proc-1')), + ({'SNOWFLAKE_SERVICE_NAME': 'JOB_A1'}, None, ('snowflake', None, 'JOB_A1')), ]) def test_scheduler_context(no_scheduler, env, processing, expected): for name, value in env.items(): @@ -202,41 +145,25 @@ def test_scheduler_context(no_scheduler, env, processing, expected): context = provenance.scheduler_context() - if expected is None: - assert context is None - else: - assert context == dict( - zip(('platform', 'job_id', 'job_run_id'), expected)) - - -def test_manifest_from_a_newer_schema_is_ignored(tmp_path): - snapshot = GraphSnapshot('snapshot', tmp_path, {}) - assert snapshot.manifest is None - - (tmp_path / SNAPSHOT_MANIFEST_FILE).write_text( - json.dumps({ - 'kind': 'graph_snapshot', - 'schema_version': 2 - })) - with pytest.warns(UserWarning, match='schema_version 2'): - assert snapshot.manifest is None + assert context == (expected and dict( + zip(('platform', 'job_id', 'job_run_id'), expected))) - for version in (True, 0): - (tmp_path / SNAPSHOT_MANIFEST_FILE).write_text( - json.dumps({'schema_version': version})) - with pytest.warns(UserWarning, match='schema_version'): - assert snapshot.manifest is None - (tmp_path / SNAPSHOT_MANIFEST_FILE).write_text('{"kind": ') - with pytest.warns(UserWarning, match='unreadable'): - assert snapshot.manifest is None +@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(), 's3://bucket-a/exports/tt', None, - 's3://bucket-a/exports/tt'), (SimpleNamespace(config=None, source_backend=None, database='DB', schema_name='SCHEMA'), 'TT', None, 'DB.SCHEMA.TT'), (None, 'synthetic', 'SELECT 1', 'sql:' + diff --git a/packages/kumo-engine/test/test_snapshot_staging_cleanup.py b/packages/kumo-engine/test/test_snapshot_staging_cleanup.py index 806b443fd..ea8abe9c0 100644 --- a/packages/kumo-engine/test/test_snapshot_staging_cleanup.py +++ b/packages/kumo-engine/test/test_snapshot_staging_cleanup.py @@ -20,9 +20,8 @@ def prepare_ingest(connector, table_name, dest_dir, *, dtypes, sql=None): return SimpleNamespace( definition=SimpleNamespace(cols=[], source_table=None), - connector=SimpleNamespace( - source_type='LOCAL_FILE', - source_backend=SimpleNamespace(prepare_ingest=prepare_ingest)), + connector=SimpleNamespace(source_backend=SimpleNamespace( + prepare_ingest=prepare_ingest)), table_name='t', ) From acb2810e06d06d753674ec8b4353184f800b249c Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Wed, 21 Oct 2026 18:22:30 +0000 Subject: [PATCH 09/80] Reuse existing helpers in the snapshot manifest The snapshot manifest carried its own atomic writer, version lookups, a DiskGraph format-version binding, warehouse name formatting and row/edge counts that ingest_stats.json in the same directory already holds. Move the atomic write to kumo_core.io.filesystem.atomic_write_bytes and use it from FilesystemArtifactStore, the prediction manifest storage and the engine manifest writers. Record the engine version via kumo_core.version.get_kumo_version (as model/config.json does) and other packages via installed_distribution_version. Drop format_versions, the manifest's formats/tables/edges and the unused diff helper; expose the counts as GraphSnapshot.counts. Warehouse sessions gain table_location next to the object stores' table_uri. Verified with the snapshot-manifest, trainer-registry, snapshot cleanup, kumo-core artifact/storage and connector source/backend suites. Co-Authored-By: Claude Opus 5.5 --- .../diskgraph/diskgraph-index/src/graph.rs | 4 +- .../python/diskgraph/__init__.py | 2 - .../python/tests/test_format_versions.py | 17 ------ .../diskgraph/diskgraph-pybind/src/lib.rs | 15 +----- .../src/kumo_bigquery/session.py | 5 ++ packages/kumo-core/src/kumo_core/artifacts.py | 35 ++---------- .../kumo-core/src/kumo_core/io/filesystem.py | 32 +++++++++++ .../src/kumo_core/train/_artifact_storage.py | 19 ++----- .../src/kumo_databricks/session.py | 5 ++ packages/kumo-engine/kumo/engine/_manifest.py | 54 ++----------------- .../kumo-engine/kumo/engine/_provenance.py | 30 +++++------ packages/kumo-engine/kumo/engine/graph.py | 47 ++++++++-------- .../kumo-engine/kumo/engine/trainer/_spec.py | 7 ++- .../test/test_snapshot_manifest.py | 53 +++++++++++------- .../kumo-engine/test/test_trainer_registry.py | 6 +-- .../src/kumo_snowflake/session.py | 5 ++ 16 files changed, 139 insertions(+), 197 deletions(-) delete mode 100644 packages/diskgraph/diskgraph-pybind/python/tests/test_format_versions.py diff --git a/packages/diskgraph/diskgraph-index/src/graph.rs b/packages/diskgraph/diskgraph-index/src/graph.rs index b75d64a4f..0fcb6c2cc 100644 --- a/packages/diskgraph/diskgraph-index/src/graph.rs +++ b/packages/diskgraph/diskgraph-index/src/graph.rs @@ -27,9 +27,9 @@ use std::path::{Path, PathBuf}; use std::sync::Arc; const TEMP_META_MAGIC: [u8; 4] = *b"DGTP"; -pub const TEMP_META_VERSION: u8 = 2; +const TEMP_META_VERSION: u8 = 2; pub const INDEX_SHARD_COMMIT_DIR: &str = "_tmp/index_commits"; -pub const INDEX_SHARD_COMMIT_VERSION: u32 = 1; +const INDEX_SHARD_COMMIT_VERSION: u32 = 1; #[derive(Debug, Deserialize, Serialize)] struct IndexShardCommit { diff --git a/packages/diskgraph/diskgraph-pybind/python/diskgraph/__init__.py b/packages/diskgraph/diskgraph-pybind/python/diskgraph/__init__.py index 5e2fbc51e..d157a5b0c 100644 --- a/packages/diskgraph/diskgraph-pybind/python/diskgraph/__init__.py +++ b/packages/diskgraph/diskgraph-pybind/python/diskgraph/__init__.py @@ -40,7 +40,6 @@ verify_snapshot_integrity, has_resumable_ingest, validated_index_resume_progress, - format_versions, compute_col_stats_for_parquet, resolve_temporal, detect_multicat_delimiter, @@ -142,7 +141,6 @@ def __getattr__(name): "verify_snapshot_integrity", "has_resumable_ingest", "validated_index_resume_progress", - "format_versions", "compute_col_stats_for_parquet", "DiskGraphPlanCompiler", "SplitSpec", diff --git a/packages/diskgraph/diskgraph-pybind/python/tests/test_format_versions.py b/packages/diskgraph/diskgraph-pybind/python/tests/test_format_versions.py deleted file mode 100644 index c8f22422f..000000000 --- a/packages/diskgraph/diskgraph-pybind/python/tests/test_format_versions.py +++ /dev/null @@ -1,17 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. -# All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -import diskgraph - - -def test_format_versions_names_every_on_disk_format(): - versions = diskgraph.format_versions() - assert set(versions) == { - "layout", - "shard_commit", - "temp_meta", - "ingest_report", - "manifest_schema", - } - assert all(type(value) is int and value > 0 for value in versions.values()) diff --git a/packages/diskgraph/diskgraph-pybind/src/lib.rs b/packages/diskgraph/diskgraph-pybind/src/lib.rs index 14387e8fa..c907055d1 100644 --- a/packages/diskgraph/diskgraph-pybind/src/lib.rs +++ b/packages/diskgraph/diskgraph-pybind/src/lib.rs @@ -45,7 +45,7 @@ create_exception!(diskgraph, DiskgraphSystemError, PyException); // ``_MANIFEST.json``. create_exception!(diskgraph, DiskgraphCorruptionError, PyException); -use std::collections::{BTreeMap, HashMap}; +use std::collections::HashMap; use std::sync::Arc; use std::thread; use tracing::{error, info}; @@ -1567,18 +1567,6 @@ fn validated_index_resume_progress(output_dir: &str) -> PyResult BTreeMap<&'static str, u32> { - BTreeMap::from([ - ("layout", diskgraph_index::layout::LAYOUT_VERSION), - ("shard_commit", diskgraph_index::graph::INDEX_SHARD_COMMIT_VERSION), - ("temp_meta", u32::from(diskgraph_index::graph::TEMP_META_VERSION)), - ("ingest_report", diskgraph_ingest::stats::INGEST_REPORT_VERSION), - ("manifest_schema", diskgraph_ingest::snapshot_mirror::MANIFEST_SCHEMA_VERSION), - ]) -} - #[pymodule] #[pyo3(name = "_c")] fn diskgraph_pybind(m: &Bound<'_, PyModule>) -> PyResult<()> { @@ -1620,7 +1608,6 @@ fn diskgraph_pybind(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_function(wrap_pyfunction!(verify_snapshot_integrity, 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!(format_versions, m)?)?; m.add_function(wrap_pyfunction!(compute_col_stats_for_parquet, m)?)?; m.add_function(wrap_pyfunction!(col_stats_export::read_col_stats, m)?)?; m.add_function(wrap_pyfunction!(col_stats_export::read_edge_stats, m)?)?; 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-core/src/kumo_core/artifacts.py b/packages/kumo-core/src/kumo_core/artifacts.py index 7a4226974..838f4c0c4 100644 --- a/packages/kumo-core/src/kumo_core/artifacts.py +++ b/packages/kumo-core/src/kumo_core/artifacts.py @@ -15,7 +15,7 @@ from pathlib import Path from typing import Any, ClassVar, Mapping, Protocol, runtime_checkable -from .io.filesystem import local_path +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 +66,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 +233,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(): 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/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-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-engine/kumo/engine/_manifest.py b/packages/kumo-engine/kumo/engine/_manifest.py index 57e4efb52..ca75b7dcb 100644 --- a/packages/kumo-engine/kumo/engine/_manifest.py +++ b/packages/kumo-engine/kumo/engine/_manifest.py @@ -4,37 +4,14 @@ """Shared JSON manifest format for engine artifacts.""" import json -import os -import tempfile import warnings from datetime import datetime, timezone from pathlib import Path -from typing import Any, Dict, Mapping, Optional, Tuple - -SCHEMA_VERSION = 1 +from typing import Any, Dict, Mapping, Optional +from kumo_core.io.filesystem import atomic_write_bytes -def atomic_write(path: Path, payload: bytes) -> None: - 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) +SCHEMA_VERSION = 1 def new_manifest(kind: str, id: str, parents: Mapping[str, Any], @@ -54,8 +31,8 @@ def new_manifest(kind: str, id: str, parents: Mapping[str, Any], def write_manifest(path: Path, manifest: Mapping[str, Any]) -> None: - atomic_write(path, (json.dumps(manifest, indent=2, sort_keys=True) + - '\n').encode('utf-8')) + payload = json.dumps(manifest, indent=2, sort_keys=True) + '\n' + atomic_write_bytes(path, payload.encode('utf-8')) def read_manifest(path: Path) -> Optional[Dict[str, Any]]: @@ -78,24 +55,3 @@ def read_manifest(path: Path) -> Optional[Dict[str, Any]]: ) return None return manifest - - -def _flatten(value: Mapping[str, Any], prefix: str = '') -> Dict[str, Any]: - flat: Dict[str, Any] = {} - for key, item in value.items(): - name = f'{prefix}{key}' - if isinstance(item, Mapping) and item: - flat.update(_flatten(item, f'{name}.')) - elif item is not None: - flat[name] = item - return flat - - -def diff_manifests(a: Mapping[str, Any], - b: Mapping[str, Any]) -> Dict[str, Tuple[Any, Any]]: - left, right = _flatten(a), _flatten(b) - return { - key: (left.get(key), right.get(key)) - for key in sorted(left.keys() | right.keys()) - if key not in left or key not in right or left[key] != right[key] - } diff --git a/packages/kumo-engine/kumo/engine/_provenance.py b/packages/kumo-engine/kumo/engine/_provenance.py index 5cbf5dc6f..9aab73b96 100644 --- a/packages/kumo-engine/kumo/engine/_provenance.py +++ b/packages/kumo-engine/kumo/engine/_provenance.py @@ -10,20 +10,12 @@ import subprocess import warnings from functools import cache -from importlib import metadata from pathlib import Path from typing import Any, Dict, Optional PROCESSING_JOB_CONFIG = Path('/opt/ml/config/processingjobconfig.json') -def _version(distribution: str) -> Optional[str]: - try: - return metadata.version(distribution) - except metadata.PackageNotFoundError: - return None - - @cache def _build_commit(version: str) -> Optional[str]: if '+' in version: @@ -40,13 +32,18 @@ def _build_commit(version: str) -> Optional[str]: def environment() -> Dict[str, Optional[str]]: + from kumo_core.version import ( + get_kumo_version, + installed_distribution_version, + ) + from kumo.engine._version import __version__ return { - 'kumo_engine': __version__, + 'kumo_engine': get_kumo_version(), 'build_commit': _build_commit(__version__), - 'kumo_core': _version('kumo-core'), - 'torch': _version('torch'), + 'kumo_core': installed_distribution_version('kumo-core'), + 'torch': installed_distribution_version('torch'), 'python': platform.python_version(), } @@ -100,13 +97,10 @@ def source_location(connector: Any, table_name: str, root = getattr(connector.config, 'root_dir', None) if root is not None: return str(Path(root, table_name)) - table_uri = getattr(connector.source_backend, 'table_uri', None) - if table_uri is not None: - return table_uri(connector, table_name) - namespace = (getattr(connector, field, None) - for field in ('database', 'catalog', 'project_id', - 'schema_name', 'dataset_id')) - return '.'.join([*filter(None, namespace), table_name]) + backend = connector.source_backend + locate = (getattr(backend, 'table_uri', None) + or getattr(backend, 'table_location', None)) + return table_name if locate is None else locate(connector, table_name) def source_table(connector: Any, table_name: str, sql: Optional[str], diff --git a/packages/kumo-engine/kumo/engine/graph.py b/packages/kumo-engine/kumo/engine/graph.py index 358d8d3dc..7c8ca4b4d 100644 --- a/packages/kumo-engine/kumo/engine/graph.py +++ b/packages/kumo-engine/kumo/engine/graph.py @@ -127,6 +127,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 @@ -150,10 +158,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] = {} @@ -225,14 +231,17 @@ def manifest(self) -> Dict[str, Any] | None: """The snapshot manifest stored next to the index, if any.""" return read_manifest(Path(self.index_path) / SNAPSHOT_MANIFEST_FILE) - -def _ingest_stats(result: Any) -> Dict[str, Any]: - try: - stats = json.loads(Path( - result.ingest_stats_path).read_text()).get('stats') - except (AttributeError, OSError, TypeError, ValueError): - return {} - return stats if isinstance(stats, dict) else {} + @property + def counts(self) -> Dict[str, Dict[str, int]]: + """Source rows per table and edges per edge type in this index.""" + from kumo_core.diskgraph.snapshot import DiskgraphIngestResult + + stats = _ingest_stats( + DiskgraphIngestResult.for_index(self.index_path).ingest_stats_path) + return { + 'rows': dict(stats.get('source_rows_by_table', {})), + 'edges': dict(stats.get('edge_counts', {})), + } def _write_snapshot_manifest(result: Any, snapshot_id: str, @@ -242,18 +251,8 @@ def _write_snapshot_manifest(result: Any, snapshot_id: str, if not fresh and path.exists(): return try: - from diskgraph import format_versions - - stats = _ingest_stats(result) - rows = stats.get('source_rows_by_table', {}) - fields = { - 'formats': format_versions(), - 'tables': {a: {'rows': rows.get(a)} for a in source_tables}, - 'edges': dict(stats.get('edge_counts', {})), - } - if not fresh: - fields.update(created_at=None, environment=None, scheduler=None, - formats=None, reconstructed=True) + fields = {} if fresh else dict(created_at=None, environment=None, + scheduler=None, reconstructed=True) write_manifest( path, new_manifest('graph_snapshot', snapshot_id, diff --git a/packages/kumo-engine/kumo/engine/trainer/_spec.py b/packages/kumo-engine/kumo/engine/trainer/_spec.py index 21cefcce4..147817c0e 100644 --- a/packages/kumo-engine/kumo/engine/trainer/_spec.py +++ b/packages/kumo-engine/kumo/engine/trainer/_spec.py @@ -15,10 +15,9 @@ from pathlib import Path from typing import Any, Dict +from kumo_core.io.filesystem import atomic_write_bytes from kumoapi.model_plan import ModelPlan, TrainingTableGenerationPlan -from kumo.engine._manifest import atomic_write - #: Bundle metadata filename and schema version. SPEC_FILE = 'training-spec.json' SPEC_VERSION = 1 @@ -71,7 +70,7 @@ def write_training_spec( training_table_generation=to_json_dict(training_table_generation), ) 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'), @@ -89,7 +88,7 @@ def write_training_spec( 'size_bytes': spec_path.stat().st_size, 'sha256': _sha256(spec_path), }] - atomic_write( + atomic_write_bytes( manifest_path, (json.dumps(manifest, indent=2, sort_keys=True) + '\n').encode('utf-8'), ) diff --git a/packages/kumo-engine/test/test_snapshot_manifest.py b/packages/kumo-engine/test/test_snapshot_manifest.py index 34525caf5..92c55d555 100644 --- a/packages/kumo-engine/test/test_snapshot_manifest.py +++ b/packages/kumo-engine/test/test_snapshot_manifest.py @@ -11,13 +11,21 @@ import numpy as np import pandas as pd import pytest -from diskgraph import format_versions +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 FileConnector, Graph, S3Connector, Table, __version__ -from kumo.engine._manifest import diff_manifests +from kumo.engine import ( + BigQueryConnector, + DatabricksConnector, + FileConnector, + Graph, + S3Connector, + SnowflakeConnector, + Table, +) +from kumo.engine.connector.identity import DatabricksIdentity from kumo.engine.graph import SNAPSHOT_MANIFEST_FILE, GraphSnapshot _SCHEDULER_ENV = ('KUMO_SCHEDULER', 'KUMO_JOB_ID', 'KUMO_JOB_RUN_ID', @@ -73,24 +81,26 @@ def test_snapshot_manifest_records_sources_and_diffs(make_snapshot, assert manifest['kind'] == 'graph_snapshot' assert manifest['id'] == before.id - assert manifest['formats'] == format_versions() 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 manifest['tables'] == {'ents': {'rows': 5}, 'events': {'rows': 30}} - assert manifest['edges']['events_ENT_ID_to_ents_ID'] == 30 - assert manifest['environment']['kumo_engine'] == __version__ + assert 'tables' not in manifest and 'edges' not in manifest + assert before.counts['rows'] == {'ents': 5, 'events': 30} + assert before.counts['edges']['events_ENT_ID_to_ents_ID'] == 30 + assert manifest['environment']['kumo_engine'] == get_kumo_version() assert manifest['scheduler'] is None assert datetime.fromisoformat(manifest['created_at']).tzinfo is not None assert str(tmp_path / 'scratch') not in json.dumps(manifest) - diff = diff_manifests(manifest, make_snapshot(40).manifest) - assert diff['tables.events.rows'] == (30, 40) - assert 'parents.source_tables.events.version' in diff - assert 'parents.source_tables.ents.version' not in diff + after = make_snapshot(40) + assert after.counts['rows']['events'] == 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( @@ -106,11 +116,10 @@ def test_index_without_manifest_is_reused_and_backfilled( reused = make_snapshot() assert reused == snapshot - changed = diff_manifests(original, reused.manifest) - assert changed['reconstructed'] == (None, True) - assert all(key.split('.')[0] in ('created_at', 'reconstructed', - 'environment', 'formats') - for key in changed) + assert reused.manifest == { + **original, 'created_at': None, 'environment': None, + 'scheduler': None, 'reconstructed': True + } def test_manifest_write_failure_leaves_snapshot_unchanged( @@ -164,8 +173,16 @@ def test_unsupported_manifest_is_ignored(tmp_path, content): @pytest.mark.parametrize('connector, name, sql, location', [ (S3Connector(root_dir='s3://bucket-a/data/'), 'exports/tt/', None, 's3://bucket-a/data/exports/tt'), - (SimpleNamespace(config=None, source_backend=None, database='DB', - schema_name='SCHEMA'), 'TT', None, 'DB.SCHEMA.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(config=None, source_backend=None), 'TT', None, 'TT'), (None, 'synthetic', 'SELECT 1', 'sql:' + hashlib.sha256(b'SELECT 1').hexdigest()), ]) diff --git a/packages/kumo-engine/test/test_trainer_registry.py b/packages/kumo-engine/test/test_trainer_registry.py index c0cb2b52b..9900d5ff0 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._manifest as manifest_module + import kumo_core.io.filesystem as fs_module bundle = _fake_published_bundle(tmp_path) original_manifest = (bundle / MANIFEST_FILE).read_bytes() - real_replace = manifest_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(manifest_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', diff --git a/packages/kumo-snowflake/src/kumo_snowflake/session.py b/packages/kumo-snowflake/src/kumo_snowflake/session.py index 0e237c868..a896821df 100644 --- a/packages/kumo-snowflake/src/kumo_snowflake/session.py +++ b/packages/kumo-snowflake/src/kumo_snowflake/session.py @@ -461,6 +461,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, table_name) + @staticmethod def _qualified(connector: SourceConnector, table_name: str) -> str: return (f'{_quoted(resolved(_target_config(connector).database))}.' From 11ec73ff0df643ca99bdbc489268c8628003c274 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Wed, 21 Oct 2026 21:00:28 +0000 Subject: [PATCH 10/80] Record DiskGraph's version and build commit in manifests Manifests recorded the engine, kumo-core and torch versions but not the DiskGraph build that wrote the index, so two builds with the same format constants could not be told apart. The pybind build stamps the git commit (DISKGRAPH_BUILD_COMMIT overrides it; "unknown" without git) and exposes diskgraph.__version__ and diskgraph.BUILD_COMMIT. environment() records both, or null for an older extension. Verified with the diskgraph export test and the snapshot manifest tests. Co-Authored-By: Claude Opus 5.5 --- packages/diskgraph/diskgraph-pybind/build.rs | 29 +++++++++++++++++ .../python/diskgraph/__init__.py | 2 ++ .../python/tests/test_exception_exports.py | 5 +++ .../diskgraph/diskgraph-pybind/src/lib.rs | 2 ++ .../kumo-engine/kumo/engine/_provenance.py | 32 +++++++++++++++---- 5 files changed, 64 insertions(+), 6 deletions(-) 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..548081716 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, 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..aa4eca8e9 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,8 @@ 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 diff --git a/packages/diskgraph/diskgraph-pybind/src/lib.rs b/packages/diskgraph/diskgraph-pybind/src/lib.rs index c907055d1..e02bda485 100644 --- a/packages/diskgraph/diskgraph-pybind/src/lib.rs +++ b/packages/diskgraph/diskgraph-pybind/src/lib.rs @@ -1593,6 +1593,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)?)?; diff --git a/packages/kumo-engine/kumo/engine/_provenance.py b/packages/kumo-engine/kumo/engine/_provenance.py index 9aab73b96..e4ecdfd4f 100644 --- a/packages/kumo-engine/kumo/engine/_provenance.py +++ b/packages/kumo-engine/kumo/engine/_provenance.py @@ -31,6 +31,14 @@ def _build_commit(version: str) -> Optional[str]: return None +def _diskgraph() -> Any: + try: + import diskgraph + return diskgraph + except ImportError: + return None + + def environment() -> Dict[str, Optional[str]]: from kumo_core.version import ( get_kumo_version, @@ -39,10 +47,13 @@ def environment() -> Dict[str, Optional[str]]: from kumo.engine._version import __version__ + diskgraph = _diskgraph() return { '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(), } @@ -94,9 +105,6 @@ 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() - root = getattr(connector.config, 'root_dir', None) - if root is not None: - return str(Path(root, table_name)) backend = connector.source_backend locate = (getattr(backend, 'table_uri', None) or getattr(backend, 'table_location', None)) @@ -104,16 +112,28 @@ def source_location(connector: Any, table_name: str, def source_table(connector: Any, table_name: str, sql: Optional[str], - version: Any) -> Dict[str, Optional[str]]: - entry: Dict[str, Optional[str]] = { + version: Any, *, + record_locations: bool = True) -> Dict[str, Any]: + from kumo_core.sources.identity import version_kind + + entry: Dict[str, Any] = { 'connector': None, 'location': None, 'version': None if version is None else str(version), + 'version_kind': version_kind(version), } + dependencies = getattr(version, 'dependencies', None) + if dependencies is not None: + entry['dependencies'] = { + name: str(value) + for name, value in dependencies.items() + } try: kind = connector.source_type entry['connector'] = str(getattr(kind, 'value', kind)) - entry['location'] = source_location(connector, table_name, sql) + entry['location'] = (source_location(connector, table_name, sql) + if record_locations or sql is not None else + table_name) except Exception as error: warnings.warn(f"Could not record a source table's location: {error}", UserWarning, stacklevel=3) From 3643180db5af0a9b5fa8e6adf8884adf5e611f88 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Wed, 21 Oct 2026 21:00:28 +0000 Subject: [PATCH 11/80] Type source versions and allow location-free snapshot manifests Lineage could not tell an exact source version from a modification time or an hour bucket, Custom SQL kept only its folded hash, and every manifest recorded full source locations. Connectors now return versions tagged exact, timestamp, hour_bucket or unknown, and Custom SQL versions keep the dependency versions they fold. Each source table entry records version_kind (and dependencies for SQL). Graph.snapshot(record_locations=False) records connector and table name only; a reused index's manifest is rewritten when the choice changes. Snowflake locations carry the account; object-store URIs keep their percent-encoding; file locations name the resolved file. Verified with the snapshot manifest and connector identity tests. Co-Authored-By: Claude Opus 5.5 --- .../src/kumo_bigquery/identity.py | 10 +++- packages/kumo-bigquery/tests/test_identity.py | 3 + .../src/kumo_core/sources/identity.py | 35 ++++++++++-- .../src/kumo_core/sources/object_store.py | 8 ++- .../kumo_core/sources/object_store_backend.py | 2 +- .../sources/object_store_sessions.py | 11 +++- .../tests/sources/unit/test_sql_identity.py | 6 ++ .../src/kumo_databricks/identity.py | 8 ++- .../kumo-databricks/tests/test_identity.py | 2 + .../kumo-engine/kumo/engine/connector/file.py | 5 ++ packages/kumo-engine/kumo/engine/graph.py | 55 ++++++++++++------ .../test/test_snapshot_manifest.py | 56 +++++++++++++++++-- .../src/kumo_snowflake/identity.py | 9 ++- .../src/kumo_snowflake/session.py | 9 ++- .../kumo-snowflake/tests/test_identity.py | 2 + 15 files changed, 176 insertions(+), 45 deletions(-) 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/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/sources/identity.py b/packages/kumo-core/src/kumo_core/sources/identity.py index ac69d6b92..989fd6629 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``, ``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..09c11a32b 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(), "timestamp") class _CsvSource(NamedTuple): 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 712df3f36..ba4f9b893 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: @@ -106,8 +108,11 @@ def _over(self, connector: SourceConnector, bucket: str, root: str, def table_uri(self, connector: SourceConnector, table_name: str) -> str: """Full ``scheme://bucket/key`` of ``table_name``, without credentials.""" - _, root, name, _ = self._placement(connector, table_name) - return f'{self.scheme}://{posixpath.join(root, name).rstrip("/")}' + 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]]: 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-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/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/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 7c8ca4b4d..428f0bbc3 100644 --- a/packages/kumo-engine/kumo/engine/graph.py +++ b/packages/kumo-engine/kumo/engine/graph.py @@ -245,21 +245,33 @@ def counts(self) -> Dict[str, Dict[str, int]]: def _write_snapshot_manifest(result: Any, snapshot_id: str, - source_tables: Dict[str, Dict[str, Any]], *, - fresh: bool) -> None: + source_tables: Dict[str, Dict[str, Any]], + record_locations: bool, *, fresh: bool) -> None: path = Path(result.index_path) / SNAPSHOT_MANIFEST_FILE - if not fresh and path.exists(): - return try: - fields = {} if fresh else dict(created_at=None, environment=None, - scheduler=None, reconstructed=True) - write_manifest( - path, - new_manifest('graph_snapshot', snapshot_id, - {'source_tables': source_tables}, **fields)) + if not fresh and path.exists(): + existing = read_manifest(path) + if existing is None or existing.get('record_locations', + True) == record_locations: + return + manifest = { + **existing, 'record_locations': record_locations, + 'parents': { + **existing['parents'], 'source_tables': source_tables + } + } + else: + fields = {} if fresh else dict(created_at=None, environment=None, + scheduler=None, reconstructed=True) + manifest = new_manifest('graph_snapshot', snapshot_id, + {'source_tables': source_tables}, + record_locations=record_locations, + **fields) + write_manifest(path, manifest) except Exception as error: if fresh or path not in _warned_backfill: - _warned_backfill.add(path) + if not fresh: + _warned_backfill.add(path) warnings.warn( f"Could not write the snapshot manifest at {path}: {error}", UserWarning, @@ -396,7 +408,7 @@ def validate(self) -> None: """ self.native_graph - def snapshot(self) -> GraphSnapshot: + def snapshot(self, *, record_locations: bool = True) -> 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` @@ -405,6 +417,11 @@ 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. + + Args: + record_locations: if False, the snapshot manifest records each + source as connector and table name only (no bucket, path or + account). """ from filelock import FileLock @@ -434,10 +451,11 @@ def snapshot(self) -> GraphSnapshot: # 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) + return self._snapshot_no_lock(graph_key, persist_dir, + record_locations) - def _snapshot_no_lock(self, graph_key: str, persist_dir: Path) \ - -> GraphSnapshot: + def _snapshot_no_lock(self, graph_key: str, persist_dir: Path, + record_locations: 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, @@ -503,8 +521,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) + source_tables[alias] = source_table( + connector, name, sql, prepared.version, + record_locations=record_locations) paths = prepared.source.get('paths') or [] if paths: @@ -555,7 +574,7 @@ def _snapshot_no_lock(self, graph_key: str, persist_dir: Path) \ snapshot_id = f'{graph_key}@{data_version[:16]}' _write_snapshot_manifest(result, snapshot_id, source_tables, - fresh=existing is None) + record_locations, fresh=existing is None) _warn_ingest_data_loss(result, graph_config) return GraphSnapshot( diff --git a/packages/kumo-engine/test/test_snapshot_manifest.py b/packages/kumo-engine/test/test_snapshot_manifest.py index 92c55d555..71c3b22b0 100644 --- a/packages/kumo-engine/test/test_snapshot_manifest.py +++ b/packages/kumo-engine/test/test_snapshot_manifest.py @@ -11,6 +11,7 @@ import numpy as np import pandas as pd import pytest +from kumo_core.sources.identity import SourceVersion from kumo_core.version import get_kumo_version from kumoapi.graph import Edge from kumoapi.typing import Stype @@ -25,7 +26,10 @@ SnowflakeConnector, Table, ) -from kumo.engine.connector.identity import DatabricksIdentity +from kumo.engine.connector.identity import ( + DatabricksIdentity, + SnowflakeIdentity, +) from kumo.engine.graph import SNAPSHOT_MANIFEST_FILE, GraphSnapshot _SCHEDULER_ENV = ('KUMO_SCHEDULER', 'KUMO_JOB_ID', 'KUMO_JOB_RUN_ID', @@ -49,7 +53,7 @@ def make_snapshot(tmp_path, no_scheduler): data.mkdir() pd.DataFrame({'ID': np.arange(5)}).to_parquet(data / 'ents.parquet') - def make(events: int | None = None) -> GraphSnapshot: + def make(events: int | None = None, **options) -> GraphSnapshot: if events is not None: pd.DataFrame({ 'EID': np.arange(events), @@ -62,14 +66,14 @@ def make(events: int | None = None) -> GraphSnapshot: return Graph( { 'ents': - Table.from_source_table(connector['ents.parquet'], + 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() + }, [Edge('events', 'ENT_ID', 'ents')]).snapshot(**options) return make @@ -87,10 +91,14 @@ def test_snapshot_manifest_records_sources_and_diffs(make_snapshot, 'ents': str(data / 'ents.parquet'), 'events': str(data / 'events.parquet'), } + assert {e['version_kind'] + for e in manifest['parents']['source_tables'].values() + } == {'timestamp'} assert 'tables' not in manifest and 'edges' not in manifest assert before.counts['rows'] == {'ents': 5, 'events': 30} assert before.counts['edges']['events_ENT_ID_to_ents_ID'] == 30 assert manifest['environment']['kumo_engine'] == get_kumo_version() + assert manifest['environment']['diskgraph'] assert manifest['scheduler'] is None assert datetime.fromisoformat(manifest['created_at']).tzinfo is not None assert str(tmp_path / 'scratch') not in json.dumps(manifest) @@ -173,6 +181,13 @@ def test_unsupported_manifest_is_ignored(tmp_path, content): @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( @@ -193,4 +208,35 @@ def test_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 == {'connector': None, 'location': None, 'version': '7'} + assert entry == {'connector': None, 'location': None, 'version': '7', + 'version_kind': 'unknown'} + + +def test_custom_sql_records_dependency_versions(): + 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_locations=False) + assert entry == { + 'connector': 'SNOWFLAKE', + 'location': 'sql:' + hashlib.sha256(b'SELECT * FROM orders').hexdigest(), + 'version': 'abc', + 'version_kind': 'hour_bucket', + 'dependencies': {'orders': '17'}, + } + + +def test_record_locations_false_keeps_table_names_only(make_snapshot, + tmp_path): + names = {'ents': 'ents', 'events': 'events.parquet'} + redacted = make_snapshot(30, record_locations=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) + assert redacted.manifest['record_locations'] is False + + full = make_snapshot(record_locations=True) + assert full == redacted + assert full.manifest['parents']['source_tables']['ents'][ + 'location'] == str((tmp_path / 'data' / 'ents.parquet').resolve()) + assert full.manifest['created_at'] == redacted.manifest['created_at'] 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 a896821df..08b8c4ebf 100644 --- a/packages/kumo-snowflake/src/kumo_snowflake/session.py +++ b/packages/kumo-snowflake/src/kumo_snowflake/session.py @@ -463,8 +463,13 @@ 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, table_name) + """``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: 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 677d15af4ea4464b2929e02f92c7e7bb1c7651dd Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Wed, 21 Oct 2026 22:11:50 +0000 Subject: [PATCH 12/80] Pick the source-location lookup by connector type source_location chose between table_uri and table_location by probing the backend's attributes. Object stores now expose table_location like the file and warehouse backends, and the manifest asks for a location only for the connector types that record one; any other connector records its table name. Verified with test_snapshot_manifest.py. Co-Authored-By: Claude Opus 5.5 --- .../kumo_core/sources/object_store_sessions.py | 3 ++- packages/kumo-engine/kumo/engine/_provenance.py | 17 +++++++++++------ .../kumo-engine/test/test_snapshot_manifest.py | 3 ++- 3 files changed, 15 insertions(+), 8 deletions(-) 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 ba4f9b893..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 @@ -106,7 +106,8 @@ def _over(self, connector: SourceConnector, bucket: str, root: str, fetch=access.fetch, object_identity=access.object_identity) - def table_uri(self, connector: SourceConnector, table_name: str) -> str: + 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: diff --git a/packages/kumo-engine/kumo/engine/_provenance.py b/packages/kumo-engine/kumo/engine/_provenance.py index e4ecdfd4f..12a3620f2 100644 --- a/packages/kumo-engine/kumo/engine/_provenance.py +++ b/packages/kumo-engine/kumo/engine/_provenance.py @@ -14,6 +14,8 @@ from typing import Any, Dict, Optional PROCESSING_JOB_CONFIG = Path('/opt/ml/config/processingjobconfig.json') +_LOCATED_CONNECTORS = frozenset( + {'LOCAL_FILE', 'S3', 'GCS', 'SNOWFLAKE', 'DATABRICKS', 'BIGQUERY'}) @cache @@ -105,10 +107,14 @@ 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() - backend = connector.source_backend - locate = (getattr(backend, 'table_uri', None) - or getattr(backend, 'table_location', None)) - return table_name if locate is None else locate(connector, table_name) + 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], @@ -129,8 +135,7 @@ def source_table(connector: Any, table_name: str, sql: Optional[str], for name, value in dependencies.items() } try: - kind = connector.source_type - entry['connector'] = str(getattr(kind, 'value', kind)) + entry['connector'] = _connector_kind(connector) entry['location'] = (source_location(connector, table_name, sql) if record_locations or sql is not None else table_name) diff --git a/packages/kumo-engine/test/test_snapshot_manifest.py b/packages/kumo-engine/test/test_snapshot_manifest.py index 71c3b22b0..58d519f3a 100644 --- a/packages/kumo-engine/test/test_snapshot_manifest.py +++ b/packages/kumo-engine/test/test_snapshot_manifest.py @@ -197,7 +197,8 @@ def test_unsupported_manifest_is_ignored(tmp_path, content): 'cat.sch.tt', None, '`cat`.`sch`.`tt`'), (BigQueryConnector(project_id='data-project', dataset_id='analytics'), 'tt', None, 'data-project.analytics.tt'), - (SimpleNamespace(config=None, source_backend=None), 'TT', None, 'TT'), + (SimpleNamespace(source_type='OTHER', source_backend=None), 'TT', None, + 'TT'), (None, 'synthetic', 'SELECT 1', 'sql:' + hashlib.sha256(b'SELECT 1').hexdigest()), ]) From e1d24f5478eed34730cdd58e6a76c43f9f152e8e Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Wed, 21 Oct 2026 22:15:02 +0000 Subject: [PATCH 13/80] Record Custom SQL text by default; one record_sources switch The snapshot manifest recorded Custom SQL only as a hash, so a model's lineage could not show which query produced a source. Custom SQL entries now carry the SQL text and its dialect next to the hash and dependency versions. graph.snapshot(record_locations=False) becomes record_sources=False and turns off both: connector and table name only, and the SQL as a hash only. Verified with test_snapshot_manifest.py. Co-Authored-By: Claude Opus 5.5 --- .../kumo-engine/kumo/engine/_provenance.py | 7 ++++-- packages/kumo-engine/kumo/engine/graph.py | 24 +++++++++---------- .../test/test_snapshot_manifest.py | 20 +++++++++------- 3 files changed, 29 insertions(+), 22 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/_provenance.py b/packages/kumo-engine/kumo/engine/_provenance.py index 12a3620f2..c4b2a1d19 100644 --- a/packages/kumo-engine/kumo/engine/_provenance.py +++ b/packages/kumo-engine/kumo/engine/_provenance.py @@ -119,7 +119,7 @@ def _connector_kind(connector: Any) -> str: def source_table(connector: Any, table_name: str, sql: Optional[str], version: Any, *, - record_locations: bool = True) -> Dict[str, Any]: + record_sources: bool = True) -> Dict[str, Any]: from kumo_core.sources.identity import version_kind entry: Dict[str, Any] = { @@ -137,8 +137,11 @@ def source_table(connector: Any, table_name: str, sql: Optional[str], try: entry['connector'] = _connector_kind(connector) entry['location'] = (source_location(connector, table_name, sql) - if record_locations or sql is not None else + if record_sources or sql is not None else table_name) + if sql is not None and record_sources: + entry['sql'] = sql + entry['dialect'] = entry['connector'].lower() except Exception as error: warnings.warn(f"Could not record a source table's location: {error}", UserWarning, stacklevel=3) diff --git a/packages/kumo-engine/kumo/engine/graph.py b/packages/kumo-engine/kumo/engine/graph.py index 428f0bbc3..d8775ee81 100644 --- a/packages/kumo-engine/kumo/engine/graph.py +++ b/packages/kumo-engine/kumo/engine/graph.py @@ -246,16 +246,16 @@ def counts(self) -> Dict[str, Dict[str, int]]: def _write_snapshot_manifest(result: Any, snapshot_id: str, source_tables: Dict[str, Dict[str, Any]], - record_locations: bool, *, fresh: bool) -> None: + 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) - if existing is None or existing.get('record_locations', - True) == record_locations: + if existing is None or existing.get('record_sources', + True) == record_sources: return manifest = { - **existing, 'record_locations': record_locations, + **existing, 'record_sources': record_sources, 'parents': { **existing['parents'], 'source_tables': source_tables } @@ -265,7 +265,7 @@ def _write_snapshot_manifest(result: Any, snapshot_id: str, scheduler=None, reconstructed=True) manifest = new_manifest('graph_snapshot', snapshot_id, {'source_tables': source_tables}, - record_locations=record_locations, + record_sources=record_sources, **fields) write_manifest(path, manifest) except Exception as error: @@ -408,7 +408,7 @@ def validate(self) -> None: """ self.native_graph - def snapshot(self, *, record_locations: bool = True) -> GraphSnapshot: + def snapshot(self, *, record_sources: bool = True) -> 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` @@ -419,9 +419,9 @@ def snapshot(self, *, record_locations: bool = True) -> GraphSnapshot: it's reused as is. Args: - record_locations: if False, the snapshot manifest records each + record_sources: if False, the snapshot manifest records each source as connector and table name only (no bucket, path or - account). + account) and Custom SQL as a hash only (no SQL text). """ from filelock import FileLock @@ -452,10 +452,10 @@ def snapshot(self, *, record_locations: bool = True) -> GraphSnapshot: # lock behind after a crash. with FileLock(f'{persist_dir}.lock'): return self._snapshot_no_lock(graph_key, persist_dir, - record_locations) + record_sources) def _snapshot_no_lock(self, graph_key: str, persist_dir: Path, - record_locations: bool = True) -> GraphSnapshot: + 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, @@ -523,7 +523,7 @@ def _snapshot_no_lock(self, graph_key: str, persist_dir: Path, table_data_versions[alias] = prepared.version source_tables[alias] = source_table( connector, name, sql, prepared.version, - record_locations=record_locations) + record_sources=record_sources) paths = prepared.source.get('paths') or [] if paths: @@ -574,7 +574,7 @@ def _snapshot_no_lock(self, graph_key: str, persist_dir: Path, snapshot_id = f'{graph_key}@{data_version[:16]}' _write_snapshot_manifest(result, snapshot_id, source_tables, - record_locations, fresh=existing is None) + record_sources, fresh=existing is None) _warn_ingest_data_loss(result, graph_config) return GraphSnapshot( diff --git a/packages/kumo-engine/test/test_snapshot_manifest.py b/packages/kumo-engine/test/test_snapshot_manifest.py index 58d519f3a..bffad8017 100644 --- a/packages/kumo-engine/test/test_snapshot_manifest.py +++ b/packages/kumo-engine/test/test_snapshot_manifest.py @@ -213,30 +213,34 @@ def test_unrecordable_location_is_null(): 'version_kind': 'unknown'} -def test_custom_sql_records_dependency_versions(): +@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_locations=False) - assert entry == { + 'q', 'SELECT * FROM orders', version, record_sources=record_sources) + 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') + assert entry == expected -def test_record_locations_false_keeps_table_names_only(make_snapshot, - tmp_path): +def test_record_sources_false_keeps_table_names_only(make_snapshot, + tmp_path): names = {'ents': 'ents', 'events': 'events.parquet'} - redacted = make_snapshot(30, record_locations=False) + 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) - assert redacted.manifest['record_locations'] is False + assert redacted.manifest['record_sources'] is False - full = make_snapshot(record_locations=True) + full = make_snapshot(record_sources=True) assert full == redacted assert full.manifest['parents']['source_tables']['ents'][ 'location'] == str((tmp_path / 'data' / 'ents.parquet').resolve()) From 3fae7ca8559d6138cd87c6818c79523295c94c52 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Wed, 21 Oct 2026 22:27:19 +0000 Subject: [PATCH 14/80] Stamp the engine build commit into wheels; type directory versions Manifests from an installed kumo-engine wheel recorded build_commit null: only a "+sha" local version or a git checkout supplied it. The wheel build now appends BUILD_COMMIT to kumo/engine/_version.py in the build directory (KUMO_ENGINE_BUILD_COMMIT, else git, else None), and environment() reads it. The source file is unchanged. A Parquet directory's version is a sha256 fingerprint of its parts' paths, sizes and mtimes, but was typed "timestamp" like a single file's mtime. It is now "listing". Verified with test_snapshot_manifest.py, kumo-core tests/sources/parquet/test_object_store.py and a build_py run with and without the variable. Co-Authored-By: Claude Opus 5.5 --- .../src/kumo_core/sources/identity.py | 2 +- .../src/kumo_core/sources/object_store.py | 2 +- .../sources/parquet/test_object_store.py | 4 +-- .../kumo-engine/kumo/engine/_provenance.py | 4 +++ packages/kumo-engine/setup.py | 30 +++++++++++++++++++ 5 files changed, 38 insertions(+), 4 deletions(-) diff --git a/packages/kumo-core/src/kumo_core/sources/identity.py b/packages/kumo-core/src/kumo_core/sources/identity.py index 989fd6629..0e3d8e664 100644 --- a/packages/kumo-core/src/kumo_core/sources/identity.py +++ b/packages/kumo-core/src/kumo_core/sources/identity.py @@ -22,7 +22,7 @@ class SourceVersion(str): """A version string that also says what kind of evidence it is: - ``exact``, ``timestamp``, ``hour_bucket`` or ``unknown``.""" + ``exact``, ``timestamp``, ``listing``, ``hour_bucket`` or ``unknown``.""" 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 09c11a32b..828600080 100644 --- a/packages/kumo-core/src/kumo_core/sources/object_store.py +++ b/packages/kumo-core/src/kumo_core/sources/object_store.py @@ -110,7 +110,7 @@ def _source_version(path: str, entries: Sequence[Any]) -> str: for entry in sorted(entries, key=lambda entry: entry.path) ] return SourceVersion( - hashlib.sha256(json.dumps(metadata).encode()).hexdigest(), "timestamp") + hashlib.sha256(json.dumps(metadata).encode()).hexdigest(), "listing") class _CsvSource(NamedTuple): 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-engine/kumo/engine/_provenance.py b/packages/kumo-engine/kumo/engine/_provenance.py index c4b2a1d19..2bd279372 100644 --- a/packages/kumo-engine/kumo/engine/_provenance.py +++ b/packages/kumo-engine/kumo/engine/_provenance.py @@ -22,6 +22,10 @@ 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 diff --git a/packages/kumo-engine/setup.py b/packages/kumo-engine/setup.py index 2c174c3c3..be32d63a4 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", From 5656515681b366f682b6b0346c19410e364664e3 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Thu, 22 Oct 2026 01:10:17 +0000 Subject: [PATCH 15/80] Carry the record_sources choice on GraphSnapshot A custom training table built against a snapshot read the snapshot's record_sources choice back from its manifest file, defaulting to recording full source locations when that file was missing or unreadable. GraphSnapshot.record_sources now holds the choice the snapshot was taken with (None when unknown, e.g. a hand-built or loaded snapshot). It is not part of equality: the same index taken with either choice is the same snapshot. Verified with test_snapshot_manifest.py. Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/graph.py | 4 ++++ packages/kumo-engine/test/test_snapshot_manifest.py | 2 ++ 2 files changed, 6 insertions(+) diff --git a/packages/kumo-engine/kumo/engine/graph.py b/packages/kumo-engine/kumo/engine/graph.py index d8775ee81..c10efb40c 100644 --- a/packages/kumo-engine/kumo/engine/graph.py +++ b/packages/kumo-engine/kumo/engine/graph.py @@ -220,11 +220,14 @@ 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 manifest(self) -> Dict[str, Any] | None: @@ -581,6 +584,7 @@ def _snapshot_no_lock(self, graph_key: str, persist_dir: Path, 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/test/test_snapshot_manifest.py b/packages/kumo-engine/test/test_snapshot_manifest.py index bffad8017..bd8f10564 100644 --- a/packages/kumo-engine/test/test_snapshot_manifest.py +++ b/packages/kumo-engine/test/test_snapshot_manifest.py @@ -239,9 +239,11 @@ def test_record_sources_false_keeps_table_names_only(make_snapshot, assert {alias: e['location'] for alias, e in sources.items()} == names assert str(tmp_path) not in json.dumps(redacted.manifest) 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'] From 72308cf4c01d2aa3247e8f5bca0fe2c77acd4d22 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Thu, 22 Oct 2026 01:11:05 +0000 Subject: [PATCH 16/80] Record Custom SQL's namespace and normalized fingerprint A Custom SQL source recorded its exact text, its dialect and a hash of the exact text. Unqualified table names in that text resolve in the connector's database/schema (catalog/schema, project/dataset), which was not recorded, and a whitespace- or comment-only edit changed the only fingerprint. The source record now adds: - namespace: the connector's [database, schema] the query resolved in (source detail, omitted with record_sources=False); - normalized_sql_fingerprint: the existing Custom SQL identity (custom_sql_synthetic_name, the name the query's table gets), which ignores formatting; kept with record_sources=False as it is a hash. Verified with test_snapshot_manifest.py. Co-Authored-By: Claude Opus 5.5 --- .../kumo-engine/kumo/engine/_provenance.py | 24 ++++++++++++++++--- .../test/test_snapshot_manifest.py | 12 +++++++++- 2 files changed, 32 insertions(+), 4 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/_provenance.py b/packages/kumo-engine/kumo/engine/_provenance.py index 2bd279372..500ff8c69 100644 --- a/packages/kumo-engine/kumo/engine/_provenance.py +++ b/packages/kumo-engine/kumo/engine/_provenance.py @@ -143,10 +143,28 @@ def source_table(connector: Any, table_name: str, sql: Optional[str], entry['location'] = (source_location(connector, table_name, sql) if record_sources or sql is not None else table_name) - if sql is not None and record_sources: - entry['sql'] = sql - entry['dialect'] = entry['connector'].lower() + 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 entry + + +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/test/test_snapshot_manifest.py b/packages/kumo-engine/test/test_snapshot_manifest.py index bd8f10564..2c879fa60 100644 --- a/packages/kumo-engine/test/test_snapshot_manifest.py +++ b/packages/kumo-engine/test/test_snapshot_manifest.py @@ -11,7 +11,9 @@ 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 @@ -227,7 +229,15 @@ def test_custom_sql_records_text_unless_sources_are_off(record_sources): 'dependencies': {'orders': '17'}, } if record_sources: - expected.update(sql='SELECT * FROM orders', dialect='snowflake') + 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 From 7b6ad7905c33df9c409e20cfddb6bd9584f0db61 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Thu, 22 Oct 2026 18:40:25 +0000 Subject: [PATCH 17/80] Make the snapshot manifest a frozen dataclass Manifests were plain dicts, so readers indexed string keys with no names, completion or type checking. kumo.engine._manifest now defines frozen dataclasses for the shared envelope (Manifest), Environment, SchedulerContext, SourceTableEntry and SnapshotManifest, each with to_dict()/from_dict(). Keys a newer kumo-engine writes are kept in `extra` and written back, so a newer file round-trips byte for byte; the on-disk JSON is unchanged. GraphSnapshot.manifest returns a SnapshotManifest, and environment(), scheduler_context() and source_table() return the dataclasses. Verified with test_snapshot_manifest.py. Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/_manifest.py | 170 ++++++++++++++++-- .../kumo-engine/kumo/engine/_provenance.py | 68 +++---- packages/kumo-engine/kumo/engine/graph.py | 48 ++--- .../test/test_snapshot_manifest.py | 93 ++++++---- 4 files changed, 275 insertions(+), 104 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/_manifest.py b/packages/kumo-engine/kumo/engine/_manifest.py index ca75b7dcb..e4466f3a9 100644 --- a/packages/kumo-engine/kumo/engine/_manifest.py +++ b/packages/kumo-engine/kumo/engine/_manifest.py @@ -4,38 +4,169 @@ """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 from datetime import datetime, timezone from pathlib import Path -from typing import Any, Dict, Mapping, Optional +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 new_manifest(kind: str, id: str, parents: Mapping[str, Any], - **fields: Any) -> Dict[str, Any]: + +def omit_none() -> Any: + """A field written only when it is not None.""" + 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'} + 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 + }) + + +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, 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 SchedulerContext(Record): + platform: str + job_id: Optional[str] + job_run_id: 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() + normalized_sql_fingerprint: Optional[str] = omit_none() + sql: Optional[str] = omit_none() + dialect: Optional[str] = omit_none() + namespace: Optional[List[Any]] = omit_none() + + +@dataclass(frozen=True, kw_only=True) +class Manifest(Record): + kind: str + schema_version: int = SCHEMA_VERSION + id: str + created_at: Optional[str] + environment: Optional[Environment] + scheduler: Optional[SchedulerContext] + reconstructed: Optional[bool] = omit_none() + + +def stamp() -> Dict[str, Any]: + """Envelope fields for a manifest written now.""" from kumo.engine._provenance import environment, scheduler_context - return { - 'kind': kind, - 'schema_version': SCHEMA_VERSION, - 'id': id, - 'parents': dict(parents), - 'created_at': datetime.now(timezone.utc).isoformat(timespec='seconds'), - 'environment': environment(), - 'scheduler': scheduler_context(), - **fields, - } + return dict( + created_at=datetime.now(timezone.utc).isoformat(timespec='seconds'), + environment=environment(), scheduler=scheduler_context()) + + +#: Envelope fields for a manifest rebuilt after the fact. +RECONSTRUCTED: Dict[str, Any] = dict(created_at=None, environment=None, + scheduler=None, reconstructed=True) + +@dataclass(frozen=True, kw_only=True) +class SnapshotParents(Record): + source_tables: Dict[str, SourceTableEntry] -def write_manifest(path: Path, manifest: Mapping[str, Any]) -> None: - payload = json.dumps(manifest, indent=2, sort_keys=True) + '\n' + +@dataclass(frozen=True, kw_only=True) +class SnapshotManifest(Manifest): + kind: str = 'graph_snapshot' + parents: SnapshotParents + record_sources: Optional[bool] = omit_none() + + +def write_manifest(path: Path, manifest: Manifest) -> None: + payload = json.dumps(manifest.to_dict(), indent=2, sort_keys=True) + '\n' atomic_write_bytes(path, payload.encode('utf-8')) -def read_manifest(path: Path) -> Optional[Dict[str, Any]]: +def read_manifest(path: Path, kind: Type[M]) -> Optional[M]: try: manifest = json.loads(path.read_text(encoding='utf-8')) except FileNotFoundError: @@ -54,4 +185,9 @@ def read_manifest(path: Path) -> Optional[Dict[str, Any]]: stacklevel=3, ) return None - return manifest + 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/_provenance.py b/packages/kumo-engine/kumo/engine/_provenance.py index 500ff8c69..e463b68e6 100644 --- a/packages/kumo-engine/kumo/engine/_provenance.py +++ b/packages/kumo-engine/kumo/engine/_provenance.py @@ -13,6 +13,12 @@ from pathlib import Path from typing import Any, Dict, Optional +from kumo.engine._manifest import ( + Environment, + SchedulerContext, + SourceTableEntry, +) + PROCESSING_JOB_CONFIG = Path('/opt/ml/config/processingjobconfig.json') _LOCATED_CONNECTORS = frozenset( {'LOCAL_FILE', 'S3', 'GCS', 'SNOWFLAKE', 'DATABRICKS', 'BIGQUERY'}) @@ -45,7 +51,7 @@ def _diskgraph() -> Any: return None -def environment() -> Dict[str, Optional[str]]: +def environment() -> Environment: from kumo_core.version import ( get_kumo_version, installed_distribution_version, @@ -54,24 +60,24 @@ def environment() -> Dict[str, Optional[str]]: from kumo.engine._version import __version__ diskgraph = _diskgraph() - return { - '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(), - } + 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 _context(platform_name: str, job_id: Optional[str], - job_run_id: Optional[str]) -> Dict[str, Any]: - return {'platform': platform_name, 'job_id': job_id, - 'job_run_id': job_run_id} + job_run_id: Optional[str]) -> SchedulerContext: + return SchedulerContext(platform=platform_name, job_id=job_id, + job_run_id=job_run_id) -def _detect() -> Optional[Dict[str, Any]]: +def _detect() -> Optional[SchedulerContext]: env = os.environ if env.get('TRAINING_JOB_NAME'): return _context('sagemaker', env['TRAINING_JOB_NAME'], @@ -91,7 +97,7 @@ def _detect() -> Optional[Dict[str, Any]]: return None -def scheduler_context() -> Optional[Dict[str, Any]]: +def scheduler_context() -> Optional[SchedulerContext]: env = os.environ scheduler = env.get('KUMO_SCHEDULER') job_id = env.get('KUMO_JOB_ID') @@ -99,10 +105,10 @@ def scheduler_context() -> Optional[Dict[str, Any]]: detected = _detect() if scheduler or job_id or job_run_id: if detected is not None and scheduler in (None, '', - detected['platform']): - scheduler = detected['platform'] - job_id = job_id or detected['job_id'] - job_run_id = job_run_id or detected['job_run_id'] + detected.platform): + scheduler = detected.platform + job_id = job_id or detected.job_id + job_run_id = job_run_id or detected.job_run_id return _context(scheduler or 'other', job_id, job_run_id) return detected @@ -123,21 +129,11 @@ def _connector_kind(connector: Any) -> str: def source_table(connector: Any, table_name: str, sql: Optional[str], version: Any, *, - record_sources: bool = True) -> Dict[str, Any]: + record_sources: bool = True) -> SourceTableEntry: from kumo_core.sources.identity import version_kind - entry: Dict[str, Any] = { - 'connector': None, - 'location': None, - 'version': None if version is None else str(version), - 'version_kind': version_kind(version), - } + entry: Dict[str, Any] = {'connector': None, 'location': None} dependencies = getattr(version, 'dependencies', None) - if dependencies is not None: - entry['dependencies'] = { - name: str(value) - for name, value in dependencies.items() - } try: entry['connector'] = _connector_kind(connector) entry['location'] = (source_location(connector, table_name, sql) @@ -148,7 +144,15 @@ def source_table(connector: Any, table_name: str, sql: Optional[str], except Exception as error: warnings.warn(f"Could not record a source table's location: {error}", UserWarning, stacklevel=3) - return entry + 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 _record_sql(entry: Dict[str, Any], connector: Any, sql: str, diff --git a/packages/kumo-engine/kumo/engine/graph.py b/packages/kumo-engine/kumo/engine/graph.py index c10efb40c..1e9274581 100644 --- a/packages/kumo-engine/kumo/engine/graph.py +++ b/packages/kumo-engine/kumo/engine/graph.py @@ -61,7 +61,7 @@ import json import warnings -from dataclasses import dataclass, field +from dataclasses import dataclass, field, replace from functools import cached_property from pathlib import Path from typing import Any, Dict, Mapping, Sequence @@ -69,7 +69,15 @@ from kumo_core.pquery import NativeGraph, normalize_graph from kumoapi.graph import Edge -from kumo.engine._manifest import new_manifest, read_manifest, write_manifest +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, @@ -230,9 +238,10 @@ class GraphSnapshot: record_sources: bool | None = field(default=None, compare=False) @property - def manifest(self) -> Dict[str, Any] | None: + 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) / SNAPSHOT_MANIFEST_FILE, + SnapshotManifest) @property def counts(self) -> Dict[str, Dict[str, int]]: @@ -248,28 +257,25 @@ def counts(self) -> Dict[str, Dict[str, int]]: def _write_snapshot_manifest(result: Any, snapshot_id: str, - source_tables: Dict[str, Dict[str, Any]], + 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) - if existing is None or existing.get('record_sources', - True) == record_sources: + existing = read_manifest(path, SnapshotManifest) + if existing is None or (existing.record_sources + is not False) == record_sources: return - manifest = { - **existing, 'record_sources': record_sources, - 'parents': { - **existing['parents'], 'source_tables': source_tables - } - } + manifest = replace( + existing, record_sources=record_sources, + parents=replace(existing.parents, + source_tables=source_tables)) else: - fields = {} if fresh else dict(created_at=None, environment=None, - scheduler=None, reconstructed=True) - manifest = new_manifest('graph_snapshot', snapshot_id, - {'source_tables': source_tables}, - record_sources=record_sources, - **fields) + 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: @@ -490,7 +496,7 @@ def _snapshot_no_lock(self, graph_key: str, persist_dir: Path, table_data_versions: Dict[str, str] = {} source_paths: Dict[str, str] = {} - source_tables: Dict[str, Dict[str, Any]] = {} + source_tables: Dict[str, SourceTableEntry] = {} try: for alias, table in self._named_tables.items(): diff --git a/packages/kumo-engine/test/test_snapshot_manifest.py b/packages/kumo-engine/test/test_snapshot_manifest.py index 2c879fa60..10a6e163e 100644 --- a/packages/kumo-engine/test/test_snapshot_manifest.py +++ b/packages/kumo-engine/test/test_snapshot_manifest.py @@ -5,6 +5,7 @@ import hashlib import json import os +from dataclasses import replace from datetime import datetime from types import SimpleNamespace @@ -28,6 +29,13 @@ SnowflakeConnector, Table, ) +from kumo.engine._manifest import ( + SchedulerContext, + SnapshotManifest, + SourceTableEntry, + read_manifest, + write_manifest, +) from kumo.engine.connector.identity import ( DatabricksIdentity, SnowflakeIdentity, @@ -85,32 +93,32 @@ def test_snapshot_manifest_records_sources_and_diffs(make_snapshot, before = make_snapshot(30) manifest = before.manifest - assert manifest['kind'] == 'graph_snapshot' - assert manifest['id'] == before.id + assert manifest.kind == 'graph_snapshot' + assert manifest.id == before.id data = (tmp_path / 'data').resolve() - assert {a: e['location'] - for a, e in manifest['parents']['source_tables'].items()} == { + 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() + assert {e.version_kind + for e in manifest.parents.source_tables.values() } == {'timestamp'} - assert 'tables' not in manifest and 'edges' not in manifest + assert manifest.extra == {} assert before.counts['rows'] == {'ents': 5, 'events': 30} assert before.counts['edges']['events_ENT_ID_to_ents_ID'] == 30 - assert manifest['environment']['kumo_engine'] == get_kumo_version() - assert manifest['environment']['diskgraph'] - assert manifest['scheduler'] is None - assert datetime.fromisoformat(manifest['created_at']).tzinfo is not None - assert str(tmp_path / 'scratch') not in json.dumps(manifest) + assert manifest.environment.kumo_engine == get_kumo_version() + assert manifest.environment.diskgraph + assert manifest.scheduler is None + 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) assert after.counts['rows']['events'] == 40 - sources = manifest['parents']['source_tables'] - new_sources = after.manifest['parents']['source_tables'] + 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'] + assert new_sources['events'].version != sources['events'].version def test_index_without_manifest_is_reused_and_backfilled( @@ -126,10 +134,9 @@ def test_index_without_manifest_is_reused_and_backfilled( reused = make_snapshot() assert reused == snapshot - assert reused.manifest == { - **original, 'created_at': None, 'environment': None, - 'scheduler': None, 'reconstructed': True - } + assert reused.manifest == replace(original, created_at=None, + environment=None, scheduler=None, + reconstructed=True) def test_manifest_write_failure_leaves_snapshot_unchanged( @@ -164,8 +171,8 @@ def test_scheduler_context(no_scheduler, env, processing, expected): context = provenance.scheduler_context() - assert context == (expected and dict( - zip(('platform', 'job_id', 'job_run_id'), expected))) + assert context == (expected and SchedulerContext( + **dict(zip(('platform', 'job_id', 'job_run_id'), expected)))) @pytest.mark.parametrize('content', [ @@ -211,8 +218,8 @@ def test_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 == {'connector': None, 'location': None, 'version': '7', - 'version_kind': 'unknown'} + assert entry == SourceTableEntry(connector=None, location=None, + version='7', version_kind='unknown') @pytest.mark.parametrize('record_sources', [True, False]) @@ -220,7 +227,8 @@ 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) + 'q', 'SELECT * FROM orders', version, + record_sources=record_sources).to_dict() expected = { 'connector': 'SNOWFLAKE', 'location': 'sql:' + hashlib.sha256(b'SELECT * FROM orders').hexdigest(), @@ -234,10 +242,10 @@ def test_custom_sql_records_text_unless_sources_are_off(record_sources): 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( + 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 reformatted.location != entry['location'] assert entry == expected @@ -245,15 +253,32 @@ 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) - assert redacted.manifest['record_sources'] is 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'] + 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_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'} From 27996c2c90f8fa4eb90334ed61ef31dc12063c6a Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sat, 24 Oct 2026 17:22:03 +0000 Subject: [PATCH 18/80] Add the MLOps switch; stub parents; drop scheduler context from manifests The series turns MLOps on per block or per call rather than through environment variables, every manifest file gets one name, and platforms record their own job ids. - kumo.MLOps(artifact_root, tracking=, registry=, export=) and the kumo.mlops(cfg) block (a contextvar: the nearest block, replaced by a call's mlops=; per thread and asyncio task). The root is a URI through connector_for_root (moved here) or a (connector, path) pair whose sink holds files there; a missing root or unknown export kind raises. - Manifests: kinds snapshot/training_table/model/prediction; a parent manifest with no record (or a bare id) reads as a stub with only kind, id and reconstructed set; embedded() drops a manifest's statistics and mutable fields where a child embeds it. The snapshot manifest is _provenance.json (kumo_core.artifacts.PROVENANCE_FILE). - Removed SchedulerContext, Manifest.scheduler and scheduler detection. GraphSnapshot.counts is private. - docs/mlops.md holds the series docs; the README links it. Verified: test_mlops.py (root forms, precedence, nesting, threads and tasks, stubs), test_snapshot_manifest.py, registry and import suites. Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/README.md | 1 + packages/kumo-engine/docs/mlops.md | 25 ++++ packages/kumo-engine/kumo/engine/__init__.py | 5 + packages/kumo-engine/kumo/engine/_manifest.py | 78 +++++++++--- packages/kumo-engine/kumo/engine/_mlops.py | 116 ++++++++++++++++++ .../kumo-engine/kumo/engine/_provenance.py | 53 +------- .../kumo/engine/connector/__init__.py | 19 +++ packages/kumo-engine/kumo/engine/graph.py | 8 +- packages/kumo-engine/test/test_mlops.py | 87 +++++++++++++ .../test/test_snapshot_manifest.py | 54 ++------ 10 files changed, 329 insertions(+), 117 deletions(-) create mode 100644 packages/kumo-engine/docs/mlops.md create mode 100644 packages/kumo-engine/kumo/engine/_mlops.py create mode 100644 packages/kumo-engine/test/test_mlops.py diff --git a/packages/kumo-engine/README.md b/packages/kumo-engine/README.md index 38f193819..0326c9d6c 100644 --- a/packages/kumo-engine/README.md +++ b/packages/kumo-engine/README.md @@ -92,6 +92,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..ffc175768 --- /dev/null +++ b/packages/kumo-engine/docs/mlops.md @@ -0,0 +1,25 @@ +# MLOps with Kumo Engine + +## Turning MLOps on + +`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(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. +- **`export`** names the artifact kinds stored under the root when they are produced: `snapshot`, `training_table`, `model` and `prediction` (default: all). +- **`tracking`** and **`registry`** record fits and prediction runs and register models. +- **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. + +## 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, and keeps Custom SQL only as a hash. diff --git a/packages/kumo-engine/kumo/engine/__init__.py b/packages/kumo-engine/kumo/engine/__init__.py index 4d89af512..64d4df37b 100644 --- a/packages/kumo-engine/kumo/engine/__init__.py +++ b/packages/kumo-engine/kumo/engine/__init__.py @@ -25,6 +25,8 @@ from kumoapi.plugin import ModelRegistryLocation from kumoapi.table import Column +from kumo.engine._manifest import SnapshotManifest +from kumo.engine._mlops import MLOps, mlops from kumo.engine._version import __version__ from kumo.engine.config import configure from kumo.engine.connector import ( @@ -83,6 +85,9 @@ "Edge", "Graph", "GraphSnapshot", + "MLOps", + "mlops", + "SnapshotManifest", "PredictiveQuery", "TrainingTable", "PredictionTable", diff --git a/packages/kumo-engine/kumo/engine/_manifest.py b/packages/kumo-engine/kumo/engine/_manifest.py index e4466f3a9..0b6c306f6 100644 --- a/packages/kumo-engine/kumo/engine/_manifest.py +++ b/packages/kumo-engine/kumo/engine/_manifest.py @@ -7,7 +7,7 @@ import types import warnings from collections.abc import Mapping as AbcMapping -from dataclasses import dataclass, field, fields +from dataclasses import dataclass, field, fields, replace from datetime import datetime, timezone from pathlib import Path from typing import ( @@ -84,6 +84,8 @@ def _typed(hint: Any, value: Any) -> Any: 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: @@ -104,13 +106,6 @@ class Environment(Record): python: Optional[str] -@dataclass(frozen=True, kw_only=True) -class SchedulerContext(Record): - platform: str - job_id: Optional[str] - job_run_id: Optional[str] - - @dataclass(frozen=True, kw_only=True) class SourceTableEntry(Record): connector: Optional[str] @@ -124,29 +119,75 @@ class SourceTableEntry(Record): namespace: Optional[List[Any]] = omit_none() +#: The artifact kinds, as manifests and the artifact root name them. +KINDS = ('snapshot', 'training_table', 'model', 'prediction') + + @dataclass(frozen=True, kw_only=True) class Manifest(Record): - kind: str + """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] - environment: Optional[Environment] - scheduler: Optional[SchedulerContext] + 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, scheduler_context + from kumo.engine._provenance import environment return dict( created_at=datetime.now(timezone.utc).isoformat(timespec='seconds'), - environment=environment(), scheduler=scheduler_context()) + environment=environment()) #: Envelope fields for a manifest rebuilt after the fact. RECONSTRUCTED: Dict[str, Any] = dict(created_at=None, environment=None, - scheduler=None, reconstructed=True) + 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: + if isinstance(value, str): + return hint(id=value, reconstructed=True) + 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) @@ -156,8 +197,11 @@ class SnapshotParents(Record): @dataclass(frozen=True, kw_only=True) class SnapshotManifest(Manifest): - kind: str = 'graph_snapshot' - parents: SnapshotParents + """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() diff --git a/packages/kumo-engine/kumo/engine/_mlops.py b/packages/kumo-engine/kumo/engine/_mlops.py new file mode 100644 index 000000000..20831531a --- /dev/null +++ b/packages/kumo-engine/kumo/engine/_mlops.py @@ -0,0 +1,116 @@ +# 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 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: a ``kumo.engine.tracking.Tracking`` that records fits and + prediction runs. + registry: where fits register models: a ``ModelOutputConfig`` with a + registry ``location``, or a ``ModelRegistryLocation`` registered + through the MLflow tracker. + export: the artifact kinds stored under the root when produced + (``snapshot``, ``training_table``, ``model``, ``prediction``); + ``True`` for all, ``False`` for none. + """ + + artifact_root: Union[str, Tuple[Any, str]] + tracking: Any = None + registry: Any = None + export: Union[bool, Iterable[str]] = 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 + + +#: The settings a resolved call runs with. +Settings = MLOps + +_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, model and prediction run is stored + under ``settings.artifact_root``, and fits and prediction runs are + tracked and registered as ``settings`` says. 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 index e463b68e6..0314818b8 100644 --- a/packages/kumo-engine/kumo/engine/_provenance.py +++ b/packages/kumo-engine/kumo/engine/_provenance.py @@ -1,11 +1,9 @@ # SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. # All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Runtime versions and scheduler job identity recorded in manifests.""" +"""Runtime versions and source tables recorded in manifests.""" import hashlib -import json -import os import platform import subprocess import warnings @@ -13,13 +11,8 @@ from pathlib import Path from typing import Any, Dict, Optional -from kumo.engine._manifest import ( - Environment, - SchedulerContext, - SourceTableEntry, -) +from kumo.engine._manifest import Environment, SourceTableEntry -PROCESSING_JOB_CONFIG = Path('/opt/ml/config/processingjobconfig.json') _LOCATED_CONNECTORS = frozenset( {'LOCAL_FILE', 'S3', 'GCS', 'SNOWFLAKE', 'DATABRICKS', 'BIGQUERY'}) @@ -71,48 +64,6 @@ def environment() -> Environment: ) -def _context(platform_name: str, job_id: Optional[str], - job_run_id: Optional[str]) -> SchedulerContext: - return SchedulerContext(platform=platform_name, job_id=job_id, - job_run_id=job_run_id) - - -def _detect() -> Optional[SchedulerContext]: - env = os.environ - if env.get('TRAINING_JOB_NAME'): - return _context('sagemaker', env['TRAINING_JOB_NAME'], - env.get('TRAINING_JOB_ARN')) - try: - processing = json.loads(PROCESSING_JOB_CONFIG.read_text()) - except (OSError, ValueError): - processing = None - if isinstance(processing, dict) and processing.get('ProcessingJobName'): - return _context('sagemaker', processing['ProcessingJobName'], - processing.get('ProcessingJobArn')) - service = env.get('SNOWFLAKE_SERVICE_NAME') - if service: - return _context('snowflake', None, service) - if env.get('DATABRICKS_RUNTIME_VERSION'): - return _context('databricks', None, None) - return None - - -def scheduler_context() -> Optional[SchedulerContext]: - env = os.environ - scheduler = env.get('KUMO_SCHEDULER') - job_id = env.get('KUMO_JOB_ID') - job_run_id = env.get('KUMO_JOB_RUN_ID') - detected = _detect() - if scheduler or job_id or job_run_id: - if detected is not None and scheduler in (None, '', - detected.platform): - scheduler = detected.platform - job_id = job_id or detected.job_id - job_run_id = job_run_id or detected.job_run_id - return _context(scheduler or 'other', job_id, job_run_id) - return detected - - def source_location(connector: Any, table_name: str, sql: Optional[str] = None) -> str: if sql is not None: diff --git a/packages/kumo-engine/kumo/engine/connector/__init__.py b/packages/kumo-engine/kumo/engine/connector/__init__.py index 64c869506..c3a21f4be 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 @@ -77,4 +95,5 @@ def file_store(connector: Any, path: str) -> Optional[Tuple[Any, str]]: "GCSIdentity", "BigQueryIdentity", "DatabricksIdentity", + "connector_for_root", ] diff --git a/packages/kumo-engine/kumo/engine/graph.py b/packages/kumo-engine/kumo/engine/graph.py index 1e9274581..6b25908ee 100644 --- a/packages/kumo-engine/kumo/engine/graph.py +++ b/packages/kumo-engine/kumo/engine/graph.py @@ -66,6 +66,7 @@ from pathlib import Path from typing import Any, Dict, Mapping, Sequence +from kumo_core.artifacts import PROVENANCE_FILE from kumo_core.pquery import NativeGraph, normalize_graph from kumoapi.graph import Edge @@ -87,7 +88,7 @@ from kumo.engine._storage import network_storage from kumo.engine.table import Table -SNAPSHOT_MANIFEST_FILE = 'snapshot-manifest.json' +SNAPSHOT_MANIFEST_FILE = PROVENANCE_FILE _warned_scratch: set[Path] = set() _warned_backfill: set[Path] = set() @@ -244,7 +245,7 @@ def manifest(self) -> SnapshotManifest | None: SnapshotManifest) @property - def counts(self) -> Dict[str, Dict[str, int]]: + def _counts(self) -> Dict[str, Dict[str, int]]: """Source rows per table and edges per edge type in this index.""" from kumo_core.diskgraph.snapshot import DiskgraphIngestResult @@ -268,8 +269,7 @@ def _write_snapshot_manifest(result: Any, snapshot_id: str, return manifest = replace( existing, record_sources=record_sources, - parents=replace(existing.parents, - source_tables=source_tables)) + parents=SnapshotParents(source_tables=source_tables)) else: manifest = SnapshotManifest( id=snapshot_id, diff --git a/packages/kumo-engine/test/test_mlops.py b/packages/kumo-engine/test/test_mlops.py new file mode 100644 index 000000000..9a76770c7 --- /dev/null +++ b/packages/kumo-engine/test/test_mlops.py @@ -0,0 +1,87 @@ +# 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', + 'prediction'} + + 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_snapshot_manifest.py b/packages/kumo-engine/test/test_snapshot_manifest.py index 10a6e163e..46c411771 100644 --- a/packages/kumo-engine/test/test_snapshot_manifest.py +++ b/packages/kumo-engine/test/test_snapshot_manifest.py @@ -30,7 +30,6 @@ Table, ) from kumo.engine._manifest import ( - SchedulerContext, SnapshotManifest, SourceTableEntry, read_manifest, @@ -42,23 +41,10 @@ ) from kumo.engine.graph import SNAPSHOT_MANIFEST_FILE, GraphSnapshot -_SCHEDULER_ENV = ('KUMO_SCHEDULER', 'KUMO_JOB_ID', 'KUMO_JOB_RUN_ID', - 'TRAINING_JOB_NAME', 'TRAINING_JOB_ARN', - 'SNOWFLAKE_SERVICE_NAME', 'DATABRICKS_RUNTIME_VERSION') - @pytest.fixture -def no_scheduler(tmp_path, monkeypatch): - for name in _SCHEDULER_ENV: - monkeypatch.delenv(name, raising=False) - monkeypatch.setattr(provenance, 'PROCESSING_JOB_CONFIG', - tmp_path / 'processingjobconfig.json') - return monkeypatch - - -@pytest.fixture -def make_snapshot(tmp_path, no_scheduler): - no_scheduler.setattr('kumo.engine._scratch._root', tmp_path / 'scratch') +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') @@ -93,7 +79,8 @@ def test_snapshot_manifest_records_sources_and_diffs(make_snapshot, before = make_snapshot(30) manifest = before.manifest - assert manifest.kind == 'graph_snapshot' + assert manifest.kind == 'snapshot' + assert SNAPSHOT_MANIFEST_FILE == '_provenance.json' assert manifest.id == before.id data = (tmp_path / 'data').resolve() assert {a: e.location @@ -105,16 +92,16 @@ def test_snapshot_manifest_records_sources_and_diffs(make_snapshot, for e in manifest.parents.source_tables.values() } == {'timestamp'} assert manifest.extra == {} - assert before.counts['rows'] == {'ents': 5, 'events': 30} - assert before.counts['edges']['events_ENT_ID_to_ents_ID'] == 30 + assert before._counts['rows'] == {'ents': 5, 'events': 30} + assert before._counts['edges']['events_ENT_ID_to_ents_ID'] == 30 assert manifest.environment.kumo_engine == get_kumo_version() assert manifest.environment.diskgraph - assert manifest.scheduler is None + 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) - assert after.counts['rows']['events'] == 40 + assert after._counts['rows']['events'] == 40 sources = manifest.parents.source_tables new_sources = after.manifest.parents.source_tables assert new_sources['ents'] == sources['ents'] @@ -135,7 +122,7 @@ def test_index_without_manifest_is_reused_and_backfilled( assert reused == snapshot assert reused.manifest == replace(original, created_at=None, - environment=None, scheduler=None, + environment=None, reconstructed=True) @@ -152,29 +139,6 @@ def fail(path, manifest): assert (snapshot.index_path / '_SUCCESS').is_file() -@pytest.mark.parametrize('env, processing, expected', [ - ({}, None, None), - ({'KUMO_SCHEDULER': 'databricks', 'KUMO_JOB_ID': '11', - 'KUMO_JOB_RUN_ID': '22', 'TRAINING_JOB_NAME': 'ignored'}, None, - ('databricks', '11', '22')), - ({'KUMO_JOB_RUN_ID': 'run-9', 'TRAINING_JOB_NAME': 'train-1'}, None, - ('sagemaker', 'train-1', 'run-9')), - ({}, {'ProcessingJobName': 'proc-1', 'ProcessingJobArn': 'arn:proc-1'}, - ('sagemaker', 'proc-1', 'arn:proc-1')), - ({'SNOWFLAKE_SERVICE_NAME': 'JOB_A1'}, None, ('snowflake', None, 'JOB_A1')), -]) -def test_scheduler_context(no_scheduler, env, processing, expected): - for name, value in env.items(): - no_scheduler.setenv(name, value) - if processing is not None: - provenance.PROCESSING_JOB_CONFIG.write_text(json.dumps(processing)) - - context = provenance.scheduler_context() - - assert context == (expected and SchedulerContext( - **dict(zip(('platform', 'job_id', 'job_run_id'), expected)))) - - @pytest.mark.parametrize('content', [ None, '{"schema_version": 2}', '{"schema_version": true}', '{"kind": ' ]) From 9d3ccf74d62fca95a92b6f0410b99676acc83803 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Tue, 20 Oct 2026 20:28:23 +0000 Subject: [PATCH 19/80] Record training-table lineage across export and reload A training table exported with TrainingTable.export and reloaded with from_source_table becomes an anonymous custom table: its tt_id, snapshot id, query and generation plan are lost, so a model trained on an exported or edited table cannot say where its rows came from. Generation (PredictiveQuery.generate_training_table and TrainingTable.from_source_table) now writes training-table-manifest.json into the staging directory: tt_id, parent snapshot id, query, generation plan, task type, rows per split, effective split seed, source (pql or custom), environment and scheduler. TrainingTable.tt_id and TrainingTable.manifest read it back (None for older tables). Directory exports to file and object-store sinks add _kumo_manifest.json, the manifest plus a content digest of the parts and the export identity, through a new sidecars argument of ObjectStoreSinkBackend.write_files that export_selected_files writes before _SUCCESS. Reloading that exact directory, with a matching _KUMO_EXPORT marker, records derived_from and whether the staged rows differ from the exported digest. All new writes and reads are best-effort: failures warn, nothing new raises, and exported part files are unchanged. Verified with new tests for the manifest fields, the sidecar contents and idempotent re-export, byte-identical parts versus the plain exporter, reload of an untouched, edited, sidecar-less, garbled or foreign export, a failing manifest write, a warehouse sink receiving no sidecars, and sidecar ordering in export_selected_files; plus the existing training-table export, custom-source, file-connector, snapshot-manifest, prediction export, registry, artifact-export and S3 suites. Co-Authored-By: Claude Opus 5.5 --- .../src/kumo_core/artifact_export.py | 13 +- .../src/kumo_core/object_store_sink.py | 29 +++- .../kumo-core/tests/test_artifact_export.py | 29 ++++ .../kumo/engine/pquery/predictive_query.py | 6 +- .../kumo/engine/pquery/training_table.py | 147 +++++++++++++++++- .../test/test_training_table_export.py | 144 +++++++++++++++++ 6 files changed, 356 insertions(+), 12 deletions(-) diff --git a/packages/kumo-core/src/kumo_core/artifact_export.py b/packages/kumo-core/src/kumo_core/artifact_export.py index 0d73620ca..8abd316d6 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) are written into a directory destination + after the data and before ``_SUCCESS``; ``fs`` then needs ``write_bytes``. """ 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,9 @@ 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(): + fs.write_bytes(osp.join(output_path, name), payload) return publish_export_destination(output_path, fs=fs) 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..066d6218e 100644 --- a/packages/kumo-core/src/kumo_core/object_store_sink.py +++ b/packages/kumo-core/src/kumo_core/object_store_sink.py @@ -130,8 +130,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): - """Export selected Parquet without collecting their rows.""" + export_id=None, extra_columns=None, + sidecars: Mapping[str, bytes] = {}): + """Export selected Parquet without collecting their rows. + + ``sidecars`` are written into a directory destination before + ``_SUCCESS``. + """ del project_id if write_mode is not WriteMode.OVERWRITE: raise NotImplementedError( @@ -144,9 +149,22 @@ def write_files(self, connector, table_name, files, _, 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 read_sidecar(self, connector: SourceConnector, table_name: str, + name: str) -> Optional[bytes]: + """Return the file ``name`` inside directory ``table_name``, if any.""" + from pyarrow.fs import FileType + + filesystem, path = self._destination(connector, table_name) + target = f"{path.rstrip('/')}/{name}" + if filesystem.get_file_info(target).type is not FileType.File: + return None + with filesystem.open_input_stream(target) as reader: + return reader.readall() + def publish_model( self, connector: SourceConnector, @@ -342,8 +360,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/tests/test_artifact_export.py b/packages/kumo-core/tests/test_artifact_export.py index 39aba03b2..726d6847c 100644 --- a/packages/kumo-core/tests/test_artifact_export.py +++ b/packages/kumo-core/tests/test_artifact_export.py @@ -222,3 +222,32 @@ def interrupted(source, destination, **kwargs): pd.testing.assert_frame_equal(pd.read_csv(target / 'part-00000.csv'), expected) assert (target / '_SUCCESS').exists() + + +@pytest.mark.parametrize('directory', [True, False]) +def test_sidecars_are_written_before_success(tmp_path, monkeypatch, + directory): + 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) + touched = fs.touch + + def touch(path): + assert (target / 'meta.json').read_bytes() == b'{}' + touched(path) + + monkeypatch.setattr(fs, 'touch', touch) + export_selected_files([str(source)], + str(target) + ('/' if directory else ''), fs=fs, + identity='job-1', sidecars={'meta.json': b'{}'}) + if directory: + assert (target / '_SUCCESS').exists() + else: + assert not (target / 'meta.json').exists() + assert pq.read_table(tmp_path / 'output.parquet').num_rows == 1 diff --git a/packages/kumo-engine/kumo/engine/pquery/predictive_query.py b/packages/kumo-engine/kumo/engine/pquery/predictive_query.py index 20fb51a9e..32736e86d 100644 --- a/packages/kumo-engine/kumo/engine/pquery/predictive_query.py +++ b/packages/kumo-engine/kumo/engine/pquery/predictive_query.py @@ -90,7 +90,7 @@ PredictionTable, _compile_prediction_table, ) -from kumo.engine.pquery.training_table import TrainingTable +from kumo.engine.pquery.training_table import TrainingTable, _record_manifest logger = logging.getLogger(__name__) @@ -308,7 +308,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 +316,8 @@ def generate_training_table( required_stages=task_plan.task_dataset.required_stages, generation_plan=generation_plan, ) + _record_manifest(table, snapshot, task_plan.train_table, 'pql') + return table 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 ae58aac63..9fab4ba7d 100644 --- a/packages/kumo-engine/kumo/engine/pquery/training_table.py +++ b/packages/kumo-engine/kumo/engine/pquery/training_table.py @@ -2,21 +2,32 @@ # All rights reserved. # SPDX-License-Identifier: Apache-2.0 +import hashlib +import json +import logging +import warnings from dataclasses import dataclass, replace from pathlib import Path -from typing import TYPE_CHECKING, Optional +from typing import TYPE_CHECKING, Any, Dict, Optional, Sequence, Tuple from kumo_core.model_plan import TrainingTableGenerationPlan from kumoapi.model_plan import \ TrainingTableGenerationPlan as PublicGenerationPlan from kumoapi.train import TrainingTableSpec +from kumo.engine._manifest import new_manifest, read_manifest, 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 +logger = logging.getLogger(__name__) + +STAGING_MANIFEST_FILE = 'training-table-manifest.json' +EXPORT_MANIFEST_FILE = '_kumo_manifest.json' +_RANDOM_SPLITS = ('RandomSplit', 'RandomLinkSplit') + @dataclass(frozen=True) class TrainingTable: @@ -38,6 +49,18 @@ class TrainingTable: #: Effective generation settings carried forward into ``Trainer.fit``. generation_plan: TrainingTableGenerationPlan + @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[Dict[str, Any]]: + """Lineage recorded at generation; None for older tables.""" + return read_manifest(self.staging_dir / STAGING_MANIFEST_FILE) + def export(self, output_config: "TrainingTableOutputConfig") -> str: r"""Write the staged rows with a ``SPLIT`` column (``0`` train, ``1`` validation, ``2`` test) for :meth:`from_source_table`. Returns the @@ -62,8 +85,40 @@ def export(self, output_config: "TrainingTableOutputConfig") -> str: root.mkdir(parents=True, exist_ok=True) with tempfile.TemporaryDirectory(dir=root) as output: files = export_train_table(str(self.staging_dir), output) + payload = self._export_manifest(connector, output_config, files, + export_id) + extra = ({} if payload is None else { + 'sidecars': { + EXPORT_MANIFEST_FILE: payload + } + }) 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, + **extra) + + def _export_manifest(self, connector, output_config, files: Sequence[str], + export_id: str) -> Optional[bytes]: + from kumo_core.object_store_sink import ObjectStoreSinkBackend + + if (not isinstance(connector.sink_backend, ObjectStoreSinkBackend) + or not (output_config.table_name.endswith('/') + or len(files) > 1)): + logger.debug('Training-table lineage is only recorded for ' + 'directory exports to files or object stores.') + return None + try: + manifest = self.manifest + if manifest is None: + return None + manifest.update(content_sha256=_content_digest(files), + export_identity=export_id) + return (json.dumps(manifest, indent=2, sort_keys=True) + + '\n').encode('utf-8') + except Exception as error: + warnings.warn( + f"Exporting without a training-table manifest: {error}", + UserWarning, stacklevel=3) + return None @classmethod def from_source_table( @@ -127,6 +182,7 @@ 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") + derived_from, modified = _export_lineage(source_table, rows_path) generation_plan = from_json_dict( to_json_dict(public_plan), TrainingTableGenerationPlan, @@ -183,7 +239,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 +247,88 @@ def from_source_table( required_stages=task_plan.task_dataset.required_stages, generation_plan=generation_plan, ) + _record_manifest(table, snapshot, task_plan.train_table, 'custom', + derived_from, modified=modified) + return table + + +def _record_manifest(table: TrainingTable, snapshot: GraphSnapshot, + request: Any, source: str, + derived_from: Optional[str] = None, + **fields: Any) -> None: + path = table.staging_dir / STAGING_MANIFEST_FILE + try: + from kumo_core.io.json import to_json_dict + from kumo_core.training.export import read_spec + + spec = read_spec(str(table.staging_dir)) + seed = None + if request.split_family in _RANDOM_SPLITS: + seed = (request.split_config or {}).get('rnd_seed') or 0 + parents = { + 'snapshot': snapshot.id, + 'derived_from': derived_from, + } + write_manifest( + path, + new_manifest( + 'training_table', spec['tt_id'], parents, + query=table.query, + generation_plan=to_json_dict(table.generation_plan), + task_type=request.task_type, + rows={ + stage['key']: stage['row_count'] + for stage in spec.get('split_stages') or [] + }, + split_seed=seed, + source=source, + **fields, + )) + except Exception as error: + warnings.warn( + f"Could not write the training-table manifest at {path}: " + f"{error}", UserWarning, stacklevel=3) + + +def _content_digest(paths: Sequence[Any]) -> str: + digests = [] + for path in paths: + digest = hashlib.sha256() + with open(path, 'rb') as handle: + for chunk in iter(lambda: handle.read(1 << 20), b''): + digest.update(chunk) + digests.append(digest.hexdigest()) + return hashlib.sha256(''.join(sorted(digests)).encode()).hexdigest() + + +def _export_lineage(source_table: SourceTable, + parts: Path) -> Tuple[Optional[str], Optional[bool]]: + from kumo_core.artifact_export import EXPORT_MARKER_NAME, export_marker + from kumo_core.object_store_sink import ObjectStoreSinkBackend + + connector = source_table.connector + try: + backend = connector.sink_backend + except NotImplementedError: + return None, None + if source_table.sql is not None or not isinstance( + backend, ObjectStoreSinkBackend): + return None, None + name = source_table.table_name + try: + payload = backend.read_sidecar(connector, name, EXPORT_MANIFEST_FILE) + if payload is None: + return None, None + exported = json.loads(payload) + marker = backend.read_sidecar(connector, name, EXPORT_MARKER_NAME) + if marker is None or marker.decode('utf-8') != export_marker( + exported['export_identity']): + logger.debug('%s does not belong to the export it records', name) + return None, None + digest = _content_digest(sorted(parts.glob('*.parquet'))) + return exported['id'], digest != exported['content_sha256'] + except Exception as error: + warnings.warn( + f"Ignoring the training-table manifest of {name!r}: {error}", + UserWarning, stacklevel=3) + return None, 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..296a3fcd7 100644 --- a/packages/kumo-engine/test/test_training_table_export.py +++ b/packages/kumo-engine/test/test_training_table_export.py @@ -2,13 +2,17 @@ # All rights reserved. # SPDX-License-Identifier: Apache-2.0 +import hashlib import json +import shutil from collections import Counter from functools import partial +from types import SimpleNamespace import numpy as np import pandas as pd import pyarrow as pa +import pyarrow.compute as pc import pyarrow.parquet as pq import pytest from kumoapi.graph import Edge @@ -292,3 +296,143 @@ 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 _digest(paths): + return hashlib.sha256(''.join( + sorted(hashlib.sha256(p.read_bytes()).hexdigest() + for p in paths)).encode()).hexdigest() + + +def _reload(graph, name, pquery): + _, snapshot, connector, _ = graph + return kumo.TrainingTable.from_source_table( + connector[name], TrainingTableSpec(split_col='SPLIT'), pq=pquery, + snapshot=snapshot) + + +def test_manifest_lineage_across_export_and_reload(graph, tmp_path): + from kumo_core.training.export import export_train_table, identity + + _, snapshot, connector, out = graph + query = 'PREDICT ents.LABEL FOR EACH ents.ID' + pquery, table, _ = _export(graph, query, 'lineage') + manifest = table.manifest + assert manifest['kind'] == 'training_table' + assert manifest['id'] == table.tt_id == _spec(table)['tt_id'] + assert manifest['parents'] == { + 'snapshot': snapshot.id, + 'derived_from': None + } + assert (manifest['query'], manifest['source']) == (query, 'pql') + assert manifest['task_type'] == 'binary_classification' + assert manifest['rows'] == _seeds(table)[1] + assert manifest['split_seed'] == 0 + assert manifest['generation_plan']['split'] == 'inferred' + + directory = out / 'lineage' + parts = sorted(directory.glob('*.parquet')) + baseline = tmp_path / 'baseline' + expected = export_train_table(str(table.staging_dir), str(baseline)) + assert [p.name for p in parts] == [p.split('/')[-1] for p in expected] + assert [p.read_bytes() for p in parts] == [ + (baseline / p.name).read_bytes() for p in parts + ] + assert sorted(p.name for p in directory.iterdir()) == sorted( + [p.name for p in parts] + + ['_KUMO_EXPORT', '_SUCCESS', '_kumo_manifest.json']) + sidecar = (directory / '_kumo_manifest.json').read_bytes() + assert json.loads(sidecar) == { + **manifest, 'content_sha256': _digest(parts), + 'export_identity': identity(_spec(table)) + } + + table.export(TrainingTableOutputConfig(connector, 'lineage')) + assert (directory / '_kumo_manifest.json').read_bytes() == sidecar + assert [p.read_bytes() for p in parts] == [ + (baseline / p.name).read_bytes() for p in parts + ] + + reloaded = _reload(graph, 'lineage', pquery).manifest + assert reloaded['parents'] == { + 'snapshot': snapshot.id, + 'derived_from': table.tt_id + } + assert (reloaded['source'], reloaded['modified']) == ('custom', False) + assert reloaded['id'] != table.tt_id + + edited = out / 'lineage_edited' + shutil.copytree(directory, edited) + rows = pq.read_table(parts[0]) + flipped = pc.invert(rows['TARGET1']) + pq.write_table(rows.set_column(rows.schema.get_field_index('TARGET1'), + 'TARGET1', flipped), + edited / parts[0].name) + edited_manifest = _reload(graph, 'lineage_edited', pquery).manifest + assert edited_manifest['parents']['derived_from'] == table.tt_id + assert edited_manifest['modified'] is True + + +@pytest.mark.parametrize('damage', ['missing', 'garbled', 'foreign']) +def test_reload_without_matching_sidecar_records_no_lineage(graph, damage): + _, snapshot, connector, out = graph + pquery, table, _ = _export(graph, 'PREDICT ents.LABEL FOR EACH ents.ID', + f'unlinked_base_{damage}') + directory = out / f'unlinked_{damage}' + shutil.copytree(out / f'unlinked_base_{damage}', directory) + sidecar = directory / '_kumo_manifest.json' + if damage == 'missing': + sidecar.unlink() + elif damage == 'garbled': + sidecar.write_text('{not json') + else: + (directory / '_KUMO_EXPORT').write_text('another export') + + if damage == 'garbled': + with pytest.warns(UserWarning, match='training-table manifest'): + reloaded = _reload(graph, directory.name, pquery) + else: + reloaded = _reload(graph, directory.name, pquery) + assert reloaded.manifest['parents'] == { + 'snapshot': snapshot.id, + 'derived_from': None + } + assert reloaded.manifest['modified'] is None + assert _seeds(reloaded)[1] == _seeds(table)[1] + + +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'): + pquery, table, exported = _export( + graph, 'PREDICT ents.LABEL FOR EACH ents.ID', 'no_manifest') + assert table.manifest is None + assert not (out / 'no_manifest' / '_kumo_manifest.json').exists() + assert (out / 'no_manifest' / '_SUCCESS').exists() + assert exported.num_rows == sum(_seeds(table)[1].values()) + + +def test_warehouse_sink_receives_no_sidecars(graph): + graph_, snapshot, _, _ = 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) + return table_name + + connector = SimpleNamespace( + source_type='SNOWFLAKE', + sink_backend=SimpleNamespace(write_files=write_files)) + table.export(TrainingTableOutputConfig(connector, 'DB.SCHEMA.TT')) + assert calls == [{'export_id': f'{table.tt_id}:tt-export-v1'}] From 5f0178f21d5ec502b1da7b7c3fc0d28524d0da5a Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Tue, 20 Oct 2026 20:40:33 +0000 Subject: [PATCH 20/80] Digest exported parts at the source for training-table lineage Reload compared the export digest against the staged copy, but staging re-encodes some Parquet (narrow integer keys, list columns, dictionary columns), so an untouched export could read as modified. The digest is now read from the export directory itself: ObjectStoreSinkBackend.content_digest streams the visible Parquet parts through the connector filesystem, and artifact_export.content_sha256 is the one formula export and reload share. Sink lookup moves inside the best-effort guard, a missing sidecar or destination logs at debug, CSV directory exports get no sidecar, and _record_manifest lineage arguments are keyword-only. Verified with a reload of an export carrying a list weight column (staged bytes differ, modified stays false), a CSV directory export, and the training-table export, custom-source and artifact-export suites. Co-Authored-By: Claude Opus 5.5 --- .../src/kumo_core/artifact_export.py | 5 ++ .../src/kumo_core/object_store_sink.py | 24 +++++++++ .../kumo/engine/pquery/training_table.py | 49 ++++++++++--------- .../test/test_training_table_export.py | 46 +++++++++++++++++ 4 files changed, 102 insertions(+), 22 deletions(-) diff --git a/packages/kumo-core/src/kumo_core/artifact_export.py b/packages/kumo-core/src/kumo_core/artifact_export.py index 8abd316d6..bb34aed00 100644 --- a/packages/kumo-core/src/kumo_core/artifact_export.py +++ b/packages/kumo-core/src/kumo_core/artifact_export.py @@ -139,6 +139,11 @@ def _renderable(column: "pd.Series") -> bool: or pa.types.is_fixed_size_list(arrow)) +def content_sha256(part_digests) -> str: + """Digest of an export's parts from their sha256 hex digests, any order.""" + return hashlib.sha256(''.join(sorted(part_digests)).encode()).hexdigest() + + class ArtifactDestinationConflict(ValueError): """A directory belongs to another export.""" 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 066d6218e..658da3609 100644 --- a/packages/kumo-core/src/kumo_core/object_store_sink.py +++ b/packages/kumo-core/src/kumo_core/object_store_sink.py @@ -153,6 +153,30 @@ def write_files(self, connector, table_name, files, sidecars=sidecars) return f'{self.scheme}://{output}' + def content_digest(self, connector: SourceConnector, + table_name: str) -> str: + """``content_sha256`` of the visible Parquet parts in ``table_name``.""" + import hashlib + + from kumo_core.artifact_export import content_sha256 + from kumo_core.sources import layout + from pyarrow.fs import FileSelector, FileType + + filesystem, path = self._destination(connector, table_name) + digests = [] + for info in filesystem.get_file_info( + FileSelector(path.rstrip('/'), allow_not_found=True)): + if (info.type is not FileType.File + or layout.is_hidden(info.base_name) + or layout.file_type(info.base_name) != layout.PARQUET): + continue + digest = hashlib.sha256() + with filesystem.open_input_stream(info.path) as reader: + while chunk := reader.read(_CHUNK): + digest.update(chunk) + digests.append(digest.hexdigest()) + return content_sha256(digests) + def read_sidecar(self, connector: SourceConnector, table_name: str, name: str) -> Optional[bytes]: """Return the file ``name`` inside directory ``table_name``, if any.""" diff --git a/packages/kumo-engine/kumo/engine/pquery/training_table.py b/packages/kumo-engine/kumo/engine/pquery/training_table.py index 9fab4ba7d..b9e652857 100644 --- a/packages/kumo-engine/kumo/engine/pquery/training_table.py +++ b/packages/kumo-engine/kumo/engine/pquery/training_table.py @@ -2,7 +2,6 @@ # All rights reserved. # SPDX-License-Identifier: Apache-2.0 -import hashlib import json import logging import warnings @@ -99,10 +98,13 @@ def export(self, output_config: "TrainingTableOutputConfig") -> str: def _export_manifest(self, connector, output_config, files: Sequence[str], export_id: str) -> Optional[bytes]: from kumo_core.object_store_sink import ObjectStoreSinkBackend + from kumo_core.sources import layout + name = output_config.table_name if (not isinstance(connector.sink_backend, ObjectStoreSinkBackend) - or not (output_config.table_name.endswith('/') - or len(files) > 1)): + or not (name.endswith('/') or len(files) > 1) + or layout.file_type(name.rstrip('/')) not in ( + None, layout.PARQUET)): logger.debug('Training-table lineage is only recorded for ' 'directory exports to files or object stores.') return None @@ -182,7 +184,7 @@ 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") - derived_from, modified = _export_lineage(source_table, rows_path) + derived_from, modified = _export_lineage(source_table) generation_plan = from_json_dict( to_json_dict(public_plan), TrainingTableGenerationPlan, @@ -248,12 +250,12 @@ def from_source_table( generation_plan=generation_plan, ) _record_manifest(table, snapshot, task_plan.train_table, 'custom', - derived_from, modified=modified) + derived_from=derived_from, modified=modified) return table def _record_manifest(table: TrainingTable, snapshot: GraphSnapshot, - request: Any, source: str, + request: Any, source: str, *, derived_from: Optional[str] = None, **fields: Any) -> None: path = table.staging_dir / STAGING_MANIFEST_FILE @@ -290,42 +292,45 @@ def _record_manifest(table: TrainingTable, snapshot: GraphSnapshot, f"{error}", UserWarning, stacklevel=3) -def _content_digest(paths: Sequence[Any]) -> str: +def _content_digest(paths: Sequence[str]) -> str: + import hashlib + + from kumo_core.artifact_export import content_sha256 + digests = [] for path in paths: - digest = hashlib.sha256() with open(path, 'rb') as handle: - for chunk in iter(lambda: handle.read(1 << 20), b''): - digest.update(chunk) - digests.append(digest.hexdigest()) - return hashlib.sha256(''.join(sorted(digests)).encode()).hexdigest() + digests.append(hashlib.file_digest(handle, 'sha256').hexdigest()) + return content_sha256(digests) -def _export_lineage(source_table: SourceTable, - parts: Path) -> Tuple[Optional[str], Optional[bool]]: +def _export_lineage( + source_table: SourceTable) -> Tuple[Optional[str], Optional[bool]]: from kumo_core.artifact_export import EXPORT_MARKER_NAME, export_marker from kumo_core.object_store_sink import ObjectStoreSinkBackend connector = source_table.connector - try: - backend = connector.sink_backend - except NotImplementedError: - return None, None - if source_table.sql is not None or not isinstance( - backend, ObjectStoreSinkBackend): - return None, None name = source_table.table_name try: + backend = connector.sink_backend + if source_table.sql is not None or not isinstance( + backend, ObjectStoreSinkBackend): + return None, None payload = backend.read_sidecar(connector, name, EXPORT_MANIFEST_FILE) if payload is None: + logger.debug('%s has no training-table manifest', name) return None, None + except Exception as error: + logger.debug('No training-table manifest for %s: %s', name, error) + return None, None + try: exported = json.loads(payload) marker = backend.read_sidecar(connector, name, EXPORT_MARKER_NAME) if marker is None or marker.decode('utf-8') != export_marker( exported['export_identity']): logger.debug('%s does not belong to the export it records', name) return None, None - digest = _content_digest(sorted(parts.glob('*.parquet'))) + digest = backend.content_digest(connector, name) return exported['id'], digest != exported['content_sha256'] except Exception as error: warnings.warn( diff --git a/packages/kumo-engine/test/test_training_table_export.py b/packages/kumo-engine/test/test_training_table_export.py index 296a3fcd7..4fc947f67 100644 --- a/packages/kumo-engine/test/test_training_table_export.py +++ b/packages/kumo-engine/test/test_training_table_export.py @@ -436,3 +436,49 @@ def write_files(connector, table_name, files, write_mode, **kwargs): sink_backend=SimpleNamespace(write_files=write_files)) table.export(TrainingTableOutputConfig(connector, 'DB.SCHEMA.TT')) assert calls == [{'export_id': f'{table.tt_id}:tt-export-v1'}] + + +def test_reencoded_staging_does_not_mark_export_modified(graph, monkeypatch): + from kumo.engine.pquery import _custom_source + + _, snapshot, connector, out = graph + query = ('PREDICT LIST_DISTINCT(events.ITEM_ID, 0, 7, days) ' + 'RANK TOP 3 FOR EACH ents.ID') + pquery, table, _ = _export(graph, query, 'lineage_weighted') + directory = out / 'lineage_weighted' + parts = sorted(directory.glob('*.parquet')) + for part in parts: + rows = pq.read_table(part) + weights = pa.array([[1.0] * len(values) + for values in rows['TARGET1'].to_pylist()], + pa.list_(pa.float64())) + pq.write_table(rows.append_column('WEIGHT', weights), part) + sidecar = directory / '_kumo_manifest.json' + sidecar.write_text(json.dumps({ + **json.loads(sidecar.read_text()), 'content_sha256': _digest(parts) + })) + staged = [] + stage_source_table = _custom_source.stage_source_table + + def capture(*args, **kwargs): + staged.append(stage_source_table(*args, **kwargs)) + return staged[-1] + + monkeypatch.setattr(_custom_source, 'stage_source_table', capture) + manifest = kumo.TrainingTable.from_source_table( + connector['lineage_weighted'], + TrainingTableSpec(weight_col='WEIGHT', split_col='SPLIT'), + pq=pquery, snapshot=snapshot).manifest + assert _digest(sorted(staged[0].glob('*.parquet'))) != _digest(parts) + assert manifest['parents']['derived_from'] == table.tt_id + assert manifest['modified'] is False + + +def test_csv_directory_export_gets_no_sidecar(graph): + graph_, snapshot, connector, out = graph + table = kumo.PredictiveQuery( + graph=graph_, query='PREDICT ents.LABEL FOR EACH ents.ID' + ).generate_training_table(snapshot) + table.export(TrainingTableOutputConfig(connector, 'rows.csv/')) + names = {p.name for p in (out / 'rows.csv').iterdir()} + assert '_SUCCESS' in names and '_kumo_manifest.json' not in names From f1b261a2604f79504144e6130d043a3d360f5680 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Wed, 21 Oct 2026 06:56:27 +0000 Subject: [PATCH 21/80] Record a custom table's source instead of re-reading its export Reload lineage (derived_from, modified) read the export sidecar back and re-downloaded and hashed every part to compare against content_sha256: a second full remote read on top of staging, for a flag that could only be recorded. A custom table now records where its rows came from in parents.source_table: the full location (absolute path for a FileConnector, the full s3:// or gs:// URI for object stores, the fully qualified name for warehouse tables, a sha256 of the SQL for Custom SQL) and the connector's get_data_version stamp. Either is null with a warning when it cannot be determined. The export sidecar stays write-only and now holds just the training-table manifest; derived_from, modified, content_sha256, export_identity, ObjectStoreSinkBackend.read_sidecar and content_digest are removed. No engine code reads the sidecar. Verified with the training-table export, custom-source, snapshot-manifest and kumo-core artifact-export suites, including a FileConnector reload recording its path and version, failing location and version lookups recording nulls, and S3, warehouse and SQL location rendering. Co-Authored-By: Claude Opus 5.5 --- .../src/kumo_core/artifact_export.py | 5 - .../src/kumo_core/object_store_sink.py | 36 ---- .../kumo/engine/pquery/training_table.py | 124 ++++++-------- .../test/test_training_table_export.py | 162 +++++++----------- 4 files changed, 120 insertions(+), 207 deletions(-) diff --git a/packages/kumo-core/src/kumo_core/artifact_export.py b/packages/kumo-core/src/kumo_core/artifact_export.py index bb34aed00..8abd316d6 100644 --- a/packages/kumo-core/src/kumo_core/artifact_export.py +++ b/packages/kumo-core/src/kumo_core/artifact_export.py @@ -139,11 +139,6 @@ def _renderable(column: "pd.Series") -> bool: or pa.types.is_fixed_size_list(arrow)) -def content_sha256(part_digests) -> str: - """Digest of an export's parts from their sha256 hex digests, any order.""" - return hashlib.sha256(''.join(sorted(part_digests)).encode()).hexdigest() - - class ArtifactDestinationConflict(ValueError): """A directory belongs to another export.""" 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 658da3609..48f31e2d9 100644 --- a/packages/kumo-core/src/kumo_core/object_store_sink.py +++ b/packages/kumo-core/src/kumo_core/object_store_sink.py @@ -153,42 +153,6 @@ def write_files(self, connector, table_name, files, sidecars=sidecars) return f'{self.scheme}://{output}' - def content_digest(self, connector: SourceConnector, - table_name: str) -> str: - """``content_sha256`` of the visible Parquet parts in ``table_name``.""" - import hashlib - - from kumo_core.artifact_export import content_sha256 - from kumo_core.sources import layout - from pyarrow.fs import FileSelector, FileType - - filesystem, path = self._destination(connector, table_name) - digests = [] - for info in filesystem.get_file_info( - FileSelector(path.rstrip('/'), allow_not_found=True)): - if (info.type is not FileType.File - or layout.is_hidden(info.base_name) - or layout.file_type(info.base_name) != layout.PARQUET): - continue - digest = hashlib.sha256() - with filesystem.open_input_stream(info.path) as reader: - while chunk := reader.read(_CHUNK): - digest.update(chunk) - digests.append(digest.hexdigest()) - return content_sha256(digests) - - def read_sidecar(self, connector: SourceConnector, table_name: str, - name: str) -> Optional[bytes]: - """Return the file ``name`` inside directory ``table_name``, if any.""" - from pyarrow.fs import FileType - - filesystem, path = self._destination(connector, table_name) - target = f"{path.rstrip('/')}/{name}" - if filesystem.get_file_info(target).type is not FileType.File: - return None - with filesystem.open_input_stream(target) as reader: - return reader.readall() - def publish_model( self, connector: SourceConnector, diff --git a/packages/kumo-engine/kumo/engine/pquery/training_table.py b/packages/kumo-engine/kumo/engine/pquery/training_table.py index b9e652857..924ef06ca 100644 --- a/packages/kumo-engine/kumo/engine/pquery/training_table.py +++ b/packages/kumo-engine/kumo/engine/pquery/training_table.py @@ -7,7 +7,7 @@ import warnings from dataclasses import dataclass, replace from pathlib import Path -from typing import TYPE_CHECKING, Any, Dict, Optional, Sequence, Tuple +from typing import TYPE_CHECKING, Any, Dict, Optional, Sequence from kumo_core.model_plan import TrainingTableGenerationPlan from kumoapi.model_plan import \ @@ -84,8 +84,7 @@ def export(self, output_config: "TrainingTableOutputConfig") -> str: root.mkdir(parents=True, exist_ok=True) with tempfile.TemporaryDirectory(dir=root) as output: files = export_train_table(str(self.staging_dir), output) - payload = self._export_manifest(connector, output_config, files, - export_id) + payload = self._export_manifest(connector, output_config, files) extra = ({} if payload is None else { 'sidecars': { EXPORT_MANIFEST_FILE: payload @@ -95,8 +94,8 @@ def export(self, output_config: "TrainingTableOutputConfig") -> str: output_config.write_mode, export_id=export_id, **extra) - def _export_manifest(self, connector, output_config, files: Sequence[str], - export_id: str) -> Optional[bytes]: + def _export_manifest(self, connector, output_config, + files: Sequence[str]) -> Optional[bytes]: from kumo_core.object_store_sink import ObjectStoreSinkBackend from kumo_core.sources import layout @@ -105,22 +104,14 @@ def _export_manifest(self, connector, output_config, files: Sequence[str], or not (name.endswith('/') or len(files) > 1) or layout.file_type(name.rstrip('/')) not in ( None, layout.PARQUET)): - logger.debug('Training-table lineage is only recorded for ' - 'directory exports to files or object stores.') + logger.debug('Training-table manifests are only exported to ' + 'Parquet directories on files or object stores.') return None - try: - manifest = self.manifest - if manifest is None: - return None - manifest.update(content_sha256=_content_digest(files), - export_identity=export_id) - return (json.dumps(manifest, indent=2, sort_keys=True) + - '\n').encode('utf-8') - except Exception as error: - warnings.warn( - f"Exporting without a training-table manifest: {error}", - UserWarning, stacklevel=3) + manifest = self.manifest + if manifest is None: return None + return (json.dumps(manifest, indent=2, sort_keys=True) + + '\n').encode('utf-8') @classmethod def from_source_table( @@ -184,7 +175,7 @@ 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") - derived_from, modified = _export_lineage(source_table) + lineage = _source_lineage(source_table) generation_plan = from_json_dict( to_json_dict(public_plan), TrainingTableGenerationPlan, @@ -250,14 +241,13 @@ def from_source_table( generation_plan=generation_plan, ) _record_manifest(table, snapshot, task_plan.train_table, 'custom', - derived_from=derived_from, modified=modified) + source_table=lineage) return table def _record_manifest(table: TrainingTable, snapshot: GraphSnapshot, request: Any, source: str, *, - derived_from: Optional[str] = None, - **fields: Any) -> None: + source_table: Optional[Dict[str, Any]] = None) -> None: path = table.staging_dir / STAGING_MANIFEST_FILE try: from kumo_core.io.json import to_json_dict @@ -267,10 +257,9 @@ def _record_manifest(table: TrainingTable, snapshot: GraphSnapshot, seed = None if request.split_family in _RANDOM_SPLITS: seed = (request.split_config or {}).get('rnd_seed') or 0 - parents = { - 'snapshot': snapshot.id, - 'derived_from': derived_from, - } + parents: Dict[str, Any] = {'snapshot': snapshot.id} + if source_table is not None: + parents['source_table'] = source_table write_manifest( path, new_manifest( @@ -284,7 +273,6 @@ def _record_manifest(table: TrainingTable, snapshot: GraphSnapshot, }, split_seed=seed, source=source, - **fields, )) except Exception as error: warnings.warn( @@ -292,48 +280,44 @@ def _record_manifest(table: TrainingTable, snapshot: GraphSnapshot, f"{error}", UserWarning, stacklevel=3) -def _content_digest(paths: Sequence[str]) -> str: - import hashlib - - from kumo_core.artifact_export import content_sha256 - - digests = [] - for path in paths: - with open(path, 'rb') as handle: - digests.append(hashlib.file_digest(handle, 'sha256').hexdigest()) - return content_sha256(digests) - - -def _export_lineage( - source_table: SourceTable) -> Tuple[Optional[str], Optional[bool]]: - from kumo_core.artifact_export import EXPORT_MARKER_NAME, export_marker - from kumo_core.object_store_sink import ObjectStoreSinkBackend - - connector = source_table.connector - name = source_table.table_name +def _source_lineage(source_table: SourceTable) -> Dict[str, Optional[str]]: + connector, name, sql = (source_table.connector, source_table.table_name, + source_table.sql) + lineage: Dict[str, Optional[str]] = {'location': None, 'version': None} try: - backend = connector.sink_backend - if source_table.sql is not None or not isinstance( - backend, ObjectStoreSinkBackend): - return None, None - payload = backend.read_sidecar(connector, name, EXPORT_MANIFEST_FILE) - if payload is None: - logger.debug('%s has no training-table manifest', name) - return None, None - except Exception as error: - logger.debug('No training-table manifest for %s: %s', name, error) - return None, None - try: - exported = json.loads(payload) - marker = backend.read_sidecar(connector, name, EXPORT_MARKER_NAME) - if marker is None or marker.decode('utf-8') != export_marker( - exported['export_identity']): - logger.debug('%s does not belong to the export it records', name) - return None, None - digest = backend.content_digest(connector, name) - return exported['id'], digest != exported['content_sha256'] + lineage['location'] = _location(connector, name, sql) except Exception as error: warnings.warn( - f"Ignoring the training-table manifest of {name!r}: {error}", - UserWarning, stacklevel=3) - return None, None + f"Could not record where the custom table was loaded from: " + f"{error}", UserWarning, stacklevel=3) + if sql is None: + try: + version = connector.source_backend.get_data_version( + connector, name) + lineage['version'] = None if version is None else str(version) + except Exception as error: + warnings.warn( + f"Could not record the custom table's data version: {error}", + UserWarning, stacklevel=3) + return lineage + + +def _location(connector: Any, name: str, sql: Optional[str]) -> str: + import hashlib + import posixpath + + from kumo_core.sources.object_store_sessions import ObjectStoreSessions + + if sql is not None: + return 'sql:' + hashlib.sha256(sql.encode()).hexdigest() + root = getattr(connector.config, 'root_dir', None) + if root is not None: + return str(Path(root, name)) + backend = connector.source_backend + if isinstance(backend, ObjectStoreSessions): + _, prefix, key, _ = backend._placement(connector, name) + return f'{backend.scheme}://{posixpath.join(prefix, key).rstrip("/")}' + namespace = (getattr(connector, field, None) + for field in ('database', 'catalog', 'project_id', + 'schema_name', 'dataset_id')) + return '.'.join([*filter(None, namespace), name]) diff --git a/packages/kumo-engine/test/test_training_table_export.py b/packages/kumo-engine/test/test_training_table_export.py index 4fc947f67..0ad521558 100644 --- a/packages/kumo-engine/test/test_training_table_export.py +++ b/packages/kumo-engine/test/test_training_table_export.py @@ -4,7 +4,6 @@ import hashlib import json -import shutil from collections import Counter from functools import partial from types import SimpleNamespace @@ -12,7 +11,6 @@ import numpy as np import pandas as pd import pyarrow as pa -import pyarrow.compute as pc import pyarrow.parquet as pq import pytest from kumoapi.graph import Edge @@ -298,12 +296,6 @@ def test_custom_training_table_declares_integer_timestamp_unit(graph, unit, assert _rows(seeds[stage], columns) == _rows(original_seeds[stage], columns) -def _digest(paths): - return hashlib.sha256(''.join( - sorted(hashlib.sha256(p.read_bytes()).hexdigest() - for p in paths)).encode()).hexdigest() - - def _reload(graph, name, pquery): _, snapshot, connector, _ = graph return kumo.TrainingTable.from_source_table( @@ -311,8 +303,8 @@ def _reload(graph, name, pquery): snapshot=snapshot) -def test_manifest_lineage_across_export_and_reload(graph, tmp_path): - from kumo_core.training.export import export_train_table, identity +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' @@ -320,10 +312,7 @@ def test_manifest_lineage_across_export_and_reload(graph, tmp_path): manifest = table.manifest assert manifest['kind'] == 'training_table' assert manifest['id'] == table.tt_id == _spec(table)['tt_id'] - assert manifest['parents'] == { - 'snapshot': snapshot.id, - 'derived_from': None - } + assert manifest['parents'] == {'snapshot': snapshot.id} assert (manifest['query'], manifest['source']) == (query, 'pql') assert manifest['task_type'] == 'binary_classification' assert manifest['rows'] == _seeds(table)[1] @@ -342,65 +331,82 @@ def test_manifest_lineage_across_export_and_reload(graph, tmp_path): [p.name for p in parts] + ['_KUMO_EXPORT', '_SUCCESS', '_kumo_manifest.json']) sidecar = (directory / '_kumo_manifest.json').read_bytes() - assert json.loads(sidecar) == { - **manifest, 'content_sha256': _digest(parts), - 'export_identity': identity(_spec(table)) - } + assert json.loads(sidecar) == manifest table.export(TrainingTableOutputConfig(connector, 'lineage')) assert (directory / '_kumo_manifest.json').read_bytes() == sidecar - assert [p.read_bytes() for p in parts] == [ - (baseline / p.name).read_bytes() for p in parts - ] - reloaded = _reload(graph, 'lineage', pquery).manifest - assert reloaded['parents'] == { + +def test_custom_table_records_its_source(graph, recwarn): + _, snapshot, connector, out = graph + pquery, table, _ = _export(graph, 'PREDICT ents.LABEL FOR EACH ents.ID', + 'reloaded') + (out / 'reloaded' / '_kumo_manifest.json').write_text('{not json') + reloaded = _reload(graph, 'reloaded', pquery) + manifest = reloaded.manifest + assert manifest['source'] == 'custom' + assert manifest['id'] != table.tt_id + assert manifest['parents'] == { 'snapshot': snapshot.id, - 'derived_from': table.tt_id + 'source_table': { + 'location': str(out / 'reloaded'), + 'version': connector.source_backend.get_data_version( + connector, 'reloaded'), + }, } - assert (reloaded['source'], reloaded['modified']) == ('custom', False) - assert reloaded['id'] != table.tt_id - - edited = out / 'lineage_edited' - shutil.copytree(directory, edited) - rows = pq.read_table(parts[0]) - flipped = pc.invert(rows['TARGET1']) - pq.write_table(rows.set_column(rows.schema.get_field_index('TARGET1'), - 'TARGET1', flipped), - edited / parts[0].name) - edited_manifest = _reload(graph, 'lineage_edited', pquery).manifest - assert edited_manifest['parents']['derived_from'] == table.tt_id - assert edited_manifest['modified'] is True - - -@pytest.mark.parametrize('damage', ['missing', 'garbled', 'foreign']) -def test_reload_without_matching_sidecar_records_no_lineage(graph, damage): + assert manifest['parents']['source_table']['version'] + assert not [w for w in recwarn if 'manifest' in str(w.message)] + assert _seeds(reloaded)[1] == _seeds(table)[1] + + +@pytest.mark.parametrize('failing', ['_location', 'get_data_version']) +def test_unrecordable_source_is_null(graph, monkeypatch, failing): + from kumo.engine.pquery import training_table + _, snapshot, connector, out = graph pquery, table, _ = _export(graph, 'PREDICT ents.LABEL FOR EACH ents.ID', - f'unlinked_base_{damage}') - directory = out / f'unlinked_{damage}' - shutil.copytree(out / f'unlinked_base_{damage}', directory) - sidecar = directory / '_kumo_manifest.json' - if damage == 'missing': - sidecar.unlink() - elif damage == 'garbled': - sidecar.write_text('{not json') - else: - (directory / '_KUMO_EXPORT').write_text('another export') + f'unrecorded_{failing}') - if damage == 'garbled': - with pytest.warns(UserWarning, match='training-table manifest'): - reloaded = _reload(graph, directory.name, pquery) + def fail(*args, **kwargs): + raise OSError('listing denied') + + target = (training_table if failing == '_location' else + type(connector.source_backend)) + monkeypatch.setattr(target, failing, fail) + with pytest.warns(UserWarning, match='listing denied'): + reloaded = _reload(graph, f'unrecorded_{failing}', pquery) + recorded = reloaded.manifest['parents']['source_table'] + if failing == '_location': + assert recorded['location'] is None and recorded['version'] else: - reloaded = _reload(graph, directory.name, pquery) - assert reloaded.manifest['parents'] == { - 'snapshot': snapshot.id, - 'derived_from': None - } - assert reloaded.manifest['modified'] is None + assert recorded == { + 'location': str(out / f'unrecorded_{failing}'), + 'version': None + } assert _seeds(reloaded)[1] == _seeds(table)[1] +@pytest.mark.parametrize('connector, name, location', [ + (kumo.S3Connector(root_dir='s3://bucket/data/'), 'exports/tt/', + 's3://bucket/data/exports/tt'), + (kumo.S3Connector(), 's3://bucket/exports/tt', 's3://bucket/exports/tt'), + (SimpleNamespace(config=None, source_backend=None, database='DB', + schema_name='SCHEMA'), 'TT', 'DB.SCHEMA.TT'), +]) +def test_source_location(connector, name, location): + from kumo.engine.pquery.training_table import _location + + assert _location(connector, name, None) == location + + +def test_sql_location_is_a_digest(): + from kumo.engine.pquery.training_table import _location + + sql = 'SELECT * FROM labels' + assert _location(None, 'synthetic', sql) == ( + 'sql:' + hashlib.sha256(sql.encode()).hexdigest()) + + def test_manifest_write_failure_keeps_generation_and_export( graph, monkeypatch): from kumo.engine.pquery import training_table @@ -438,42 +444,6 @@ def write_files(connector, table_name, files, write_mode, **kwargs): assert calls == [{'export_id': f'{table.tt_id}:tt-export-v1'}] -def test_reencoded_staging_does_not_mark_export_modified(graph, monkeypatch): - from kumo.engine.pquery import _custom_source - - _, snapshot, connector, out = graph - query = ('PREDICT LIST_DISTINCT(events.ITEM_ID, 0, 7, days) ' - 'RANK TOP 3 FOR EACH ents.ID') - pquery, table, _ = _export(graph, query, 'lineage_weighted') - directory = out / 'lineage_weighted' - parts = sorted(directory.glob('*.parquet')) - for part in parts: - rows = pq.read_table(part) - weights = pa.array([[1.0] * len(values) - for values in rows['TARGET1'].to_pylist()], - pa.list_(pa.float64())) - pq.write_table(rows.append_column('WEIGHT', weights), part) - sidecar = directory / '_kumo_manifest.json' - sidecar.write_text(json.dumps({ - **json.loads(sidecar.read_text()), 'content_sha256': _digest(parts) - })) - staged = [] - stage_source_table = _custom_source.stage_source_table - - def capture(*args, **kwargs): - staged.append(stage_source_table(*args, **kwargs)) - return staged[-1] - - monkeypatch.setattr(_custom_source, 'stage_source_table', capture) - manifest = kumo.TrainingTable.from_source_table( - connector['lineage_weighted'], - TrainingTableSpec(weight_col='WEIGHT', split_col='SPLIT'), - pq=pquery, snapshot=snapshot).manifest - assert _digest(sorted(staged[0].glob('*.parquet'))) != _digest(parts) - assert manifest['parents']['derived_from'] == table.tt_id - assert manifest['modified'] is False - - def test_csv_directory_export_gets_no_sidecar(graph): graph_, snapshot, connector, out = graph table = kumo.PredictiveQuery( From e25b68858e313c074d076c50477e8452fdf56314 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Wed, 21 Oct 2026 07:46:37 +0000 Subject: [PATCH 22/80] Record a custom table's source in the snapshot manifest's shape parents.source_table now uses the shared source_table helper from the snapshot manifest: {"connector", "location", "version"}, with the same location rules, so a custom table and a snapshot source can be matched on connector and location. The version is still the connector's get_data_version stamp, null with a warning when it cannot be read. Verified with test_training_table_export.py. Co-Authored-By: Claude Opus 5.5 --- .../kumo/engine/pquery/training_table.py | 34 +++---------------- .../test/test_training_table_export.py | 32 ++++------------- 2 files changed, 10 insertions(+), 56 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/pquery/training_table.py b/packages/kumo-engine/kumo/engine/pquery/training_table.py index 924ef06ca..a63f19a7d 100644 --- a/packages/kumo-engine/kumo/engine/pquery/training_table.py +++ b/packages/kumo-engine/kumo/engine/pquery/training_table.py @@ -281,43 +281,17 @@ def _record_manifest(table: TrainingTable, snapshot: GraphSnapshot, def _source_lineage(source_table: SourceTable) -> Dict[str, Optional[str]]: + from kumo.engine._provenance import source_table as lineage + connector, name, sql = (source_table.connector, source_table.table_name, source_table.sql) - lineage: Dict[str, Optional[str]] = {'location': None, 'version': None} - try: - lineage['location'] = _location(connector, name, sql) - except Exception as error: - warnings.warn( - f"Could not record where the custom table was loaded from: " - f"{error}", UserWarning, stacklevel=3) + version = None if sql is None: try: version = connector.source_backend.get_data_version( connector, name) - lineage['version'] = None if version is None else str(version) except Exception as error: warnings.warn( f"Could not record the custom table's data version: {error}", UserWarning, stacklevel=3) - return lineage - - -def _location(connector: Any, name: str, sql: Optional[str]) -> str: - import hashlib - import posixpath - - from kumo_core.sources.object_store_sessions import ObjectStoreSessions - - if sql is not None: - return 'sql:' + hashlib.sha256(sql.encode()).hexdigest() - root = getattr(connector.config, 'root_dir', None) - if root is not None: - return str(Path(root, name)) - backend = connector.source_backend - if isinstance(backend, ObjectStoreSessions): - _, prefix, key, _ = backend._placement(connector, name) - return f'{backend.scheme}://{posixpath.join(prefix, key).rstrip("/")}' - namespace = (getattr(connector, field, None) - for field in ('database', 'catalog', 'project_id', - 'schema_name', 'dataset_id')) - return '.'.join([*filter(None, namespace), name]) + return lineage(connector, name, sql, version) diff --git a/packages/kumo-engine/test/test_training_table_export.py b/packages/kumo-engine/test/test_training_table_export.py index 0ad521558..8a9285fbc 100644 --- a/packages/kumo-engine/test/test_training_table_export.py +++ b/packages/kumo-engine/test/test_training_table_export.py @@ -2,7 +2,6 @@ # All rights reserved. # SPDX-License-Identifier: Apache-2.0 -import hashlib import json from collections import Counter from functools import partial @@ -349,6 +348,7 @@ def test_custom_table_records_its_source(graph, recwarn): assert manifest['parents'] == { 'snapshot': snapshot.id, 'source_table': { + 'connector': 'LOCAL_FILE', 'location': str(out / 'reloaded'), 'version': connector.source_backend.get_data_version( connector, 'reloaded'), @@ -359,9 +359,9 @@ def test_custom_table_records_its_source(graph, recwarn): assert _seeds(reloaded)[1] == _seeds(table)[1] -@pytest.mark.parametrize('failing', ['_location', 'get_data_version']) +@pytest.mark.parametrize('failing', ['source_location', 'get_data_version']) def test_unrecordable_source_is_null(graph, monkeypatch, failing): - from kumo.engine.pquery import training_table + from kumo.engine import _provenance _, snapshot, connector, out = graph pquery, table, _ = _export(graph, 'PREDICT ents.LABEL FOR EACH ents.ID', @@ -370,43 +370,23 @@ def test_unrecordable_source_is_null(graph, monkeypatch, failing): def fail(*args, **kwargs): raise OSError('listing denied') - target = (training_table if failing == '_location' else + target = (_provenance if failing == 'source_location' else type(connector.source_backend)) monkeypatch.setattr(target, failing, fail) with pytest.warns(UserWarning, match='listing denied'): reloaded = _reload(graph, f'unrecorded_{failing}', pquery) recorded = reloaded.manifest['parents']['source_table'] - if failing == '_location': + if failing == 'source_location': assert recorded['location'] is None and recorded['version'] else: assert recorded == { + 'connector': 'LOCAL_FILE', 'location': str(out / f'unrecorded_{failing}'), 'version': None } assert _seeds(reloaded)[1] == _seeds(table)[1] -@pytest.mark.parametrize('connector, name, location', [ - (kumo.S3Connector(root_dir='s3://bucket/data/'), 'exports/tt/', - 's3://bucket/data/exports/tt'), - (kumo.S3Connector(), 's3://bucket/exports/tt', 's3://bucket/exports/tt'), - (SimpleNamespace(config=None, source_backend=None, database='DB', - schema_name='SCHEMA'), 'TT', 'DB.SCHEMA.TT'), -]) -def test_source_location(connector, name, location): - from kumo.engine.pquery.training_table import _location - - assert _location(connector, name, None) == location - - -def test_sql_location_is_a_digest(): - from kumo.engine.pquery.training_table import _location - - sql = 'SELECT * FROM labels' - assert _location(None, 'synthetic', sql) == ( - 'sql:' + hashlib.sha256(sql.encode()).hexdigest()) - - def test_manifest_write_failure_keeps_generation_and_export( graph, monkeypatch): from kumo.engine.pquery import training_table From 8de3f80629853dfe9c703062e07063b45a7c8d4d Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Wed, 21 Oct 2026 08:12:42 +0000 Subject: [PATCH 23/80] Keep sidecar writes best-effort; version custom tables from their read A failing _kumo_manifest.json write (a bucket policy, a transient error) raised out of export_selected_files before _SUCCESS, so an export that works today failed. Each sidecar write now warns and continues, and a failure building the training-table manifest exports without it. A custom table's version came from a second get_data_version call after staging, which could describe a later state of the source and was null for Custom SQL. stage_source_table now returns the version prepare_ingest reports for the rows it staged, the same stamp the snapshot manifest records. Verified with test_training_table_export.py, connector/test_custom_source.py, test_custom_sql.py and kumo-core test_artifact_export.py. Co-Authored-By: Claude Opus 5.5 --- .../src/kumo_core/artifact_export.py | 7 ++++- .../kumo-core/tests/test_artifact_export.py | 24 +++++++++++++++ .../kumo/engine/pquery/_custom_source.py | 9 ++++-- .../kumo/engine/pquery/prediction_table.py | 3 +- .../kumo/engine/pquery/training_table.py | 30 +++++++++---------- .../test/connector/test_custom_source.py | 17 ++++++----- packages/kumo-engine/test/test_custom_sql.py | 2 +- .../kumo-engine/test/test_predictive_query.py | 6 ++-- .../test/test_training_table_export.py | 25 +++++----------- 9 files changed, 72 insertions(+), 51 deletions(-) diff --git a/packages/kumo-core/src/kumo_core/artifact_export.py b/packages/kumo-core/src/kumo_core/artifact_export.py index 8abd316d6..25df7dc7f 100644 --- a/packages/kumo-core/src/kumo_core/artifact_export.py +++ b/packages/kumo-core/src/kumo_core/artifact_export.py @@ -236,7 +236,12 @@ def export_selected_files(files, output_path, *, fs, identity, if output_is_dir: for name, payload in sidecars.items(): - fs.write_bytes(osp.join(output_path, name), payload) + 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/tests/test_artifact_export.py b/packages/kumo-core/tests/test_artifact_export.py index 726d6847c..e22acd5d8 100644 --- a/packages/kumo-core/tests/test_artifact_export.py +++ b/packages/kumo-core/tests/test_artifact_export.py @@ -251,3 +251,27 @@ def touch(path): else: assert not (target / 'meta.json').exists() assert pq.read_table(tmp_path / 'output.parquet').num_rows == 1 + + +def test_failed_sidecar_write_still_publishes(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) + fs = _ArrowExportFiles(LocalFileSystem(), None) + write_bytes = fs.write_bytes + + def denied(path, value): + if path.endswith('meta.json'): + raise PermissionError('denied') + write_bytes(path, value) + + fs.write_bytes = denied + target = tmp_path / 'output' + with pytest.warns(UserWarning, match='denied'): + export_selected_files([str(source)], f'{target}/', fs=fs, + identity='job-1', sidecars={'meta.json': b'{}'}) + assert (target / '_SUCCESS').exists() diff --git a/packages/kumo-engine/kumo/engine/pquery/_custom_source.py b/packages/kumo-engine/kumo/engine/pquery/_custom_source.py index a98c74975..7650d7904 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,10 @@ 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]]: + """Read ``source_table`` through its connector and retain its parquet + parts; also returns the connector's version of the rows read.""" 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 +65,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/training_table.py b/packages/kumo-engine/kumo/engine/pquery/training_table.py index a63f19a7d..8958e537e 100644 --- a/packages/kumo-engine/kumo/engine/pquery/training_table.py +++ b/packages/kumo-engine/kumo/engine/pquery/training_table.py @@ -84,7 +84,13 @@ def export(self, output_config: "TrainingTableOutputConfig") -> str: root.mkdir(parents=True, exist_ok=True) with tempfile.TemporaryDirectory(dir=root) as output: files = export_train_table(str(self.staging_dir), output) - payload = self._export_manifest(connector, output_config, files) + try: + payload = self._export_manifest(connector, output_config, + files) + except Exception as error: + warnings.warn(f"Exporting without a training-table manifest: " + f"{error}", UserWarning, stacklevel=2) + payload = None extra = ({} if payload is None else { 'sidecars': { EXPORT_MANIFEST_FILE: payload @@ -174,8 +180,9 @@ 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") - lineage = _source_lineage(source_table) + rows_path, version = stage_source_table(source_table, + run_dir / "source-table") + lineage = _source_lineage(source_table, version) generation_plan = from_json_dict( to_json_dict(public_plan), TrainingTableGenerationPlan, @@ -280,18 +287,9 @@ def _record_manifest(table: TrainingTable, snapshot: GraphSnapshot, f"{error}", UserWarning, stacklevel=3) -def _source_lineage(source_table: SourceTable) -> Dict[str, Optional[str]]: +def _source_lineage(source_table: SourceTable, + version: Optional[str]) -> Dict[str, Optional[str]]: from kumo.engine._provenance import source_table as lineage - connector, name, sql = (source_table.connector, source_table.table_name, - source_table.sql) - version = None - if sql is None: - try: - version = connector.source_backend.get_data_version( - connector, name) - except Exception as error: - warnings.warn( - f"Could not record the custom table's data version: {error}", - UserWarning, stacklevel=3) - return lineage(connector, name, sql, version) + return lineage(source_table.connector, source_table.table_name, + source_table.sql, version) 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_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_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_training_table_export.py b/packages/kumo-engine/test/test_training_table_export.py index 8a9285fbc..6337fc179 100644 --- a/packages/kumo-engine/test/test_training_table_export.py +++ b/packages/kumo-engine/test/test_training_table_export.py @@ -226,9 +226,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) @@ -359,31 +359,20 @@ def test_custom_table_records_its_source(graph, recwarn): assert _seeds(reloaded)[1] == _seeds(table)[1] -@pytest.mark.parametrize('failing', ['source_location', 'get_data_version']) -def test_unrecordable_source_is_null(graph, monkeypatch, failing): +def test_unrecordable_location_is_null(graph, monkeypatch): from kumo.engine import _provenance - _, snapshot, connector, out = graph pquery, table, _ = _export(graph, 'PREDICT ents.LABEL FOR EACH ents.ID', - f'unrecorded_{failing}') + 'unrecorded') def fail(*args, **kwargs): raise OSError('listing denied') - target = (_provenance if failing == 'source_location' else - type(connector.source_backend)) - monkeypatch.setattr(target, failing, fail) + monkeypatch.setattr(_provenance, 'source_location', fail) with pytest.warns(UserWarning, match='listing denied'): - reloaded = _reload(graph, f'unrecorded_{failing}', pquery) + reloaded = _reload(graph, 'unrecorded', pquery) recorded = reloaded.manifest['parents']['source_table'] - if failing == 'source_location': - assert recorded['location'] is None and recorded['version'] - else: - assert recorded == { - 'connector': 'LOCAL_FILE', - 'location': str(out / f'unrecorded_{failing}'), - 'version': None - } + assert recorded['location'] is None and recorded['version'] assert _seeds(reloaded)[1] == _seeds(table)[1] From ea6fe84ac6cff41cb32e39d37723126c2af06d5a Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Wed, 21 Oct 2026 09:33:25 +0000 Subject: [PATCH 24/80] Trim training-table lineage to essentials The export sidecar path re-serialised the staged manifest behind a redundant directory check and a catch-all, and the tests restated fields and covered helpers owned by the snapshot-manifest PR. Copy the staged manifest bytes as the sidecar, inline the single-use source-lineage helper, shorten docstrings, and merge the lineage tests into one export/reload test plus best-effort and no-sidecar cases. Behaviour is unchanged. Verified: test_training_table_export, test_custom_source, test_custom_sql, test_predictive_query -k custom and kumo-core test_artifact_export pass. Co-Authored-By: Claude Opus 5.5 --- .../src/kumo_core/artifact_export.py | 4 +- .../src/kumo_core/object_store_sink.py | 6 +- .../kumo-core/tests/test_artifact_export.py | 47 +++------ .../kumo/engine/pquery/_custom_source.py | 3 +- .../kumo/engine/pquery/training_table.py | 59 ++++-------- .../test/test_training_table_export.py | 96 ++++--------------- 6 files changed, 53 insertions(+), 162 deletions(-) diff --git a/packages/kumo-core/src/kumo_core/artifact_export.py b/packages/kumo-core/src/kumo_core/artifact_export.py index 25df7dc7f..311b972bc 100644 --- a/packages/kumo-core/src/kumo_core/artifact_export.py +++ b/packages/kumo-core/src/kumo_core/artifact_export.py @@ -205,8 +205,8 @@ def export_selected_files(files, output_path, *, fs, identity, Filesystem/authentication and retry wrappers stay with the caller. The directory ownership, output names and CSV retry decisions are shared. - ``sidecars`` (name to bytes) are written into a directory destination - after the data and before ``_SUCCESS``; ``fs`` then needs ``write_bytes``. + ``sidecars`` (name to bytes) land in a directory destination before + ``_SUCCESS``. """ import os.path as osp import tempfile 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 48f31e2d9..93eed57ce 100644 --- a/packages/kumo-core/src/kumo_core/object_store_sink.py +++ b/packages/kumo-core/src/kumo_core/object_store_sink.py @@ -132,11 +132,7 @@ def write_files(self, connector, table_name, files, write_mode=WriteMode.OVERWRITE, project_id=None, *, export_id=None, extra_columns=None, sidecars: Mapping[str, bytes] = {}): - """Export selected Parquet without collecting their rows. - - ``sidecars`` are written into a directory destination before - ``_SUCCESS``. - """ + """Export selected Parquet without collecting their rows.""" del project_id if write_mode is not WriteMode.OVERWRITE: raise NotImplementedError( diff --git a/packages/kumo-core/tests/test_artifact_export.py b/packages/kumo-core/tests/test_artifact_export.py index e22acd5d8..b3701f00f 100644 --- a/packages/kumo-core/tests/test_artifact_export.py +++ b/packages/kumo-core/tests/test_artifact_export.py @@ -224,36 +224,8 @@ def interrupted(source, destination, **kwargs): assert (target / '_SUCCESS').exists() -@pytest.mark.parametrize('directory', [True, False]) -def test_sidecars_are_written_before_success(tmp_path, monkeypatch, - directory): - 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) - touched = fs.touch - def touch(path): - assert (target / 'meta.json').read_bytes() == b'{}' - touched(path) - - monkeypatch.setattr(fs, 'touch', touch) - export_selected_files([str(source)], - str(target) + ('/' if directory else ''), fs=fs, - identity='job-1', sidecars={'meta.json': b'{}'}) - if directory: - assert (target / '_SUCCESS').exists() - else: - assert not (target / 'meta.json').exists() - assert pq.read_table(tmp_path / 'output.parquet').num_rows == 1 - - -def test_failed_sidecar_write_still_publishes(tmp_path): +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 @@ -261,17 +233,24 @@ def test_failed_sidecar_write_still_publishes(tmp_path): source = tmp_path / 'part-0.parquet' pq.write_table(pa.table({'id': [1]}), source) + target = tmp_path / 'output' fs = _ArrowExportFiles(LocalFileSystem(), None) - write_bytes = fs.write_bytes + write_bytes, touch = fs.write_bytes, fs.touch def denied(path, value): - if path.endswith('meta.json'): + if path.endswith('bad.json'): raise PermissionError('denied') write_bytes(path, value) - fs.write_bytes = denied - target = tmp_path / 'output' + 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'{}'}) + identity='job-1', sidecars={ + 'meta.json': b'{}', + 'bad.json': b'{}' + }) assert (target / '_SUCCESS').exists() diff --git a/packages/kumo-engine/kumo/engine/pquery/_custom_source.py b/packages/kumo-engine/kumo/engine/pquery/_custom_source.py index 7650d7904..74ea2e9d9 100644 --- a/packages/kumo-engine/kumo/engine/pquery/_custom_source.py +++ b/packages/kumo-engine/kumo/engine/pquery/_custom_source.py @@ -18,8 +18,7 @@ def stage_source_table(source_table: SourceTable, destination: Path) -> Tuple[Path, Optional[str]]: - """Read ``source_table`` through its connector and retain its parquet - parts; also returns the connector's version of the rows read.""" + """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) diff --git a/packages/kumo-engine/kumo/engine/pquery/training_table.py b/packages/kumo-engine/kumo/engine/pquery/training_table.py index 8958e537e..165eab0c1 100644 --- a/packages/kumo-engine/kumo/engine/pquery/training_table.py +++ b/packages/kumo-engine/kumo/engine/pquery/training_table.py @@ -2,18 +2,17 @@ # All rights reserved. # SPDX-License-Identifier: Apache-2.0 -import json -import logging import warnings from dataclasses import dataclass, replace from pathlib import Path -from typing import TYPE_CHECKING, Any, Dict, Optional, Sequence +from typing import TYPE_CHECKING, Any, Dict, Optional 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 new_manifest, read_manifest, write_manifest from kumo.engine.connector.source_table import SourceTable from kumo.engine.graph import GraphSnapshot @@ -21,8 +20,6 @@ if TYPE_CHECKING: from kumo.engine.trainer.output_config import TrainingTableOutputConfig -logger = logging.getLogger(__name__) - STAGING_MANIFEST_FILE = 'training-table-manifest.json' EXPORT_MANIFEST_FILE = '_kumo_manifest.json' _RANDOM_SPLITS = ('RandomSplit', 'RandomLinkSplit') @@ -84,40 +81,22 @@ def export(self, output_config: "TrainingTableOutputConfig") -> str: root.mkdir(parents=True, exist_ok=True) with tempfile.TemporaryDirectory(dir=root) as output: files = export_train_table(str(self.staging_dir), output) - try: - payload = self._export_manifest(connector, output_config, - files) - except Exception as error: - warnings.warn(f"Exporting without a training-table manifest: " - f"{error}", UserWarning, stacklevel=2) - payload = None - extra = ({} if payload is None else { - 'sidecars': { - EXPORT_MANIFEST_FILE: payload - } - }) return write_files(connector, output_config.table_name, files, output_config.write_mode, export_id=export_id, - **extra) + **self._export_sidecars(connector, + output_config)) - def _export_manifest(self, connector, output_config, - files: Sequence[str]) -> Optional[bytes]: + def _export_sidecars(self, connector, output_config) -> Dict[str, Any]: from kumo_core.object_store_sink import ObjectStoreSinkBackend from kumo_core.sources import layout - name = output_config.table_name + manifest = self.staging_dir / STAGING_MANIFEST_FILE + name = output_config.table_name.rstrip('/') if (not isinstance(connector.sink_backend, ObjectStoreSinkBackend) - or not (name.endswith('/') or len(files) > 1) - or layout.file_type(name.rstrip('/')) not in ( - None, layout.PARQUET)): - logger.debug('Training-table manifests are only exported to ' - 'Parquet directories on files or object stores.') - return None - manifest = self.manifest - if manifest is None: - return None - return (json.dumps(manifest, indent=2, sort_keys=True) + - '\n').encode('utf-8') + or layout.file_type(name) == layout.CSV + or not manifest.is_file()): + return {} + return {'sidecars': {EXPORT_MANIFEST_FILE: manifest.read_bytes()}} @classmethod def from_source_table( @@ -182,7 +161,6 @@ def from_source_table( run_dir = train_table_scratch_dir(run_id) rows_path, version = stage_source_table(source_table, run_dir / "source-table") - lineage = _source_lineage(source_table, version) generation_plan = from_json_dict( to_json_dict(public_plan), TrainingTableGenerationPlan, @@ -247,8 +225,11 @@ def from_source_table( required_stages=task_plan.task_dataset.required_stages, generation_plan=generation_plan, ) - _record_manifest(table, snapshot, task_plan.train_table, 'custom', - source_table=lineage) + _record_manifest( + table, snapshot, task_plan.train_table, 'custom', + source_table=provenance.source_table(source_table.connector, + source_table.table_name, + source_table.sql, version)) return table @@ -285,11 +266,3 @@ def _record_manifest(table: TrainingTable, snapshot: GraphSnapshot, warnings.warn( f"Could not write the training-table manifest at {path}: " f"{error}", UserWarning, stacklevel=3) - - -def _source_lineage(source_table: SourceTable, - version: Optional[str]) -> Dict[str, Optional[str]]: - from kumo.engine._provenance import source_table as lineage - - return lineage(source_table.connector, source_table.table_name, - source_table.sql, version) diff --git a/packages/kumo-engine/test/test_training_table_export.py b/packages/kumo-engine/test/test_training_table_export.py index 6337fc179..119e1be1b 100644 --- a/packages/kumo-engine/test/test_training_table_export.py +++ b/packages/kumo-engine/test/test_training_table_export.py @@ -295,13 +295,6 @@ def test_custom_training_table_declares_integer_timestamp_unit(graph, unit, assert _rows(seeds[stage], columns) == _rows(original_seeds[stage], columns) -def _reload(graph, name, pquery): - _, snapshot, connector, _ = graph - return kumo.TrainingTable.from_source_table( - connector[name], TrainingTableSpec(split_col='SPLIT'), pq=pquery, - snapshot=snapshot) - - def test_manifest_and_export_sidecar(graph, tmp_path): from kumo_core.training.export import export_train_table @@ -309,70 +302,31 @@ def test_manifest_and_export_sidecar(graph, tmp_path): query = 'PREDICT ents.LABEL FOR EACH ents.ID' pquery, table, _ = _export(graph, query, 'lineage') manifest = table.manifest - assert manifest['kind'] == 'training_table' assert manifest['id'] == table.tt_id == _spec(table)['tt_id'] assert manifest['parents'] == {'snapshot': snapshot.id} assert (manifest['query'], manifest['source']) == (query, 'pql') - assert manifest['task_type'] == 'binary_classification' assert manifest['rows'] == _seeds(table)[1] assert manifest['split_seed'] == 0 - assert manifest['generation_plan']['split'] == 'inferred' directory = out / 'lineage' - parts = sorted(directory.glob('*.parquet')) baseline = tmp_path / 'baseline' - expected = export_train_table(str(table.staging_dir), str(baseline)) - assert [p.name for p in parts] == [p.split('/')[-1] for p in expected] - assert [p.read_bytes() for p in parts] == [ - (baseline / p.name).read_bytes() for p in parts - ] - assert sorted(p.name for p in directory.iterdir()) == sorted( - [p.name for p in parts] + - ['_KUMO_EXPORT', '_SUCCESS', '_kumo_manifest.json']) - sidecar = (directory / '_kumo_manifest.json').read_bytes() - assert json.loads(sidecar) == manifest - - table.export(TrainingTableOutputConfig(connector, 'lineage')) - assert (directory / '_kumo_manifest.json').read_bytes() == sidecar - - -def test_custom_table_records_its_source(graph, recwarn): - _, snapshot, connector, out = graph - pquery, table, _ = _export(graph, 'PREDICT ents.LABEL FOR EACH ents.ID', - 'reloaded') - (out / 'reloaded' / '_kumo_manifest.json').write_text('{not json') - reloaded = _reload(graph, 'reloaded', pquery) + 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 / '_kumo_manifest.json').read_bytes()) == manifest + + reloaded = kumo.TrainingTable.from_source_table( + connector['lineage'], TrainingTableSpec(split_col='SPLIT'), pq=pquery, + snapshot=snapshot) manifest = reloaded.manifest assert manifest['source'] == 'custom' - assert manifest['id'] != table.tt_id - assert manifest['parents'] == { - 'snapshot': snapshot.id, - 'source_table': { - 'connector': 'LOCAL_FILE', - 'location': str(out / 'reloaded'), - 'version': connector.source_backend.get_data_version( - connector, 'reloaded'), - }, + assert manifest['parents']['source_table'] == { + 'connector': 'LOCAL_FILE', + 'location': str(directory), + 'version': connector.source_backend.get_data_version( + connector, 'lineage'), } - assert manifest['parents']['source_table']['version'] - assert not [w for w in recwarn if 'manifest' in str(w.message)] - assert _seeds(reloaded)[1] == _seeds(table)[1] - - -def test_unrecordable_location_is_null(graph, monkeypatch): - from kumo.engine import _provenance - - pquery, table, _ = _export(graph, 'PREDICT ents.LABEL FOR EACH ents.ID', - 'unrecorded') - - def fail(*args, **kwargs): - raise OSError('listing denied') - - monkeypatch.setattr(_provenance, 'source_location', fail) - with pytest.warns(UserWarning, match='listing denied'): - reloaded = _reload(graph, 'unrecorded', pquery) - recorded = reloaded.manifest['parents']['source_table'] - assert recorded['location'] is None and recorded['version'] assert _seeds(reloaded)[1] == _seeds(table)[1] @@ -387,16 +341,15 @@ def fail(*args, **kwargs): monkeypatch.setattr(training_table, 'write_manifest', fail) with pytest.warns(UserWarning, match='read-only scratch'): - pquery, table, exported = _export( + _, table, exported = _export( graph, 'PREDICT ents.LABEL FOR EACH ents.ID', 'no_manifest') assert table.manifest is None assert not (out / 'no_manifest' / '_kumo_manifest.json').exists() - assert (out / 'no_manifest' / '_SUCCESS').exists() assert exported.num_rows == sum(_seeds(table)[1].values()) -def test_warehouse_sink_receives_no_sidecars(graph): - graph_, snapshot, _, _ = 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) @@ -404,20 +357,11 @@ def test_warehouse_sink_receives_no_sidecars(graph): def write_files(connector, table_name, files, write_mode, **kwargs): calls.append(kwargs) - return table_name - connector = SimpleNamespace( + warehouse = SimpleNamespace( source_type='SNOWFLAKE', sink_backend=SimpleNamespace(write_files=write_files)) - table.export(TrainingTableOutputConfig(connector, 'DB.SCHEMA.TT')) + table.export(TrainingTableOutputConfig(warehouse, 'DB.SCHEMA.TT')) assert calls == [{'export_id': f'{table.tt_id}:tt-export-v1'}] - - -def test_csv_directory_export_gets_no_sidecar(graph): - graph_, snapshot, connector, out = graph - table = kumo.PredictiveQuery( - graph=graph_, query='PREDICT ents.LABEL FOR EACH ents.ID' - ).generate_training_table(snapshot) table.export(TrainingTableOutputConfig(connector, 'rows.csv/')) - names = {p.name for p in (out / 'rows.csv').iterdir()} - assert '_SUCCESS' in names and '_kumo_manifest.json' not in names + assert not (out / 'rows.csv' / '_kumo_manifest.json').exists() From a494a099fb67de23883d67f71f6e8519d02fc7bb Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Wed, 21 Oct 2026 18:24:22 +0000 Subject: [PATCH 25/80] Read row counts and source mode from DiskGraph's manifest training-table-manifest.json repeated the per-stage row counts and the pql/custom source that DiskGraph already records in the table's manifest.json (stage_counts, source_mode) and tt_ingest_spec.json. Drop rows and source from the training-table manifest and expose TrainingTable.source_mode from DiskGraph's manifest; a custom table is still identified by parents.source_table. Verified with the training-table export, custom-source and predictive query suites. Co-Authored-By: Claude Opus 5.5 --- .../kumo/engine/pquery/predictive_query.py | 2 +- .../kumo/engine/pquery/training_table.py | 22 ++++++++++--------- .../test/test_training_table_export.py | 7 +++--- 3 files changed, 17 insertions(+), 14 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/pquery/predictive_query.py b/packages/kumo-engine/kumo/engine/pquery/predictive_query.py index 32736e86d..01a425fe2 100644 --- a/packages/kumo-engine/kumo/engine/pquery/predictive_query.py +++ b/packages/kumo-engine/kumo/engine/pquery/predictive_query.py @@ -316,7 +316,7 @@ def generate_training_table( required_stages=task_plan.task_dataset.required_stages, generation_plan=generation_plan, ) - _record_manifest(table, snapshot, task_plan.train_table, 'pql') + _record_manifest(table, snapshot, task_plan.train_table) return table def generate_prediction_table( diff --git a/packages/kumo-engine/kumo/engine/pquery/training_table.py b/packages/kumo-engine/kumo/engine/pquery/training_table.py index 165eab0c1..6f62b17dc 100644 --- a/packages/kumo-engine/kumo/engine/pquery/training_table.py +++ b/packages/kumo-engine/kumo/engine/pquery/training_table.py @@ -2,6 +2,7 @@ # All rights reserved. # SPDX-License-Identifier: Apache-2.0 +import json import warnings from dataclasses import dataclass, replace from pathlib import Path @@ -57,6 +58,14 @@ def manifest(self) -> Optional[Dict[str, Any]]: """Lineage recorded at generation; None for older tables.""" return read_manifest(self.staging_dir / STAGING_MANIFEST_FILE) + @property + def source_mode(self) -> Optional[str]: + """``pquery`` or ``custom_tt``, as DiskGraph recorded at generation.""" + try: + return json.loads(self.manifest_path.read_text())['source_mode'] + except (OSError, ValueError, KeyError, TypeError): + return None + def export(self, output_config: "TrainingTableOutputConfig") -> str: r"""Write the staged rows with a ``SPLIT`` column (``0`` train, ``1`` validation, ``2`` test) for :meth:`from_source_table`. Returns the @@ -226,7 +235,7 @@ def from_source_table( generation_plan=generation_plan, ) _record_manifest( - table, snapshot, task_plan.train_table, 'custom', + table, snapshot, task_plan.train_table, source_table=provenance.source_table(source_table.connector, source_table.table_name, source_table.sql, version)) @@ -234,14 +243,12 @@ def from_source_table( def _record_manifest(table: TrainingTable, snapshot: GraphSnapshot, - request: Any, source: str, *, + request: Any, *, source_table: Optional[Dict[str, Any]] = None) -> None: path = table.staging_dir / STAGING_MANIFEST_FILE try: from kumo_core.io.json import to_json_dict - from kumo_core.training.export import read_spec - spec = read_spec(str(table.staging_dir)) seed = None if request.split_family in _RANDOM_SPLITS: seed = (request.split_config or {}).get('rnd_seed') or 0 @@ -251,16 +258,11 @@ def _record_manifest(table: TrainingTable, snapshot: GraphSnapshot, write_manifest( path, new_manifest( - 'training_table', spec['tt_id'], parents, + 'training_table', table.tt_id, parents, query=table.query, generation_plan=to_json_dict(table.generation_plan), task_type=request.task_type, - rows={ - stage['key']: stage['row_count'] - for stage in spec.get('split_stages') or [] - }, split_seed=seed, - source=source, )) except Exception as error: warnings.warn( diff --git a/packages/kumo-engine/test/test_training_table_export.py b/packages/kumo-engine/test/test_training_table_export.py index 119e1be1b..64791b751 100644 --- a/packages/kumo-engine/test/test_training_table_export.py +++ b/packages/kumo-engine/test/test_training_table_export.py @@ -304,8 +304,9 @@ def test_manifest_and_export_sidecar(graph, tmp_path): manifest = table.manifest assert manifest['id'] == table.tt_id == _spec(table)['tt_id'] assert manifest['parents'] == {'snapshot': snapshot.id} - assert (manifest['query'], manifest['source']) == (query, 'pql') - assert manifest['rows'] == _seeds(table)[1] + assert manifest['query'] == query + assert 'rows' not in manifest and 'source' not in manifest + assert table.source_mode == 'pquery' assert manifest['split_seed'] == 0 directory = out / 'lineage' @@ -320,7 +321,7 @@ def test_manifest_and_export_sidecar(graph, tmp_path): connector['lineage'], TrainingTableSpec(split_col='SPLIT'), pq=pquery, snapshot=snapshot) manifest = reloaded.manifest - assert manifest['source'] == 'custom' + assert reloaded.source_mode == 'custom_tt' assert manifest['parents']['source_table'] == { 'connector': 'LOCAL_FILE', 'location': str(directory), From fc4120285146d6dc0b37e034d550d1ea4ebae897 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Wed, 21 Oct 2026 21:07:01 +0000 Subject: [PATCH 26/80] Record a custom table's source as the snapshot's manifest does A custom training table always recorded its source's full location and no version kind, even when its snapshot was taken with record_locations=False. The source entry now carries version_kind and follows the snapshot manifest's record_locations choice. Verified with the training-table export tests. Co-Authored-By: Claude Opus 5.5 --- .../kumo/engine/pquery/training_table.py | 8 +++++--- .../test/test_training_table_export.py | 17 ++++++++++++++++- 2 files changed, 21 insertions(+), 4 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/pquery/training_table.py b/packages/kumo-engine/kumo/engine/pquery/training_table.py index 6f62b17dc..26c4cf000 100644 --- a/packages/kumo-engine/kumo/engine/pquery/training_table.py +++ b/packages/kumo-engine/kumo/engine/pquery/training_table.py @@ -236,9 +236,11 @@ def from_source_table( ) _record_manifest( table, snapshot, task_plan.train_table, - source_table=provenance.source_table(source_table.connector, - source_table.table_name, - source_table.sql, version)) + source_table=provenance.source_table( + source_table.connector, source_table.table_name, + source_table.sql, version, + record_locations=(snapshot.manifest + or {}).get('record_locations', True))) return table diff --git a/packages/kumo-engine/test/test_training_table_export.py b/packages/kumo-engine/test/test_training_table_export.py index 64791b751..f9a6ec05b 100644 --- a/packages/kumo-engine/test/test_training_table_export.py +++ b/packages/kumo-engine/test/test_training_table_export.py @@ -324,13 +324,28 @@ def test_manifest_and_export_sidecar(graph, tmp_path): assert reloaded.source_mode == 'custom_tt' assert manifest['parents']['source_table'] == { 'connector': 'LOCAL_FILE', - 'location': str(directory), + 'location': str(directory.resolve()), 'version': connector.source_backend.get_data_version( connector, 'lineage'), + 'version_kind': 'timestamp', } assert _seeds(reloaded)[1] == _seeds(table)[1] +def test_custom_table_follows_the_snapshots_record_locations(graph): + graph_, _, connector, _ = graph + pquery, _, _ = _export(graph, 'PREDICT ents.LABEL FOR EACH ents.ID', + 'unlocated') + try: + redacted = graph_.snapshot(record_locations=False) + table = kumo.TrainingTable.from_source_table( + connector['unlocated'], TrainingTableSpec(split_col='SPLIT'), + pq=pquery, snapshot=redacted) + finally: + graph_.snapshot() + assert table.manifest['parents']['source_table']['location'] == 'unlocated' + + def test_manifest_write_failure_keeps_generation_and_export( graph, monkeypatch): from kumo.engine.pquery import training_table From d9743a821ab10b9b4d3811a130f38c3c5e60f1f0 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Wed, 21 Oct 2026 22:13:59 +0000 Subject: [PATCH 27/80] Export a training table even when its manifest sidecar fails A failure while preparing _kumo_manifest.json (an unreadable staging manifest, an unexpected sink) failed the whole export. Preparing the sidecar is now best-effort: the rows are exported without it and a warning names the file. Verified with test_training_table_export.py. Co-Authored-By: Claude Opus 5.5 --- .../kumo/engine/pquery/training_table.py | 14 ++++++++++---- .../test/test_training_table_export.py | 16 ++++++++++++++++ 2 files changed, 26 insertions(+), 4 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/pquery/training_table.py b/packages/kumo-engine/kumo/engine/pquery/training_table.py index 26c4cf000..b05aefc44 100644 --- a/packages/kumo-engine/kumo/engine/pquery/training_table.py +++ b/packages/kumo-engine/kumo/engine/pquery/training_table.py @@ -101,11 +101,17 @@ def _export_sidecars(self, connector, output_config) -> Dict[str, Any]: manifest = self.staging_dir / STAGING_MANIFEST_FILE name = output_config.table_name.rstrip('/') - if (not isinstance(connector.sink_backend, ObjectStoreSinkBackend) - or layout.file_type(name) == layout.CSV - or not manifest.is_file()): + try: + if (not isinstance(connector.sink_backend, ObjectStoreSinkBackend) + 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 {} - return {'sidecars': {EXPORT_MANIFEST_FILE: manifest.read_bytes()}} @classmethod def from_source_table( diff --git a/packages/kumo-engine/test/test_training_table_export.py b/packages/kumo-engine/test/test_training_table_export.py index f9a6ec05b..b43c508ba 100644 --- a/packages/kumo-engine/test/test_training_table_export.py +++ b/packages/kumo-engine/test/test_training_table_export.py @@ -364,6 +364,22 @@ def fail(*args, **kwargs): 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 / 'training-table-manifest.json' + manifest.chmod(0) + try: + with pytest.warns(UserWarning, match='_kumo_manifest.json'): + table.export(TrainingTableOutputConfig(connector, 'unreadable')) + finally: + manifest.chmod(0o644) + assert list((out / 'unreadable').glob('*.parquet')) + assert not (out / 'unreadable' / '_kumo_manifest.json').exists() + + def test_no_sidecar_for_warehouse_or_csv_exports(graph): graph_, snapshot, connector, out = graph table = kumo.PredictiveQuery( From 20fe503a0d4b9e6e9dbe8ad7bf8d345847f78a95 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Wed, 21 Oct 2026 22:17:08 +0000 Subject: [PATCH 28/80] Follow the snapshot's record_sources choice graph.snapshot(record_locations=False) was renamed record_sources=False in the snapshot manifest; a custom training table built against such a snapshot reads the new field and records its source the same way. Verified with test_training_table_export.py and test_snapshot_manifest.py. Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/pquery/training_table.py | 4 ++-- packages/kumo-engine/test/test_training_table_export.py | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/pquery/training_table.py b/packages/kumo-engine/kumo/engine/pquery/training_table.py index b05aefc44..61996bdb2 100644 --- a/packages/kumo-engine/kumo/engine/pquery/training_table.py +++ b/packages/kumo-engine/kumo/engine/pquery/training_table.py @@ -245,8 +245,8 @@ def from_source_table( source_table=provenance.source_table( source_table.connector, source_table.table_name, source_table.sql, version, - record_locations=(snapshot.manifest - or {}).get('record_locations', True))) + record_sources=(snapshot.manifest + or {}).get('record_sources', True))) return table diff --git a/packages/kumo-engine/test/test_training_table_export.py b/packages/kumo-engine/test/test_training_table_export.py index b43c508ba..95b8755ec 100644 --- a/packages/kumo-engine/test/test_training_table_export.py +++ b/packages/kumo-engine/test/test_training_table_export.py @@ -332,12 +332,12 @@ def test_manifest_and_export_sidecar(graph, tmp_path): assert _seeds(reloaded)[1] == _seeds(table)[1] -def test_custom_table_follows_the_snapshots_record_locations(graph): +def test_custom_table_follows_the_snapshots_record_sources(graph): graph_, _, connector, _ = graph pquery, _, _ = _export(graph, 'PREDICT ents.LABEL FOR EACH ents.ID', 'unlocated') try: - redacted = graph_.snapshot(record_locations=False) + redacted = graph_.snapshot(record_sources=False) table = kumo.TrainingTable.from_source_table( connector['unlocated'], TrainingTableSpec(split_col='SPLIT'), pq=pquery, snapshot=redacted) From d88a447eac828cfacabb18e59499340f7155f7c9 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Wed, 21 Oct 2026 22:29:02 +0000 Subject: [PATCH 29/80] Expect the listing kind for a reloaded export directory #50 types a directory's listing fingerprint as "listing"; the reloaded custom table records that kind. Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/test/test_training_table_export.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/kumo-engine/test/test_training_table_export.py b/packages/kumo-engine/test/test_training_table_export.py index 95b8755ec..68dd9983a 100644 --- a/packages/kumo-engine/test/test_training_table_export.py +++ b/packages/kumo-engine/test/test_training_table_export.py @@ -327,7 +327,7 @@ def test_manifest_and_export_sidecar(graph, tmp_path): 'location': str(directory.resolve()), 'version': connector.source_backend.get_data_version( connector, 'lineage'), - 'version_kind': 'timestamp', + 'version_kind': 'listing', } assert _seeds(reloaded)[1] == _seeds(table)[1] From 918d0f2d122a69bba8ce639c0c0b4a29eef2c016 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Thu, 22 Oct 2026 01:12:53 +0000 Subject: [PATCH 30/80] Record a custom table's source only when the snapshot did A custom training table read the snapshot's record_sources choice from the snapshot manifest and recorded full source locations when that file was missing or unreadable, so a snapshot taken with record_sources=False could still leak locations through the table's manifest. The table now follows GraphSnapshot.record_sources, then the manifest's recorded choice, and records connector and table name only when neither says the snapshot recorded sources. Verified with test_training_table_export.py and test_snapshot_manifest.py. Co-Authored-By: Claude Opus 5.5 --- .../kumo/engine/pquery/training_table.py | 9 +++++++-- .../test/test_training_table_export.py | 20 +++++++++++++------ 2 files changed, 21 insertions(+), 8 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/pquery/training_table.py b/packages/kumo-engine/kumo/engine/pquery/training_table.py index 61996bdb2..f212d665a 100644 --- a/packages/kumo-engine/kumo/engine/pquery/training_table.py +++ b/packages/kumo-engine/kumo/engine/pquery/training_table.py @@ -245,11 +245,16 @@ def from_source_table( source_table=provenance.source_table( source_table.connector, source_table.table_name, source_table.sql, version, - record_sources=(snapshot.manifest - or {}).get('record_sources', True))) + record_sources=_record_sources(snapshot))) return table +def _record_sources(snapshot: GraphSnapshot) -> bool: + if snapshot.record_sources is not None: + return snapshot.record_sources + return (snapshot.manifest or {}).get('record_sources') is True + + def _record_manifest(table: TrainingTable, snapshot: GraphSnapshot, request: Any, *, source_table: Optional[Dict[str, Any]] = None) -> None: diff --git a/packages/kumo-engine/test/test_training_table_export.py b/packages/kumo-engine/test/test_training_table_export.py index 68dd9983a..47c6b08a0 100644 --- a/packages/kumo-engine/test/test_training_table_export.py +++ b/packages/kumo-engine/test/test_training_table_export.py @@ -4,6 +4,7 @@ import json from collections import Counter +from dataclasses import replace from functools import partial from types import SimpleNamespace @@ -332,18 +333,25 @@ def test_manifest_and_export_sidecar(graph, tmp_path): assert _seeds(reloaded)[1] == _seeds(table)[1] -def test_custom_table_follows_the_snapshots_record_sources(graph): - graph_, _, connector, _ = graph +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') - try: - redacted = graph_.snapshot(record_sources=False) + + def location(snapshot): table = kumo.TrainingTable.from_source_table( connector['unlocated'], TrainingTableSpec(split_col='SPLIT'), - pq=pquery, snapshot=redacted) + pq=pquery, snapshot=snapshot) + return table.manifest['parents']['source_table']['location'] + + try: + assert location(graph_.snapshot(record_sources=False)) == 'unlocated' finally: graph_.snapshot() - assert table.manifest['parents']['source_table']['location'] == 'unlocated' + 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( From 27fc36d030ce6595c268e768fa160c238d0e9344 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Thu, 22 Oct 2026 18:42:54 +0000 Subject: [PATCH 31/80] Make the training-table manifest a frozen dataclass TrainingTable.manifest returned a dict like the snapshot manifest did. It now returns a TrainingTableManifest (parents: snapshot and, for a custom table, its source table) built on the shared manifest dataclasses; the file on disk and the exported _kumo_manifest.json are unchanged. Verified with test_training_table_export.py. Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/_manifest.py | 16 ++++++ .../kumo/engine/pquery/training_table.py | 29 +++++++---- .../test/test_training_table_export.py | 51 ++++++++++++++----- 3 files changed, 73 insertions(+), 23 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/_manifest.py b/packages/kumo-engine/kumo/engine/_manifest.py index 0b6c306f6..57adce548 100644 --- a/packages/kumo-engine/kumo/engine/_manifest.py +++ b/packages/kumo-engine/kumo/engine/_manifest.py @@ -205,6 +205,22 @@ class SnapshotManifest(Manifest): record_sources: Optional[bool] = omit_none() +@dataclass(frozen=True, kw_only=True) +class TrainingTableParents(Record): + snapshot: str + source_table: Optional[SourceTableEntry] = omit_none() + + +@dataclass(frozen=True, kw_only=True) +class TrainingTableManifest(Manifest): + kind: str = 'training_table' + parents: TrainingTableParents + query: str + generation_plan: Dict[str, Any] + task_type: str + split_seed: Optional[int] + + def write_manifest(path: Path, manifest: Manifest) -> None: payload = json.dumps(manifest.to_dict(), indent=2, sort_keys=True) + '\n' atomic_write_bytes(path, payload.encode('utf-8')) diff --git a/packages/kumo-engine/kumo/engine/pquery/training_table.py b/packages/kumo-engine/kumo/engine/pquery/training_table.py index f212d665a..990b85dcb 100644 --- a/packages/kumo-engine/kumo/engine/pquery/training_table.py +++ b/packages/kumo-engine/kumo/engine/pquery/training_table.py @@ -14,7 +14,14 @@ from kumoapi.train import TrainingTableSpec from kumo.engine import _provenance as provenance -from kumo.engine._manifest import new_manifest, read_manifest, write_manifest +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 @@ -54,9 +61,10 @@ def tt_id(self) -> str: return read_spec(str(self.staging_dir))['tt_id'] @property - def manifest(self) -> Optional[Dict[str, Any]]: + 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 / STAGING_MANIFEST_FILE, + TrainingTableManifest) @property def source_mode(self) -> Optional[str]: @@ -252,12 +260,13 @@ def from_source_table( def _record_sources(snapshot: GraphSnapshot) -> bool: if snapshot.record_sources is not None: return snapshot.record_sources - return (snapshot.manifest or {}).get('record_sources') is True + manifest = snapshot.manifest + return manifest is not None and manifest.record_sources is True def _record_manifest(table: TrainingTable, snapshot: GraphSnapshot, request: Any, *, - source_table: Optional[Dict[str, Any]] = None) -> None: + source_table: Optional[SourceTableEntry] = None) -> None: path = table.staging_dir / STAGING_MANIFEST_FILE try: from kumo_core.io.json import to_json_dict @@ -265,17 +274,17 @@ def _record_manifest(table: TrainingTable, snapshot: GraphSnapshot, seed = None if request.split_family in _RANDOM_SPLITS: seed = (request.split_config or {}).get('rnd_seed') or 0 - parents: Dict[str, Any] = {'snapshot': snapshot.id} - if source_table is not None: - parents['source_table'] = source_table write_manifest( path, - new_manifest( - 'training_table', table.tt_id, parents, + TrainingTableManifest( + id=table.tt_id, + parents=TrainingTableParents(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( diff --git a/packages/kumo-engine/test/test_training_table_export.py b/packages/kumo-engine/test/test_training_table_export.py index 47c6b08a0..7bd2e20f5 100644 --- a/packages/kumo-engine/test/test_training_table_export.py +++ b/packages/kumo-engine/test/test_training_table_export.py @@ -19,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 @@ -303,12 +310,12 @@ def test_manifest_and_export_sidecar(graph, tmp_path): query = 'PREDICT ents.LABEL FOR EACH ents.ID' pquery, table, _ = _export(graph, query, 'lineage') manifest = table.manifest - assert manifest['id'] == table.tt_id == _spec(table)['tt_id'] - assert manifest['parents'] == {'snapshot': snapshot.id} - assert manifest['query'] == query - assert 'rows' not in manifest and 'source' not in manifest + assert manifest.id == table.tt_id == _spec(table)['tt_id'] + assert manifest.parents == TrainingTableParents(snapshot=snapshot.id) + assert manifest.query == query + assert manifest.extra == {} assert table.source_mode == 'pquery' - assert manifest['split_seed'] == 0 + assert manifest.split_seed == 0 directory = out / 'lineage' baseline = tmp_path / 'baseline' @@ -316,23 +323,41 @@ def test_manifest_and_export_sidecar(graph, tmp_path): for part in baseline.glob('*.parquet'): assert (directory / part.name).read_bytes() == part.read_bytes() assert json.loads( - (directory / '_kumo_manifest.json').read_bytes()) == manifest + (directory / '_kumo_manifest.json').read_bytes()) == manifest.to_dict() reloaded = kumo.TrainingTable.from_source_table( connector['lineage'], TrainingTableSpec(split_col='SPLIT'), pq=pquery, snapshot=snapshot) manifest = reloaded.manifest assert reloaded.source_mode == 'custom_tt' - assert manifest['parents']['source_table'] == { - 'connector': 'LOCAL_FILE', - 'location': str(directory.resolve()), - 'version': connector.source_backend.get_data_version( + 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', - } + 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 / 'training-table-manifest.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 @@ -343,7 +368,7 @@ 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'] + return table.manifest.parents.source_table.location try: assert location(graph_.snapshot(record_sources=False)) == 'unlocated' From 7858aceb75b699cc9274d89047b48e14df47855d Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sat, 24 Oct 2026 17:28:44 +0000 Subject: [PATCH 32/80] Give training tables a deterministic id, embedded snapshot, root store The training-table manifest named its snapshot by id only and its own id was DiskGraph's random tt id, so the same table generated twice got two ids and no artifact store could recognise it. - TrainingTable.id: tt- plus a hash of the snapshot id, query, generation plan, drawn split seed, the custom source's version and the engine/DiskGraph builds (random when a custom source has no version). DiskGraph's extension dir keeps its random tt id. - The manifest is _provenance.json and embeds the snapshot manifest (a stub when the snapshot has none; a bare id reads as a stub). - With MLOps on (mlops= or a kumo.mlops block), each table is stored once under /training_tables//; failures warn. - TrainingTable.load_from_connector(connector, path=, pq=, snapshot=) materializes an export again. The export sidecar follows any file destination (is_file_destination). source_mode is private. Verified: test_training_table_export.py (19: deterministic id, embedded snapshot, stored once, stub reading), custom source, Custom SQL and predictive query suites. Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/docs/mlops.md | 12 ++ packages/kumo-engine/kumo/engine/_manifest.py | 15 ++- .../kumo/engine/pquery/predictive_query.py | 9 +- .../kumo/engine/pquery/training_table.py | 115 +++++++++++++++--- .../test/test_training_table_export.py | 63 +++++++--- 5 files changed, 173 insertions(+), 41 deletions(-) diff --git a/packages/kumo-engine/docs/mlops.md b/packages/kumo-engine/docs/mlops.md index ffc175768..e1480f29e 100644 --- a/packages/kumo-engine/docs/mlops.md +++ b/packages/kumo-engine/docs/mlops.md @@ -23,3 +23,15 @@ with kumo.mlops(cfg): 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, 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//`. diff --git a/packages/kumo-engine/kumo/engine/_manifest.py b/packages/kumo-engine/kumo/engine/_manifest.py index 57adce548..a47a0737f 100644 --- a/packages/kumo-engine/kumo/engine/_manifest.py +++ b/packages/kumo-engine/kumo/engine/_manifest.py @@ -207,18 +207,21 @@ class SnapshotManifest(Manifest): @dataclass(frozen=True, kw_only=True) class TrainingTableParents(Record): - snapshot: str + 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: TrainingTableParents - query: str - generation_plan: Dict[str, Any] - task_type: str - split_seed: Optional[int] + 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 def write_manifest(path: Path, manifest: Manifest) -> None: diff --git a/packages/kumo-engine/kumo/engine/pquery/predictive_query.py b/packages/kumo-engine/kumo/engine/pquery/predictive_query.py index 01a425fe2..5cbf4771c 100644 --- a/packages/kumo-engine/kumo/engine/pquery/predictive_query.py +++ b/packages/kumo-engine/kumo/engine/pquery/predictive_query.py @@ -90,7 +90,7 @@ PredictionTable, _compile_prediction_table, ) -from kumo.engine.pquery.training_table import TrainingTable, _record_manifest +from kumo.engine.pquery.training_table import TrainingTable, _finish logger = logging.getLogger(__name__) @@ -219,6 +219,8 @@ 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 @@ -233,6 +235,8 @@ def generate_training_table( 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'): @@ -316,8 +320,7 @@ def generate_training_table( required_stages=task_plan.task_dataset.required_stages, generation_plan=generation_plan, ) - _record_manifest(table, snapshot, task_plan.train_table) - return table + return _finish(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 990b85dcb..36a045755 100644 --- a/packages/kumo-engine/kumo/engine/pquery/training_table.py +++ b/packages/kumo-engine/kumo/engine/pquery/training_table.py @@ -2,12 +2,14 @@ # All rights reserved. # SPDX-License-Identifier: Apache-2.0 +import hashlib import json import warnings -from dataclasses import dataclass, replace +from dataclasses import dataclass, field, replace from pathlib import Path 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 @@ -28,8 +30,8 @@ if TYPE_CHECKING: from kumo.engine.trainer.output_config import TrainingTableOutputConfig -STAGING_MANIFEST_FILE = 'training-table-manifest.json' -EXPORT_MANIFEST_FILE = '_kumo_manifest.json' +STAGING_MANIFEST_FILE = PROVENANCE_FILE +EXPORT_MANIFEST_FILE = PROVENANCE_FILE _RANDOM_SPLITS = ('RandomSplit', 'RandomLinkSplit') @@ -52,6 +54,9 @@ 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: @@ -67,7 +72,7 @@ def manifest(self) -> Optional[TrainingTableManifest]: TrainingTableManifest) @property - def source_mode(self) -> Optional[str]: + def _source_mode(self) -> Optional[str]: """``pquery`` or ``custom_tt``, as DiskGraph recorded at generation.""" try: return json.loads(self.manifest_path.read_text())['source_mode'] @@ -104,13 +109,14 @@ def export(self, output_config: "TrainingTableOutputConfig") -> str: output_config)) def _export_sidecars(self, connector, output_config) -> Dict[str, Any]: - from kumo_core.object_store_sink import ObjectStoreSinkBackend 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 isinstance(connector.sink_backend, ObjectStoreSinkBackend) + if (not is_file_destination(connector, output_config.table_name) or layout.file_type(name) == layout.CSV or not manifest.is_file()): return {} @@ -121,6 +127,35 @@ def _export_sidecars(self, connector, output_config) -> Dict[str, Any]: 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( cls, @@ -130,6 +165,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. @@ -142,7 +178,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. @@ -248,13 +286,13 @@ def from_source_table( required_stages=task_plan.task_dataset.required_stages, generation_plan=generation_plan, ) - _record_manifest( - table, snapshot, task_plan.train_table, + 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))) - return table + record_sources=_record_sources(snapshot)), + source_version=version) def _record_sources(snapshot: GraphSnapshot) -> bool: @@ -264,22 +302,63 @@ 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: + 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, *, + 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 - seed = None - if request.split_family in _RANDOM_SPLITS: - seed = (request.split_config or {}).get('rnd_seed') or 0 + parent = snapshot.manifest write_manifest( path, TrainingTableManifest( - id=table.tt_id, - parents=TrainingTableParents(snapshot=snapshot.id, - source_table=source_table), + 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, diff --git a/packages/kumo-engine/test/test_training_table_export.py b/packages/kumo-engine/test/test_training_table_export.py index 7bd2e20f5..181518cc1 100644 --- a/packages/kumo-engine/test/test_training_table_export.py +++ b/packages/kumo-engine/test/test_training_table_export.py @@ -310,12 +310,17 @@ def test_manifest_and_export_sidecar(graph, tmp_path): query = 'PREDICT ents.LABEL FOR EACH ents.ID' pquery, table, _ = _export(graph, query, 'lineage') manifest = table.manifest - assert manifest.id == table.tt_id == _spec(table)['tt_id'] - assert manifest.parents == TrainingTableParents(snapshot=snapshot.id) + 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 table.source_mode == 'pquery' + assert table._source_mode == 'pquery' 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' @@ -323,13 +328,13 @@ def test_manifest_and_export_sidecar(graph, tmp_path): for part in baseline.glob('*.parquet'): assert (directory / part.name).read_bytes() == part.read_bytes() assert json.loads( - (directory / '_kumo_manifest.json').read_bytes()) == manifest.to_dict() + (directory / '_provenance.json').read_bytes()) == manifest.to_dict() - reloaded = kumo.TrainingTable.from_source_table( - connector['lineage'], TrainingTableSpec(split_col='SPLIT'), pq=pquery, - snapshot=snapshot) + reloaded = kumo.TrainingTable.load_from_connector( + connector, path='lineage', pq=pquery, snapshot=snapshot) manifest = reloaded.manifest - assert reloaded.source_mode == 'custom_tt' + assert reloaded._source_mode == 'custom_tt' + assert reloaded.id != table.id assert manifest.parents.source_table == SourceTableEntry( connector='LOCAL_FILE', location=str(directory.resolve()), @@ -344,7 +349,7 @@ 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 / 'training-table-manifest.json').read_text()) + (table.staging_dir / '_provenance.json').read_text()) newer['future'] = 1 newer['parents']['label_source'] = 'x' source = tmp_path / 'newer.json' @@ -393,7 +398,7 @@ def fail(*args, **kwargs): _, table, exported = _export( graph, 'PREDICT ents.LABEL FOR EACH ents.ID', 'no_manifest') assert table.manifest is None - assert not (out / 'no_manifest' / '_kumo_manifest.json').exists() + assert not (out / 'no_manifest' / '_provenance.json').exists() assert exported.num_rows == sum(_seeds(table)[1].values()) @@ -402,15 +407,15 @@ def test_unreadable_manifest_still_exports(graph): table = kumo.PredictiveQuery( graph=graph_, query='PREDICT ents.LABEL FOR EACH ents.ID' ).generate_training_table(snapshot) - manifest = table.staging_dir / 'training-table-manifest.json' + manifest = table.staging_dir / '_provenance.json' manifest.chmod(0) try: - with pytest.warns(UserWarning, match='_kumo_manifest.json'): + 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' / '_kumo_manifest.json').exists() + assert not (out / 'unreadable' / '_provenance.json').exists() def test_no_sidecar_for_warehouse_or_csv_exports(graph): @@ -429,4 +434,34 @@ def write_files(connector, table_name, files, write_mode, **kwargs): 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' / '_kumo_manifest.json').exists() + 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() + + +def test_a_bare_snapshot_id_parent_reads_as_a_stub(tmp_path): + path = tmp_path / '_provenance.json' + path.write_text(json.dumps({ + 'kind': 'training_table', 'schema_version': 1, 'id': 'tt-1', + 'parents': {'snapshot': 'abc@0123'}})) + manifest = read_manifest(path, TrainingTableManifest) + assert manifest.parents.snapshot.id == 'abc@0123' + assert manifest.parents.snapshot.reconstructed From 732f4c3245006eb6423f544fbc9f8c781ac1081e Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sat, 24 Oct 2026 17:37:10 +0000 Subject: [PATCH 33/80] Export TrainingTableManifest Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/__init__.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/packages/kumo-engine/kumo/engine/__init__.py b/packages/kumo-engine/kumo/engine/__init__.py index 64d4df37b..52b0540ae 100644 --- a/packages/kumo-engine/kumo/engine/__init__.py +++ b/packages/kumo-engine/kumo/engine/__init__.py @@ -25,7 +25,7 @@ from kumoapi.plugin import ModelRegistryLocation from kumoapi.table import Column -from kumo.engine._manifest import SnapshotManifest +from kumo.engine._manifest import SnapshotManifest, TrainingTableManifest from kumo.engine._mlops import MLOps, mlops from kumo.engine._version import __version__ from kumo.engine.config import configure @@ -88,6 +88,7 @@ "MLOps", "mlops", "SnapshotManifest", + "TrainingTableManifest", "PredictiveQuery", "TrainingTable", "PredictionTable", From cca0d0028416949f9f3b25de4dce018be9bdb3ab Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Tue, 20 Oct 2026 20:49:35 +0000 Subject: [PATCH 34/80] Record bundle_id and fit provenance in every model bundle A fitted model bundle had no stable id and recorded no fit time, seed, trials, library versions or parent ids. verify_model_bundle already computed a 16-hex content digest that every caller discarded. verify_model_bundle now returns the 64-hex bundle_id: sha256 over the sorted [relative_path, sha256] pairs of the verified files (plus an unlisted training-spec.json, as before), excluding provenance.json. TrainingResult.bundle_id and Trainer.bundle_id carry it from fit and load; older bundles get an id at load without being rewritten. A new kumo-core add_bundle_file lists a written file in the publication manifest and replaces the engine's own training-spec amend. Fit writes provenance.json (role provenance, integrity-covered, outside the id): parents (training table id, snapshot id, warm-start bundle id), seed, fit times, selected trial, per-trial status/seconds/best epoch/val metrics/varying plan params, environment and scheduler. TrainingResult.provenance reads it only when listed. Provenance and id failures warn; the bundle still publishes. Verified with new tests for the id formula and exclusion, add_bundle_file, a two-trial fit's provenance, id stability across load and copy, a bundle without provenance loading and predicting identically, a stray unlisted provenance file, a failing provenance write and the warm-start parent id; plus the registry, compatibility, warm-start, holdout, snapshot, carry- forward, explanation, training-table export, Databricks and Snowflake registry suites. Co-Authored-By: Claude Opus 5.5 --- packages/kumo-core/src/kumo_core/artifacts.py | 56 ++++- packages/kumo-core/tests/test_artifacts.py | 42 ++++ .../kumo/engine/trainer/_provenance.py | 146 +++++++++++++ .../kumo-engine/kumo/engine/trainer/_spec.py | 43 +--- .../kumo/engine/trainer/trainer.py | 53 ++++- .../test/test_bundle_provenance.py | 194 ++++++++++++++++++ .../kumo-engine/test/test_trainer_registry.py | 10 +- packages/kumo-engine/test/test_warm_start.py | 1 + 8 files changed, 493 insertions(+), 52 deletions(-) create mode 100644 packages/kumo-engine/kumo/engine/trainer/_provenance.py create mode 100644 packages/kumo-engine/test/test_bundle_provenance.py diff --git a/packages/kumo-core/src/kumo_core/artifacts.py b/packages/kumo-core/src/kumo_core/artifacts.py index 838f4c0c4..7ea69b635 100644 --- a/packages/kumo-core/src/kumo_core/artifacts.py +++ b/packages/kumo-core/src/kumo_core/artifacts.py @@ -248,14 +248,17 @@ def _sha256(path: Path) -> str: return digest.hexdigest() +PROVENANCE_FILE = 'provenance.json' + + 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: the SHA-256 of the sorted ``[relative_path, + sha256]`` pairs of 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; its digest still participates in the id. """ from .sources import layout from .sources.errors import PermanentSourceError @@ -273,7 +276,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: @@ -325,7 +328,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: " @@ -335,11 +338,42 @@ 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] + payload = json.dumps( + sorted(pair for pair in verified 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() + ] + temporary = manifest_path.with_name(f'.{manifest_path.name}.tmp') + try: + with temporary.open('w', encoding='utf-8') as handle: + handle.write(json.dumps(manifest, indent=2, sort_keys=True) + '\n') + handle.flush() + os.fsync(handle.fileno()) + os.replace(temporary, manifest_path) + FilesystemArtifactStore._fsync_directory(bundle_dir) + finally: + temporary.unlink(missing_ok=True) @runtime_checkable diff --git a/packages/kumo-core/tests/test_artifacts.py b/packages/kumo-core/tests/test_artifacts.py index 071f4339b..e29622b96 100644 --- a/packages/kumo-core/tests/test_artifacts.py +++ b/packages/kumo-core/tests/test_artifacts.py @@ -15,7 +15,9 @@ FilesystemVitalArtifactPublisher, VitalArtifactPublicationPolicy, VitalArtifactSources, + add_bundle_file, resolve_vital_artifact_sources, + verify_model_bundle, ) from kumo_core.config import NativeRuntimeConfig from kumo_core.manifest import RunManifest, RunPaths @@ -237,3 +239,43 @@ 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_covers_listed_files_except_provenance( + tmp_path: Path) -> None: + sources, _ = _vital_sources(tmp_path) + bundle = tmp_path / "bundle" + VitalArtifactPublicationPolicy().publish(sources, bundle) + (bundle / "training-spec.json").write_text("{}") + entries = json.loads( + (bundle / "publication-manifest.json").read_text())["artifacts"] + pairs = sorted([entry["relative_path"], entry["sha256"]] + for entry in entries) + pairs = sorted(pairs + [[ + "training-spec.json", + hashlib.sha256(b"{}").hexdigest(), + ]]) + expected = hashlib.sha256( + json.dumps(pairs, separators=(",", ":")).encode()).hexdigest() + + assert verify_model_bundle(bundle) == expected + add_bundle_file(bundle, "training_spec", "training-spec.json") + assert verify_model_bundle(bundle) == expected + + (bundle / "provenance.json").write_text('{"a": 1}') + add_bundle_file(bundle, "provenance", "provenance.json") + assert verify_model_bundle(bundle) == expected + (bundle / "provenance.json").write_text('{"a": 2}') + with pytest.raises(ValueError, match="provenance.json"): + verify_model_bundle(bundle) + add_bundle_file(bundle, "provenance", "provenance.json") + assert verify_model_bundle(bundle) == expected + + 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) != expected 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..31740c462 --- /dev/null +++ b/packages/kumo-engine/kumo/engine/trainer/_provenance.py @@ -0,0 +1,146 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +r"""Descriptive fit record stored as ``provenance.json`` in a model bundle. + +It is integrity-covered by the publication manifest but excluded from the +bundle id, so its volatile fields (times, environment, scheduler) never change +which model a bundle is. +""" + +import csv +import json +import math +import re +from datetime import datetime, timezone +from pathlib import Path +from typing import Any, Dict, List, Optional + +from kumo.engine._manifest import ( + _flatten, + new_manifest, + read_manifest, + write_manifest, +) +from kumo.engine.pquery.training_table import TrainingTable +from kumo.engine.trainer._spec import ( + MANIFEST_FILE, + PROVENANCE_ROLE, + SPEC_VERSION, +) + + +def write_provenance( + bundle_dir: Path, + *, + train_table: TrainingTable, + snapshot_id: str, + warm_start_bundle_id: Optional[str], + started_at: datetime, + home_dir: Path, +) -> None: + from kumo_core.artifacts import PROVENANCE_FILE, add_bundle_file + + config = _read_json(bundle_dir / 'model' / 'config.json') or {} + finished_at = datetime.now(timezone.utc) + parents = { + 'training_table': train_table.tt_id, + 'training_table_content_sha256': + (train_table.manifest or {}).get('content_sha256'), + 'snapshot': snapshot_id, + 'warm_start': warm_start_bundle_id, + } + write_manifest( + bundle_dir / PROVENANCE_FILE, + new_manifest( + 'model_bundle', None, parents, + bundle_format_version=SPEC_VERSION, + seed=config.get('random_state_seed'), + fit={ + 'started_at': started_at.isoformat(timespec='seconds'), + 'finished_at': finished_at.isoformat(timespec='seconds'), + 'seconds': round((finished_at - started_at).total_seconds(), + 3), + }, + selected_trial=config.get('version'), + trials=_trials(home_dir, Path(config['root_dir'])), + )) + add_bundle_file(bundle_dir, PROVENANCE_ROLE, PROVENANCE_FILE) + + +def read_provenance(bundle_dir: Path) -> Optional[Dict[str, Any]]: + manifest = _read_json(bundle_dir / MANIFEST_FILE) or {} + for entry in manifest.get('artifacts') or []: + if entry.get('role') == PROVENANCE_ROLE: + return read_manifest(bundle_dir / entry['relative_path']) + return None + + +def _trials(home_dir: Path, automl_dir: Path) -> List[Dict[str, Any]]: + status = { + entry['trial']: entry + for entry in (_read_json(home_dir / 'status.json') or {}).get( + 'trials') or [] + } + versions = { + int(match.group(1)) + for path in automl_dir.iterdir() + if (match := re.fullmatch(r'version_(\d+)', path.name)) + } + trials, plans = [], {} + for trial in sorted(status.keys() | versions): + log_dir = automl_dir / f'version_{trial}' + config = _read_json(log_dir / 'config.json') or {} + progress = _read_json(log_dir / 'train_progress.json') or {} + best_epoch = (config.get('refit_cfg') or {}).get('best_val_epoch', -1) + succeeded = status.get(trial, {}).get('succeed') + trials.append({ + 'trial': trial, + 'status': (None if succeeded is None else + 'succeeded' if succeeded else 'failed'), + 'seconds': progress.get('elapsed_time'), + 'best_epoch': best_epoch if best_epoch >= 0 else None, + 'metrics': _val_metrics(log_dir / 'val_result.csv'), + }) + plans[trial] = { + key: json.dumps(value, sort_keys=True) + for key, value in _flatten( + _read_json(log_dir / 'model_plan.json') or {}).items() + } + varying = sorted( + key for key in {key + for plan in plans.values() for key in plan} + if len({plan.get(key) + for plan in plans.values()}) > 1) + for entry in trials: + plan = plans[entry['trial']] + entry['params'] = { + key: json.loads(plan[key]) if key in plan else None + for key in varying + } + return trials + + +def _val_metrics(path: Path) -> Dict[str, float]: + try: + with open(path, newline='', encoding='utf-8') as handle: + row = next(csv.DictReader(handle), None) or {} + except OSError: + return {} + metrics = {} + for key, value in row.items(): + try: + number = float(value) + except (TypeError, ValueError): + continue + if key.startswith('val_') and math.isfinite(number): + metrics[key] = number + return metrics + + +def _read_json(path: Path) -> Optional[Dict[str, Any]]: + try: + data = json.loads(path.read_text(encoding='utf-8')) + except (OSError, ValueError): + return None + return data if isinstance(data, dict) else None diff --git a/packages/kumo-engine/kumo/engine/trainer/_spec.py b/packages/kumo-engine/kumo/engine/trainer/_spec.py index 147817c0e..637714d45 100644 --- a/packages/kumo-engine/kumo/engine/trainer/_spec.py +++ b/packages/kumo-engine/kumo/engine/trainer/_spec.py @@ -9,11 +9,10 @@ separate and must be supplied through a compatible prediction snapshot. """ -import hashlib import json from dataclasses import asdict, dataclass from pathlib import Path -from typing import Any, Dict +from typing import Any, Dict, Tuple from kumo_core.io.filesystem import atomic_write_bytes from kumoapi.model_plan import ModelPlan, TrainingTableGenerationPlan @@ -23,10 +22,10 @@ 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' @dataclass(frozen=True) @@ -45,12 +44,6 @@ class TrainingSpec: training_table_generation: Dict[str, Any] -def _sha256(path: Path) -> str: - digest = hashlib.sha256() - digest.update(path.read_bytes()) - return digest.hexdigest() - - def write_training_spec( bundle_dir: Path, *, @@ -60,6 +53,7 @@ def write_training_spec( training_table_generation: TrainingTableGenerationPlan, ) -> 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( @@ -75,23 +69,7 @@ def write_training_spec( (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_bytes( - 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: @@ -125,11 +103,12 @@ def read_training_spec(bundle_dir: Path) -> TrainingSpec: ) -def read_bundle_artifact_paths(bundle_dir: Path) -> Dict[str, str]: - r"""Rebuild the verified ``role -> path`` artifact mapping.""" +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( @@ -150,5 +129,5 @@ 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) + }, bundle_id diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index 28c3156a3..bc0824125 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -83,6 +83,9 @@ class TrainingResult: 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. + + `bundle_id` is the content id of the model bundle's scoring files, + computed when the bundle is written or loaded. """ metrics: Dict[str, float] @@ -91,6 +94,14 @@ class TrainingResult: published_model: Optional[PublishedModel] = None holdout_path: Optional[Path] = None holdout_truncated: Optional[bool] = False + bundle_id: Optional[str] = None + + @property + def provenance(self) -> Optional[Dict[str, Any]]: + """The bundle's recorded fit provenance; None for older bundles.""" + from kumo.engine.trainer._provenance import read_provenance + + return read_provenance(self.bundle_dir) def holdout_df(self) -> pd.DataFrame: r"""Read the holdout predictions as a pandas DataFrame. @@ -299,6 +310,13 @@ def training_table_generation_plan( def scratch_dir(self) -> Path: return trainer_scratch_dir(self._id) + @property + def bundle_id(self) -> Optional[str]: + """Content id of the fitted or loaded model bundle.""" + if self._training_result is None: + return None + return self._training_result.bundle_id + @_without_torch_scatter_hint def fit( self, @@ -363,18 +381,21 @@ def fit( raise NotImplementedError( f'{type(export_config.connector).__name__} cannot publish ' f'models') + 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'], @@ -524,6 +545,10 @@ def fit( training_table_generation=train_table.generation_plan, export_config=export_config, model_plan=effective_plan, + provenance=dict(train_table=train_table, + snapshot_id=snapshot.id, + warm_start_bundle_id=warm_start_bundle_id, + started_at=started_at, home_dir=run_dir), ) def _finish_fit( @@ -535,8 +560,13 @@ def _finish_fit( training_table_generation: TrainingTableGenerationPlan, export_config: Optional[ModelOutputConfig], model_plan: Optional[ModelPlan] = None, + provenance: Optional[Mapping[str, Any]] = None, ) -> TrainingResult: """Commit durable fit outputs before making this trainer usable.""" + from kumo_core.artifacts import verify_model_bundle + + 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, @@ -545,6 +575,20 @@ def _finish_fit( snapshot_id=snapshot_id, training_table_generation=training_table_generation, ) + if provenance is not None: + try: + write_provenance(training_result.bundle_dir, **provenance) + except Exception as error: + warnings.warn( + f"Publishing the model bundle without provenance.json: " + f"{error}", UserWarning, stacklevel=3) + try: + training_result = replace( + training_result, + bundle_id=verify_model_bundle(training_result.bundle_dir)) + except (OSError, ValueError) as error: + warnings.warn(f"Could not compute the bundle id: {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( @@ -837,7 +881,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) @@ -849,5 +893,6 @@ def _load_bundle( 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) + trainer._training_result = TrainingResult(metrics, paths, bundle_dir, + bundle_id=bundle_id) return trainer 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..443c90742 --- /dev/null +++ b/packages/kumo-engine/test/test_bundle_provenance.py @@ -0,0 +1,194 @@ +# 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 + +import numpy as np +import pandas as pd +import pytest +from kumo_core.artifacts import add_bundle_file +from kumoapi.graph import Edge +from kumoapi.model_plan import ModelPlan, TrainingTableGenerationPlan +from kumoapi.typing import Stype + +import kumo.engine as kumo +from kumo.engine._scratch import get_scratch_dir +from kumo.engine.trainer._spec import MANIFEST_FILE +from kumo.engine.trainer.trainer import TrainingResult + + +@pytest.fixture(scope='module') +def two_trials(tmp_path_factory): + root = tmp_path_factory.mktemp('provenance') + previous = get_scratch_dir() + kumo.configure(scratch_dir=root / 'scratch') + try: + rng = np.random.default_rng(0) + data = root / 'data' + data.mkdir() + pd.DataFrame({'ID': np.arange(50)}).to_parquet(data / 'ents.parquet') + pd.DataFrame({ + 'EID': + np.arange(3000), + 'ENT_ID': + rng.integers(0, 50, 3000), + 'TS': + pd.Timestamp('2024-01-01') + + pd.to_timedelta(rng.integers(0, 120, 3000), unit='D'), + }).to_parquet(data / 'events.parquet') + connector = kumo.FileConnector(root_dir=str(data)) + events = kumo.Table.from_source_table( + connector['events.parquet'], primary_key='EID', + time_column='TS').infer_metadata().with_column( + 'ENT_ID', stype=Stype.ID) + graph = kumo.Graph( + tables={ + 'ents': + kumo.Table.from_source_table(connector['ents.parquet'], + primary_key='ID').infer_metadata(), + 'events': + events, + }, + edges=[Edge(src_table='events', fkey='ENT_ID', dst_table='ents')]) + snapshot = graph.snapshot() + pquery = kumo.PredictiveQuery( + graph=graph, query='PREDICT COUNT(events.*, 0, 7, days) > 0 ' + 'FOR EACH ents.ID') + plan = ModelPlan() + plan.training_job.disable_explain = True + plan.training_job.num_experiments = 2 + plan.optimization.max_epochs = 1 + plan.optimization.max_steps_per_epoch = 2 + table = pquery.generate_training_table(snapshot) + trainer = kumo.Trainer(model_plan=plan) + result = trainer.fit(table, snapshot=snapshot) + yield (trainer, result, table, snapshot, + pquery.generate_prediction_table(snapshot)) + finally: + kumo.configure(scratch_dir=previous) + + +def _predict(trainer, table, snapshot): + result = trainer.predict(table, snapshot=snapshot, + binary_classification_threshold=.5) + return result.frame.sort_values('ENTITY').reset_index(drop=True) + + +def _roles(bundle): + manifest = json.loads((bundle / MANIFEST_FILE).read_text()) + return {entry['role']: entry for entry in manifest['artifacts']} + + +def _strip_provenance(bundle): + manifest = json.loads((bundle / MANIFEST_FILE).read_text()) + manifest['artifacts'] = [ + entry for entry in manifest['artifacts'] + if entry['role'] != 'provenance' + ] + (bundle / MANIFEST_FILE).write_text(json.dumps(manifest)) + (bundle / 'provenance.json').unlink() + + +def test_fit_records_listed_provenance(two_trials): + trainer, result, table, snapshot, _ = two_trials + assert _roles(result.bundle_dir)['provenance']['relative_path'] == ( + 'provenance.json') + provenance = result.provenance + assert provenance['kind'] == 'model_bundle' + assert provenance['bundle_format_version'] == 1 + assert provenance['parents'] == { + 'training_table': table.tt_id, + 'training_table_content_sha256': None, + 'snapshot': snapshot.id, + 'warm_start': None, + } + assert provenance['seed'] is None + assert provenance['fit']['seconds'] >= 0 + assert [trial['trial'] for trial in provenance['trials']] == [0, 1] + assert {trial['status'] for trial in provenance['trials']} == {'succeeded'} + assert provenance['selected_trial'] in (0, 1) + for trial in provenance['trials']: + assert trial['best_epoch'] == 1 + assert trial['seconds'] > 0 + assert 'val_auroc' in trial['metrics'] + params = [trial['params'] for trial in provenance['trials']] + assert params[0] and params[0].keys() == params[1].keys() + assert all(params[0][key] != params[1][key] for key in params[0]) + assert provenance['environment']['kumo_engine'] == kumo.__version__ + assert trainer.bundle_id == result.bundle_id + + +def test_bundle_id_is_content_identity(two_trials, tmp_path): + trainer, result, *_ = two_trials + assert re.fullmatch('[0-9a-f]{64}', result.bundle_id) + assert kumo.Trainer.load(result.bundle_dir).bundle_id == result.bundle_id + copy = tmp_path / 'copy' + shutil.copytree(result.bundle_dir, copy) + assert kumo.Trainer.load(copy).bundle_id == result.bundle_id + + provenance = json.loads((copy / 'provenance.json').read_text()) + provenance['fit']['seconds'] = -1 + (copy / 'provenance.json').write_text(json.dumps(provenance)) + add_bundle_file(copy, 'provenance', 'provenance.json') + reloaded = kumo.Trainer.load(copy) + assert reloaded.bundle_id == result.bundle_id + assert reloaded._training_result.provenance['fit']['seconds'] == -1 + + config = copy / 'model' / 'config.json' + config.write_text(config.read_text() + ' ') + add_bundle_file(copy, 'trainer_config', 'model/config.json') + assert kumo.Trainer.load(copy).bundle_id != result.bundle_id + + +def test_bundle_without_provenance_loads_and_predicts_identically( + two_trials, tmp_path): + trainer, result, _, snapshot, prediction_table = two_trials + old = tmp_path / 'old' + shutil.copytree(result.bundle_dir, old) + _strip_provenance(old) + + restored = kumo.Trainer.load(old) + assert restored.bundle_id == result.bundle_id + assert restored._training_result.provenance is None + pd.testing.assert_frame_equal(_predict(restored, prediction_table, + snapshot), + _predict(trainer, prediction_table, snapshot)) + + (old / 'provenance.json').write_text(json.dumps(result.provenance)) + stray = kumo.Trainer.load(old) + assert stray._training_result.provenance is None + assert stray.bundle_id == result.bundle_id + + +def test_provenance_failure_warns_and_still_publishes(tmp_path, monkeypatch): + from kumo.engine.trainer import _provenance + + bundle = tmp_path / 'bundle' + (bundle / 'model').mkdir(parents=True) + (bundle / 'model' / 'config.json').write_text('{"root_dir": "automl"}') + (bundle / MANIFEST_FILE).write_text(json.dumps({'artifacts': []})) + add_bundle_file(bundle, 'trainer_config', 'model/config.json') + + def fail(*args): + raise RuntimeError('unreadable trials') + + monkeypatch.setattr(_provenance, '_trials', fail) + table = type('Table', (), {'tt_id': 'tt-1', 'manifest': None})() + trainer = kumo.Trainer(ModelPlan()) + with pytest.warns(UserWarning, match='unreadable trials'): + 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, snapshot_id='snap', + warm_start_bundle_id=None, + started_at=datetime.now(timezone.utc), + home_dir=tmp_path)) + assert 'provenance' not in _roles(bundle) + assert result.provenance is None + assert re.fullmatch('[0-9a-f]{64}', result.bundle_id) + assert trainer._training_result is result diff --git a/packages/kumo-engine/test/test_trainer_registry.py b/packages/kumo-engine/test/test_trainer_registry.py index 9900d5ff0..e668544b2 100644 --- a/packages/kumo-engine/test/test_trainer_registry.py +++ b/packages/kumo-engine/test/test_trainer_registry.py @@ -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).bundle_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_warm_start.py b/packages/kumo-engine/test/test_warm_start.py index e79130045..d60a645cb 100644 --- a/packages/kumo-engine/test/test_warm_start.py +++ b/packages/kumo-engine/test/test_warm_start.py @@ -182,6 +182,7 @@ 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.provenance['parents']['warm_start'] == source.bundle_id assert destination.model_plan == configured_plan assert _hash_bundle(durable) == original_hashes assert _hash_bundle(source_bundle) == original_hashes From 1b676995074f2b6feaf7864edfc7a492510e38cc Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Tue, 20 Oct 2026 21:03:27 +0000 Subject: [PATCH 35/80] Address bundle provenance review feedback Any bundle-id failure now warns instead of only OSError/ValueError. Provenance drops the always-null training-table content digest and records a training_table object (id, derived_from, modified) copied from the training-table manifest. Missing trial logs yield an empty trials list so parents, seed, times and environment are still written. Single-element trial params are unwrapped. new_manifest accepts a null id. Verified with the bundle provenance (incl. a new missing-trial-logs test), trainer registry, warm start and kumo-core artifact suites. Co-Authored-By: Claude Opus 5.5 --- .../kumo/engine/trainer/_provenance.py | 26 ++++++++--- .../kumo/engine/trainer/trainer.py | 2 +- .../test/test_bundle_provenance.py | 45 ++++++++++++------- 3 files changed, 51 insertions(+), 22 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/trainer/_provenance.py b/packages/kumo-engine/kumo/engine/trainer/_provenance.py index 31740c462..6c6ae586e 100644 --- a/packages/kumo-engine/kumo/engine/trainer/_provenance.py +++ b/packages/kumo-engine/kumo/engine/trainer/_provenance.py @@ -43,10 +43,10 @@ def write_provenance( config = _read_json(bundle_dir / 'model' / 'config.json') or {} finished_at = datetime.now(timezone.utc) + tt_id = train_table.tt_id + table_manifest = train_table.manifest or {} parents = { - 'training_table': train_table.tt_id, - 'training_table_content_sha256': - (train_table.manifest or {}).get('content_sha256'), + 'training_table': tt_id, 'snapshot': snapshot_id, 'warm_start': warm_start_bundle_id, } @@ -62,8 +62,14 @@ def write_provenance( 'seconds': round((finished_at - started_at).total_seconds(), 3), }, + training_table={ + 'id': tt_id, + 'derived_from': + (table_manifest.get('parents') or {}).get('derived_from'), + 'modified': table_manifest.get('modified'), + }, selected_trial=config.get('version'), - trials=_trials(home_dir, Path(config['root_dir'])), + trials=_trials(home_dir, config.get('root_dir')), )) add_bundle_file(bundle_dir, PROVENANCE_ROLE, PROVENANCE_FILE) @@ -76,7 +82,11 @@ def read_provenance(bundle_dir: Path) -> Optional[Dict[str, Any]]: return None -def _trials(home_dir: Path, automl_dir: Path) -> List[Dict[str, Any]]: +def _trials(home_dir: Path, + root_dir: Optional[str]) -> List[Dict[str, Any]]: + automl_dir = Path(root_dir) if root_dir else None + if automl_dir is None or not automl_dir.is_dir(): + return [] status = { entry['trial']: entry for entry in (_read_json(home_dir / 'status.json') or {}).get( @@ -115,12 +125,16 @@ def _trials(home_dir: Path, automl_dir: Path) -> List[Dict[str, Any]]: for entry in trials: plan = plans[entry['trial']] entry['params'] = { - key: json.loads(plan[key]) if key in plan else None + key: _unwrap(json.loads(plan[key])) if key in plan else None for key in varying } return trials +def _unwrap(value: Any) -> Any: + return value[0] if isinstance(value, list) and len(value) == 1 else value + + def _val_metrics(path: Path) -> Dict[str, float]: try: with open(path, newline='', encoding='utf-8') as handle: diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index bc0824125..b7ec406c1 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -586,7 +586,7 @@ def _finish_fit( training_result = replace( training_result, bundle_id=verify_model_bundle(training_result.bundle_dir)) - except (OSError, ValueError) as error: + except Exception as error: warnings.warn(f"Could not compute the bundle id: {error}", UserWarning, stacklevel=3) if export_config is not None: diff --git a/packages/kumo-engine/test/test_bundle_provenance.py b/packages/kumo-engine/test/test_bundle_provenance.py index 443c90742..ba3e383de 100644 --- a/packages/kumo-engine/test/test_bundle_provenance.py +++ b/packages/kumo-engine/test/test_bundle_provenance.py @@ -102,10 +102,14 @@ def test_fit_records_listed_provenance(two_trials): assert provenance['bundle_format_version'] == 1 assert provenance['parents'] == { 'training_table': table.tt_id, - 'training_table_content_sha256': None, 'snapshot': snapshot.id, 'warm_start': None, } + assert provenance['training_table'] == { + 'id': table.tt_id, + 'derived_from': None, + 'modified': None, + } assert provenance['seed'] is None assert provenance['fit']['seconds'] >= 0 assert [trial['trial'] for trial in provenance['trials']] == [0, 1] @@ -118,6 +122,8 @@ def test_fit_records_listed_provenance(two_trials): params = [trial['params'] for trial in provenance['trials']] assert params[0] and params[0].keys() == params[1].keys() assert all(params[0][key] != params[1][key] for key in params[0]) + assert not any(isinstance(value, list) and len(value) == 1 + for value in params[0].values()) assert provenance['environment']['kumo_engine'] == kumo.__version__ assert trainer.bundle_id == result.bundle_id @@ -164,31 +170,40 @@ def test_bundle_without_provenance_loads_and_predicts_identically( assert stray.bundle_id == result.bundle_id -def test_provenance_failure_warns_and_still_publishes(tmp_path, monkeypatch): - from kumo.engine.trainer import _provenance - +def _finish_stub_fit(tmp_path, trainer): bundle = tmp_path / 'bundle' (bundle / 'model').mkdir(parents=True) - (bundle / 'model' / 'config.json').write_text('{"root_dir": "automl"}') + (bundle / 'model' / 'config.json').write_text('{}') (bundle / MANIFEST_FILE).write_text(json.dumps({'artifacts': []})) add_bundle_file(bundle, 'trainer_config', 'model/config.json') + table = type('Table', (), {'tt_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, snapshot_id='snap', + warm_start_bundle_id=None, + started_at=datetime.now(timezone.utc), + home_dir=tmp_path)) + + +def test_provenance_without_trial_logs_still_records_parents(tmp_path): + result = _finish_stub_fit(tmp_path, kumo.Trainer(ModelPlan())) + assert result.provenance['trials'] == [] + assert result.provenance['parents']['training_table'] == 'tt-1' + + +def test_provenance_failure_warns_and_still_publishes(tmp_path, monkeypatch): + from kumo.engine.trainer import _provenance def fail(*args): raise RuntimeError('unreadable trials') monkeypatch.setattr(_provenance, '_trials', fail) - table = type('Table', (), {'tt_id': 'tt-1', 'manifest': None})() trainer = kumo.Trainer(ModelPlan()) with pytest.warns(UserWarning, match='unreadable trials'): - 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, snapshot_id='snap', - warm_start_bundle_id=None, - started_at=datetime.now(timezone.utc), - home_dir=tmp_path)) - assert 'provenance' not in _roles(bundle) + result = _finish_stub_fit(tmp_path, trainer) + assert 'provenance' not in _roles(result.bundle_dir) assert result.provenance is None assert re.fullmatch('[0-9a-f]{64}', result.bundle_id) assert trainer._training_result is result From d15d97e367407e73d1973219fafa53751ef2df18 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Tue, 20 Oct 2026 22:35:09 +0000 Subject: [PATCH 36/80] Read Trainer.bundle_id defensively Training results built without a bundle_id field (test doubles) made the property raise; return None for them. Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/trainer/trainer.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index b7ec406c1..28596f615 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -313,9 +313,7 @@ def scratch_dir(self) -> Path: @property def bundle_id(self) -> Optional[str]: """Content id of the fitted or loaded model bundle.""" - if self._training_result is None: - return None - return self._training_result.bundle_id + return getattr(self._training_result, 'bundle_id', None) @_without_torch_scatter_hint def fit( From ffbf2065700eab97046efd42cc6071fb6c43db4c Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Tue, 20 Oct 2026 23:56:18 +0000 Subject: [PATCH 37/80] Record the published model's fit mode in provenance.json metrics.json holds test_* (the fit model) and test_refit_* (the refit_trainval model) and nothing says which, if any, belongs to the shipped checkpoint; a refit_full model has none of its own. Record published_model ('fit', 'refit_trainval' or 'refit_full') from the best trainer's refit flags, the order kumo-core publishes in; null when unknown. metrics.json is unchanged. Verified with a two-trial fit and each combination of refit flags. Co-Authored-By: Claude Opus 5.5 --- .../kumo/engine/trainer/_provenance.py | 2 ++ .../kumo/engine/trainer/trainer.py | 20 +++++++++++++++++-- .../test/test_bundle_provenance.py | 20 +++++++++++++++++-- 3 files changed, 38 insertions(+), 4 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/trainer/_provenance.py b/packages/kumo-engine/kumo/engine/trainer/_provenance.py index 6c6ae586e..222fc931c 100644 --- a/packages/kumo-engine/kumo/engine/trainer/_provenance.py +++ b/packages/kumo-engine/kumo/engine/trainer/_provenance.py @@ -38,6 +38,7 @@ def write_provenance( warm_start_bundle_id: Optional[str], started_at: datetime, home_dir: Path, + published_model: Optional[str], ) -> None: from kumo_core.artifacts import PROVENANCE_FILE, add_bundle_file @@ -69,6 +70,7 @@ def write_provenance( 'modified': table_manifest.get('modified'), }, selected_trial=config.get('version'), + published_model=published_model, trials=_trials(home_dir, config.get('root_dir')), )) add_bundle_file(bundle_dir, PROVENANCE_ROLE, PROVENANCE_FILE) diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index 28596f615..4794169ca 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -98,7 +98,11 @@ class TrainingResult: @property def provenance(self) -> Optional[Dict[str, Any]]: - """The bundle's recorded fit provenance; None for older bundles.""" + """The bundle's recorded fit provenance; None for older bundles. + + Its `published_model` is the shipped checkpoint's fit mode, unlike + the registry record in `TrainingResult.published_model`. + """ from kumo.engine.trainer._provenance import read_provenance return read_provenance(self.bundle_dir) @@ -210,6 +214,16 @@ def _published_holdout(trainer: Any) -> tuple[Optional[Path], Optional[bool]]: return _find_holdout(trainer.get_log_dir(fit_mode=published)) +def _published_fit_mode(trainer: Any) -> Optional[str]: + refit = getattr(trainer, 'refit_cfg', None) + for mode, flag in (('refit_full', 'refitted_full'), + ('refit_trainval', 'refitted_trainval'), + ('fit', 'fitted')): + if getattr(refit, flag, False) is True: + return mode + return None + + 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. @@ -527,6 +541,7 @@ def fit( holdout_path, holdout_truncated = _published_holdout( result.best_trainer) + published_model = _published_fit_mode(result.best_trainer) training_result = TrainingResult( metrics=dict(result.metrics), artifact_paths=dict(result.artifact_paths), @@ -546,7 +561,8 @@ def fit( provenance=dict(train_table=train_table, snapshot_id=snapshot.id, warm_start_bundle_id=warm_start_bundle_id, - started_at=started_at, home_dir=run_dir), + started_at=started_at, home_dir=run_dir, + published_model=published_model), ) def _finish_fit( diff --git a/packages/kumo-engine/test/test_bundle_provenance.py b/packages/kumo-engine/test/test_bundle_provenance.py index ba3e383de..9c4363303 100644 --- a/packages/kumo-engine/test/test_bundle_provenance.py +++ b/packages/kumo-engine/test/test_bundle_provenance.py @@ -6,6 +6,7 @@ import re import shutil from datetime import datetime, timezone +from types import SimpleNamespace import numpy as np import pandas as pd @@ -18,7 +19,7 @@ import kumo.engine as kumo from kumo.engine._scratch import get_scratch_dir from kumo.engine.trainer._spec import MANIFEST_FILE -from kumo.engine.trainer.trainer import TrainingResult +from kumo.engine.trainer.trainer import TrainingResult, _published_fit_mode @pytest.fixture(scope='module') @@ -115,6 +116,7 @@ def test_fit_records_listed_provenance(two_trials): assert [trial['trial'] for trial in provenance['trials']] == [0, 1] assert {trial['status'] for trial in provenance['trials']} == {'succeeded'} assert provenance['selected_trial'] in (0, 1) + assert provenance['published_model'] == 'refit_trainval' for trial in provenance['trials']: assert trial['best_epoch'] == 1 assert trial['seconds'] > 0 @@ -184,13 +186,27 @@ def _finish_stub_fit(tmp_path, trainer): provenance=dict(train_table=table, snapshot_id='snap', warm_start_bundle_id=None, started_at=datetime.now(timezone.utc), - home_dir=tmp_path)) + home_dir=tmp_path, published_model=None)) def test_provenance_without_trial_logs_still_records_parents(tmp_path): result = _finish_stub_fit(tmp_path, kumo.Trainer(ModelPlan())) assert result.provenance['trials'] == [] assert result.provenance['parents']['training_table'] == 'tt-1' + assert result.provenance['published_model'] is None + + +@pytest.mark.parametrize('flags,mode', [ + ({}, None), + ({'fitted': True}, 'fit'), + ({'fitted': True, 'refitted_trainval': True}, 'refit_trainval'), + ({'fitted': True, 'refitted_full': True}, 'refit_full'), + ({'fitted': True, 'refitted_trainval': True, 'refitted_full': True}, + 'refit_full'), +]) +def test_published_fit_mode_follows_refit_flags(flags, mode): + trainer = SimpleNamespace(refit_cfg=SimpleNamespace(**flags)) + assert _published_fit_mode(trainer) == mode def test_provenance_failure_warns_and_still_publishes(tmp_path, monkeypatch): From 2a6fea337ad914345344c037fc7ac0a95afca840 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Wed, 21 Oct 2026 06:58:17 +0000 Subject: [PATCH 38/80] Mirror the training table's source in bundle provenance #55 no longer records derived_from or modified for a reloaded table, so provenance.json's training_table block stopped having anything to copy. The block now records the table id, whether it was generated (pql) or loaded (custom), the snapshot it was generated on or validated against, and, for custom tables, the location and data version they were loaded from. The bundle stays self-contained without the scratch staging directory; query and plans stay only in training-spec.json. Verified with test_bundle_provenance.py (pql fit, custom-table stub, manifest-less stub) and test_trainer_registry.py. Co-Authored-By: Claude Opus 5.5 --- .../kumo/engine/trainer/_provenance.py | 24 ++++++++++---- .../test/test_bundle_provenance.py | 32 ++++++++++++++++--- 2 files changed, 45 insertions(+), 11 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/trainer/_provenance.py b/packages/kumo-engine/kumo/engine/trainer/_provenance.py index 222fc931c..34bc53d9e 100644 --- a/packages/kumo-engine/kumo/engine/trainer/_provenance.py +++ b/packages/kumo-engine/kumo/engine/trainer/_provenance.py @@ -45,7 +45,6 @@ def write_provenance( config = _read_json(bundle_dir / 'model' / 'config.json') or {} finished_at = datetime.now(timezone.utc) tt_id = train_table.tt_id - table_manifest = train_table.manifest or {} parents = { 'training_table': tt_id, 'snapshot': snapshot_id, @@ -63,12 +62,7 @@ def write_provenance( 'seconds': round((finished_at - started_at).total_seconds(), 3), }, - training_table={ - 'id': tt_id, - 'derived_from': - (table_manifest.get('parents') or {}).get('derived_from'), - 'modified': table_manifest.get('modified'), - }, + training_table=_training_table(tt_id, train_table.manifest), selected_trial=config.get('version'), published_model=published_model, trials=_trials(home_dir, config.get('root_dir')), @@ -76,6 +70,22 @@ def write_provenance( add_bundle_file(bundle_dir, PROVENANCE_ROLE, PROVENANCE_FILE) +def _training_table(tt_id: str, + manifest: Optional[Dict[str, Any]]) -> Dict[str, Any]: + manifest = manifest or {} + parents = manifest.get('parents') or {} + record = { + 'id': tt_id, + 'source': manifest.get('source'), + 'snapshot': parents.get('snapshot'), + } + if record['source'] == 'custom': + source_table = parents.get('source_table') or {} + record.update(location=source_table.get('location'), + version=source_table.get('version')) + return record + + def read_provenance(bundle_dir: Path) -> Optional[Dict[str, Any]]: manifest = _read_json(bundle_dir / MANIFEST_FILE) or {} for entry in manifest.get('artifacts') or []: diff --git a/packages/kumo-engine/test/test_bundle_provenance.py b/packages/kumo-engine/test/test_bundle_provenance.py index 9c4363303..b90810eed 100644 --- a/packages/kumo-engine/test/test_bundle_provenance.py +++ b/packages/kumo-engine/test/test_bundle_provenance.py @@ -108,8 +108,8 @@ def test_fit_records_listed_provenance(two_trials): } assert provenance['training_table'] == { 'id': table.tt_id, - 'derived_from': None, - 'modified': None, + 'source': 'pql', + 'snapshot': snapshot.id, } assert provenance['seed'] is None assert provenance['fit']['seconds'] >= 0 @@ -172,13 +172,13 @@ def test_bundle_without_provenance_loads_and_predicts_identically( assert stray.bundle_id == result.bundle_id -def _finish_stub_fit(tmp_path, trainer): +def _finish_stub_fit(tmp_path, trainer, manifest=None): bundle = tmp_path / 'bundle' (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 = type('Table', (), {'tt_id': 'tt-1', 'manifest': None})() + table = type('Table', (), {'tt_id': 'tt-1', 'manifest': manifest})() return trainer._finish_fit( TrainingResult({}, {}, bundle), predictive_query='PREDICT x', snapshot_id='snap', export_config=None, @@ -193,9 +193,33 @@ def test_provenance_without_trial_logs_still_records_parents(tmp_path): result = _finish_stub_fit(tmp_path, kumo.Trainer(ModelPlan())) assert result.provenance['trials'] == [] assert result.provenance['parents']['training_table'] == 'tt-1' + assert result.provenance['training_table'] == { + 'id': 'tt-1', + 'source': None, + 'snapshot': None + } assert result.provenance['published_model'] is None +def test_provenance_records_custom_table_source(tmp_path): + source_table = {'location': 's3://bucket/exports/tt', 'version': '17'} + manifest = { + 'source': 'custom', + 'query': 'PREDICT x', + 'parents': { + 'snapshot': 'snap', + 'source_table': source_table + }, + } + result = _finish_stub_fit(tmp_path, kumo.Trainer(ModelPlan()), manifest) + assert result.provenance['training_table'] == { + 'id': 'tt-1', + 'source': 'custom', + 'snapshot': 'snap', + **source_table, + } + + @pytest.mark.parametrize('flags,mode', [ ({}, None), ({'fitted': True}, 'fit'), From 68e72b91d3a5f09356feb35b34a630e4fc3dddcd Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Wed, 21 Oct 2026 07:47:45 +0000 Subject: [PATCH 39/80] Trim provenance.json to facts no other bundle file holds seed copied model/config.json, bundle_format_version restated training-spec.json's spec version, and trials scraped AutoML-internal status, metric and plan files that experiment tracking already records. All three are dropped; selected_trial stays, read from the bundle's own model/config.json, and write_provenance no longer needs the run's home directory. The training_table block of a custom table now copies connector, location and version from the table manifest's parents.source_table, the same keys the snapshot manifest uses for its sources. Verified with test_bundle_provenance.py and test_trainer_registry.py. Co-Authored-By: Claude Opus 5.5 --- .../kumo/engine/trainer/_provenance.py | 98 ++----------------- .../kumo/engine/trainer/trainer.py | 2 +- .../test/test_bundle_provenance.py | 31 +++--- 3 files changed, 19 insertions(+), 112 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/trainer/_provenance.py b/packages/kumo-engine/kumo/engine/trainer/_provenance.py index 34bc53d9e..4fbce3479 100644 --- a/packages/kumo-engine/kumo/engine/trainer/_provenance.py +++ b/packages/kumo-engine/kumo/engine/trainer/_provenance.py @@ -8,26 +8,14 @@ which model a bundle is. """ -import csv import json -import math -import re from datetime import datetime, timezone from pathlib import Path -from typing import Any, Dict, List, Optional +from typing import Any, Dict, Optional -from kumo.engine._manifest import ( - _flatten, - new_manifest, - read_manifest, - write_manifest, -) +from kumo.engine._manifest import new_manifest, read_manifest, write_manifest from kumo.engine.pquery.training_table import TrainingTable -from kumo.engine.trainer._spec import ( - MANIFEST_FILE, - PROVENANCE_ROLE, - SPEC_VERSION, -) +from kumo.engine.trainer._spec import MANIFEST_FILE, PROVENANCE_ROLE def write_provenance( @@ -37,7 +25,6 @@ def write_provenance( snapshot_id: str, warm_start_bundle_id: Optional[str], started_at: datetime, - home_dir: Path, published_model: Optional[str], ) -> None: from kumo_core.artifacts import PROVENANCE_FILE, add_bundle_file @@ -54,8 +41,6 @@ def write_provenance( bundle_dir / PROVENANCE_FILE, new_manifest( 'model_bundle', None, parents, - bundle_format_version=SPEC_VERSION, - seed=config.get('random_state_seed'), fit={ 'started_at': started_at.isoformat(timespec='seconds'), 'finished_at': finished_at.isoformat(timespec='seconds'), @@ -65,7 +50,6 @@ def write_provenance( training_table=_training_table(tt_id, train_table.manifest), selected_trial=config.get('version'), published_model=published_model, - trials=_trials(home_dir, config.get('root_dir')), )) add_bundle_file(bundle_dir, PROVENANCE_ROLE, PROVENANCE_FILE) @@ -81,8 +65,10 @@ def _training_table(tt_id: str, } if record['source'] == 'custom': source_table = parents.get('source_table') or {} - record.update(location=source_table.get('location'), - version=source_table.get('version')) + record.update({ + key: source_table.get(key) + for key in ('connector', 'location', 'version') + }) return record @@ -94,76 +80,6 @@ def read_provenance(bundle_dir: Path) -> Optional[Dict[str, Any]]: return None -def _trials(home_dir: Path, - root_dir: Optional[str]) -> List[Dict[str, Any]]: - automl_dir = Path(root_dir) if root_dir else None - if automl_dir is None or not automl_dir.is_dir(): - return [] - status = { - entry['trial']: entry - for entry in (_read_json(home_dir / 'status.json') or {}).get( - 'trials') or [] - } - versions = { - int(match.group(1)) - for path in automl_dir.iterdir() - if (match := re.fullmatch(r'version_(\d+)', path.name)) - } - trials, plans = [], {} - for trial in sorted(status.keys() | versions): - log_dir = automl_dir / f'version_{trial}' - config = _read_json(log_dir / 'config.json') or {} - progress = _read_json(log_dir / 'train_progress.json') or {} - best_epoch = (config.get('refit_cfg') or {}).get('best_val_epoch', -1) - succeeded = status.get(trial, {}).get('succeed') - trials.append({ - 'trial': trial, - 'status': (None if succeeded is None else - 'succeeded' if succeeded else 'failed'), - 'seconds': progress.get('elapsed_time'), - 'best_epoch': best_epoch if best_epoch >= 0 else None, - 'metrics': _val_metrics(log_dir / 'val_result.csv'), - }) - plans[trial] = { - key: json.dumps(value, sort_keys=True) - for key, value in _flatten( - _read_json(log_dir / 'model_plan.json') or {}).items() - } - varying = sorted( - key for key in {key - for plan in plans.values() for key in plan} - if len({plan.get(key) - for plan in plans.values()}) > 1) - for entry in trials: - plan = plans[entry['trial']] - entry['params'] = { - key: _unwrap(json.loads(plan[key])) if key in plan else None - for key in varying - } - return trials - - -def _unwrap(value: Any) -> Any: - return value[0] if isinstance(value, list) and len(value) == 1 else value - - -def _val_metrics(path: Path) -> Dict[str, float]: - try: - with open(path, newline='', encoding='utf-8') as handle: - row = next(csv.DictReader(handle), None) or {} - except OSError: - return {} - metrics = {} - for key, value in row.items(): - try: - number = float(value) - except (TypeError, ValueError): - continue - if key.startswith('val_') and math.isfinite(number): - metrics[key] = number - return metrics - - def _read_json(path: Path) -> Optional[Dict[str, Any]]: try: data = json.loads(path.read_text(encoding='utf-8')) diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index 4794169ca..07c7bbf86 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -561,7 +561,7 @@ def fit( provenance=dict(train_table=train_table, snapshot_id=snapshot.id, warm_start_bundle_id=warm_start_bundle_id, - started_at=started_at, home_dir=run_dir, + started_at=started_at, published_model=published_model), ) diff --git a/packages/kumo-engine/test/test_bundle_provenance.py b/packages/kumo-engine/test/test_bundle_provenance.py index b90810eed..999dc552e 100644 --- a/packages/kumo-engine/test/test_bundle_provenance.py +++ b/packages/kumo-engine/test/test_bundle_provenance.py @@ -100,7 +100,6 @@ def test_fit_records_listed_provenance(two_trials): 'provenance.json') provenance = result.provenance assert provenance['kind'] == 'model_bundle' - assert provenance['bundle_format_version'] == 1 assert provenance['parents'] == { 'training_table': table.tt_id, 'snapshot': snapshot.id, @@ -111,21 +110,10 @@ def test_fit_records_listed_provenance(two_trials): 'source': 'pql', 'snapshot': snapshot.id, } - assert provenance['seed'] is None + assert not {'seed', 'bundle_format_version', 'trials'} & provenance.keys() assert provenance['fit']['seconds'] >= 0 - assert [trial['trial'] for trial in provenance['trials']] == [0, 1] - assert {trial['status'] for trial in provenance['trials']} == {'succeeded'} assert provenance['selected_trial'] in (0, 1) assert provenance['published_model'] == 'refit_trainval' - for trial in provenance['trials']: - assert trial['best_epoch'] == 1 - assert trial['seconds'] > 0 - assert 'val_auroc' in trial['metrics'] - params = [trial['params'] for trial in provenance['trials']] - assert params[0] and params[0].keys() == params[1].keys() - assert all(params[0][key] != params[1][key] for key in params[0]) - assert not any(isinstance(value, list) and len(value) == 1 - for value in params[0].values()) assert provenance['environment']['kumo_engine'] == kumo.__version__ assert trainer.bundle_id == result.bundle_id @@ -186,12 +174,11 @@ def _finish_stub_fit(tmp_path, trainer, manifest=None): provenance=dict(train_table=table, snapshot_id='snap', warm_start_bundle_id=None, started_at=datetime.now(timezone.utc), - home_dir=tmp_path, published_model=None)) + published_model=None)) -def test_provenance_without_trial_logs_still_records_parents(tmp_path): +def test_provenance_without_a_table_manifest(tmp_path): result = _finish_stub_fit(tmp_path, kumo.Trainer(ModelPlan())) - assert result.provenance['trials'] == [] assert result.provenance['parents']['training_table'] == 'tt-1' assert result.provenance['training_table'] == { 'id': 'tt-1', @@ -202,7 +189,11 @@ def test_provenance_without_trial_logs_still_records_parents(tmp_path): def test_provenance_records_custom_table_source(tmp_path): - source_table = {'location': 's3://bucket/exports/tt', 'version': '17'} + source_table = { + 'connector': 'S3', + 'location': 's3://bucket/exports/tt', + 'version': '17' + } manifest = { 'source': 'custom', 'query': 'PREDICT x', @@ -237,11 +228,11 @@ def test_provenance_failure_warns_and_still_publishes(tmp_path, monkeypatch): from kumo.engine.trainer import _provenance def fail(*args): - raise RuntimeError('unreadable trials') + raise RuntimeError('unreadable manifest') - monkeypatch.setattr(_provenance, '_trials', fail) + monkeypatch.setattr(_provenance, '_training_table', fail) trainer = kumo.Trainer(ModelPlan()) - with pytest.warns(UserWarning, match='unreadable trials'): + with pytest.warns(UserWarning, match='unreadable manifest'): result = _finish_stub_fit(tmp_path, trainer) assert 'provenance' not in _roles(result.bundle_dir) assert result.provenance is None From c5fe930a858aa725190f595a9c5d011c58c3d3f2 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Wed, 21 Oct 2026 08:14:48 +0000 Subject: [PATCH 40/80] Record fit_mode and source versions in provenance; expose it on Trainer published_model in provenance.json held the shipped checkpoint's fit mode while TrainingResult.published_model is the registry record; it is now fit_mode. selected_trial indexed AutoML log directories that the bundle does not ship, numbered differently from fit events, and had no reader; it is dropped, as is training_table.id, which repeated parents.training_table. The fit snapshot's parents.source_tables (connector, location, version per table) is copied into provenance.json, so a bundle still says which source versions it was trained on after the scratch index is gone. Trainer.provenance returns the record for a fitted or loaded trainer. Verified with test_bundle_provenance.py, test_trainer_registry.py and test_warm_start.py. Co-Authored-By: Claude Opus 5.5 --- .../kumo/engine/trainer/_provenance.py | 30 ++++++++++--------- .../kumo/engine/trainer/trainer.py | 18 ++++++----- .../test/test_bundle_provenance.py | 25 ++++++++-------- 3 files changed, 38 insertions(+), 35 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/trainer/_provenance.py b/packages/kumo-engine/kumo/engine/trainer/_provenance.py index 4fbce3479..745b3826c 100644 --- a/packages/kumo-engine/kumo/engine/trainer/_provenance.py +++ b/packages/kumo-engine/kumo/engine/trainer/_provenance.py @@ -1,12 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. # All rights reserved. # SPDX-License-Identifier: Apache-2.0 -r"""Descriptive fit record stored as ``provenance.json`` in a model bundle. - -It is integrity-covered by the publication manifest but excluded from the -bundle id, so its volatile fields (times, environment, scheduler) never change -which model a bundle is. -""" +"""Fit record in ``provenance.json``; excluded from the bundle id.""" import json from datetime import datetime, timezone @@ -14,6 +9,7 @@ from typing import Any, Dict, Optional from kumo.engine._manifest import new_manifest, read_manifest, write_manifest +from kumo.engine.graph import GraphSnapshot from kumo.engine.pquery.training_table import TrainingTable from kumo.engine.trainer._spec import MANIFEST_FILE, PROVENANCE_ROLE @@ -25,11 +21,11 @@ def write_provenance( snapshot_id: str, warm_start_bundle_id: Optional[str], started_at: datetime, - published_model: Optional[str], + fit_mode: Optional[str], + snapshot: Optional[GraphSnapshot] = None, ) -> None: from kumo_core.artifacts import PROVENANCE_FILE, add_bundle_file - config = _read_json(bundle_dir / 'model' / 'config.json') or {} finished_at = datetime.now(timezone.utc) tt_id = train_table.tt_id parents = { @@ -47,19 +43,25 @@ def write_provenance( 'seconds': round((finished_at - started_at).total_seconds(), 3), }, - training_table=_training_table(tt_id, train_table.manifest), - selected_trial=config.get('version'), - published_model=published_model, + training_table=_training_table(train_table.manifest), + source_tables=_source_tables(snapshot), + fit_mode=fit_mode, )) add_bundle_file(bundle_dir, PROVENANCE_ROLE, PROVENANCE_FILE) -def _training_table(tt_id: str, - manifest: Optional[Dict[str, Any]]) -> Dict[str, Any]: +def _source_tables( + snapshot: Optional[GraphSnapshot]) -> Optional[Dict[str, Any]]: + manifest = None if snapshot is None else snapshot.manifest + if not manifest or manifest.get('reconstructed'): + return None + return (manifest.get('parents') or {}).get('source_tables') + + +def _training_table(manifest: Optional[Dict[str, Any]]) -> Dict[str, Any]: manifest = manifest or {} parents = manifest.get('parents') or {} record = { - 'id': tt_id, 'source': manifest.get('source'), 'snapshot': parents.get('snapshot'), } diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index 07c7bbf86..0423426ae 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -98,11 +98,7 @@ class TrainingResult: @property def provenance(self) -> Optional[Dict[str, Any]]: - """The bundle's recorded fit provenance; None for older bundles. - - Its `published_model` is the shipped checkpoint's fit mode, unlike - the registry record in `TrainingResult.published_model`. - """ + """The bundle's recorded fit provenance; None for older bundles.""" from kumo.engine.trainer._provenance import read_provenance return read_provenance(self.bundle_dir) @@ -329,6 +325,12 @@ def bundle_id(self) -> Optional[str]: """Content id of the fitted or loaded model bundle.""" return getattr(self._training_result, 'bundle_id', None) + @property + def provenance(self) -> Optional[Dict[str, Any]]: + """Fit provenance of the fitted or loaded bundle, if recorded.""" + result = self._training_result + return None if result is None else result.provenance + @_without_torch_scatter_hint def fit( self, @@ -541,7 +543,7 @@ def fit( holdout_path, holdout_truncated = _published_holdout( result.best_trainer) - published_model = _published_fit_mode(result.best_trainer) + fit_mode = _published_fit_mode(result.best_trainer) training_result = TrainingResult( metrics=dict(result.metrics), artifact_paths=dict(result.artifact_paths), @@ -561,8 +563,8 @@ def fit( provenance=dict(train_table=train_table, snapshot_id=snapshot.id, warm_start_bundle_id=warm_start_bundle_id, - started_at=started_at, - published_model=published_model), + started_at=started_at, fit_mode=fit_mode, + snapshot=snapshot), ) def _finish_fit( diff --git a/packages/kumo-engine/test/test_bundle_provenance.py b/packages/kumo-engine/test/test_bundle_provenance.py index 999dc552e..9b4725c6e 100644 --- a/packages/kumo-engine/test/test_bundle_provenance.py +++ b/packages/kumo-engine/test/test_bundle_provenance.py @@ -106,14 +106,13 @@ def test_fit_records_listed_provenance(two_trials): 'warm_start': None, } assert provenance['training_table'] == { - 'id': table.tt_id, 'source': 'pql', - 'snapshot': snapshot.id, + 'snapshot': snapshot.id } - assert not {'seed', 'bundle_format_version', 'trials'} & provenance.keys() + assert provenance['source_tables'] == ( + snapshot.manifest['parents']['source_tables']) assert provenance['fit']['seconds'] >= 0 - assert provenance['selected_trial'] in (0, 1) - assert provenance['published_model'] == 'refit_trainval' + assert provenance['fit_mode'] == 'refit_trainval' assert provenance['environment']['kumo_engine'] == kumo.__version__ assert trainer.bundle_id == result.bundle_id @@ -121,7 +120,9 @@ def test_fit_records_listed_provenance(two_trials): def test_bundle_id_is_content_identity(two_trials, tmp_path): trainer, result, *_ = two_trials assert re.fullmatch('[0-9a-f]{64}', result.bundle_id) - assert kumo.Trainer.load(result.bundle_dir).bundle_id == result.bundle_id + loaded = kumo.Trainer.load(result.bundle_dir) + assert loaded.bundle_id == result.bundle_id + assert loaded.provenance == result.provenance copy = tmp_path / 'copy' shutil.copytree(result.bundle_dir, copy) assert kumo.Trainer.load(copy).bundle_id == result.bundle_id @@ -132,7 +133,7 @@ def test_bundle_id_is_content_identity(two_trials, tmp_path): add_bundle_file(copy, 'provenance', 'provenance.json') reloaded = kumo.Trainer.load(copy) assert reloaded.bundle_id == result.bundle_id - assert reloaded._training_result.provenance['fit']['seconds'] == -1 + assert reloaded.provenance['fit']['seconds'] == -1 config = copy / 'model' / 'config.json' config.write_text(config.read_text() + ' ') @@ -149,14 +150,14 @@ def test_bundle_without_provenance_loads_and_predicts_identically( restored = kumo.Trainer.load(old) assert restored.bundle_id == result.bundle_id - assert restored._training_result.provenance is None + assert restored.provenance is None pd.testing.assert_frame_equal(_predict(restored, prediction_table, snapshot), _predict(trainer, prediction_table, snapshot)) (old / 'provenance.json').write_text(json.dumps(result.provenance)) stray = kumo.Trainer.load(old) - assert stray._training_result.provenance is None + assert stray.provenance is None assert stray.bundle_id == result.bundle_id @@ -174,18 +175,17 @@ def _finish_stub_fit(tmp_path, trainer, manifest=None): provenance=dict(train_table=table, snapshot_id='snap', warm_start_bundle_id=None, started_at=datetime.now(timezone.utc), - published_model=None)) + fit_mode=None)) def test_provenance_without_a_table_manifest(tmp_path): result = _finish_stub_fit(tmp_path, kumo.Trainer(ModelPlan())) assert result.provenance['parents']['training_table'] == 'tt-1' assert result.provenance['training_table'] == { - 'id': 'tt-1', 'source': None, 'snapshot': None } - assert result.provenance['published_model'] is None + assert result.provenance['source_tables'] is None def test_provenance_records_custom_table_source(tmp_path): @@ -204,7 +204,6 @@ def test_provenance_records_custom_table_source(tmp_path): } result = _finish_stub_fit(tmp_path, kumo.Trainer(ModelPlan()), manifest) assert result.provenance['training_table'] == { - 'id': 'tt-1', 'source': 'custom', 'snapshot': 'snap', **source_table, From d3de46ac2b1815681a15736857da3143624bdbb1 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Wed, 21 Oct 2026 09:35:07 +0000 Subject: [PATCH 41/80] Trim bundle provenance to essentials The provenance tests fitted their own copy of the shared fitted-model fixture, restated manifest fields, tested the private fit-mode helper case by case and spread one back-compat scenario over three tests. Reuse the conftest fitted fixture, merge the load/back-compat tests, fold the missing-manifest case into the custom-source stub test, reduce the kumo-core bundle-id test to the provenance-exclusion contract, drop the fit-mode helper test, find a listed bundle file through one listed_file helper, and shorten docstrings. Behaviour is unchanged. Verified: test_bundle_provenance, test_trainer_registry and kumo-core test_artifacts pass. Co-Authored-By: Claude Opus 5.5 --- packages/kumo-core/src/kumo_core/artifacts.py | 9 +- packages/kumo-core/tests/test_artifacts.py | 36 +--- .../kumo/engine/trainer/_provenance.py | 21 +- .../kumo-engine/kumo/engine/trainer/_spec.py | 13 +- .../kumo/engine/trainer/trainer.py | 3 - .../test/test_bundle_provenance.py | 195 ++++-------------- 6 files changed, 69 insertions(+), 208 deletions(-) diff --git a/packages/kumo-core/src/kumo_core/artifacts.py b/packages/kumo-core/src/kumo_core/artifacts.py index 7ea69b635..7fb91d1ef 100644 --- a/packages/kumo-core/src/kumo_core/artifacts.py +++ b/packages/kumo-core/src/kumo_core/artifacts.py @@ -254,11 +254,10 @@ 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: the SHA-256 of the sorted ``[relative_path, - sha256]`` pairs of 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; its digest still participates in the id. + 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 .sources import layout from .sources.errors import PermanentSourceError diff --git a/packages/kumo-core/tests/test_artifacts.py b/packages/kumo-core/tests/test_artifacts.py index e29622b96..9061d7ee7 100644 --- a/packages/kumo-core/tests/test_artifacts.py +++ b/packages/kumo-core/tests/test_artifacts.py @@ -241,41 +241,23 @@ def best_ckpt(self): assert json.loads(sources.metrics.read_text()) == {"test_accuracy": 0.8} -def test_bundle_id_covers_listed_files_except_provenance( - tmp_path: Path) -> None: + +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 / "training-spec.json").write_text("{}") - entries = json.loads( - (bundle / "publication-manifest.json").read_text())["artifacts"] - pairs = sorted([entry["relative_path"], entry["sha256"]] - for entry in entries) - pairs = sorted(pairs + [[ - "training-spec.json", - hashlib.sha256(b"{}").hexdigest(), - ]]) - expected = hashlib.sha256( - json.dumps(pairs, separators=(",", ":")).encode()).hexdigest() - - assert verify_model_bundle(bundle) == expected - add_bundle_file(bundle, "training_spec", "training-spec.json") - assert verify_model_bundle(bundle) == expected - - (bundle / "provenance.json").write_text('{"a": 1}') - add_bundle_file(bundle, "provenance", "provenance.json") - assert verify_model_bundle(bundle) == expected - (bundle / "provenance.json").write_text('{"a": 2}') - with pytest.raises(ValueError, match="provenance.json"): - verify_model_bundle(bundle) - add_bundle_file(bundle, "provenance", "provenance.json") - assert verify_model_bundle(bundle) == expected + bundle_id = verify_model_bundle(bundle) + for payload in ('{"a": 1}', '{"a": 2}'): + (bundle / "provenance.json").write_text(payload) + add_bundle_file(bundle, "provenance", "provenance.json") + assert verify_model_bundle(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) != expected + assert verify_model_bundle(bundle) != bundle_id diff --git a/packages/kumo-engine/kumo/engine/trainer/_provenance.py b/packages/kumo-engine/kumo/engine/trainer/_provenance.py index 745b3826c..231ee13f8 100644 --- a/packages/kumo-engine/kumo/engine/trainer/_provenance.py +++ b/packages/kumo-engine/kumo/engine/trainer/_provenance.py @@ -3,7 +3,6 @@ # SPDX-License-Identifier: Apache-2.0 """Fit record in ``provenance.json``; excluded from the bundle id.""" -import json from datetime import datetime, timezone from pathlib import Path from typing import Any, Dict, Optional @@ -11,7 +10,7 @@ from kumo.engine._manifest import new_manifest, read_manifest, write_manifest from kumo.engine.graph import GraphSnapshot from kumo.engine.pquery.training_table import TrainingTable -from kumo.engine.trainer._spec import MANIFEST_FILE, PROVENANCE_ROLE +from kumo.engine.trainer._spec import PROVENANCE_ROLE, listed_file def write_provenance( @@ -27,9 +26,8 @@ def write_provenance( from kumo_core.artifacts import PROVENANCE_FILE, add_bundle_file finished_at = datetime.now(timezone.utc) - tt_id = train_table.tt_id parents = { - 'training_table': tt_id, + 'training_table': train_table.tt_id, 'snapshot': snapshot_id, 'warm_start': warm_start_bundle_id, } @@ -75,16 +73,5 @@ def _training_table(manifest: Optional[Dict[str, Any]]) -> Dict[str, Any]: def read_provenance(bundle_dir: Path) -> Optional[Dict[str, Any]]: - manifest = _read_json(bundle_dir / MANIFEST_FILE) or {} - for entry in manifest.get('artifacts') or []: - if entry.get('role') == PROVENANCE_ROLE: - return read_manifest(bundle_dir / entry['relative_path']) - return None - - -def _read_json(path: Path) -> Optional[Dict[str, Any]]: - try: - data = json.loads(path.read_text(encoding='utf-8')) - except (OSError, ValueError): - return None - return data if isinstance(data, dict) else None + path = listed_file(bundle_dir, PROVENANCE_ROLE) + return None if path is None else read_manifest(path) diff --git a/packages/kumo-engine/kumo/engine/trainer/_spec.py b/packages/kumo-engine/kumo/engine/trainer/_spec.py index 637714d45..2c755fab4 100644 --- a/packages/kumo-engine/kumo/engine/trainer/_spec.py +++ b/packages/kumo-engine/kumo/engine/trainer/_spec.py @@ -12,7 +12,7 @@ import json from dataclasses import asdict, dataclass from pathlib import Path -from typing import Any, Dict, Tuple +from typing import Any, Dict, Optional, Tuple from kumo_core.io.filesystem import atomic_write_bytes from kumoapi.model_plan import ModelPlan, TrainingTableGenerationPlan @@ -103,6 +103,17 @@ def read_training_spec(bundle_dir: Path) -> TrainingSpec: ) +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.""" diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index 0423426ae..623e7918a 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -83,9 +83,6 @@ class TrainingResult: 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. - - `bundle_id` is the content id of the model bundle's scoring files, - computed when the bundle is written or loaded. """ metrics: Dict[str, float] diff --git a/packages/kumo-engine/test/test_bundle_provenance.py b/packages/kumo-engine/test/test_bundle_provenance.py index 9b4725c6e..fff92cd3f 100644 --- a/packages/kumo-engine/test/test_bundle_provenance.py +++ b/packages/kumo-engine/test/test_bundle_provenance.py @@ -8,166 +8,70 @@ from datetime import datetime, timezone from types import SimpleNamespace -import numpy as np import pandas as pd import pytest from kumo_core.artifacts import add_bundle_file -from kumoapi.graph import Edge from kumoapi.model_plan import ModelPlan, TrainingTableGenerationPlan -from kumoapi.typing import Stype import kumo.engine as kumo -from kumo.engine._scratch import get_scratch_dir from kumo.engine.trainer._spec import MANIFEST_FILE -from kumo.engine.trainer.trainer import TrainingResult, _published_fit_mode - - -@pytest.fixture(scope='module') -def two_trials(tmp_path_factory): - root = tmp_path_factory.mktemp('provenance') - previous = get_scratch_dir() - kumo.configure(scratch_dir=root / 'scratch') - try: - rng = np.random.default_rng(0) - data = root / 'data' - data.mkdir() - pd.DataFrame({'ID': np.arange(50)}).to_parquet(data / 'ents.parquet') - pd.DataFrame({ - 'EID': - np.arange(3000), - 'ENT_ID': - rng.integers(0, 50, 3000), - 'TS': - pd.Timestamp('2024-01-01') + - pd.to_timedelta(rng.integers(0, 120, 3000), unit='D'), - }).to_parquet(data / 'events.parquet') - connector = kumo.FileConnector(root_dir=str(data)) - events = kumo.Table.from_source_table( - connector['events.parquet'], primary_key='EID', - time_column='TS').infer_metadata().with_column( - 'ENT_ID', stype=Stype.ID) - graph = kumo.Graph( - tables={ - 'ents': - kumo.Table.from_source_table(connector['ents.parquet'], - primary_key='ID').infer_metadata(), - 'events': - events, - }, - edges=[Edge(src_table='events', fkey='ENT_ID', dst_table='ents')]) - snapshot = graph.snapshot() - pquery = kumo.PredictiveQuery( - graph=graph, query='PREDICT COUNT(events.*, 0, 7, days) > 0 ' - 'FOR EACH ents.ID') - plan = ModelPlan() - plan.training_job.disable_explain = True - plan.training_job.num_experiments = 2 - plan.optimization.max_epochs = 1 - plan.optimization.max_steps_per_epoch = 2 - table = pquery.generate_training_table(snapshot) - trainer = kumo.Trainer(model_plan=plan) - result = trainer.fit(table, snapshot=snapshot) - yield (trainer, result, table, snapshot, - pquery.generate_prediction_table(snapshot)) - finally: - kumo.configure(scratch_dir=previous) - - -def _predict(trainer, table, snapshot): - result = trainer.predict(table, snapshot=snapshot, - binary_classification_threshold=.5) - return result.frame.sort_values('ENTITY').reset_index(drop=True) - - -def _roles(bundle): - manifest = json.loads((bundle / MANIFEST_FILE).read_text()) - return {entry['role']: entry for entry in manifest['artifacts']} - - -def _strip_provenance(bundle): - manifest = json.loads((bundle / MANIFEST_FILE).read_text()) - manifest['artifacts'] = [ - entry for entry in manifest['artifacts'] - if entry['role'] != 'provenance' - ] - (bundle / MANIFEST_FILE).write_text(json.dumps(manifest)) - (bundle / 'provenance.json').unlink() +from kumo.engine.trainer.trainer import TrainingResult -def test_fit_records_listed_provenance(two_trials): - trainer, result, table, snapshot, _ = two_trials - assert _roles(result.bundle_dir)['provenance']['relative_path'] == ( - 'provenance.json') +def test_fit_records_provenance(fitted): + trainer, _, snapshot = fitted + result = trainer._training_result provenance = result.provenance - assert provenance['kind'] == 'model_bundle' - assert provenance['parents'] == { - 'training_table': table.tt_id, - 'snapshot': snapshot.id, - 'warm_start': None, - } + parents = provenance['parents'] + assert (parents['snapshot'], parents['warm_start']) == (snapshot.id, None) + assert parents['training_table'] assert provenance['training_table'] == { 'source': 'pql', 'snapshot': snapshot.id } assert provenance['source_tables'] == ( snapshot.manifest['parents']['source_tables']) - assert provenance['fit']['seconds'] >= 0 assert provenance['fit_mode'] == 'refit_trainval' - assert provenance['environment']['kumo_engine'] == kumo.__version__ + assert re.fullmatch('[0-9a-f]{64}', trainer.bundle_id) assert trainer.bundle_id == result.bundle_id -def test_bundle_id_is_content_identity(two_trials, tmp_path): - trainer, result, *_ = two_trials - assert re.fullmatch('[0-9a-f]{64}', result.bundle_id) +def test_bundle_without_provenance_loads_and_predicts_identically( + fitted, tmp_path): + trainer, prediction_table, snapshot = fitted + result = trainer._training_result loaded = kumo.Trainer.load(result.bundle_dir) - assert loaded.bundle_id == result.bundle_id - assert loaded.provenance == result.provenance - copy = tmp_path / 'copy' - shutil.copytree(result.bundle_dir, copy) - assert kumo.Trainer.load(copy).bundle_id == result.bundle_id - - provenance = json.loads((copy / 'provenance.json').read_text()) - provenance['fit']['seconds'] = -1 - (copy / 'provenance.json').write_text(json.dumps(provenance)) - add_bundle_file(copy, 'provenance', 'provenance.json') - reloaded = kumo.Trainer.load(copy) - assert reloaded.bundle_id == result.bundle_id - assert reloaded.provenance['fit']['seconds'] == -1 - - config = copy / 'model' / 'config.json' - config.write_text(config.read_text() + ' ') - add_bundle_file(copy, 'trainer_config', 'model/config.json') - assert kumo.Trainer.load(copy).bundle_id != result.bundle_id + assert (loaded.bundle_id, loaded.provenance) == (result.bundle_id, + result.provenance) - -def test_bundle_without_provenance_loads_and_predicts_identically( - two_trials, tmp_path): - trainer, result, _, snapshot, prediction_table = two_trials old = tmp_path / 'old' shutil.copytree(result.bundle_dir, old) - _strip_provenance(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.bundle_id == result.bundle_id assert restored.provenance is None - pd.testing.assert_frame_equal(_predict(restored, prediction_table, - snapshot), - _predict(trainer, prediction_table, snapshot)) - (old / 'provenance.json').write_text(json.dumps(result.provenance)) - stray = kumo.Trainer.load(old) - assert stray.provenance is None - assert stray.bundle_id == result.bundle_id + 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 _finish_stub_fit(tmp_path, trainer, manifest=None): - bundle = tmp_path / 'bundle' +def _finish_stub_fit(bundle, trainer, manifest=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 = type('Table', (), {'tt_id': 'tt-1', 'manifest': manifest})() + table = SimpleNamespace(tt_id='tt-1', manifest=manifest) return trainer._finish_fit( TrainingResult({}, {}, bundle), predictive_query='PREDICT x', snapshot_id='snap', export_config=None, @@ -178,17 +82,7 @@ def _finish_stub_fit(tmp_path, trainer, manifest=None): fit_mode=None)) -def test_provenance_without_a_table_manifest(tmp_path): - result = _finish_stub_fit(tmp_path, kumo.Trainer(ModelPlan())) - assert result.provenance['parents']['training_table'] == 'tt-1' - assert result.provenance['training_table'] == { - 'source': None, - 'snapshot': None - } - assert result.provenance['source_tables'] is None - - -def test_provenance_records_custom_table_source(tmp_path): +def test_provenance_records_custom_or_missing_table_source(tmp_path): source_table = { 'connector': 'S3', 'location': 's3://bucket/exports/tt', @@ -196,31 +90,23 @@ def test_provenance_records_custom_table_source(tmp_path): } manifest = { 'source': 'custom', - 'query': 'PREDICT x', 'parents': { 'snapshot': 'snap', 'source_table': source_table }, } - result = _finish_stub_fit(tmp_path, kumo.Trainer(ModelPlan()), manifest) - assert result.provenance['training_table'] == { + custom = _finish_stub_fit(tmp_path / 'custom', kumo.Trainer(ModelPlan()), + manifest) + assert custom.provenance['training_table'] == { 'source': 'custom', 'snapshot': 'snap', **source_table, } - - -@pytest.mark.parametrize('flags,mode', [ - ({}, None), - ({'fitted': True}, 'fit'), - ({'fitted': True, 'refitted_trainval': True}, 'refit_trainval'), - ({'fitted': True, 'refitted_full': True}, 'refit_full'), - ({'fitted': True, 'refitted_trainval': True, 'refitted_full': True}, - 'refit_full'), -]) -def test_published_fit_mode_follows_refit_flags(flags, mode): - trainer = SimpleNamespace(refit_cfg=SimpleNamespace(**flags)) - assert _published_fit_mode(trainer) == mode + older = _finish_stub_fit(tmp_path / 'older', kumo.Trainer(ModelPlan())) + assert older.provenance['training_table'] == { + 'source': None, + 'snapshot': None + } def test_provenance_failure_warns_and_still_publishes(tmp_path, monkeypatch): @@ -232,8 +118,7 @@ def fail(*args): monkeypatch.setattr(_provenance, '_training_table', fail) trainer = kumo.Trainer(ModelPlan()) with pytest.warns(UserWarning, match='unreadable manifest'): - result = _finish_stub_fit(tmp_path, trainer) - assert 'provenance' not in _roles(result.bundle_dir) + result = _finish_stub_fit(tmp_path / 'bundle', trainer) assert result.provenance is None assert re.fullmatch('[0-9a-f]{64}', result.bundle_id) assert trainer._training_result is result From efaeced434717e1c8a120b5adb13629c23f82019 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Wed, 21 Oct 2026 18:29:44 +0000 Subject: [PATCH 42/80] Reuse shared helpers for bundle provenance add_bundle_file inlined its own atomic write, the engine re-walked the trainer's refit flags to name the published fit mode, and provenance.json repeated the fit snapshot that training-spec.json already records as snapshot_id. add_bundle_file now uses kumo_core.io.filesystem.atomic_write_bytes. kumo_core's Trainer gains published_fit_mode(), the order _set_fit_mode_to_use already used, and fit records it as fit_mode. provenance.json drops parents.snapshot; training_table.snapshot stays because a table may be generated on a different snapshot than the fit's. training_table.source is DiskGraph's source_mode (pquery or custom_tt). Verified with the bundle provenance, trainer registry, warm start, training-table export, snapshot manifest and kumo-core artifact/training suites. Co-Authored-By: Claude Opus 5.5 --- packages/kumo-core/src/kumo_core/artifacts.py | 12 ++---------- .../kumo/engine/trainer/_provenance.py | 16 ++++++---------- .../kumo/engine/trainer/trainer.py | 14 ++------------ .../test/test_bundle_provenance.py | 19 +++++++++---------- .../test/test_semantic_carry_forward.py | 7 +++---- 5 files changed, 22 insertions(+), 46 deletions(-) diff --git a/packages/kumo-core/src/kumo_core/artifacts.py b/packages/kumo-core/src/kumo_core/artifacts.py index 7fb91d1ef..68b22e2f8 100644 --- a/packages/kumo-core/src/kumo_core/artifacts.py +++ b/packages/kumo-core/src/kumo_core/artifacts.py @@ -363,16 +363,8 @@ def add_bundle_file(bundle_dir: Path, role: str, relative_path: str) -> None: sha256=VitalArtifactPublicationPolicy._sha256(path), ).to_dict() ] - temporary = manifest_path.with_name(f'.{manifest_path.name}.tmp') - try: - with temporary.open('w', encoding='utf-8') as handle: - handle.write(json.dumps(manifest, indent=2, sort_keys=True) + '\n') - handle.flush() - os.fsync(handle.fileno()) - os.replace(temporary, manifest_path) - FilesystemArtifactStore._fsync_directory(bundle_dir) - finally: - temporary.unlink(missing_ok=True) + payload = json.dumps(manifest, indent=2, sort_keys=True) + '\n' + atomic_write_bytes(manifest_path, payload.encode('utf-8')) @runtime_checkable diff --git a/packages/kumo-engine/kumo/engine/trainer/_provenance.py b/packages/kumo-engine/kumo/engine/trainer/_provenance.py index 231ee13f8..62826acf5 100644 --- a/packages/kumo-engine/kumo/engine/trainer/_provenance.py +++ b/packages/kumo-engine/kumo/engine/trainer/_provenance.py @@ -17,7 +17,6 @@ def write_provenance( bundle_dir: Path, *, train_table: TrainingTable, - snapshot_id: str, warm_start_bundle_id: Optional[str], started_at: datetime, fit_mode: Optional[str], @@ -28,7 +27,6 @@ def write_provenance( finished_at = datetime.now(timezone.utc) parents = { 'training_table': train_table.tt_id, - 'snapshot': snapshot_id, 'warm_start': warm_start_bundle_id, } write_manifest( @@ -41,7 +39,7 @@ def write_provenance( 'seconds': round((finished_at - started_at).total_seconds(), 3), }, - training_table=_training_table(train_table.manifest), + training_table=_training_table(train_table), source_tables=_source_tables(snapshot), fit_mode=fit_mode, )) @@ -56,17 +54,15 @@ def _source_tables( return (manifest.get('parents') or {}).get('source_tables') -def _training_table(manifest: Optional[Dict[str, Any]]) -> Dict[str, Any]: - manifest = manifest or {} - parents = manifest.get('parents') or {} +def _training_table(train_table: TrainingTable) -> Dict[str, Any]: + parents = (train_table.manifest or {}).get('parents') or {} record = { - 'source': manifest.get('source'), + 'source': train_table.source_mode, 'snapshot': parents.get('snapshot'), } - if record['source'] == 'custom': - source_table = parents.get('source_table') or {} + if 'source_table' in parents: record.update({ - key: source_table.get(key) + key: (parents['source_table'] or {}).get(key) for key in ('connector', 'location', 'version') }) return record diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index 623e7918a..f2a3c95a8 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -207,16 +207,6 @@ def _published_holdout(trainer: Any) -> tuple[Optional[Path], Optional[bool]]: return _find_holdout(trainer.get_log_dir(fit_mode=published)) -def _published_fit_mode(trainer: Any) -> Optional[str]: - refit = getattr(trainer, 'refit_cfg', None) - for mode, flag in (('refit_full', 'refitted_full'), - ('refit_trainval', 'refitted_trainval'), - ('fit', 'fitted')): - if getattr(refit, flag, False) is True: - return mode - return None - - 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. @@ -540,7 +530,8 @@ def fit( holdout_path, holdout_truncated = _published_holdout( result.best_trainer) - fit_mode = _published_fit_mode(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), @@ -558,7 +549,6 @@ def fit( export_config=export_config, model_plan=effective_plan, provenance=dict(train_table=train_table, - snapshot_id=snapshot.id, warm_start_bundle_id=warm_start_bundle_id, started_at=started_at, fit_mode=fit_mode, snapshot=snapshot), diff --git a/packages/kumo-engine/test/test_bundle_provenance.py b/packages/kumo-engine/test/test_bundle_provenance.py index fff92cd3f..5c891c58c 100644 --- a/packages/kumo-engine/test/test_bundle_provenance.py +++ b/packages/kumo-engine/test/test_bundle_provenance.py @@ -23,10 +23,10 @@ def test_fit_records_provenance(fitted): result = trainer._training_result provenance = result.provenance parents = provenance['parents'] - assert (parents['snapshot'], parents['warm_start']) == (snapshot.id, None) - assert parents['training_table'] + assert parents.keys() == {'training_table', 'warm_start'} + assert parents['training_table'] and parents['warm_start'] is None assert provenance['training_table'] == { - 'source': 'pql', + 'source': 'pquery', 'snapshot': snapshot.id } assert provenance['source_tables'] == ( @@ -66,18 +66,18 @@ def predict(model): pd.testing.assert_frame_equal(predict(restored), predict(trainer)) -def _finish_stub_fit(bundle, trainer, manifest=None): +def _finish_stub_fit(bundle, trainer, manifest=None, source_mode=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-1', manifest=manifest) + table = SimpleNamespace(tt_id='tt-1', manifest=manifest, + source_mode=source_mode) 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, snapshot_id='snap', - warm_start_bundle_id=None, + provenance=dict(train_table=table, warm_start_bundle_id=None, started_at=datetime.now(timezone.utc), fit_mode=None)) @@ -89,16 +89,15 @@ def test_provenance_records_custom_or_missing_table_source(tmp_path): 'version': '17' } manifest = { - 'source': 'custom', 'parents': { 'snapshot': 'snap', 'source_table': source_table }, } custom = _finish_stub_fit(tmp_path / 'custom', kumo.Trainer(ModelPlan()), - manifest) + manifest, 'custom_tt') assert custom.provenance['training_table'] == { - 'source': 'custom', + 'source': 'custom_tt', 'snapshot': 'snap', **source_table, } 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 c7288937b843b252786805b85ef156892ef098b9 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Wed, 21 Oct 2026 18:31:54 +0000 Subject: [PATCH 43/80] Compute a fresh bundle's id from its publication manifest Fit re-read and re-hashed every bundle file to compute bundle_id right after publication and add_bundle_file had hashed them; on a bundle with a 1 GiB checkpoint that is about 0.76 s of SHA-256 per fit. kumo-core gains bundle_id_from_manifest, the same id formula over the digests the publication manifest already lists; verify_model_bundle uses the shared formula and still fully verifies bundles on load and fetch. Verified with the kumo-core artifact suite (the two ids agree) and the bundle provenance and warm-start suites. Co-Authored-By: Claude Opus 5.5 --- packages/kumo-core/src/kumo_core/artifacts.py | 23 +++++++++++++++++-- packages/kumo-core/tests/test_artifacts.py | 2 ++ .../kumo/engine/trainer/trainer.py | 4 ++-- 3 files changed, 25 insertions(+), 4 deletions(-) diff --git a/packages/kumo-core/src/kumo_core/artifacts.py b/packages/kumo-core/src/kumo_core/artifacts.py index 68b22e2f8..c5730c55f 100644 --- a/packages/kumo-core/src/kumo_core/artifacts.py +++ b/packages/kumo-core/src/kumo_core/artifacts.py @@ -13,7 +13,14 @@ import tempfile from dataclasses import dataclass, field from pathlib import Path -from typing import Any, ClassVar, Mapping, Protocol, runtime_checkable +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 @@ -339,8 +346,20 @@ def verify_model_bundle(bundle: Path) -> str: if 'training-spec.json' not in seen and spec_path.is_file(): verified.append(('training-spec.json', VitalArtifactPublicationPolicy._sha256(spec_path))) + 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 verified if pair[0] != PROVENANCE_FILE), + sorted(pair for pair in pairs if pair[0] != PROVENANCE_FILE), separators=(',', ':')) return hashlib.sha256(payload.encode('utf-8')).hexdigest() diff --git a/packages/kumo-core/tests/test_artifacts.py b/packages/kumo-core/tests/test_artifacts.py index 9061d7ee7..eadaf82e1 100644 --- a/packages/kumo-core/tests/test_artifacts.py +++ b/packages/kumo-core/tests/test_artifacts.py @@ -16,6 +16,7 @@ VitalArtifactPublicationPolicy, VitalArtifactSources, add_bundle_file, + bundle_id_from_manifest, resolve_vital_artifact_sources, verify_model_bundle, ) @@ -252,6 +253,7 @@ def test_bundle_id_ignores_provenance_only(tmp_path: Path) -> None: (bundle / "provenance.json").write_text(payload) add_bundle_file(bundle, "provenance", "provenance.json") 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"] diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index f2a3c95a8..7f3183295 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -566,7 +566,7 @@ def _finish_fit( provenance: Optional[Mapping[str, Any]] = None, ) -> TrainingResult: """Commit durable fit outputs before making this trainer usable.""" - from kumo_core.artifacts import verify_model_bundle + from kumo_core.artifacts import bundle_id_from_manifest from kumo.engine.trainer._provenance import write_provenance @@ -588,7 +588,7 @@ def _finish_fit( try: training_result = replace( training_result, - bundle_id=verify_model_bundle(training_result.bundle_dir)) + 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) From 58b80d391746c253aac29b4d51771fc461287855 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Wed, 21 Oct 2026 21:10:20 +0000 Subject: [PATCH 44/80] Copy a custom table's full source entry into provenance provenance.json picked connector, location and version from the training table's source entry and dropped version_kind and Custom SQL dependencies that the training-table manifest now records. The whole entry is copied, so provenance mirrors the training-table manifest the way source_tables mirror the snapshot manifest. Verified with the bundle provenance and training-table export tests. Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/trainer/_provenance.py | 5 +---- packages/kumo-engine/test/test_bundle_provenance.py | 3 ++- 2 files changed, 3 insertions(+), 5 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/trainer/_provenance.py b/packages/kumo-engine/kumo/engine/trainer/_provenance.py index 62826acf5..12a15ade9 100644 --- a/packages/kumo-engine/kumo/engine/trainer/_provenance.py +++ b/packages/kumo-engine/kumo/engine/trainer/_provenance.py @@ -61,10 +61,7 @@ def _training_table(train_table: TrainingTable) -> Dict[str, Any]: 'snapshot': parents.get('snapshot'), } if 'source_table' in parents: - record.update({ - key: (parents['source_table'] or {}).get(key) - for key in ('connector', 'location', 'version') - }) + record.update(parents['source_table'] or {}) return record diff --git a/packages/kumo-engine/test/test_bundle_provenance.py b/packages/kumo-engine/test/test_bundle_provenance.py index 5c891c58c..d5c2c194c 100644 --- a/packages/kumo-engine/test/test_bundle_provenance.py +++ b/packages/kumo-engine/test/test_bundle_provenance.py @@ -86,7 +86,8 @@ def test_provenance_records_custom_or_missing_table_source(tmp_path): source_table = { 'connector': 'S3', 'location': 's3://bucket/exports/tt', - 'version': '17' + 'version': '17', + 'version_kind': 'timestamp', } manifest = { 'parents': { From f70ee4d5481797de3350ee26822f0ad00ba734ce Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Wed, 21 Oct 2026 22:19:39 +0000 Subject: [PATCH 45/80] Store the bundle id in provenance.json provenance.json recorded "id": null, so a provenance file copied out of its bundle (logged to a tracker or registry) could not be matched to the bundle it describes. The id excludes provenance.json, so it is now computed first and written as the provenance id. Verified with test_bundle_provenance.py. Co-Authored-By: Claude Opus 5.5 --- .../kumo/engine/trainer/_provenance.py | 3 ++- .../kumo-engine/kumo/engine/trainer/trainer.py | 16 +++++++++------- .../kumo-engine/test/test_bundle_provenance.py | 3 ++- 3 files changed, 13 insertions(+), 9 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/trainer/_provenance.py b/packages/kumo-engine/kumo/engine/trainer/_provenance.py index 12a15ade9..f59d9cb98 100644 --- a/packages/kumo-engine/kumo/engine/trainer/_provenance.py +++ b/packages/kumo-engine/kumo/engine/trainer/_provenance.py @@ -16,6 +16,7 @@ def write_provenance( bundle_dir: Path, *, + bundle_id: Optional[str], train_table: TrainingTable, warm_start_bundle_id: Optional[str], started_at: datetime, @@ -32,7 +33,7 @@ def write_provenance( write_manifest( bundle_dir / PROVENANCE_FILE, new_manifest( - 'model_bundle', None, parents, + 'model_bundle', bundle_id, parents, fit={ 'started_at': started_at.isoformat(timespec='seconds'), 'finished_at': finished_at.isoformat(timespec='seconds'), diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index 7f3183295..557ea6f23 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -578,13 +578,6 @@ def _finish_fit( snapshot_id=snapshot_id, training_table_generation=training_table_generation, ) - if provenance is not None: - try: - write_provenance(training_result.bundle_dir, **provenance) - except Exception as error: - warnings.warn( - f"Publishing the model bundle without provenance.json: " - f"{error}", UserWarning, stacklevel=3) try: training_result = replace( training_result, @@ -592,6 +585,15 @@ def _finish_fit( except Exception as error: warnings.warn(f"Could not compute the bundle id: {error}", UserWarning, stacklevel=3) + if provenance 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) published = export_config.connector.sink_backend.publish_model( diff --git a/packages/kumo-engine/test/test_bundle_provenance.py b/packages/kumo-engine/test/test_bundle_provenance.py index d5c2c194c..c20cd54dd 100644 --- a/packages/kumo-engine/test/test_bundle_provenance.py +++ b/packages/kumo-engine/test/test_bundle_provenance.py @@ -33,7 +33,8 @@ def test_fit_records_provenance(fitted): snapshot.manifest['parents']['source_tables']) assert provenance['fit_mode'] == 'refit_trainval' assert re.fullmatch('[0-9a-f]{64}', trainer.bundle_id) - assert trainer.bundle_id == result.bundle_id + assert trainer.bundle_id == result.bundle_id == provenance['id'] + assert kumo.Trainer.load(result.bundle_dir).bundle_id == provenance['id'] def test_bundle_without_provenance_loads_and_predicts_identically( From 42171445e48bb9a2a9ea52dbca01ef6acb989d73 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Thu, 22 Oct 2026 18:47:02 +0000 Subject: [PATCH 46/80] Make bundle provenance a frozen dataclass TrainingResult.provenance and Trainer.provenance returned a dict. They now return a BundleProvenance (parents, fit timing, the training table record and source tables) built on the shared manifest dataclasses. provenance.json on disk and the bundle id are unchanged; the training table record still writes its source table's keys inline. Verified with test_bundle_provenance.py and test_warm_start.py. Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/_manifest.py | 51 +++++++++++ .../kumo/engine/trainer/_provenance.py | 62 +++++++------ .../kumo/engine/trainer/trainer.py | 5 +- .../test/test_bundle_provenance.py | 88 +++++++++++-------- packages/kumo-engine/test/test_warm_start.py | 2 +- 5 files changed, 142 insertions(+), 66 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/_manifest.py b/packages/kumo-engine/kumo/engine/_manifest.py index a47a0737f..8b1c7f860 100644 --- a/packages/kumo-engine/kumo/engine/_manifest.py +++ b/packages/kumo-engine/kumo/engine/_manifest.py @@ -16,6 +16,7 @@ List, Mapping, Optional, + Self, Type, TypeVar, Union, @@ -224,6 +225,56 @@ class TrainingTableManifest(Manifest): split_seed: Optional[int] = None +@dataclass(frozen=True, kw_only=True) +class BundleParents(Record): + training_table: str + warm_start: Optional[str] + + +@dataclass(frozen=True, kw_only=True) +class FitTiming(Record): + started_at: str + finished_at: str + seconds: float + + +@dataclass(frozen=True, kw_only=True) +class BundleTrainingTable(Record): + """Written with ``source_table``'s keys inline.""" + + source: Optional[str] + snapshot: Optional[str] + source_table: Optional[SourceTableEntry] = None + + def to_dict(self) -> Dict[str, Any]: + data = {**self.extra, 'source': self.source, 'snapshot': self.snapshot} + if self.source_table is not None: + data.update(self.source_table.to_dict()) + return data + + @classmethod + def from_dict(cls, data: Mapping[str, Any]) -> Self: + rest = { + key: value + for key, value in data.items() if key not in ('source', 'snapshot') + } + if 'version_kind' not in rest: + return cls(source=data.get('source'), snapshot=data.get('snapshot'), + extra=rest) + return cls(source=data.get('source'), snapshot=data.get('snapshot'), + source_table=SourceTableEntry.from_dict(rest)) + + +@dataclass(frozen=True, kw_only=True) +class BundleProvenance(Manifest): + kind: str = 'model_bundle' + parents: BundleParents + fit: FitTiming + training_table: BundleTrainingTable + source_tables: Optional[Dict[str, SourceTableEntry]] + fit_mode: Optional[str] + + def write_manifest(path: Path, manifest: Manifest) -> None: payload = json.dumps(manifest.to_dict(), indent=2, sort_keys=True) + '\n' atomic_write_bytes(path, payload.encode('utf-8')) diff --git a/packages/kumo-engine/kumo/engine/trainer/_provenance.py b/packages/kumo-engine/kumo/engine/trainer/_provenance.py index f59d9cb98..3287db604 100644 --- a/packages/kumo-engine/kumo/engine/trainer/_provenance.py +++ b/packages/kumo-engine/kumo/engine/trainer/_provenance.py @@ -5,9 +5,18 @@ from datetime import datetime, timezone from pathlib import Path -from typing import Any, Dict, Optional +from typing import Dict, Optional -from kumo.engine._manifest import new_manifest, read_manifest, write_manifest +from kumo.engine._manifest import ( + BundleParents, + BundleProvenance, + BundleTrainingTable, + FitTiming, + SourceTableEntry, + read_manifest, + stamp, + write_manifest, +) from kumo.engine.graph import GraphSnapshot from kumo.engine.pquery.training_table import TrainingTable from kumo.engine.trainer._spec import PROVENANCE_ROLE, listed_file @@ -26,46 +35,43 @@ def write_provenance( from kumo_core.artifacts import PROVENANCE_FILE, add_bundle_file finished_at = datetime.now(timezone.utc) - parents = { - 'training_table': train_table.tt_id, - 'warm_start': warm_start_bundle_id, - } write_manifest( bundle_dir / PROVENANCE_FILE, - new_manifest( - 'model_bundle', bundle_id, parents, - fit={ - 'started_at': started_at.isoformat(timespec='seconds'), - 'finished_at': finished_at.isoformat(timespec='seconds'), - 'seconds': round((finished_at - started_at).total_seconds(), - 3), - }, + BundleProvenance( + id=bundle_id, + parents=BundleParents(training_table=train_table.tt_id, + warm_start=warm_start_bundle_id), + 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), + ), training_table=_training_table(train_table), source_tables=_source_tables(snapshot), fit_mode=fit_mode, + **stamp(), )) add_bundle_file(bundle_dir, PROVENANCE_ROLE, PROVENANCE_FILE) def _source_tables( - snapshot: Optional[GraphSnapshot]) -> Optional[Dict[str, Any]]: + snapshot: Optional[GraphSnapshot] +) -> Optional[Dict[str, SourceTableEntry]]: manifest = None if snapshot is None else snapshot.manifest - if not manifest or manifest.get('reconstructed'): + if manifest is None or manifest.reconstructed: return None - return (manifest.get('parents') or {}).get('source_tables') + return manifest.parents.source_tables -def _training_table(train_table: TrainingTable) -> Dict[str, Any]: - parents = (train_table.manifest or {}).get('parents') or {} - record = { - 'source': train_table.source_mode, - 'snapshot': parents.get('snapshot'), - } - if 'source_table' in parents: - record.update(parents['source_table'] or {}) - return record +def _training_table(train_table: TrainingTable) -> BundleTrainingTable: + manifest = train_table.manifest + return BundleTrainingTable( + source=train_table.source_mode, + snapshot=None if manifest is None else manifest.parents.snapshot, + source_table=None + if manifest is None else manifest.parents.source_table) -def read_provenance(bundle_dir: Path) -> Optional[Dict[str, Any]]: +def read_provenance(bundle_dir: Path) -> Optional[BundleProvenance]: path = listed_file(bundle_dir, PROVENANCE_ROLE) - return None if path is None else read_manifest(path) + return None if path is None else read_manifest(path, BundleProvenance) diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index 557ea6f23..c318590f0 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -52,6 +52,7 @@ require_embedding_models, require_trained_embedding_models, ) +from kumo.engine._manifest import BundleProvenance from kumo.engine._scratch import new_run_id, trainer_scratch_dir from kumo.engine.connector.base import Connector from kumo.engine.graph import GraphSnapshot @@ -94,7 +95,7 @@ class TrainingResult: bundle_id: Optional[str] = None @property - def provenance(self) -> Optional[Dict[str, Any]]: + def provenance(self) -> Optional[BundleProvenance]: """The bundle's recorded fit provenance; None for older bundles.""" from kumo.engine.trainer._provenance import read_provenance @@ -313,7 +314,7 @@ def bundle_id(self) -> Optional[str]: return getattr(self._training_result, 'bundle_id', None) @property - def provenance(self) -> Optional[Dict[str, Any]]: + def provenance(self) -> Optional[BundleProvenance]: """Fit provenance of the fitted or loaded bundle, if recorded.""" result = self._training_result return None if result is None else result.provenance diff --git a/packages/kumo-engine/test/test_bundle_provenance.py b/packages/kumo-engine/test/test_bundle_provenance.py index c20cd54dd..608564460 100644 --- a/packages/kumo-engine/test/test_bundle_provenance.py +++ b/packages/kumo-engine/test/test_bundle_provenance.py @@ -14,6 +14,14 @@ from kumoapi.model_plan import ModelPlan, TrainingTableGenerationPlan import kumo.engine as kumo +from kumo.engine._manifest import ( + BundleProvenance, + BundleTrainingTable, + SourceTableEntry, + TrainingTableParents, + read_manifest, + write_manifest, +) from kumo.engine.trainer._spec import MANIFEST_FILE from kumo.engine.trainer.trainer import TrainingResult @@ -22,19 +30,36 @@ def test_fit_records_provenance(fitted): trainer, _, snapshot = fitted result = trainer._training_result provenance = result.provenance - parents = provenance['parents'] - assert parents.keys() == {'training_table', 'warm_start'} - assert parents['training_table'] and parents['warm_start'] is None - assert provenance['training_table'] == { - 'source': 'pquery', - 'snapshot': snapshot.id - } - assert provenance['source_tables'] == ( - snapshot.manifest['parents']['source_tables']) - assert provenance['fit_mode'] == 'refit_trainval' + parents = provenance.parents + assert parents.training_table and parents.warm_start is None + assert parents.extra == {} + assert provenance.training_table == BundleTrainingTable( + source='pquery', snapshot=snapshot.id) + assert provenance.source_tables == snapshot.manifest.parents.source_tables + assert provenance.fit_mode == 'refit_trainval' assert re.fullmatch('[0-9a-f]{64}', trainer.bundle_id) - assert trainer.bundle_id == result.bundle_id == provenance['id'] - assert kumo.Trainer.load(result.bundle_dir).bundle_id == provenance['id'] + assert trainer.bundle_id == result.bundle_id == provenance.id + assert kumo.Trainer.load(result.bundle_dir).bundle_id == provenance.id + + +def test_provenance_round_trips_unchanged(fitted, tmp_path): + trainer, _, _ = fitted + newer = json.loads( + (trainer._training_result.bundle_dir / 'provenance.json').read_text()) + newer['future'] = True + newer['fit']['gpu_seconds'] = 1.5 + newer['training_table'].update(connector='S3', location='s3://b/tt', + version='1', version_kind='exact') + source = tmp_path / 'newer.json' + source.write_text(json.dumps(newer, indent=2, sort_keys=True) + '\n') + + provenance = read_manifest(source, BundleProvenance) + write_manifest(tmp_path / 'out.json', provenance) + + assert (tmp_path / 'out.json').read_bytes() == source.read_bytes() + assert provenance.extra == {'future': True} + assert provenance.fit.extra == {'gpu_seconds': 1.5} + assert provenance.training_table.source_table.location == 's3://b/tt' def test_bundle_without_provenance_loads_and_predicts_identically( @@ -79,35 +104,28 @@ def _finish_stub_fit(bundle, trainer, manifest=None, source_mode=None): snapshot_id='snap', export_config=None, training_table_generation=TrainingTableGenerationPlan(), provenance=dict(train_table=table, warm_start_bundle_id=None, - started_at=datetime.now(timezone.utc), - fit_mode=None)) + started_at=datetime.now(timezone.utc), fit_mode=None)) def test_provenance_records_custom_or_missing_table_source(tmp_path): - source_table = { - 'connector': 'S3', - 'location': 's3://bucket/exports/tt', - 'version': '17', - 'version_kind': 'timestamp', - } - manifest = { - 'parents': { - 'snapshot': 'snap', - 'source_table': source_table - }, - } + source_table = SourceTableEntry(connector='S3', + location='s3://bucket/exports/tt', + version='17', version_kind='timestamp') + manifest = SimpleNamespace(parents=TrainingTableParents( + snapshot='snap', source_table=source_table)) custom = _finish_stub_fit(tmp_path / 'custom', kumo.Trainer(ModelPlan()), manifest, 'custom_tt') - assert custom.provenance['training_table'] == { - 'source': 'custom_tt', - 'snapshot': 'snap', - **source_table, - } + assert custom.provenance.training_table == BundleTrainingTable( + source='custom_tt', snapshot='snap', source_table=source_table) + assert json.loads((tmp_path / 'custom' / + 'provenance.json').read_text())['training_table'] == { + 'source': 'custom_tt', + 'snapshot': 'snap', + **source_table.to_dict() + } older = _finish_stub_fit(tmp_path / 'older', kumo.Trainer(ModelPlan())) - assert older.provenance['training_table'] == { - 'source': None, - 'snapshot': None - } + assert older.provenance.training_table == BundleTrainingTable( + source=None, snapshot=None) def test_provenance_failure_warns_and_still_publishes(tmp_path, monkeypatch): diff --git a/packages/kumo-engine/test/test_warm_start.py b/packages/kumo-engine/test/test_warm_start.py index d60a645cb..3370f632f 100644 --- a/packages/kumo-engine/test/test_warm_start.py +++ b/packages/kumo-engine/test/test_warm_start.py @@ -182,7 +182,7 @@ 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.provenance['parents']['warm_start'] == source.bundle_id + assert second.provenance.parents.warm_start == source.bundle_id assert destination.model_plan == configured_plan assert _hash_bundle(durable) == original_hashes assert _hash_bundle(source_bundle) == original_hashes From 13209487d238adf9a2c97db97e7f687c32984851 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sat, 24 Oct 2026 17:37:42 +0000 Subject: [PATCH 47/80] Embed the training table in a self-contained model manifest; store models Bundle provenance copied parts of the training table and snapshot into bespoke fields (training_table source/snapshot, source_tables) and named the warm-start model by id only. - ModelManifest (was BundleProvenance), kind model, file _provenance.json: parents.training_table embeds the training-table manifest (with its snapshot), parents.warm_start the warm-start model's manifest with its own model parents cut to stubs, so retrain chains stay flat; a missing parent is a stub. The bundle id still leaves the manifest out (old and new names). - fit(mlops=): with MLOps on, the bundle is stored under /models//; a complete copy is kept, a partial one replaced; failures warn. Each fit returns a new TrainingResult (.id, .manifest); Trainer.result and Trainer.manifest expose it. Verified: test_bundle_provenance.py (manifest contents and round trip, stubs, flat warm-start chains, new result per fit, root store and repair, write failure), warm start, registry, holdout and kumo-core artifact suites. Co-Authored-By: Claude Opus 5.5 --- packages/kumo-core/src/kumo_core/artifacts.py | 6 +- packages/kumo-core/tests/test_artifacts.py | 9 +- packages/kumo-engine/docs/mlops.md | 11 ++ packages/kumo-engine/kumo/engine/__init__.py | 7 +- packages/kumo-engine/kumo/engine/_manifest.py | 51 ++---- .../kumo/engine/trainer/_provenance.py | 72 +++++---- .../kumo/engine/trainer/trainer.py | 85 ++++++++-- .../test/test_bundle_provenance.py | 153 ++++++++++-------- packages/kumo-engine/test/test_warm_start.py | 3 +- 9 files changed, 242 insertions(+), 155 deletions(-) diff --git a/packages/kumo-core/src/kumo_core/artifacts.py b/packages/kumo-core/src/kumo_core/artifacts.py index c5730c55f..c8e45e4a4 100644 --- a/packages/kumo-core/src/kumo_core/artifacts.py +++ b/packages/kumo-core/src/kumo_core/artifacts.py @@ -255,7 +255,8 @@ def _sha256(path: Path) -> str: return digest.hexdigest() -PROVENANCE_FILE = 'provenance.json' +#: Earlier name of the bundle manifest; also kept out of the bundle id. +_OLD_PROVENANCE_FILE = 'provenance.json' def verify_model_bundle(bundle: Path) -> str: @@ -359,7 +360,8 @@ def bundle_id_from_manifest(bundle: Path) -> str: def _bundle_id(pairs: Iterable[tuple[str, str]]) -> str: payload = json.dumps( - sorted(pair for pair in pairs if pair[0] != PROVENANCE_FILE), + sorted(pair for pair in pairs + if pair[0] not in (PROVENANCE_FILE, _OLD_PROVENANCE_FILE)), separators=(',', ':')) return hashlib.sha256(payload.encode('utf-8')).hexdigest() diff --git a/packages/kumo-core/tests/test_artifacts.py b/packages/kumo-core/tests/test_artifacts.py index eadaf82e1..8828667ae 100644 --- a/packages/kumo-core/tests/test_artifacts.py +++ b/packages/kumo-core/tests/test_artifacts.py @@ -242,16 +242,17 @@ def best_ckpt(self): 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 payload in ('{"a": 1}', '{"a": 2}'): - (bundle / "provenance.json").write_text(payload) - add_bundle_file(bundle, "provenance", "provenance.json") + for name, payload in (("_provenance.json", '{"a": 1}'), + ("_provenance.json", '{"a": 2}'), + ("provenance.json", '{"a": 3}')): + (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 = [ diff --git a/packages/kumo-engine/docs/mlops.md b/packages/kumo-engine/docs/mlops.md index e1480f29e..5c93f0f1f 100644 --- a/packages/kumo-engine/docs/mlops.md +++ b/packages/kumo-engine/docs/mlops.md @@ -35,3 +35,14 @@ A generated or caller-supplied training table carries `_provenance.json` in its `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` and `Trainer.manifest` describe the latest fit or the loaded bundle. + +With MLOps on, each bundle is stored under `/models//`. A complete copy that is already there is kept; a partial one is replaced. diff --git a/packages/kumo-engine/kumo/engine/__init__.py b/packages/kumo-engine/kumo/engine/__init__.py index 52b0540ae..e2d3c1500 100644 --- a/packages/kumo-engine/kumo/engine/__init__.py +++ b/packages/kumo-engine/kumo/engine/__init__.py @@ -25,7 +25,11 @@ from kumoapi.plugin import ModelRegistryLocation from kumoapi.table import Column -from kumo.engine._manifest import SnapshotManifest, TrainingTableManifest +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 @@ -89,6 +93,7 @@ "mlops", "SnapshotManifest", "TrainingTableManifest", + "ModelManifest", "PredictiveQuery", "TrainingTable", "PredictionTable", diff --git a/packages/kumo-engine/kumo/engine/_manifest.py b/packages/kumo-engine/kumo/engine/_manifest.py index 8b1c7f860..81265178f 100644 --- a/packages/kumo-engine/kumo/engine/_manifest.py +++ b/packages/kumo-engine/kumo/engine/_manifest.py @@ -16,7 +16,6 @@ List, Mapping, Optional, - Self, Type, TypeVar, Union, @@ -226,9 +225,10 @@ class TrainingTableManifest(Manifest): @dataclass(frozen=True, kw_only=True) -class BundleParents(Record): - training_table: str - warm_start: Optional[str] +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) @@ -239,40 +239,15 @@ class FitTiming(Record): @dataclass(frozen=True, kw_only=True) -class BundleTrainingTable(Record): - """Written with ``source_table``'s keys inline.""" - - source: Optional[str] - snapshot: Optional[str] - source_table: Optional[SourceTableEntry] = None - - def to_dict(self) -> Dict[str, Any]: - data = {**self.extra, 'source': self.source, 'snapshot': self.snapshot} - if self.source_table is not None: - data.update(self.source_table.to_dict()) - return data - - @classmethod - def from_dict(cls, data: Mapping[str, Any]) -> Self: - rest = { - key: value - for key, value in data.items() if key not in ('source', 'snapshot') - } - if 'version_kind' not in rest: - return cls(source=data.get('source'), snapshot=data.get('snapshot'), - extra=rest) - return cls(source=data.get('source'), snapshot=data.get('snapshot'), - source_table=SourceTableEntry.from_dict(rest)) - - -@dataclass(frozen=True, kw_only=True) -class BundleProvenance(Manifest): - kind: str = 'model_bundle' - parents: BundleParents - fit: FitTiming - training_table: BundleTrainingTable - source_tables: Optional[Dict[str, SourceTableEntry]] - fit_mode: Optional[str] +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 write_manifest(path: Path, manifest: Manifest) -> None: diff --git a/packages/kumo-engine/kumo/engine/trainer/_provenance.py b/packages/kumo-engine/kumo/engine/trainer/_provenance.py index 3287db604..4429186dd 100644 --- a/packages/kumo-engine/kumo/engine/trainer/_provenance.py +++ b/packages/kumo-engine/kumo/engine/trainer/_provenance.py @@ -1,23 +1,23 @@ # SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. # All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Fit record in ``provenance.json``; excluded from the bundle id.""" +"""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 Dict, Optional +from typing import Optional from kumo.engine._manifest import ( - BundleParents, - BundleProvenance, - BundleTrainingTable, FitTiming, - SourceTableEntry, + ModelManifest, + ModelParents, + embedded, read_manifest, stamp, + stub, write_manifest, ) -from kumo.engine.graph import GraphSnapshot from kumo.engine.pquery.training_table import TrainingTable from kumo.engine.trainer._spec import PROVENANCE_ROLE, listed_file @@ -25,53 +25,63 @@ def write_provenance( bundle_dir: Path, *, - bundle_id: Optional[str], + bundle_id: str, train_table: TrainingTable, - warm_start_bundle_id: Optional[str], + warm_start: Optional[ModelManifest], started_at: datetime, fit_mode: Optional[str], - snapshot: Optional[GraphSnapshot] = None, ) -> 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, - BundleProvenance( + ModelManifest( id=bundle_id, - parents=BundleParents(training_table=train_table.tt_id, - warm_start=warm_start_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), ), - training_table=_training_table(train_table), - source_tables=_source_tables(snapshot), fit_mode=fit_mode, **stamp(), )) add_bundle_file(bundle_dir, PROVENANCE_ROLE, PROVENANCE_FILE) -def _source_tables( - snapshot: Optional[GraphSnapshot] -) -> Optional[Dict[str, SourceTableEntry]]: - manifest = None if snapshot is None else snapshot.manifest - if manifest is None or manifest.reconstructed: - return None - return manifest.parents.source_tables +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) -def _training_table(train_table: TrainingTable) -> BundleTrainingTable: - manifest = train_table.manifest - return BundleTrainingTable( - source=train_table.source_mode, - snapshot=None if manifest is None else manifest.parents.snapshot, - source_table=None - if manifest is None else manifest.parents.source_table) + return replace(copy, parents=replace( + parents, warm_start=cut(parents.warm_start), + derived_from=cut(parents.derived_from))) -def read_provenance(bundle_dir: Path) -> Optional[BundleProvenance]: +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, BundleProvenance) + return None if path is None else read_manifest(path, ModelManifest) diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index c318590f0..bdf62243e 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -52,7 +52,7 @@ require_embedding_models, require_trained_embedding_models, ) -from kumo.engine._manifest import BundleProvenance +from kumo.engine._manifest import ModelManifest from kumo.engine._scratch import new_run_id, trainer_scratch_dir from kumo.engine.connector.base import Connector from kumo.engine.graph import GraphSnapshot @@ -95,8 +95,13 @@ class TrainingResult: bundle_id: Optional[str] = None @property - def provenance(self) -> Optional[BundleProvenance]: - """The bundle's recorded fit provenance; None for older bundles.""" + 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) @@ -197,6 +202,41 @@ def _find_holdout(log_dir: str) -> tuple[Optional[Path], Optional[bool]]: 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 upload_bundle + from pyarrow.fs import FileSelector, FileType + + 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}')) + local = { + path.relative_to(result.bundle_dir).as_posix(): path.stat().st_size + for path in Path(result.bundle_dir).rglob('*') if path.is_file() + } + 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 + } + if remote == local: + return + if remote: + 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 @@ -314,10 +354,15 @@ def bundle_id(self) -> Optional[str]: return getattr(self._training_result, 'bundle_id', None) @property - def provenance(self) -> Optional[BundleProvenance]: - """Fit provenance of the fitted or loaded bundle, if recorded.""" + def result(self) -> Optional[TrainingResult]: + """The last fit's result, or the loaded bundle's.""" + return self._training_result + + @property + def manifest(self) -> Optional[ModelManifest]: + """The fitted or loaded bundle's ``_provenance.json``, if any.""" result = self._training_result - return None if result is None else result.provenance + return None if result is None else result.manifest @_without_torch_scatter_hint def fit( @@ -327,6 +372,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`. @@ -355,6 +401,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, @@ -383,6 +434,10 @@ 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 @@ -549,10 +604,13 @@ def fit( training_table_generation=train_table.generation_plan, export_config=export_config, model_plan=effective_plan, - provenance=dict(train_table=train_table, - warm_start_bundle_id=warm_start_bundle_id, - started_at=started_at, fit_mode=fit_mode, - snapshot=snapshot), + 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( @@ -565,6 +623,7 @@ def _finish_fit( export_config: Optional[ModelOutputConfig], 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 @@ -586,15 +645,17 @@ def _finish_fit( except Exception as error: warnings.warn(f"Could not compute the bundle id: {error}", UserWarning, stacklevel=3) - if provenance is not None: + 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"Publishing the model bundle without _provenance.json: " f"{error}", UserWarning, stacklevel=3) + if settings is not None and settings.wants('model'): + _store_model(training_result, settings) if export_config is not None: spec = read_training_spec(training_result.bundle_dir) published = export_config.connector.sink_backend.publish_model( diff --git a/packages/kumo-engine/test/test_bundle_provenance.py b/packages/kumo-engine/test/test_bundle_provenance.py index 608564460..5b44d3153 100644 --- a/packages/kumo-engine/test/test_bundle_provenance.py +++ b/packages/kumo-engine/test/test_bundle_provenance.py @@ -15,61 +15,56 @@ import kumo.engine as kumo from kumo.engine._manifest import ( - BundleProvenance, - BundleTrainingTable, - SourceTableEntry, - TrainingTableParents, + 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_provenance(fitted): +def test_fit_records_a_self_contained_manifest(fitted): trainer, _, snapshot = fitted - result = trainer._training_result - provenance = result.provenance - parents = provenance.parents - assert parents.training_table and parents.warm_start is None - assert parents.extra == {} - assert provenance.training_table == BundleTrainingTable( - source='pquery', snapshot=snapshot.id) - assert provenance.source_tables == snapshot.manifest.parents.source_tables - assert provenance.fit_mode == 'refit_trainval' + result = trainer.result + manifest = result.manifest + assert manifest.kind == 'model' assert re.fullmatch('[0-9a-f]{64}', trainer.bundle_id) - assert trainer.bundle_id == result.bundle_id == provenance.id - assert kumo.Trainer.load(result.bundle_dir).bundle_id == provenance.id + assert trainer.bundle_id == result.id == manifest.id + assert trainer.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_provenance_round_trips_unchanged(fitted, tmp_path): +def test_manifest_round_trips_unchanged(fitted, tmp_path): trainer, _, _ = fitted newer = json.loads( - (trainer._training_result.bundle_dir / 'provenance.json').read_text()) + (trainer.result.bundle_dir / '_provenance.json').read_text()) newer['future'] = True newer['fit']['gpu_seconds'] = 1.5 - newer['training_table'].update(connector='S3', location='s3://b/tt', - version='1', version_kind='exact') source = tmp_path / 'newer.json' source.write_text(json.dumps(newer, indent=2, sort_keys=True) + '\n') - provenance = read_manifest(source, BundleProvenance) - write_manifest(tmp_path / 'out.json', provenance) + manifest = read_manifest(source, ModelManifest) + write_manifest(tmp_path / 'out.json', manifest) assert (tmp_path / 'out.json').read_bytes() == source.read_bytes() - assert provenance.extra == {'future': True} - assert provenance.fit.extra == {'gpu_seconds': 1.5} - assert provenance.training_table.source_table.location == 's3://b/tt' + assert manifest.extra == {'future': True} + assert manifest.fit.extra == {'gpu_seconds': 1.5} -def test_bundle_without_provenance_loads_and_predicts_identically( +def test_bundle_without_a_manifest_loads_and_predicts_identically( fitted, tmp_path): trainer, prediction_table, snapshot = fitted - result = trainer._training_result - loaded = kumo.Trainer.load(result.bundle_dir) - assert (loaded.bundle_id, loaded.provenance) == (result.bundle_id, - result.provenance) - + result = trainer.result old = tmp_path / 'old' shutil.copytree(result.bundle_dir, old) manifest = json.loads((old / MANIFEST_FILE).read_text()) @@ -78,10 +73,10 @@ def test_bundle_without_provenance_loads_and_predicts_identically( if entry['role'] != 'provenance' ] (old / MANIFEST_FILE).write_text(json.dumps(manifest)) - (old / 'provenance.json').unlink() + (old / '_provenance.json').unlink() restored = kumo.Trainer.load(old) assert restored.bundle_id == result.bundle_id - assert restored.provenance is None + assert restored.manifest is None def predict(model): return model.predict( @@ -92,52 +87,78 @@ def predict(model): pd.testing.assert_frame_equal(predict(restored), predict(trainer)) -def _finish_stub_fit(bundle, trainer, manifest=None, source_mode=None): +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-1', manifest=manifest, - source_mode=source_mode) + 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_bundle_id=None, - started_at=datetime.now(timezone.utc), fit_mode=None)) - - -def test_provenance_records_custom_or_missing_table_source(tmp_path): - source_table = SourceTableEntry(connector='S3', - location='s3://bucket/exports/tt', - version='17', version_kind='timestamp') - manifest = SimpleNamespace(parents=TrainingTableParents( - snapshot='snap', source_table=source_table)) - custom = _finish_stub_fit(tmp_path / 'custom', kumo.Trainer(ModelPlan()), - manifest, 'custom_tt') - assert custom.provenance.training_table == BundleTrainingTable( - source='custom_tt', snapshot='snap', source_table=source_table) - assert json.loads((tmp_path / 'custom' / - 'provenance.json').read_text())['training_table'] == { - 'source': 'custom_tt', - 'snapshot': 'snap', - **source_table.to_dict() - } - older = _finish_stub_fit(tmp_path / 'older', kumo.Trainer(ModelPlan())) - assert older.provenance.training_table == BundleTrainingTable( - source=None, snapshot=None) - - -def test_provenance_failure_warns_and_still_publishes(tmp_path, monkeypatch): + 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).bundle_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): + def fail(*args, **kwargs): raise RuntimeError('unreadable manifest') - monkeypatch.setattr(_provenance, '_training_table', fail) + 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.provenance is None + assert result.manifest is None assert re.fullmatch('[0-9a-f]{64}', result.bundle_id) - assert trainer._training_result is result + assert trainer.result is result diff --git a/packages/kumo-engine/test/test_warm_start.py b/packages/kumo-engine/test/test_warm_start.py index 3370f632f..25a109489 100644 --- a/packages/kumo-engine/test/test_warm_start.py +++ b/packages/kumo-engine/test/test_warm_start.py @@ -182,7 +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.provenance.parents.warm_start == source.bundle_id + assert second.manifest.parents.warm_start.id == source.bundle_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 From 456b97d11103e5ceee1b47f3b82e8a43e7aee7ef Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sat, 24 Oct 2026 17:41:09 +0000 Subject: [PATCH 48/80] Store the model under the root after the plain export Order per the series: a plain export_config publish first, then the artifact-root store. Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/trainer/trainer.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index bdf62243e..1acc5e100 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -654,8 +654,6 @@ def _finish_fit( warnings.warn( f"Publishing the model bundle without _provenance.json: " f"{error}", UserWarning, stacklevel=3) - if settings is not None and settings.wants('model'): - _store_model(training_result, settings) if export_config is not None: spec = read_training_spec(training_result.bundle_dir) published = export_config.connector.sink_backend.publish_model( @@ -672,6 +670,8 @@ 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) self._training_result = training_result self._snapshot_id = snapshot_id From 231a309b1bf704b619d993d945fdf253e7cd7bb3 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Tue, 20 Oct 2026 20:30:39 +0000 Subject: [PATCH 49/80] Add GraphSnapshot export and load through object storage A graph index can only be used on the machine that built it, so building on a CPU box and training on a GPU box means re-reading every source. GraphSnapshot.export(connector, path) copies the index to /// through any S3, GCS or file connector and returns that location; GraphSnapshot.load_from_connector (connector, path=...) copies it into the local snapshot store. Transfers run in parallel with a CRC32 per range, an interrupted export resumes, a load stages and renames so it completes or leaves nothing, and identical racing exports both succeed. Re-running an export over a committed one re-uploads any file whose size no longer matches. kumo_snapshot.json, written last, embeds the snapshot manifest; load writes it back with the loading machine's formats under `loaded`. Load compares no versions: a graph mismatch or unknown layout format warns, and an index DiskGraph cannot open fails naming the recorded and installed formats. Loading again reuses the local copy unless it fails the DiskGraph integrity check. Graph.snapshot() reuse is unchanged. Warehouse connectors raise NotImplementedError. Ported from the platform's snapshot export, with fixes: free space must cover the full size, malformed manifests raise ValueError, the load path always runs the integrity check, and a racing export cannot fail on the marker delete. Verified with the export/load tests (round trip, resume, corruption, hostile and malformed manifests, free space, concurrent and late identical exports, repair, recorded differences, unchanged in-place reuse, object-store fakes), the S3 sink test, and the existing snapshot, kumo-s3 and kumo-core object store suites. Live S3/GCS tests are env-gated and were not run. Co-Authored-By: Claude Opus 5.5 --- .../src/kumo_core/object_store_sink.py | 10 +- .../src/kumo_core/sources/object_store.py | 8 +- .../kumo-engine/kumo/engine/_snapshot_io.py | 473 ++++++++++++++++ packages/kumo-engine/kumo/engine/graph.py | 91 ++- packages/kumo-engine/test/conftest.py | 28 + .../kumo-engine/test/test_snapshot_export.py | 526 ++++++++++++++++++ .../test/test_snapshot_export_live.py | 56 ++ packages/kumo-s3/src/kumo_s3/backend.py | 33 +- packages/kumo-s3/src/kumo_s3/sink.py | 9 + packages/kumo-s3/tests/test_sink.py | 49 ++ 10 files changed, 1256 insertions(+), 27 deletions(-) create mode 100644 packages/kumo-engine/kumo/engine/_snapshot_io.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-s3/tests/test_sink.py 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 93eed57ce..4843916fd 100644 --- a/packages/kumo-core/src/kumo_core/object_store_sink.py +++ b/packages/kumo-core/src/kumo_core/object_store_sink.py @@ -100,9 +100,13 @@ 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``.""" + def filesystem(self, connector: SourceConnector, path: str, *, + parallel_uploads: bool = False) -> Tuple[FileSystem, str]: + """The authenticated filesystem and the path ``path`` resolves to. + ``parallel_uploads`` tunes S3 for many concurrent uploads (one + buffered part each, small files as one PUT); other stores ignore it. + """ + del parallel_uploads filesystem, resolved = self._destination(connector, path) return filesystem, resolved.rstrip('/') 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 828600080..bf1058d07 100644 --- a/packages/kumo-core/src/kumo_core/sources/object_store.py +++ b/packages/kumo-core/src/kumo_core/sources/object_store.py @@ -491,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 @@ -521,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: @@ -816,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-engine/kumo/engine/_snapshot_io.py b/packages/kumo-engine/kumo/engine/_snapshot_io.py new file mode 100644 index 000000000..ef38e39dc --- /dev/null +++ b/packages/kumo-engine/kumo/engine/_snapshot_io.py @@ -0,0 +1,473 @@ +# 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 shutil +import threading +import time +import warnings +import zlib +from concurrent.futures import FIRST_EXCEPTION, ThreadPoolExecutor, wait +from functools import partial +from pathlib import Path +from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple + +from pyarrow.fs import FileSelector, FileSystem, FileType + +_FORMAT = 1 +_MANIFEST = 'kumo_snapshot.json' +_EXPORTING = '_EXPORTING' +_RANGE_BYTES = 64 << 20 +_CHUNK = 8 << 20 +_INDEX_MANIFEST = '_MANIFEST.json' +_ATTEMPTS = 3 +_KEY = re.compile('[0-9a-f]{64}') +_PROGRESS_SECONDS = 30 +_UNSUPPORTED = { + 'SNOWFLAKE': 'Snowflake stages are not supported', + 'DATABRICKS': ('for a Unity Catalog Volume, use an S3Connector or ' + "GCSConnector on the volume's backing cloud path"), + 'BIGQUERY': 'BigQuery cannot hold files', +} + +logger = logging.getLogger(__name__) + + +def _location(connector: Any, path: str) -> Tuple[FileSystem, str]: + from kumo_core.object_store_sink import ObjectStoreSinkBackend + + reason = _UNSUPPORTED.get(getattr(connector.source_type, 'name', None)) + sink = None if reason else connector.sink_backend + if not isinstance(sink, ObjectStoreSinkBackend): + raise NotImplementedError( + f'{type(connector).__name__} cannot hold a snapshot export' + f"{f' ({reason})' if reason else ''}; use an S3Connector, " + 'GCSConnector or FileConnector.') + return sink.filesystem(connector, path, parallel_uploads=True) + + +def _check_key(graph_key: Any, data_version: Any, snapshot_id: Any) -> None: + if not (isinstance(graph_key, str) and isinstance(data_version, str) + and _KEY.fullmatch(graph_key) and _KEY.fullmatch(data_version) + and snapshot_id == f'{graph_key}@{data_version[:16]}'): + raise ValueError( + f'Snapshot {snapshot_id!r} does not name a graph key and data ' + 'version this package produced.') + + +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() + + +def _parallel(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 + + stop = threading.Event() + total = sum(sizes) + with copies(items) as workers, ThreadPoolExecutor(workers) as pool: + # Created after `copies` raises Arrow's IO pool, which sizes the + # client's connection pool. + client = _Client(connector, path) + logger.info('Snapshot %s of %s: %d files, %d bytes, %d workers', + action, path, len(items), total, workers) + started = time.monotonic() + futures = { + pool.submit(_retried, client, work, item, stop): size + for item, size in zip(items, sizes) + } + try: + pending, moved = set(futures), 0 + while pending: + done, pending = wait(pending, timeout=_PROGRESS_SECONDS, + return_when=FIRST_EXCEPTION) + for future in done: + if future.exception() is not None: + raise future.exception() + moved += futures[future] + elapsed = time.monotonic() - started + logger.info('Snapshot %s of %s: %d of %d bytes, %.2f GB/s', + action, path, moved, total, + moved / max(elapsed, 1e-9) / 1e9) + return [future.result() for future in futures] + except BaseException: + # Leaving the pool joins the workers, which notice `stop`. + stop.set() + for future in futures: + future.cancel() + raise + + +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 _stream(source: Path, size: int, stop: threading.Event, + write: Optional[Callable[[bytes], Any]] = None) -> List[int]: + crcs: List[int] = [] + with open(source, 'rb', buffering=0) as reader: + 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} changed while it was exported') + if offset % _RANGE_BYTES == 0: + crcs.append(0) + crcs[-1] = zlib.crc32(chunk, crcs[-1]) + if write is not None: + write(chunk) + return crcs or [0] + + +def _read_json(filesystem: FileSystem, path: str) -> Optional[dict]: + try: + with filesystem.open_input_stream(path) as reader: + return json.loads(reader.readall()) + except FileNotFoundError: + return None + + +def _open_output(filesystem: FileSystem, path: str) -> Any: + if filesystem.type_name not in ('s3', 'gcs'): + filesystem.create_dir(path.rsplit('/', 1)[0], recursive=True) + return filesystem.open_output_stream(path) + + +def _write_json(filesystem: FileSystem, path: str, value: dict) -> None: + with _open_output(filesystem, path) as writer: + writer.write(json.dumps(value).encode()) + + +def _committed(filesystem: FileSystem, root: str, location: str, + snapshot_id: str, fingerprint: str) -> bool: + existing = _read_json(filesystem, f'{root}/{_MANIFEST}') + if existing is None: + return False + if not isinstance(existing, dict): + existing = {} + if (existing.get('snapshot_id') == snapshot_id + and existing.get('index_manifest') == fingerprint): + return True + raise ValueError( + f'{location} already holds an export of snapshot ' + f"{existing.get('snapshot_id')!r} that differs from this one.") + + +def export(snapshot: Any, connector: Any, path: str) -> str: + from kumo.engine._provenance import environment + + source = Path(snapshot.index_path) + graph_key, data_version = source.parent.name, source.name + _check_key(graph_key, data_version, snapshot.id) + location = f"{path.rstrip('/')}/{graph_key}/{data_version[:16]}" + index_manifest = (source / _INDEX_MANIFEST).read_bytes() + files = sorted([(entry['rel_path'], int(entry['size_bytes'])) + for entry in json.loads(index_manifest)['files']] + + [(_INDEX_MANIFEST, len(index_manifest))], + key=lambda file: -file[1]) + fingerprint = hashlib.sha256(index_manifest).hexdigest() + identity = {'snapshot_id': snapshot.id, 'index_manifest': fingerprint} + filesystem, root = _location(connector, location) + + def committed() -> bool: + return _committed(filesystem, root, location, snapshot.id, + fingerprint) + + def present() -> Dict[str, int]: + return { + info.path[len(root) + len('/index/'):]: info.size + for info in filesystem.get_file_info( + FileSelector(f'{root}/index', recursive=True, + allow_not_found=True)) + if info.type is FileType.File + } + + def upload(client: FileSystem, file: Tuple[str, int], + stop: threading.Event) -> List[int]: + name, size = file + if uploaded.get(name) == size: + return _stream(source / name, size, stop) + with _open_output(client, f'{root}/index/{name}') as writer: + return _stream(source / name, size, stop, writer.write) + + if committed(): + # Re-running repairs files a racing or interrupted export rewrote. + uploaded = present() + stale = [file for file in files if uploaded.get(file[0]) != file[1]] + if stale: + _parallel(connector, location, stale, + [size for _, size in stale], upload, 'repair') + return location + marker = _read_json(filesystem, f'{root}/{_EXPORTING}') + if not isinstance(marker, (dict, type(None))): + marker = {} + if marker is None and filesystem.get_file_info( + FileSelector(root, allow_not_found=True)): + raise ValueError(f'{location} is not empty; remove it and re-export.') + if marker is not None and marker.get('snapshot_id') != snapshot.id: + raise ValueError( + f"{location} holds a partial export of snapshot " + f"{marker.get('snapshot_id')!r}; remove it and re-export.") + resume = (marker or {}).get('index_manifest') == fingerprint + _write_json(filesystem, f'{root}/{_EXPORTING}', identity) + uploaded = present() if resume else {} + crcs = _parallel(connector, location, files, [size for _, size in files], + upload, 'export') + filesystem, root = _location(connector, location) + # An identical export racing this one uploads the same bytes, so + # whichever finishes first commits and the other finds it committed. + if committed(): + return location + if (_read_json(filesystem, f'{root}/{_EXPORTING}') != identity + and not committed()): + raise RuntimeError( + f'Another export to {location} started while this one ran; ' + 'remove it and re-export.') + _write_json( + filesystem, f'{root}/{_MANIFEST}', { + 'format': _FORMAT, + 'snapshot_id': snapshot.id, + 'graph_key': graph_key, + 'data_version': data_version, + 'index_manifest': fingerprint, + 'graph_config': snapshot.graph_config, + 'snapshot': snapshot.manifest, + 'environment': environment(), + 'range_bytes': _RANGE_BYTES, + 'files': [{ + 'path': name, + 'size': size, + 'crc32': file_crcs + } for (name, size), file_crcs in zip(files, crcs)], + }) + try: + filesystem.delete_file(f'{root}/{_EXPORTING}') + except FileNotFoundError: + pass + return location + + +def _format_note(manifest: Any) -> str: + if not isinstance(manifest, dict) or manifest.get('format') == _FORMAT: + return '' + return (f" (export layout format {manifest.get('format')!r}; this " + f'engine reads {_FORMAT})') + + +def _malformed(path: str, detail: str, manifest: Any) -> ValueError: + return ValueError(f'The snapshot manifest at {path} {detail}; the export ' + f'is corrupt{_format_note(manifest)}.') + + +def read(connector: Any, path: str) -> Tuple[str, dict]: + filesystem, root = _location(connector, path) + manifest = _read_json(filesystem, f'{root}/{_MANIFEST}') + if manifest is None: + raise FileNotFoundError( + f'No complete snapshot export at {path}: {_MANIFEST} is missing. ' + 'Re-run GraphSnapshot.export to finish it.') + if not isinstance(manifest, dict): + raise _malformed(path, 'is not a JSON object', manifest) + if manifest.get('format') != _FORMAT: + warnings.warn( + f"Snapshot export at {path} has layout format " + f"{manifest.get('format')!r}; reading it as format {_FORMAT}.", + UserWarning, stacklevel=3) + _check_key(manifest.get('graph_key'), manifest.get('data_version'), + manifest.get('snapshot_id')) + span, files = manifest.get('range_bytes'), manifest.get('files') + if type(span) is not int or span <= 0: + raise _malformed(path, 'has an invalid range_bytes', manifest) + if not isinstance(files, list): + raise _malformed(path, 'has no file list', manifest) + for key in ('snapshot', 'environment'): + if not isinstance(manifest.get(key, {}), (dict, type(None))): + raise _malformed(path, f'has a non-object {key}', manifest) + if not isinstance(manifest.get('graph_config'), dict): + raise _malformed(path, 'has no graph_config', manifest) + seen = set() + for entry in files: + name, size, crcs = ((entry.get('path'), entry.get('size'), + entry.get('crc32')) + if isinstance(entry, dict) else (None, None, None)) + if (not isinstance(name, str) or name.startswith('/') or '\\' in name + or any(part in ('', '.', '..') for part in name.split('/')) + or name in seen or type(size) is not int or size < 0 + or not isinstance(crcs, list) + or any(type(crc) is not int for crc in crcs) + or len(crcs) != max(-(-size // span), 1)): + raise _malformed( + path, f'lists {name!r} with an invalid or repeated path, size ' + 'or checksums', manifest) + seen.add(name) + return root, manifest + + +def _download(root: str, staging: Path, client: FileSystem, + item: Tuple[str, int, int, int], stop: threading.Event) -> None: + name, offset, length, expected = item + crc = 0 + # 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'Snapshot file {name} in the export is shorter than ' + 'its manifest records; the export is corrupt.') + crc = zlib.crc32(chunk, crc) + os.pwrite(fd, chunk, start) + # Flush while other ranges download, not all at the end. + getattr(os, 'fdatasync', os.fsync)(fd) + finally: + os.close(fd) + if crc != expected: + raise ValueError( + f'Snapshot file {name} failed its checksum at offset {offset}; ' + 'the export is corrupt.') + + +def load_into(connector: Any, path: str, root: str, manifest: dict, + target: Path) -> None: + from diskgraph import format_versions, verify_snapshot_integrity + from kumo_core.diskgraph.snapshot import ( + DiskgraphIngestResult, + write_diskgraph_success_marker, + ) + + from kumo.engine._manifest import write_manifest + from kumo.engine._provenance import environment + from kumo.engine.graph import SNAPSHOT_MANIFEST_FILE + + span = manifest['range_bytes'] + files = [(entry['path'], entry['size'], entry['crc32']) + for entry in manifest['files']] + staging = target.with_name(f'{target.name}.loading') + shutil.rmtree(staging, ignore_errors=True) + try: + staging.mkdir(parents=True) + # Any old target stays until the rename, so the whole size is needed. + 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} bytes in {staging}, ' + f'which has {disk.f_bavail * disk.f_frsize} free.') + 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: + os.ftruncate(fd, size) + finally: + os.close(fd) + ranges = sorted( + ((name, offset, min(span, size - offset), crcs[index]) + for name, size, crcs in files + for index, offset in enumerate(range(0, size, span))), + key=lambda item: -item[2]) + try: + _parallel(connector, path, ranges, [item[2] for item in ranges], + partial(_download, root, staging), 'load') + except ValueError as error: + if _format_note(manifest): + raise ValueError(f'{error}{_format_note(manifest)}') from error + raise + write_diskgraph_success_marker(str(staging), + DiskgraphIngestResult.for_index(target)) + verify_snapshot_integrity(str(staging)) + _check_readable(staging, manifest) + if manifest.get('snapshot') is not None: + write_manifest( + staging / SNAPSHOT_MANIFEST_FILE, { + **manifest['snapshot'], 'loaded': { + 'formats': format_versions(), + 'environment': environment(), + } + }) + for directory in sorted({(staging / name).parent + for name, _, _ in files} | {staging}): + _fsync(directory) + shutil.rmtree(target, ignore_errors=True) + os.rename(staging, target) + _fsync(target.parent) + except BaseException: + shutil.rmtree(staging, ignore_errors=True) + raise + + +def _built_with(manifest: Dict[str, Any]) -> Tuple[Any, Any]: + recorded = manifest.get('snapshot') or {} + environment = (recorded.get('environment') or manifest.get('environment') + or {}) + return recorded.get('formats'), (environment.get('kumo_engine') + if isinstance(environment, dict) else None) + + +def _check_readable(index: Path, manifest: dict) -> None: + from diskgraph import SamplerEngine, format_versions + + from kumo.engine._version import __version__ + + try: + SamplerEngine(str(index), num_threads=1) + except Exception as error: + formats, engine = _built_with(manifest) + running = format_versions() + advice = ('The on-disk formats differ: rebuild the snapshot with ' + 'this kumo-engine, or load it with the exporting one.' + if formats is not None and formats != running else + 'Re-export with a matching kumo-engine.') + raise RuntimeError( + "This machine's DiskGraph cannot open the snapshot. It was " + f"built with DiskGraph formats {formats} by kumo-engine {engine}; " + f"installed are formats {running} and kumo-engine " + f"{__version__}{_format_note(manifest)}. {advice}") from error + + +def _fsync(path: Path) -> None: + fd = os.open(path, os.O_RDONLY) + try: + os.fsync(fd) + finally: + os.close(fd) diff --git a/packages/kumo-engine/kumo/engine/graph.py b/packages/kumo-engine/kumo/engine/graph.py index 6b25908ee..84760a51a 100644 --- a/packages/kumo-engine/kumo/engine/graph.py +++ b/packages/kumo-engine/kumo/engine/graph.py @@ -60,11 +60,12 @@ """ import json +import os import warnings from dataclasses import dataclass, field, replace from functools import cached_property from pathlib import Path -from typing import Any, Dict, Mapping, Sequence +from typing import Any, Dict, Mapping, Optional, Sequence from kumo_core.artifacts import PROVENANCE_FILE from kumo_core.pquery import NativeGraph, normalize_graph @@ -86,6 +87,7 @@ snapshot_scratch_dir, ) from kumo.engine._storage import network_storage +from kumo.engine.connector.base import Connector from kumo.engine.table import Table SNAPSHOT_MANIFEST_FILE = PROVENANCE_FILE @@ -203,6 +205,42 @@ def _warn_ingest_data_loss(result: Any, graph_config: Dict[str, Any]) -> None: ) +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(path: str) -> Optional[Any]: + r"""The index a previous load left intact at `path`, if any.""" + from diskgraph import DiskgraphCorruptionError, verify_snapshot_integrity + from kumo_core.diskgraph.snapshot import check_diskgraph_output + + marker = check_diskgraph_output(path) + if marker is None or (os.path.realpath(marker.index_path) + != os.path.realpath(path)): + return None + try: + verify_snapshot_integrity(path) + except DiskgraphCorruptionError as error: + warnings.warn(f"Reloading the damaged graph index at {path}: {error}", + UserWarning, stacklevel=3) + return None + return marker + + @dataclass(frozen=True) class GraphSnapshot: r"""The materialized artifact produced by `Graph.snapshot()`: an on-disk @@ -238,6 +276,46 @@ class GraphSnapshot: graph_config: Dict[str, Any] record_sources: bool | None = field(default=None, compare=False) + def export(self, connector: Connector, path: str) -> str: + r"""Copy the index to `//` under an + S3, GCS or file connector and return that location. Re-running + resumes.""" + from kumo.engine._snapshot_io import export + + return export(self, connector, path) + + @classmethod + def load_from_connector(cls, connector: Connector, *, path: str, + graph: Optional['Graph'] = None) -> 'GraphSnapshot': + r"""Copy the snapshot exported at `path` (the location `export` + returned) into the local snapshot store.""" + from kumo_core.diskgraph.snapshot import DiskgraphSnapshotHandle + + from kumo.engine._snapshot_io import load_into, read + + root, manifest = read(connector, path) + graph_key = manifest['graph_key'] + if graph is not None and graph.id != graph_key: + warnings.warn( + f"Snapshot at {path} was taken of graph {graph_key!r}, not of " + f"the given graph {graph.id!r}; tables, edges, schemas or " + "connectors differ from the exported graph.", + UserWarning, + stacklevel=2, + ) + target = _store().path( + DiskgraphSnapshotHandle( + graph_key=graph_key, + graph_data_version=manifest['data_version'])) + with _lock(graph_key): + result = _loaded(target) + if result is None: + load_into(connector, path, root, manifest, Path(target)) + result = _loaded(target) + return cls(id=manifest['snapshot_id'], + index_path=Path(result.index_path), + graph_config=manifest['graph_config']) + @property def manifest(self) -> SnapshotManifest | None: """The snapshot manifest stored next to the index, if any.""" @@ -432,8 +510,6 @@ def snapshot(self, *, record_sources: bool = True) -> GraphSnapshot: source as connector and table name only (no bucket, path or account) and Custom SQL as a hash only (no SQL text). """ - from filelock import FileLock - self.validate() for table in self._named_tables.values(): table.connector.validate_scratch_dir(get_scratch_dir()) @@ -441,7 +517,6 @@ def snapshot(self, *, record_sources: bool = True) -> 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 @@ -459,7 +534,7 @@ def snapshot(self, *, record_sources: bool = True) -> 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'): + with _lock(graph_key): return self._snapshot_no_lock(graph_key, persist_dir, record_sources) @@ -478,9 +553,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 @@ -490,9 +563,7 @@ 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] = {} 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/test_snapshot_export.py b/packages/kumo-engine/test/test_snapshot_export.py new file mode 100644 index 000000000..39f0549b4 --- /dev/null +++ b/packages/kumo-engine/test/test_snapshot_export.py @@ -0,0 +1,526 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import json +import shutil +import threading +import warnings +import zlib +from contextlib import contextmanager +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +from diskgraph import format_versions +from kumoapi.data_source import DataSourceType +from pyarrow.fs import LocalFileSystem, SubTreeFileSystem + +import kumo.engine as kumo +import kumo.engine._snapshot_io as snapshot_io +from kumo.engine import GraphSnapshot +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) + + +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', 'snapshot-manifest.json') + } + + +def _exported_files(snapshot): + return { + name: data + for name, data in _index_files(snapshot.index_path).items() + if not name.startswith('train_tables/') + } + + +@pytest.fixture +def exports(tmp_path): + return kumo.FileConnector(root_dir=str(tmp_path / 'exports')) + + +def _export(snapshot, exports, tmp_path): + location = snapshot.export(exports, 'one') + return location, tmp_path / 'exports' / location + + +def _edit_manifest(exported, **changes): + path = exported / 'kumo_snapshot.json' + path.write_text(json.dumps({**json.loads(path.read_text()), **changes})) + + +@pytest.fixture +def workers(monkeypatch): + @contextmanager + def four(found, worker_ceiling=None): + yield 4 + + monkeypatch.setattr('kumo_core.sources.object_store.copies', four) + + +def _no_ingest(monkeypatch): + monkeypatch.setattr( + 'kumo_core.diskgraph.ingest.execute_diskgraph_ingest', + Mock(side_effect=AssertionError('ingest must not run'))) + + +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(exports, 'one/') + graph_id, data_version = snapshot.id.split('@') + assert location == f'one/{graph_id}/{data_version}' + assert snapshot.export(exports, 'one') == location + exported = tmp_path / 'exports' / location + assert _index_files(exported / 'index') == _exported_files(snapshot) + manifest = json.loads((exported / 'kumo_snapshot.json').read_text()) + assert manifest['snapshot_id'] == snapshot.id + assert manifest['snapshot'] == snapshot.manifest + + shutil.rmtree(tmp_path / 'cpu') + scratch('gpu') + loaded = GraphSnapshot.load_from_connector(exports, path=location, + graph=graph) + assert loaded.id == snapshot.id + assert loaded.graph_config == snapshot.graph_config + assert loaded.index_path.is_relative_to(tmp_path / 'gpu') + assert _index_files(loaded.index_path) == _index_files(exported / 'index') + marker = json.loads((loaded.index_path / '_SUCCESS').read_text()) + assert marker['index_path'] == str(loaded.index_path) + recorded = loaded.manifest + assert recorded.pop('loaded') == { + 'formats': format_versions(), + 'environment': manifest['environment'], + } + assert recorded == manifest['snapshot'] + + _no_ingest(monkeypatch) + monkeypatch.setattr(snapshot_io, 'load_into', + Mock(side_effect=AssertionError('already loaded'))) + assert GraphSnapshot.load_from_connector(exports, path=location) == loaded + assert graph.snapshot() == loaded + + +@pytest.mark.parametrize('damage', [False, True]) +def test_in_place_index_reuse_is_unchanged(scratch, graph, monkeypatch, + damage): + built = graph.snapshot() + marker = (built.index_path / '_SUCCESS').read_bytes() + if damage: + stats = built.index_path / 'col_stats.json' + stats.write_bytes(stats.read_bytes()[:-1]) + (built.index_path / 'extra.bin').write_bytes(b'x') + _no_ingest(monkeypatch) + with warnings.catch_warnings(): + warnings.simplefilter('error') + assert graph.snapshot() == built + assert (built.index_path / '_SUCCESS').read_bytes() == marker + + +def test_moved_index_is_rebuilt_as_before(tmp_path, scratch, graph): + built = graph.snapshot() + shutil.move(tmp_path / 'cpu', tmp_path / 'moved') + scratch('moved') + rebuilt = graph.snapshot() + assert rebuilt.id == built.id + marker = json.loads((rebuilt.index_path / '_SUCCESS').read_text()) + assert marker['index_path'] == str(rebuilt.index_path) + + +def test_load_replaces_a_damaged_local_copy( + tmp_path, scratch, graph, exports): + snapshot = graph.snapshot() + location, _ = _export(snapshot, exports, tmp_path) + scratch('gpu') + loaded = GraphSnapshot.load_from_connector(exports, path=location) + stats = loaded.index_path / 'col_stats.json' + stats.write_bytes(stats.read_bytes()[:-1]) + with pytest.warns(UserWarning, match='Reloading the damaged'): + again = GraphSnapshot.load_from_connector(exports, path=location) + assert _index_files(again.index_path) == _exported_files(snapshot) + + +def _interrupt_export(snapshot, exports, monkeypatch): + stream = snapshot_io._stream + finished = threading.Event() + + def interrupted(source, size, stop, write=None): + if write is not None and source.name == 'diskgraph_config.json': + finished.wait(10) + raise OSError('connection reset') + crcs = stream(source, size, stop, write) + if write is not None and source.name == 'col_stats.json': + finished.set() + return crcs + + monkeypatch.setattr(snapshot_io, '_stream', interrupted) + monkeypatch.setattr(snapshot_io, '_ATTEMPTS', 1) + with pytest.raises(OSError, match='connection reset'): + snapshot.export(exports, 'one') + monkeypatch.setattr(snapshot_io, '_stream', stream) + monkeypatch.setattr(snapshot_io, '_ATTEMPTS', 3) + + +@pytest.mark.parametrize('change', [False, True]) +def test_interrupted_export_is_not_loadable_and_resumes( + tmp_path, scratch, graph, exports, monkeypatch, workers, change): + snapshot = graph.snapshot() + _interrupt_export(snapshot, exports, monkeypatch) + graph_id, data_version = snapshot.id.split('@') + location = f'one/{graph_id}/{data_version}' + with pytest.raises(FileNotFoundError, match='kumo_snapshot.json'): + GraphSnapshot.load_from_connector(exports, path=location) + if change: + stats = snapshot.index_path / 'col_stats.json' + data = stats.read_bytes() + stats.write_bytes(data[:-1] + bytes([data[-1] ^ 1])) + index_manifest = snapshot.index_path / '_MANIFEST.json' + index_manifest.write_bytes(index_manifest.read_bytes() + b' ') + + assert snapshot.export(exports, 'one') == location + scratch('gpu') + loaded = GraphSnapshot.load_from_connector(exports, path=location) + assert (loaded.index_path / 'col_stats.json').read_bytes() == ( + snapshot.index_path / 'col_stats.json').read_bytes() + + +def test_concurrent_identical_exports_both_succeed(tmp_path, scratch, graph, + exports, monkeypatch): + snapshot = graph.snapshot() + parallel = snapshot_io._parallel + calls = [] + + def racing(*args): + calls.append(args) + if len(calls) == 1: + assert snapshot.export(exports, 'one') == location + return parallel(*args) + + graph_id, data_version = snapshot.id.split('@') + location = f'one/{graph_id}/{data_version}' + monkeypatch.setattr(snapshot_io, '_parallel', racing) + assert snapshot.export(exports, 'one') == location + assert len(calls) == 2 + scratch('gpu') + loaded = GraphSnapshot.load_from_connector(exports, path=location) + assert _index_files(loaded.index_path) == _exported_files(snapshot) + + +def test_export_committed_by_a_late_identical_export_succeeds( + tmp_path, scratch, graph, exports, monkeypatch): + snapshot = graph.snapshot() + write = snapshot_io._write_json + raced = [] + + def racing(filesystem, path, value): + if path.endswith('kumo_snapshot.json') and not raced: + raced.append(None) + raced[0] = snapshot.export(exports, 'one') + return write(filesystem, path, value) + + monkeypatch.setattr(snapshot_io, '_write_json', racing) + location = snapshot.export(exports, 'one') + assert raced == [location] + assert not (tmp_path / 'exports' / location / '_EXPORTING').exists() + + +def test_rerunning_export_repairs_a_committed_export( + tmp_path, scratch, graph, exports): + snapshot = graph.snapshot() + location, exported = _export(snapshot, exports, tmp_path) + stats = exported / 'index/col_stats.json' + original = stats.read_bytes() + stats.write_bytes(original[:3]) + (exported / 'index/index_metadata.json').unlink() + assert snapshot.export(exports, 'one') == location + assert stats.read_bytes() == original + scratch('gpu') + loaded = GraphSnapshot.load_from_connector(exports, path=location) + assert _index_files(loaded.index_path) == _exported_files(snapshot) + + +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 = 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='checksum'): + GraphSnapshot.load_from_connector(exports, path=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() + (graph_dir / f'{snapshot.index_path.name}.loading/stale').write_text('x') + loaded = GraphSnapshot.load_from_connector(exports, path=location) + assert _index_files(loaded.index_path) == _exported_files(snapshot) + assert [ + path.name for path in graph_dir.iterdir() if path.suffix != '.lock' + ] == [snapshot.index_path.name] + + +def test_load_needs_the_full_size_while_an_old_copy_remains( + tmp_path, scratch, graph, exports, monkeypatch): + snapshot = graph.snapshot() + location, exported = _export(snapshot, exports, tmp_path) + total = sum(entry['size'] for entry in json.loads( + (exported / 'kumo_snapshot.json').read_text())['files']) + scratch('gpu') + target = (tmp_path / 'gpu/snapshots' / snapshot.index_path.parent.name / + snapshot.index_path.name) + target.mkdir(parents=True) + (target / 'stale').write_bytes(b'x' * total) + monkeypatch.setattr( + snapshot_io.os, 'statvfs', + lambda path: SimpleNamespace(f_frsize=1, f_bavail=total - 1)) + with pytest.raises(OSError, match=f'needs {total} bytes'): + GraphSnapshot.load_from_connector(exports, path=location) + assert (target / 'stale').is_file() + + +@pytest.mark.parametrize('field, value', [ + ('data_version', '../../victim'), + ('data_version', '/tmp/victim'), + ('graph_key', '..'), + ('files', [{ + 'path': '../../victim/x', + 'size': 1, + 'crc32': [0] + }]), +]) +def test_load_rejects_a_hostile_manifest(tmp_path, scratch, graph, exports, + field, value): + snapshot = graph.snapshot() + location, exported = _export(snapshot, exports, tmp_path) + victim = tmp_path / 'victim' + victim.mkdir() + (victim / 'keep').write_text('x') + _edit_manifest(exported, **{field: value}) + scratch('gpu/a/b') + with pytest.raises(ValueError): + GraphSnapshot.load_from_connector(exports, path=location) + assert (victim / 'keep').read_text() == 'x' + + +@pytest.mark.parametrize('edit', [ + lambda manifest: manifest.pop('files'), + lambda manifest: manifest.pop('range_bytes'), + lambda manifest: manifest.pop('graph_config'), + lambda manifest: manifest.update(snapshot=[1]), + lambda manifest: manifest['files'].append('index_metadata.json'), + lambda manifest: manifest['files'][0].pop('crc32'), + lambda manifest: manifest['files'].append(manifest['files'][0]), +]) +def test_load_rejects_a_malformed_manifest_with_value_error( + tmp_path, scratch, graph, exports, edit): + snapshot = graph.snapshot() + location, exported = _export(snapshot, exports, tmp_path) + path = exported / 'kumo_snapshot.json' + manifest = json.loads(path.read_text()) + edit(manifest) + path.write_text(json.dumps(manifest)) + scratch('gpu') + with pytest.raises(ValueError, match='export is corrupt'): + GraphSnapshot.load_from_connector(exports, path=location) + path.write_text('[]') + with pytest.raises(ValueError, match='not a JSON object'): + GraphSnapshot.load_from_connector(exports, path=location) + + +def test_load_records_differences_instead_of_refusing( + tmp_path, scratch, graph, exports, monkeypatch): + exporter_formats = {**format_versions(), 'layout': -1} + monkeypatch.setattr('diskgraph.format_versions', lambda: exporter_formats) + snapshot = graph.snapshot() + monkeypatch.undo() + location, exported = _export(snapshot, exports, tmp_path) + _edit_manifest(exported, format=2) + other = kumo.Graph(tables={'users': graph.tables['users']}) + scratch('gpu') + with pytest.warns(UserWarning) as caught: + loaded = GraphSnapshot.load_from_connector(exports, path=location, + graph=other) + messages = ' '.join(str(warning.message) for warning in caught) + assert 'layout format 2' in messages and 'was taken of graph' in messages + assert loaded.id == snapshot.id + assert loaded.manifest['formats'] == exporter_formats + assert loaded.manifest['loaded']['formats'] == format_versions() + + +def test_export_refuses_a_path_holding_other_data(tmp_path, scratch, graph, + exports): + snapshot = graph.snapshot() + graph_id, data_version = snapshot.id.split('@') + occupied = tmp_path / 'exports/occupied' / graph_id / data_version + occupied.mkdir(parents=True) + (occupied / 'other.parquet').write_bytes(b'x') + with pytest.raises(ValueError, match='not empty'): + snapshot.export(exports, 'occupied') + + _, exported = _export(snapshot, exports, tmp_path) + _edit_manifest(exported, index_manifest='other') + with pytest.raises(ValueError, match='differs from this one'): + snapshot.export(exports, 'one') + assert (exported / 'index/col_stats.json').is_file() + + +@pytest.mark.parametrize('kind, reason', [ + ('SNOWFLAKE', 'Snowflake stages'), + ('DATABRICKS', 'Unity Catalog Volume'), + ('BIGQUERY', 'BigQuery'), +]) +def test_warehouse_connectors_are_not_supported(scratch, graph, kind, reason): + connector = type(f'{kind.title()}Connector', (), { + 'source_type': DataSourceType[kind], + })() + with pytest.raises(NotImplementedError, match=reason): + graph.snapshot().export(connector, 'x') + with pytest.raises(NotImplementedError, match=reason): + GraphSnapshot.load_from_connector(connector, path='x') + + +def test_interrupted_load_joins_workers_before_cleaning_up( + tmp_path, scratch, graph, exports, monkeypatch, workers): + snapshot = graph.snapshot() + location, _ = _export(snapshot, exports, tmp_path) + download = snapshot_io._download + running = threading.Event() + + def interrupted(root, staging, client, item, stop): + if item[0] == 'col_stats.json': + running.wait(10) + raise KeyboardInterrupt + running.set() + stop.wait(1) + download(root, staging, client, item, stop) + + monkeypatch.setattr(snapshot_io, '_download', interrupted) + scratch('gpu') + with pytest.raises(KeyboardInterrupt): + GraphSnapshot.load_from_connector(exports, path=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'] + + +def test_load_explains_an_index_the_installed_diskgraph_cannot_open( + tmp_path, scratch, graph, exports): + snapshot = graph.snapshot() + location, exported = _export(snapshot, exports, tmp_path) + metadata = exported / 'index/index_metadata.json' + garbage = b'{' * metadata.stat().st_size + metadata.write_bytes(garbage) + manifest = json.loads((exported / 'kumo_snapshot.json').read_text()) + for entry in manifest['files']: + if entry['path'] == 'index_metadata.json': + entry['crc32'] = [zlib.crc32(garbage)] + manifest['snapshot']['formats'] = {'layout': -1} + manifest['snapshot']['environment']['kumo_engine'] = '0.0.1' + manifest['format'] = 2 + (exported / 'kumo_snapshot.json').write_text(json.dumps(manifest)) + + scratch('gpu') + with pytest.raises(RuntimeError, match=r"formats \{'layout': -1\} by " + r'kumo-engine 0\.0\.1; installed are formats .*' + r'export layout format 2; this engine reads 1\)\. ' + 'The on-disk formats differ: rebuild'): + GraphSnapshot.load_from_connector(exports, path=location) + + +def test_load_ignores_unknown_keys_and_a_missing_snapshot_manifest( + tmp_path, scratch, graph, exports): + snapshot = graph.snapshot() + location, exported = _export(snapshot, exports, tmp_path) + path = exported / 'kumo_snapshot.json' + manifest = json.loads(path.read_text()) + manifest['added_later'] = {'x': 1} + manifest['snapshot'] = None + for entry in manifest['files']: + entry['added_later'] = 1 + path.write_text(json.dumps(manifest)) + scratch('gpu') + loaded = GraphSnapshot.load_from_connector(exports, path=location) + assert loaded.id == snapshot.id + assert loaded.manifest is None + + +@pytest.mark.parametrize('scheme', ['s3', 'gs']) +def test_object_store_round_trip(tmp_path, scratch, graph, monkeypatch, + scheme): + storage = tmp_path / 'storage' + storage.mkdir() + bucket = SubTreeFileSystem(str(storage), LocalFileSystem()) + options = [] + + def s3(bucket_name, *, options=None, **kwargs): + options_seen.append(options) + return bucket + + options_seen = options + monkeypatch.setattr('kumo_s3.filesystem', s3) + + class Gcs(SubTreeFileSystem): + def __init__(self): + super().__init__(str(storage), LocalFileSystem()) + + monkeypatch.setattr('kumo_gcs.backend.fs.GcsFileSystem', Gcs) + + snapshot = graph.snapshot() + connector = kumo.S3Connector() if scheme == 's3' else kumo.GCSConnector() + location = snapshot.export(connector, f'{scheme}://bucket/snapshots/') + graph_id, data_version = snapshot.id.split('@') + assert location == f'{scheme}://bucket/snapshots/{graph_id}/{data_version}' + assert (storage / 'bucket/snapshots' / graph_id / data_version / + 'kumo_snapshot.json').is_file() + scratch('gpu') + loaded = GraphSnapshot.load_from_connector(connector, path=location, + graph=graph) + assert loaded.id == snapshot.id + assert _index_files(loaded.index_path) == _exported_files(snapshot) + if scheme == 's3': + assert options and all( + option['allow_delayed_open'] and not option['background_writes'] + for option in options) + + +def test_transfer_renews_an_expired_client(tmp_path, scratch, graph, exports, + monkeypatch): + snapshot = graph.snapshot() + location, _ = _export(snapshot, exports, tmp_path) + download = snapshot_io._download + failed = [] + + def expiring(root, staging, client, item, stop): + if item[0] == 'col_stats.json': + if not failed: + failed.append(client) + raise OSError('AWS Error ACCESS_DENIED: ExpiredToken') + assert client is not failed[0] + return download(root, staging, client, item, stop) + + monkeypatch.setattr(snapshot_io, '_download', expiring) + scratch('gpu') + GraphSnapshot.load_from_connector(exports, path=location) + assert failed 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..9628cedbf --- /dev/null +++ b/packages/kumo-engine/test/test_snapshot_export_live.py @@ -0,0 +1,56 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Snapshot export and load through real S3 and GCS; each store is skipped +unless its writable root is set.""" + +import contextlib +import os +import uuid + +import pytest +from pyarrow.fs import FileSelector + +import kumo.engine as kumo +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(f'LIVE SNAPSHOT {scheme.upper()} NOT RUN: set ' + f"KUMO_SNAPSHOT_LIVE_{'S3' if scheme == 's3' else 'GCS'}" + '_ROOT to a writable prefix') + 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(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 not filesystem.get_file_info( + FileSelector(prefix, allow_not_found=True)) + assert loaded.id == snapshot.id + for file in snapshot.index_path.rglob('*'): + relative = file.relative_to(snapshot.index_path) + if (file.is_file() + and file.name not in ('_SUCCESS', 'snapshot-manifest.json') + and relative.parts[0] != 'train_tables'): + assert (loaded.index_path / relative).read_bytes() == \ + file.read_bytes() diff --git a/packages/kumo-s3/src/kumo_s3/backend.py b/packages/kumo-s3/src/kumo_s3/backend.py index dfd31df9d..2b47fdfde 100644 --- a/packages/kumo-s3/src/kumo_s3/backend.py +++ b/packages/kumo-s3/src/kumo_s3/backend.py @@ -1,7 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. # All rights reserved. # SPDX-License-Identifier: Apache-2.0 -from typing import Any, Optional, Union +from typing import Any, Mapping, Optional, Union import kumo_s3 from kumo_core.sources.errors import ReconnectRequiredError @@ -13,7 +13,11 @@ ) from kumo_core.sources.principals import IdentityRef, Principal from kumo_core.sources.types import ReadMode -from kumo_s3.session import anonymous_filesystem, arrow_connection_options +from kumo_s3.session import ( + anonymous_filesystem, + arrow_connection_options, + retry_strategy, +) from kumoapi.data_source import ( DataSourceType, S3AccessKeyAuthentication, @@ -100,16 +104,25 @@ def _source_access(self, connector: Any, bucket: str, if connection else None) def _filesystem(self, connector: Any, bucket: str, - region: Optional[str]) -> S3FileSystem: + region: Optional[str], + options: Optional[Mapping[str, Any]] = None + ) -> S3FileSystem: r"""A ready filesystem is taken as it is: it may address a service that is not AWS at all, and probing its bucket name against AWS would - answer for somebody else's bucket. + answer for somebody else's bucket. ``options`` are extra + ``S3FileSystem`` keywords. """ session, allow_anonymous, connection = self._resolved_session( connector, bucket) if isinstance(session, S3FileSystem): return session + arrow_options = { + **(arrow_connection_options(connection.endpoint_url, + connection.addressing_style) + if connection else {}), + **(options or {}) + } with kumo_s3.translated(f'Signing for {bucket}', table=bucket): _require_credentials(session) # Resolve through the package so applications/tests replacing its @@ -119,9 +132,7 @@ def _filesystem(self, connector: Any, bucket: str, default_region=(connection.region or region) if connection else region, endpoint_url=connection.endpoint_url if connection else None, - options=arrow_connection_options(connection.endpoint_url, - connection.addressing_style) - if connection else None) + options=arrow_options or None) if resolved is not None: return resolved if not allow_anonymous: @@ -141,11 +152,13 @@ def _filesystem(self, connector: Any, bucket: str, if connection is not None: return anonymous_filesystem( bucket, region=connection.region or region - or session.region_name, - options=arrow_connection_options(connection.endpoint_url, - connection.addressing_style)) + or session.region_name, options=arrow_options) from pyarrow.fs import FileSystem filesystem, _ = FileSystem.from_uri(f's3://{bucket}') + if options: + return S3FileSystem(region=filesystem.region, + retry_strategy=retry_strategy(), + **options) return filesystem diff --git a/packages/kumo-s3/src/kumo_s3/sink.py b/packages/kumo-s3/src/kumo_s3/sink.py index 5e2a5a2d8..eded2e8a7 100644 --- a/packages/kumo-s3/src/kumo_s3/sink.py +++ b/packages/kumo-s3/src/kumo_s3/sink.py @@ -14,5 +14,14 @@ class S3SinkBackend(S3Sessions, ObjectStoreSinkBackend): """ source_type = DataSourceType.S3 + def filesystem(self, connector, path, *, parallel_uploads=False): + bucket, resolved, region = self._resolved(connector, path) + options = ({ + 'allow_delayed_open': True, + 'background_writes': False + } if parallel_uploads else None) + return (self._filesystem(connector, bucket, region, options), + resolved.rstrip('/')) + def __repr__(self) -> str: return f'{type(self).__name__}(identity={self._identity!r})' diff --git a/packages/kumo-s3/tests/test_sink.py b/packages/kumo-s3/tests/test_sink.py new file mode 100644 index 000000000..a0d45db50 --- /dev/null +++ b/packages/kumo-s3/tests/test_sink.py @@ -0,0 +1,49 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from types import SimpleNamespace + +import kumo_s3 +import pytest +from fakes import FakeSession +from kumo_s3 import backend +from kumo_s3.sink import S3SinkBackend +from kumoapi.data_source import ( + S3AnyURIConnectorTarget, + S3ConnectorResourceConfig, +) + + +@pytest.mark.parametrize('credentials', ['default', None]) +def test_parallel_uploads_reach_every_filesystem_construction( + monkeypatch, credentials): + built = [] + + class Record: + def __init__(self, **kwargs): + built.append(kwargs) + + monkeypatch.setattr('pyarrow.fs.S3FileSystem', Record) + monkeypatch.setattr(backend, 'S3FileSystem', Record) + monkeypatch.setattr(kumo_s3, 'ambient_session', + lambda: FakeSession(credentials=credentials)) + monkeypatch.setattr('kumo_s3.session.bucket_region', lambda bucket: 'x') + monkeypatch.setattr( + 'pyarrow.fs.FileSystem', + SimpleNamespace(from_uri=lambda uri: + (SimpleNamespace(region='us-west-2'), ''))) + connector = SimpleNamespace( + id='connector', identity=None, identity_id=None, + config=S3ConnectorResourceConfig(target=S3AnyURIConnectorTarget())) + + filesystem, path = S3SinkBackend().filesystem( + connector, 's3://bucket/prefix/', parallel_uploads=True) + + assert isinstance(filesystem, Record) + assert path == 'bucket/prefix' + assert built[-1]['allow_delayed_open'] is True + assert built[-1]['background_writes'] is False + assert 'retry_strategy' in built[-1] From 3fd23a158493decc3922495aa83ea3b1c902af3f Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Tue, 20 Oct 2026 23:07:08 +0000 Subject: [PATCH 50/80] Ship the object-store sink in the public connector image The connector image's build context is an allowlist that left out kumo_core/object_store_sink.py, so the new kumo-s3 sink test failed to import it and the s3 suite stopped at collection. Admit the module. Verified by running the s3 suite against a tree filtered by the allowlist: 49 passed (collection error reproduced without the entry). Co-Authored-By: Claude Opus 5.5 --- tests/public/Dockerfile.dockerignore | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/public/Dockerfile.dockerignore b/tests/public/Dockerfile.dockerignore index 4bedfe32e..03a5e01fb 100644 --- a/tests/public/Dockerfile.dockerignore +++ b/tests/public/Dockerfile.dockerignore @@ -12,6 +12,8 @@ !packages/kumo-core/src/kumo_core/csv_rules.py !packages/kumo-core/src/kumo_core/custom_sql.py !packages/kumo-core/src/kumo_core/artifact_staging.py +# The S3 sink tests import the shared object-store sink. +!packages/kumo-core/src/kumo_core/object_store_sink.py # Model registries validate bundles through the portable artifact contract. !packages/kumo-core/src/kumo_core/artifacts.py !packages/kumo-core/src/kumo_core/config.py From faabba580049f88365bc6d8e7213f26abd16df31 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Tue, 20 Oct 2026 23:46:37 +0000 Subject: [PATCH 51/80] Report source changes since a snapshot and warn on stale loads A loaded snapshot carried the source versions it was built from, but a scoring job had no cheap way to tell whether its sources had moved on. GraphSnapshot.source_changes(graph) asks each table's connector for its metadata-only version stamp (the value Graph.snapshot() records) and compares it with the manifest: unchanged, changed, or unknown for view and fallback hour stamps, Custom SQL, missing or reconstructed manifests and read errors. It never stages data and never raises. load_from_connector(..., graph=...) warns once naming changed tables. Docs add export routes for Databricks Volumes, Snowflake stage volumes and external volumes or stages, and the freshness rule. Verified with new tests (unchanged/changed, unknown cases, no staging calls, warning only with graph) plus the export and manifest suites. Co-Authored-By: Claude Opus 5.5 --- .../kumo-engine/kumo/engine/_snapshot_io.py | 10 +- packages/kumo-engine/kumo/engine/graph.py | 94 ++++++++++++++- .../kumo-engine/test/test_snapshot_export.py | 107 ++++++++++++++++++ 3 files changed, 203 insertions(+), 8 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/_snapshot_io.py b/packages/kumo-engine/kumo/engine/_snapshot_io.py index ef38e39dc..39f5a9709 100644 --- a/packages/kumo-engine/kumo/engine/_snapshot_io.py +++ b/packages/kumo-engine/kumo/engine/_snapshot_io.py @@ -29,9 +29,13 @@ _KEY = re.compile('[0-9a-f]{64}') _PROGRESS_SECONDS = 30 _UNSUPPORTED = { - 'SNOWFLAKE': 'Snowflake stages are not supported', - 'DATABRICKS': ('for a Unity Catalog Volume, use an S3Connector or ' - "GCSConnector on the volume's backing cloud path"), + 'SNOWFLAKE': ('Snowflake stages are not supported; in SPCS services, or ' + 'ML Jobs where a stage volume is mounted, use a ' + 'FileConnector on that volume'), + 'DATABRICKS': ('for a Unity Catalog Volume, use a FileConnector on ' + '/Volumes/// from Databricks ' + "compute, or an S3Connector or GCSConnector on the " + "volume's backing cloud path"), 'BIGQUERY': 'BigQuery cannot hold files', } diff --git a/packages/kumo-engine/kumo/engine/graph.py b/packages/kumo-engine/kumo/engine/graph.py index 84760a51a..88c650a3d 100644 --- a/packages/kumo-engine/kumo/engine/graph.py +++ b/packages/kumo-engine/kumo/engine/graph.py @@ -241,6 +241,39 @@ def _loaded(path: str) -> Optional[Any]: return marker +_HOUR_SECONDS = 3600 + + +def _comparable(stamp: Any) -> bool: + r"""Hour-aligned epoch seconds are what `hour_version()` stamps views and + unversioned sources with; `-1` marks an undated source.""" + return (isinstance(stamp, str) and stamp != '-1' + and not (stamp.isdigit() and int(stamp) % _HOUR_SECONDS == 0)) + + +def _source_change(table: Table, recorded: Any) -> Dict[str, Any]: + from kumoapi.source_table import CustomSQLSourceTable + + recorded = recorded if isinstance(recorded, dict) else {} + before, current, same_source = recorded.get('version'), None, False + try: + source = table.definition.source_table + if not isinstance(source, CustomSQLSourceTable): + connector = table.connector + same_source = recorded.get('source') == ( + f'{connector.source_type}:{table.table_name}') + stamp = connector.source_backend.get_data_version( + connector, table.table_name) + current = None if stamp is None else str(stamp) + except Exception: + current = None + if not (same_source and _comparable(before) and _comparable(current)): + 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 @@ -279,7 +312,15 @@ class GraphSnapshot: def export(self, connector: Connector, path: str) -> str: r"""Copy the index to `//` under an S3, GCS or file connector and return that location. Re-running - resumes.""" + resumes. + + Where there is no dedicated connector: on Databricks compute, use a + `FileConnector` rooted at `/Volumes///`; in + SPCS services, or ML Jobs where a stage volume is mounted, a + `FileConnector` on that stage volume (not yet tested with multi-GB + indexes); otherwise the S3 or + GCS connector on the bucket behind an external volume or stage. + """ from kumo.engine._snapshot_io import export return export(self, connector, path) @@ -288,7 +329,14 @@ def export(self, connector: Connector, path: str) -> str: def load_from_connector(cls, connector: Connector, *, path: str, graph: Optional['Graph'] = None) -> 'GraphSnapshot': r"""Copy the snapshot exported at `path` (the location `export` - returned) into the local snapshot store.""" + returned) into the local snapshot store. + + Given `graph`, reads each table's source version (one metadata call + per table, so this machine needs source credentials) and warns once + if any changed since the snapshot was built (see `source_changes`). A scoring job should + build a fresh snapshot or load one whose `source_changes` are all + `unchanged`, unless it deliberately scores an older snapshot. + """ from kumo_core.diskgraph.snapshot import DiskgraphSnapshotHandle from kumo.engine._snapshot_io import load_into, read @@ -312,9 +360,45 @@ def load_from_connector(cls, connector: Connector, *, path: str, if result is None: load_into(connector, path, root, manifest, Path(target)) result = _loaded(target) - return cls(id=manifest['snapshot_id'], - index_path=Path(result.index_path), - graph_config=manifest['graph_config']) + 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 the current data of " + f"source table(s) {', '.join(changed)}; see " + "`snapshot.source_changes(graph)`.", + UserWarning, + stacklevel=2, + ) + return snapshot + + def source_changes(self, graph: 'Graph') -> Dict[str, Dict[str, Any]]: + r"""Compare each of `graph`'s tables' current source version, read + from metadata only, with the version this snapshot recorded. Never + raises.""" + try: + manifest = self.manifest + recorded = ({} if not manifest or manifest.get('reconstructed') else + manifest['parents']['source_tables']) + except Exception: + recorded = {} + if not isinstance(recorded, dict): + recorded = {} + try: + tables = graph._named_tables + except Exception: + return {} + return { + alias: _source_change(table, recorded.get(alias)) + for alias, table in tables.items() + } @property def manifest(self) -> SnapshotManifest | None: diff --git a/packages/kumo-engine/test/test_snapshot_export.py b/packages/kumo-engine/test/test_snapshot_export.py index 39f0549b4..4c1f71cfa 100644 --- a/packages/kumo-engine/test/test_snapshot_export.py +++ b/packages/kumo-engine/test/test_snapshot_export.py @@ -3,6 +3,7 @@ # SPDX-License-Identifier: Apache-2.0 import json +import os import shutil import threading import warnings @@ -12,6 +13,7 @@ from types import SimpleNamespace from unittest.mock import Mock +import pandas as pd import pytest from diskgraph import format_versions from kumoapi.data_source import DataSourceType @@ -524,3 +526,108 @@ def expiring(root, staging, client, item, stop): scratch('gpu') GraphSnapshot.load_from_connector(exports, path=location) assert failed + + +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 + + +def _forbid_staging(monkeypatch): + from kumo_core.sources.object_store import ObjectStoreSourceReader + + from kumo.engine.connector.file import _FileSourceBackend + + 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) + + +def test_source_changes_reads_stamps_without_staging( + scratch, graph, dated, monkeypatch): + snapshot = graph.snapshot() + _forbid_staging(monkeypatch) + unchanged = {'status': 'unchanged', 'recorded': '1700000123', + 'current': '1700000123'} + assert snapshot.source_changes(graph) == {'users': unchanged, + 'events': unchanged} + _rewrite_events(dated) + assert snapshot.source_changes(graph) == { + 'users': unchanged, + 'events': {'status': 'changed', 'recorded': '1700000123', + 'current': '1700009123'}, + } + + +def test_source_changes_unknown_for_hour_stamps_and_read_errors( + scratch, graph, dated, monkeypatch): + from kumo.engine.connector.file import _FileSourceBackend + + _set_mtime(dated / 'users.parquet', 1_700_002_800) + snapshot = graph.snapshot() + stamp = _FileSourceBackend.get_data_version + + def failing(self, connector, table_name, project_id=None): + if table_name == 'events.parquet': + raise OSError('listing failed') + return stamp(self, connector, table_name) + + monkeypatch.setattr(_FileSourceBackend, 'get_data_version', failing) + assert snapshot.source_changes(graph) == { + 'users': {'status': 'unknown', 'recorded': '1700002800', + 'current': '1700002800'}, + 'events': {'status': 'unknown', 'recorded': '1700000123', + 'current': None}, + } + + +@pytest.mark.parametrize('edit', ['absent', 'reconstructed', 'missing', + 'malformed']) +def test_source_changes_unknown_without_a_recorded_version( + scratch, graph, dated, edit): + snapshot = graph.snapshot() + path = snapshot.index_path / 'snapshot-manifest.json' + manifest = json.loads(path.read_text()) + if edit == 'absent': + path.unlink() + elif edit == 'reconstructed': + path.write_text(json.dumps({**manifest, 'reconstructed': True})) + elif edit == 'missing': + del manifest['parents']['source_tables']['users'] + path.write_text(json.dumps(manifest)) + else: + path.write_text(json.dumps({**manifest, 'parents': []})) + changes = snapshot.source_changes(graph) + assert changes['users'] == {'status': 'unknown', 'recorded': None, + 'current': '1700000123'} + assert changes['events']['status'] == ( + 'unchanged' if edit == 'missing' else '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') + GraphSnapshot.load_from_connector(exports, path=location, graph=graph) + _rewrite_events(dated) + GraphSnapshot.load_from_connector(exports, path=location) + with pytest.warns(UserWarning) as caught: + GraphSnapshot.load_from_connector(exports, path=location, graph=graph) + [message] = [str(warning.message) for warning in caught] + assert 'source table(s) events;' in message From e31551749e93ffe6878866d18ea10f71efb87aaf Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Wed, 21 Oct 2026 07:44:14 +0000 Subject: [PATCH 52/80] Match source changes on the recorded connector and location The snapshot manifest now records each table as connector, full location and version, so source_changes matches the current table against the recorded connector type and location instead of the old ":" string. A table at a different location still reports unknown. Verified with test_snapshot_export.py and test_snapshot_manifest.py. Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/graph.py | 9 ++++++--- packages/kumo-engine/test/test_snapshot_export.py | 11 +++++++++++ 2 files changed, 17 insertions(+), 3 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/graph.py b/packages/kumo-engine/kumo/engine/graph.py index 88c650a3d..894e2e1e0 100644 --- a/packages/kumo-engine/kumo/engine/graph.py +++ b/packages/kumo-engine/kumo/engine/graph.py @@ -80,7 +80,7 @@ stamp, write_manifest, ) -from kumo.engine._provenance import source_table +from kumo.engine._provenance import source_location, source_table from kumo.engine._scratch import ( fingerprint, get_scratch_dir, @@ -260,8 +260,11 @@ def _source_change(table: Table, recorded: Any) -> Dict[str, Any]: source = table.definition.source_table if not isinstance(source, CustomSQLSourceTable): connector = table.connector - same_source = recorded.get('source') == ( - f'{connector.source_type}:{table.table_name}') + kind = connector.source_type + same_source = (recorded.get('connector'), + recorded.get('location')) == ( + str(getattr(kind, 'value', kind)), + source_location(connector, table.table_name)) stamp = connector.source_backend.get_data_version( connector, table.table_name) current = None if stamp is None else str(stamp) diff --git a/packages/kumo-engine/test/test_snapshot_export.py b/packages/kumo-engine/test/test_snapshot_export.py index 4c1f71cfa..a54ba9d0c 100644 --- a/packages/kumo-engine/test/test_snapshot_export.py +++ b/packages/kumo-engine/test/test_snapshot_export.py @@ -618,6 +618,17 @@ def test_source_changes_unknown_without_a_recorded_version( 'unchanged' if edit == 'missing' else 'unknown') +def test_source_changes_unknown_for_a_moved_table(scratch, graph, dated): + snapshot = graph.snapshot() + path = snapshot.index_path / 'snapshot-manifest.json' + manifest = json.loads(path.read_text()) + manifest['parents']['source_tables']['users']['location'] += '.old' + path.write_text(json.dumps(manifest)) + changes = snapshot.source_changes(graph) + assert changes['users']['status'] == 'unknown' + assert changes['events']['status'] == 'unchanged' + + 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) From 2c75c5059b9d6f26f32ad83fc89e0e1b72305bbc Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Wed, 21 Oct 2026 08:12:06 +0000 Subject: [PATCH 53/80] Keep an existing export when a rebuilt snapshot is re-exported A snapshot rebuilt on another machine has the same id but a different DiskGraph _MANIFEST.json (run-specific paths and timings), so exporting it to a location already holding that id raised "differs from this one". Exporting a loaded snapshot also copied its `loaded` block into the export record, and corrupt-export errors did not say how to recover. A committed export with the same snapshot id is now final: export logs at info and returns its location; the fingerprint only decides whether a partial upload resumes or a committed one is repaired. Only a different id is refused. The export record drops `loaded`; corrupt and checksum errors name the location to delete and re-export. README and source_changes document that a table now read from another connector or location is `unknown`. Verified: new rebuild-and-re-export test plus the existing export and manifest suites; the reviewer's rebuild script now re-exports and loads. Co-Authored-By: Claude Opus 5.5 --- .../kumo-engine/kumo/engine/_snapshot_io.py | 62 ++++++++++++------- packages/kumo-engine/kumo/engine/graph.py | 6 +- .../kumo-engine/test/test_snapshot_export.py | 29 +++++++-- 3 files changed, 68 insertions(+), 29 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/_snapshot_io.py b/packages/kumo-engine/kumo/engine/_snapshot_io.py index 39f5a9709..941e950b3 100644 --- a/packages/kumo-engine/kumo/engine/_snapshot_io.py +++ b/packages/kumo-engine/kumo/engine/_snapshot_io.py @@ -28,6 +28,7 @@ _ATTEMPTS = 3 _KEY = re.compile('[0-9a-f]{64}') _PROGRESS_SECONDS = 30 +_REEXPORT = 'Delete {} and run GraphSnapshot.export again.' _UNSUPPORTED = { 'SNOWFLAKE': ('Snowflake stages are not supported; in SPCS services, or ' 'ML Jobs where a stage volume is mounted, use a ' @@ -180,18 +181,16 @@ def _write_json(filesystem: FileSystem, path: str, value: dict) -> None: def _committed(filesystem: FileSystem, root: str, location: str, - snapshot_id: str, fingerprint: str) -> bool: + snapshot_id: str) -> Optional[dict]: existing = _read_json(filesystem, f'{root}/{_MANIFEST}') if existing is None: - return False - if not isinstance(existing, dict): - existing = {} - if (existing.get('snapshot_id') == snapshot_id - and existing.get('index_manifest') == fingerprint): - return True - raise ValueError( - f'{location} already holds an export of snapshot ' - f"{existing.get('snapshot_id')!r} that differs from this one.") + return None + found = existing.get('snapshot_id') if isinstance(existing, dict) else None + if found != snapshot_id: + raise ValueError( + f'{location} holds an export of snapshot {found!r}, not ' + f'{snapshot_id!r}; remove it and re-export.') + return existing def export(snapshot: Any, connector: Any, path: str) -> str: @@ -210,9 +209,16 @@ def export(snapshot: Any, connector: Any, path: str) -> str: identity = {'snapshot_id': snapshot.id, 'index_manifest': fingerprint} filesystem, root = _location(connector, location) - def committed() -> bool: - return _committed(filesystem, root, location, snapshot.id, - fingerprint) + def committed() -> Optional[dict]: + return _committed(filesystem, root, location, snapshot.id) + + def finished() -> bool: + existing = committed() + if existing and existing.get('index_manifest') != fingerprint: + raise RuntimeError( + f'Another build of snapshot {snapshot.id} finished exporting ' + f'to {location} while this one ran; remove it and re-export.') + return existing is not None def present() -> Dict[str, int]: return { @@ -231,7 +237,13 @@ def upload(client: FileSystem, file: Tuple[str, int], with _open_output(client, f'{root}/index/{name}') as writer: return _stream(source / name, size, stop, writer.write) - if committed(): + existing = committed() + if existing is not None and existing.get('index_manifest') != fingerprint: + logger.info( + '%s already holds snapshot %s from another build of it; keeping ' + 'that export.', location, snapshot.id) + return location + if existing is not None: # Re-running repairs files a racing or interrupted export rewrote. uploaded = present() stale = [file for file in files if uploaded.get(file[0]) != file[1]] @@ -257,13 +269,14 @@ def upload(client: FileSystem, file: Tuple[str, int], filesystem, root = _location(connector, location) # An identical export racing this one uploads the same bytes, so # whichever finishes first commits and the other finds it committed. - if committed(): + if finished(): return location if (_read_json(filesystem, f'{root}/{_EXPORTING}') != identity - and not committed()): + and not finished()): raise RuntimeError( f'Another export to {location} started while this one ran; ' 'remove it and re-export.') + recorded = snapshot.manifest _write_json( filesystem, f'{root}/{_MANIFEST}', { 'format': _FORMAT, @@ -272,7 +285,10 @@ def upload(client: FileSystem, file: Tuple[str, int], 'data_version': data_version, 'index_manifest': fingerprint, 'graph_config': snapshot.graph_config, - 'snapshot': snapshot.manifest, + 'snapshot': None if recorded is None else { + key: value + for key, value in recorded.items() if key != 'loaded' + }, 'environment': environment(), 'range_bytes': _RANGE_BYTES, 'files': [{ @@ -297,7 +313,7 @@ def _format_note(manifest: Any) -> str: def _malformed(path: str, detail: str, manifest: Any) -> ValueError: return ValueError(f'The snapshot manifest at {path} {detail}; the export ' - f'is corrupt{_format_note(manifest)}.') + f'is corrupt{_format_note(manifest)}. {_REEXPORT.format(path)}') def read(connector: Any, path: str) -> Tuple[str, dict]: @@ -359,7 +375,7 @@ def _download(root: str, staging: Path, client: FileSystem, if len(chunk) != want: raise ValueError( f'Snapshot file {name} in the export is shorter than ' - 'its manifest records; the export is corrupt.') + 'its manifest records; the export is corrupt') crc = zlib.crc32(chunk, crc) os.pwrite(fd, chunk, start) # Flush while other ranges download, not all at the end. @@ -369,7 +385,7 @@ def _download(root: str, staging: Path, client: FileSystem, if crc != expected: raise ValueError( f'Snapshot file {name} failed its checksum at offset {offset}; ' - 'the export is corrupt.') + 'the export is corrupt') def load_into(connector: Any, path: str, root: str, manifest: dict, @@ -414,9 +430,9 @@ def load_into(connector: Any, path: str, root: str, manifest: dict, _parallel(connector, path, ranges, [item[2] for item in ranges], partial(_download, root, staging), 'load') except ValueError as error: - if _format_note(manifest): - raise ValueError(f'{error}{_format_note(manifest)}') from error - raise + raise ValueError( + f'{error}{_format_note(manifest)}. {_REEXPORT.format(path)}' + ) from error write_diskgraph_success_marker(str(staging), DiskgraphIngestResult.for_index(target)) verify_snapshot_integrity(str(staging)) diff --git a/packages/kumo-engine/kumo/engine/graph.py b/packages/kumo-engine/kumo/engine/graph.py index 894e2e1e0..98fe89b48 100644 --- a/packages/kumo-engine/kumo/engine/graph.py +++ b/packages/kumo-engine/kumo/engine/graph.py @@ -384,8 +384,10 @@ def load_from_connector(cls, connector: Connector, *, path: str, def source_changes(self, graph: 'Graph') -> Dict[str, Dict[str, Any]]: r"""Compare each of `graph`'s tables' current source version, read - from metadata only, with the version this snapshot recorded. Never - raises.""" + from metadata only, with the version this snapshot recorded. A table + is `unknown` when the versions cannot be compared, including when it + now reads from a different connector or location than the one + recorded. Never raises.""" try: manifest = self.manifest recorded = ({} if not manifest or manifest.get('reconstructed') else diff --git a/packages/kumo-engine/test/test_snapshot_export.py b/packages/kumo-engine/test/test_snapshot_export.py index a54ba9d0c..abc251897 100644 --- a/packages/kumo-engine/test/test_snapshot_export.py +++ b/packages/kumo-engine/test/test_snapshot_export.py @@ -114,6 +114,9 @@ def test_export_load_round_trip_relocates_and_is_reused( 'environment': manifest['environment'], } assert recorded == manifest['snapshot'] + reexported = tmp_path / 'exports' / loaded.export(exports, 'two') + assert json.loads((reexported / 'kumo_snapshot.json').read_text())[ + 'snapshot'] == manifest['snapshot'] _no_ingest(monkeypatch) monkeypatch.setattr(snapshot_io, 'load_into', @@ -269,7 +272,7 @@ def test_failed_load_leaves_nothing_and_restarts_cleanly( stats.write_bytes(original[:-1] + bytes([original[-1] ^ 1])) scratch('gpu') - with pytest.raises(ValueError, match='checksum'): + with pytest.raises(ValueError, match='checksum.*corrupt. Delete one/'): GraphSnapshot.load_from_connector(exports, path=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'] @@ -345,7 +348,7 @@ def test_load_rejects_a_malformed_manifest_with_value_error( edit(manifest) path.write_text(json.dumps(manifest)) scratch('gpu') - with pytest.raises(ValueError, match='export is corrupt'): + with pytest.raises(ValueError, match='corrupt. Delete one/'): GraphSnapshot.load_from_connector(exports, path=location) path.write_text('[]') with pytest.raises(ValueError, match='not a JSON object'): @@ -383,12 +386,30 @@ def test_export_refuses_a_path_holding_other_data(tmp_path, scratch, graph, snapshot.export(exports, 'occupied') _, exported = _export(snapshot, exports, tmp_path) - _edit_manifest(exported, index_manifest='other') - with pytest.raises(ValueError, match='differs from this one'): + _edit_manifest(exported, snapshot_id='other') + with pytest.raises(ValueError, match="snapshot 'other', not"): snapshot.export(exports, 'one') assert (exported / 'index/col_stats.json').is_file() +def test_rebuilt_snapshot_reexports_to_the_existing_export( + tmp_path, scratch, graph, exports): + snapshot = graph.snapshot() + location, exported = _export(snapshot, exports, tmp_path) + committed = (exported / 'kumo_snapshot.json').read_bytes() + scratch('rebuilt/on/another/machine') + rebuilt = graph.snapshot() + assert rebuilt.id == snapshot.id + assert (rebuilt.index_path / '_MANIFEST.json').read_bytes() != ( + snapshot.index_path / '_MANIFEST.json').read_bytes() + + assert rebuilt.export(exports, 'one') == location + assert (exported / 'kumo_snapshot.json').read_bytes() == committed + scratch('gpu') + loaded = GraphSnapshot.load_from_connector(exports, path=location) + assert _index_files(loaded.index_path) == _exported_files(snapshot) + + @pytest.mark.parametrize('kind, reason', [ ('SNOWFLAKE', 'Snowflake stages'), ('DATABRICKS', 'Unity Catalog Volume'), From 2b40f15b18b9665e66fff8767a5c046f82021fcc Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Wed, 21 Oct 2026 08:12:06 +0000 Subject: [PATCH 54/80] Pass allow_delayed_open to S3 only when pyarrow supports it The sink always passed allow_delayed_open to S3FileSystem, which pyarrow releases inside the pinned range do not accept, so every S3 export and load raised TypeError there. The default no-connection filesystem also lacked the retry strategy the other branches set. The sink probes once whether S3FileSystem accepts the option and drops it otherwise; the no-connection branch always builds S3FileSystem with the retry strategy. Verified: sink test's fake now rejects unknown keywords and runs with and without the option; kumo-s3 unit tests pass. Co-Authored-By: Claude Opus 5.5 --- packages/kumo-s3/src/kumo_s3/backend.py | 7 ++----- packages/kumo-s3/src/kumo_s3/sink.py | 20 ++++++++++++++++---- packages/kumo-s3/tests/test_sink.py | 17 ++++++++++++++--- 3 files changed, 32 insertions(+), 12 deletions(-) diff --git a/packages/kumo-s3/src/kumo_s3/backend.py b/packages/kumo-s3/src/kumo_s3/backend.py index 2b47fdfde..d34effa5c 100644 --- a/packages/kumo-s3/src/kumo_s3/backend.py +++ b/packages/kumo-s3/src/kumo_s3/backend.py @@ -155,11 +155,8 @@ def _filesystem(self, connector: Any, bucket: str, or session.region_name, options=arrow_options) from pyarrow.fs import FileSystem filesystem, _ = FileSystem.from_uri(f's3://{bucket}') - if options: - return S3FileSystem(region=filesystem.region, - retry_strategy=retry_strategy(), - **options) - return filesystem + return S3FileSystem(region=filesystem.region, + retry_strategy=retry_strategy(), **(options or {})) class S3SourceBackend(S3Sessions, ObjectStoreSourceBackend): diff --git a/packages/kumo-s3/src/kumo_s3/sink.py b/packages/kumo-s3/src/kumo_s3/sink.py index eded2e8a7..557551b64 100644 --- a/packages/kumo-s3/src/kumo_s3/sink.py +++ b/packages/kumo-s3/src/kumo_s3/sink.py @@ -2,6 +2,9 @@ # All rights reserved. # SPDX-License-Identifier: Apache-2.0 +from functools import lru_cache +from typing import Dict + from kumo_core.object_store_sink import ObjectStoreSinkBackend from kumoapi.data_source import DataSourceType @@ -16,12 +19,21 @@ class S3SinkBackend(S3Sessions, ObjectStoreSinkBackend): def filesystem(self, connector, path, *, parallel_uploads=False): bucket, resolved, region = self._resolved(connector, path) - options = ({ - 'allow_delayed_open': True, - 'background_writes': False - } if parallel_uploads else None) + options = _parallel_upload_options() if parallel_uploads else None return (self._filesystem(connector, bucket, region, options), resolved.rstrip('/')) def __repr__(self) -> str: return f'{type(self).__name__}(identity={self._identity!r})' + + +@lru_cache(maxsize=None) +def _parallel_upload_options() -> Dict[str, bool]: + from pyarrow.fs import S3FileSystem + + try: + S3FileSystem(anonymous=True, region='us-east-1', + allow_delayed_open=True) + except TypeError: + return {'background_writes': False} + return {'allow_delayed_open': True, 'background_writes': False} diff --git a/packages/kumo-s3/tests/test_sink.py b/packages/kumo-s3/tests/test_sink.py index a0d45db50..c143321df 100644 --- a/packages/kumo-s3/tests/test_sink.py +++ b/packages/kumo-s3/tests/test_sink.py @@ -9,7 +9,7 @@ import kumo_s3 import pytest from fakes import FakeSession -from kumo_s3 import backend +from kumo_s3 import backend, sink from kumo_s3.sink import S3SinkBackend from kumoapi.data_source import ( S3AnyURIConnectorTarget, @@ -17,13 +17,21 @@ ) +@pytest.mark.parametrize('delayed_open', [True, False]) @pytest.mark.parametrize('credentials', ['default', None]) def test_parallel_uploads_reach_every_filesystem_construction( - monkeypatch, credentials): + monkeypatch, credentials, delayed_open): + known = { + 'access_key', 'secret_key', 'session_token', 'region', 'anonymous', + 'retry_strategy', 'background_writes', 'endpoint_override', 'scheme', + 'force_virtual_addressing' + } | ({'allow_delayed_open'} if delayed_open else set()) built = [] class Record: def __init__(self, **kwargs): + if set(kwargs) - known: + raise TypeError(f'unexpected {set(kwargs) - known}') built.append(kwargs) monkeypatch.setattr('pyarrow.fs.S3FileSystem', Record) @@ -35,15 +43,18 @@ def __init__(self, **kwargs): 'pyarrow.fs.FileSystem', SimpleNamespace(from_uri=lambda uri: (SimpleNamespace(region='us-west-2'), ''))) + sink._parallel_upload_options.cache_clear() connector = SimpleNamespace( id='connector', identity=None, identity_id=None, config=S3ConnectorResourceConfig(target=S3AnyURIConnectorTarget())) filesystem, path = S3SinkBackend().filesystem( connector, 's3://bucket/prefix/', parallel_uploads=True) + sink._parallel_upload_options.cache_clear() assert isinstance(filesystem, Record) assert path == 'bucket/prefix' - assert built[-1]['allow_delayed_open'] is True + assert built[-1].get('allow_delayed_open') is (True if delayed_open else + None) assert built[-1]['background_writes'] is False assert 'retry_strategy' in built[-1] From fa780f038074e9f5d043d686df3ddea81acbea64 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Wed, 21 Oct 2026 08:27:16 +0000 Subject: [PATCH 55/80] Trim comments and implementation-detail tests Export, load and source_changes docstrings and the README repeated the same guidance several times, and the manifest-validation tests ran duplicate cases of one check. Cut the docstrings to a line plus Args, tighten the README section, and drop redundant hostile/malformed manifest cases and a call-count assertion. test_snapshot_export(.py, _live.py): 38 passed, 2 skipped (live); kumo-s3 test_sink.py: 4 passed. Co-Authored-By: Claude Opus 5.5 --- .../src/kumo_core/object_store_sink.py | 4 +- packages/kumo-engine/kumo/engine/graph.py | 42 +++++++------------ .../kumo-engine/test/test_snapshot_export.py | 24 ++++------- 3 files changed, 26 insertions(+), 44 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 4843916fd..1984034d8 100644 --- a/packages/kumo-core/src/kumo_core/object_store_sink.py +++ b/packages/kumo-core/src/kumo_core/object_store_sink.py @@ -103,9 +103,7 @@ def holds_files(self, connector: SourceConnector, path: str) -> bool: def filesystem(self, connector: SourceConnector, path: str, *, parallel_uploads: bool = False) -> Tuple[FileSystem, str]: """The authenticated filesystem and the path ``path`` resolves to. - ``parallel_uploads`` tunes S3 for many concurrent uploads (one - buffered part each, small files as one PUT); other stores ignore it. - """ + ``parallel_uploads`` tunes S3 for many concurrent uploads.""" del parallel_uploads filesystem, resolved = self._destination(connector, path) return filesystem, resolved.rstrip('/') diff --git a/packages/kumo-engine/kumo/engine/graph.py b/packages/kumo-engine/kumo/engine/graph.py index 98fe89b48..772288340 100644 --- a/packages/kumo-engine/kumo/engine/graph.py +++ b/packages/kumo-engine/kumo/engine/graph.py @@ -224,7 +224,6 @@ def _lock(graph_key: str) -> Any: def _loaded(path: str) -> Optional[Any]: - r"""The index a previous load left intact at `path`, if any.""" from diskgraph import DiskgraphCorruptionError, verify_snapshot_integrity from kumo_core.diskgraph.snapshot import check_diskgraph_output @@ -245,8 +244,8 @@ def _loaded(path: str) -> Optional[Any]: def _comparable(stamp: Any) -> bool: - r"""Hour-aligned epoch seconds are what `hour_version()` stamps views and - unversioned sources with; `-1` marks an undated source.""" + # hour_version() stamps views and unversioned sources on the hour; -1 is + # an undated source. return (isinstance(stamp, str) and stamp != '-1' and not (stamp.isdigit() and int(stamp) % _HOUR_SECONDS == 0)) @@ -313,16 +312,12 @@ class GraphSnapshot: record_sources: bool | None = field(default=None, compare=False) def export(self, connector: Connector, path: str) -> str: - r"""Copy the index to `//` under an - S3, GCS or file connector and return that location. Re-running - resumes. - - Where there is no dedicated connector: on Databricks compute, use a - `FileConnector` rooted at `/Volumes///`; in - SPCS services, or ML Jobs where a stage volume is mounted, a - `FileConnector` on that stage volume (not yet tested with multi-GB - indexes); otherwise the S3 or - GCS connector on the bucket behind an external volume or stage. + r"""Copy the index to `//` and + return that location; re-running resumes. + + Args: + connector: an S3, GCS or file connector (on Databricks or + Snowflake, a `FileConnector` on a mounted volume). """ from kumo.engine._snapshot_io import export @@ -331,14 +326,12 @@ def export(self, connector: Connector, path: str) -> str: @classmethod def load_from_connector(cls, connector: Connector, *, path: str, graph: Optional['Graph'] = None) -> 'GraphSnapshot': - r"""Copy the snapshot exported at `path` (the location `export` - returned) into the local snapshot store. - - Given `graph`, reads each table's source version (one metadata call - per table, so this machine needs source credentials) and warns once - if any changed since the snapshot was built (see `source_changes`). A scoring job should - build a fresh snapshot or load one whose `source_changes` are all - `unchanged`, unless it deliberately scores an older snapshot. + 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 @@ -383,11 +376,8 @@ def load_from_connector(cls, connector: Connector, *, path: str, return snapshot def source_changes(self, graph: 'Graph') -> Dict[str, Dict[str, Any]]: - r"""Compare each of `graph`'s tables' current source version, read - from metadata only, with the version this snapshot recorded. A table - is `unknown` when the versions cannot be compared, including when it - now reads from a different connector or location than the one - recorded. Never raises.""" + r"""Each of `graph`'s tables' source status since this snapshot: + `unchanged`, `changed` or `unknown`. Never raises.""" try: manifest = self.manifest recorded = ({} if not manifest or manifest.get('reconstructed') else diff --git a/packages/kumo-engine/test/test_snapshot_export.py b/packages/kumo-engine/test/test_snapshot_export.py index abc251897..62822cd26 100644 --- a/packages/kumo-engine/test/test_snapshot_export.py +++ b/packages/kumo-engine/test/test_snapshot_export.py @@ -212,19 +212,19 @@ def test_concurrent_identical_exports_both_succeed(tmp_path, scratch, graph, exports, monkeypatch): snapshot = graph.snapshot() parallel = snapshot_io._parallel - calls = [] + raced = [] def racing(*args): - calls.append(args) - if len(calls) == 1: - assert snapshot.export(exports, 'one') == location + if not raced: + raced.append(None) + raced[0] = snapshot.export(exports, 'one') return parallel(*args) graph_id, data_version = snapshot.id.split('@') location = f'one/{graph_id}/{data_version}' monkeypatch.setattr(snapshot_io, '_parallel', racing) assert snapshot.export(exports, 'one') == location - assert len(calls) == 2 + assert raced == [location] scratch('gpu') loaded = GraphSnapshot.load_from_connector(exports, path=location) assert _index_files(loaded.index_path) == _exported_files(snapshot) @@ -308,8 +308,6 @@ def test_load_needs_the_full_size_while_an_old_copy_remains( @pytest.mark.parametrize('field, value', [ ('data_version', '../../victim'), - ('data_version', '/tmp/victim'), - ('graph_key', '..'), ('files', [{ 'path': '../../victim/x', 'size': 1, @@ -332,10 +330,7 @@ def test_load_rejects_a_hostile_manifest(tmp_path, scratch, graph, exports, @pytest.mark.parametrize('edit', [ lambda manifest: manifest.pop('files'), - lambda manifest: manifest.pop('range_bytes'), - lambda manifest: manifest.pop('graph_config'), lambda manifest: manifest.update(snapshot=[1]), - lambda manifest: manifest['files'].append('index_metadata.json'), lambda manifest: manifest['files'][0].pop('crc32'), lambda manifest: manifest['files'].append(manifest['files'][0]), ]) @@ -495,13 +490,12 @@ def test_object_store_round_trip(tmp_path, scratch, graph, monkeypatch, storage = tmp_path / 'storage' storage.mkdir() bucket = SubTreeFileSystem(str(storage), LocalFileSystem()) - options = [] + seen = [] def s3(bucket_name, *, options=None, **kwargs): - options_seen.append(options) + seen.append(options) return bucket - options_seen = options monkeypatch.setattr('kumo_s3.filesystem', s3) class Gcs(SubTreeFileSystem): @@ -523,9 +517,9 @@ def __init__(self): assert loaded.id == snapshot.id assert _index_files(loaded.index_path) == _exported_files(snapshot) if scheme == 's3': - assert options and all( + assert seen and all( option['allow_delayed_open'] and not option['background_writes'] - for option in options) + for option in seen) def test_transfer_renews_an_expired_client(tmp_path, scratch, graph, exports, From 6cf5f8295c09aa7e8391036bc951801aa5b29b11 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Wed, 21 Oct 2026 09:39:19 +0000 Subject: [PATCH 56/80] Trim snapshot export PR to essentials The export/load PR carried per-warehouse error essays, a format-note and corrupt-manifest message layer, defensive branches in source_changes, a long README section, and tests of internal race and retry sequences. Collapse unsupported connectors into one error naming the connectors to use, report corrupt exports through one message, simplify the unreadable-index error to the two format sets, trim the README to usage and limits, and keep tests of user-visible export, load, back-compat and source_changes behaviour. No documented field or API changed. Verified: test_snapshot_export, test_snapshot_export_live (skipped without live roots), test_snapshot_manifest, test_snapshot_cleanup, test_snapshot_staging_cleanup and kumo-s3 tests pass; flake8 and isort clean. Co-Authored-By: Claude Opus 5.5 --- .../src/kumo_core/object_store_sink.py | 3 +- .../kumo-engine/kumo/engine/_snapshot_io.py | 168 ++---- packages/kumo-engine/kumo/engine/graph.py | 36 +- .../kumo-engine/test/test_snapshot_export.py | 516 +++++------------- .../test/test_snapshot_export_live.py | 18 +- packages/kumo-s3/src/kumo_s3/backend.py | 3 +- 6 files changed, 201 insertions(+), 543 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 1984034d8..e4a42be99 100644 --- a/packages/kumo-core/src/kumo_core/object_store_sink.py +++ b/packages/kumo-core/src/kumo_core/object_store_sink.py @@ -102,8 +102,7 @@ def holds_files(self, connector: SourceConnector, path: str) -> bool: def filesystem(self, connector: SourceConnector, path: str, *, parallel_uploads: bool = False) -> Tuple[FileSystem, str]: - """The authenticated filesystem and the path ``path`` resolves to. - ``parallel_uploads`` tunes S3 for many concurrent uploads.""" + """The authenticated filesystem and resolved ``path``.""" del parallel_uploads filesystem, resolved = self._destination(connector, path) return filesystem, resolved.rstrip('/') diff --git a/packages/kumo-engine/kumo/engine/_snapshot_io.py b/packages/kumo-engine/kumo/engine/_snapshot_io.py index 941e950b3..c1e3f5fde 100644 --- a/packages/kumo-engine/kumo/engine/_snapshot_io.py +++ b/packages/kumo-engine/kumo/engine/_snapshot_io.py @@ -28,17 +28,6 @@ _ATTEMPTS = 3 _KEY = re.compile('[0-9a-f]{64}') _PROGRESS_SECONDS = 30 -_REEXPORT = 'Delete {} and run GraphSnapshot.export again.' -_UNSUPPORTED = { - 'SNOWFLAKE': ('Snowflake stages are not supported; in SPCS services, or ' - 'ML Jobs where a stage volume is mounted, use a ' - 'FileConnector on that volume'), - 'DATABRICKS': ('for a Unity Catalog Volume, use a FileConnector on ' - '/Volumes/// from Databricks ' - "compute, or an S3Connector or GCSConnector on the " - "volume's backing cloud path"), - 'BIGQUERY': 'BigQuery cannot hold files', -} logger = logging.getLogger(__name__) @@ -46,13 +35,14 @@ def _location(connector: Any, path: str) -> Tuple[FileSystem, str]: from kumo_core.object_store_sink import ObjectStoreSinkBackend - reason = _UNSUPPORTED.get(getattr(connector.source_type, 'name', None)) - sink = None if reason else connector.sink_backend + try: + sink = connector.sink_backend + except Exception: + sink = None if not isinstance(sink, ObjectStoreSinkBackend): raise NotImplementedError( - f'{type(connector).__name__} cannot hold a snapshot export' - f"{f' ({reason})' if reason else ''}; use an S3Connector, " - 'GCSConnector or FileConnector.') + f'{type(connector).__name__} cannot hold a snapshot export; use ' + 'an S3Connector, GCSConnector or a FileConnector on a volume.') return sink.filesystem(connector, path, parallel_uploads=True) @@ -60,9 +50,7 @@ def _check_key(graph_key: Any, data_version: Any, snapshot_id: Any) -> None: if not (isinstance(graph_key, str) and isinstance(data_version, str) and _KEY.fullmatch(graph_key) and _KEY.fullmatch(data_version) and snapshot_id == f'{graph_key}@{data_version[:16]}'): - raise ValueError( - f'Snapshot {snapshot_id!r} does not name a graph key and data ' - 'version this package produced.') + raise ValueError(f'{snapshot_id!r} is not a snapshot id.') class _Client: @@ -85,8 +73,7 @@ def _parallel(connector: Any, path: str, items: Sequence[Any], stop = threading.Event() total = sum(sizes) with copies(items) as workers, ThreadPoolExecutor(workers) as pool: - # Created after `copies` raises Arrow's IO pool, which sizes the - # client's connection pool. + # 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), total, workers) @@ -110,7 +97,6 @@ def _parallel(connector: Any, path: str, items: Sequence[Any], moved / max(elapsed, 1e-9) / 1e9) return [future.result() for future in futures] except BaseException: - # Leaving the pool joins the workers, which notice `stop`. stop.set() for future in futures: future.cancel() @@ -180,19 +166,6 @@ def _write_json(filesystem: FileSystem, path: str, value: dict) -> None: writer.write(json.dumps(value).encode()) -def _committed(filesystem: FileSystem, root: str, location: str, - snapshot_id: str) -> Optional[dict]: - existing = _read_json(filesystem, f'{root}/{_MANIFEST}') - if existing is None: - return None - found = existing.get('snapshot_id') if isinstance(existing, dict) else None - if found != snapshot_id: - raise ValueError( - f'{location} holds an export of snapshot {found!r}, not ' - f'{snapshot_id!r}; remove it and re-export.') - return existing - - def export(snapshot: Any, connector: Any, path: str) -> str: from kumo.engine._provenance import environment @@ -210,14 +183,19 @@ def export(snapshot: Any, connector: Any, path: str) -> str: filesystem, root = _location(connector, location) def committed() -> Optional[dict]: - return _committed(filesystem, root, location, snapshot.id) + existing = _read_json(filesystem, f'{root}/{_MANIFEST}') + found = (existing.get('snapshot_id') + if isinstance(existing, dict) else None) + if existing is not None and found != snapshot.id: + raise ValueError(f'{location} holds snapshot {found!r}, not ' + f'{snapshot.id!r}; remove it and re-export.') + return existing def finished() -> bool: existing = committed() if existing and existing.get('index_manifest') != fingerprint: - raise RuntimeError( - f'Another build of snapshot {snapshot.id} finished exporting ' - f'to {location} while this one ran; remove it and re-export.') + raise RuntimeError(f'Another build of {snapshot.id} was exported ' + f'to {location} meanwhile; re-export.') return existing is not None def present() -> Dict[str, int]: @@ -239,12 +217,10 @@ def upload(client: FileSystem, file: Tuple[str, int], existing = committed() if existing is not None and existing.get('index_manifest') != fingerprint: - logger.info( - '%s already holds snapshot %s from another build of it; keeping ' - 'that export.', location, snapshot.id) + logger.info('Keeping the existing export of %s at %s', snapshot.id, + location) return location if existing is not None: - # Re-running repairs files a racing or interrupted export rewrote. uploaded = present() stale = [file for file in files if uploaded.get(file[0]) != file[1]] if stale: @@ -258,24 +234,20 @@ def upload(client: FileSystem, file: Tuple[str, int], FileSelector(root, allow_not_found=True)): raise ValueError(f'{location} is not empty; remove it and re-export.') if marker is not None and marker.get('snapshot_id') != snapshot.id: - raise ValueError( - f"{location} holds a partial export of snapshot " - f"{marker.get('snapshot_id')!r}; remove it and re-export.") + raise ValueError(f'{location} holds a partial export of another ' + 'snapshot; remove it and re-export.') resume = (marker or {}).get('index_manifest') == fingerprint _write_json(filesystem, f'{root}/{_EXPORTING}', identity) uploaded = present() if resume else {} crcs = _parallel(connector, location, files, [size for _, size in files], upload, 'export') filesystem, root = _location(connector, location) - # An identical export racing this one uploads the same bytes, so - # whichever finishes first commits and the other finds it committed. if finished(): return location if (_read_json(filesystem, f'{root}/{_EXPORTING}') != identity and not finished()): - raise RuntimeError( - f'Another export to {location} started while this one ran; ' - 'remove it and re-export.') + raise RuntimeError(f'Another export to {location} started ' + 'meanwhile; re-export.') recorded = snapshot.manifest _write_json( filesystem, f'{root}/{_MANIFEST}', { @@ -285,17 +257,12 @@ def upload(client: FileSystem, file: Tuple[str, int], 'data_version': data_version, 'index_manifest': fingerprint, 'graph_config': snapshot.graph_config, - 'snapshot': None if recorded is None else { - key: value - for key, value in recorded.items() if key != 'loaded' - }, + 'snapshot': None if recorded is None else + {k: v for k, v in recorded.items() if k != 'loaded'}, 'environment': environment(), 'range_bytes': _RANGE_BYTES, - 'files': [{ - 'path': name, - 'size': size, - 'crc32': file_crcs - } for (name, size), file_crcs in zip(files, crcs)], + 'files': [{'path': name, 'size': size, 'crc32': file_crcs} + for (name, size), file_crcs in zip(files, crcs)], }) try: filesystem.delete_file(f'{root}/{_EXPORTING}') @@ -304,27 +271,19 @@ def upload(client: FileSystem, file: Tuple[str, int], return location -def _format_note(manifest: Any) -> str: - if not isinstance(manifest, dict) or manifest.get('format') == _FORMAT: - return '' - return (f" (export layout format {manifest.get('format')!r}; this " - f'engine reads {_FORMAT})') - - -def _malformed(path: str, detail: str, manifest: Any) -> ValueError: - return ValueError(f'The snapshot manifest at {path} {detail}; the export ' - f'is corrupt{_format_note(manifest)}. {_REEXPORT.format(path)}') +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) manifest = _read_json(filesystem, f'{root}/{_MANIFEST}') if manifest is None: - raise FileNotFoundError( - f'No complete snapshot export at {path}: {_MANIFEST} is missing. ' - 'Re-run GraphSnapshot.export to finish it.') + raise FileNotFoundError(f'No complete snapshot export at {path} ' + f'({_MANIFEST} is missing); re-run export.') if not isinstance(manifest, dict): - raise _malformed(path, 'is not a JSON object', manifest) + raise _corrupt(path, 'manifest is not a JSON object') if manifest.get('format') != _FORMAT: warnings.warn( f"Snapshot export at {path} has layout format " @@ -333,15 +292,10 @@ def read(connector: Any, path: str) -> Tuple[str, dict]: _check_key(manifest.get('graph_key'), manifest.get('data_version'), manifest.get('snapshot_id')) span, files = manifest.get('range_bytes'), manifest.get('files') - if type(span) is not int or span <= 0: - raise _malformed(path, 'has an invalid range_bytes', manifest) - if not isinstance(files, list): - raise _malformed(path, 'has no file list', manifest) - for key in ('snapshot', 'environment'): - if not isinstance(manifest.get(key, {}), (dict, type(None))): - raise _malformed(path, f'has a non-object {key}', manifest) - if not isinstance(manifest.get('graph_config'), dict): - raise _malformed(path, 'has no graph_config', manifest) + if (type(span) is not int or span <= 0 or not isinstance(files, list) + or not isinstance(manifest.get('graph_config'), dict) + or not isinstance(manifest.get('snapshot'), (dict, type(None)))): + raise _corrupt(path, 'malformed manifest') seen = set() for entry in files: name, size, crcs = ((entry.get('path'), entry.get('size'), @@ -353,9 +307,7 @@ def read(connector: Any, path: str) -> Tuple[str, dict]: or not isinstance(crcs, list) or any(type(crc) is not int for crc in crcs) or len(crcs) != max(-(-size // span), 1)): - raise _malformed( - path, f'lists {name!r} with an invalid or repeated path, size ' - 'or checksums', manifest) + raise _corrupt(path, f'bad file entry {name!r}') seen.add(name) return root, manifest @@ -373,19 +325,14 @@ def _download(root: str, staging: Path, client: FileSystem, want = min(_CHUNK, offset + length - start) chunk = reader.read_at(want, start) if len(chunk) != want: - raise ValueError( - f'Snapshot file {name} in the export is shorter than ' - 'its manifest records; the export is corrupt') + raise ValueError(f'{name} is truncated') crc = zlib.crc32(chunk, crc) os.pwrite(fd, chunk, start) - # Flush while other ranges download, not all at the end. getattr(os, 'fdatasync', os.fsync)(fd) finally: os.close(fd) if crc != expected: - raise ValueError( - f'Snapshot file {name} failed its checksum at offset {offset}; ' - 'the export is corrupt') + raise ValueError(f'{name} failed its checksum at offset {offset}') def load_into(connector: Any, path: str, root: str, manifest: dict, @@ -407,13 +354,12 @@ def load_into(connector: Any, path: str, root: str, manifest: dict, shutil.rmtree(staging, ignore_errors=True) try: staging.mkdir(parents=True) - # Any old target stays until the rename, so the whole size is needed. + # Any old target stays until the rename. 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} bytes in {staging}, ' - f'which has {disk.f_bavail * disk.f_frsize} free.') + 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) @@ -430,9 +376,7 @@ def load_into(connector: Any, path: str, root: str, manifest: dict, _parallel(connector, path, ranges, [item[2] for item in ranges], partial(_download, root, staging), 'load') except ValueError as error: - raise ValueError( - f'{error}{_format_note(manifest)}. {_REEXPORT.format(path)}' - ) from error + raise _corrupt(path, str(error)) from error write_diskgraph_success_marker(str(staging), DiskgraphIngestResult.for_index(target)) verify_snapshot_integrity(str(staging)) @@ -456,33 +400,17 @@ def load_into(connector: Any, path: str, root: str, manifest: dict, raise -def _built_with(manifest: Dict[str, Any]) -> Tuple[Any, Any]: - recorded = manifest.get('snapshot') or {} - environment = (recorded.get('environment') or manifest.get('environment') - or {}) - return recorded.get('formats'), (environment.get('kumo_engine') - if isinstance(environment, dict) else None) - - def _check_readable(index: Path, manifest: dict) -> None: from diskgraph import SamplerEngine, format_versions - from kumo.engine._version import __version__ - try: SamplerEngine(str(index), num_threads=1) except Exception as error: - formats, engine = _built_with(manifest) - running = format_versions() - advice = ('The on-disk formats differ: rebuild the snapshot with ' - 'this kumo-engine, or load it with the exporting one.' - if formats is not None and formats != running else - 'Re-export with a matching kumo-engine.') + built = (manifest.get('snapshot') or {}).get('formats') raise RuntimeError( - "This machine's DiskGraph cannot open the snapshot. It was " - f"built with DiskGraph formats {formats} by kumo-engine {engine}; " - f"installed are formats {running} and kumo-engine " - f"{__version__}{_format_note(manifest)}. {advice}") from error + f'This DiskGraph cannot open the snapshot: built with formats ' + f'{built}, installed are {format_versions()}. Rebuild it with ' + 'this kumo-engine or load it with the exporting one.') from error def _fsync(path: Path) -> None: diff --git a/packages/kumo-engine/kumo/engine/graph.py b/packages/kumo-engine/kumo/engine/graph.py index 772288340..7d0835559 100644 --- a/packages/kumo-engine/kumo/engine/graph.py +++ b/packages/kumo-engine/kumo/engine/graph.py @@ -240,14 +240,10 @@ def _loaded(path: str) -> Optional[Any]: return marker -_HOUR_SECONDS = 3600 - - def _comparable(stamp: Any) -> bool: - # hour_version() stamps views and unversioned sources on the hour; -1 is - # an undated source. + # Views and unversioned sources are stamped on the hour; -1 is undated. return (isinstance(stamp, str) and stamp != '-1' - and not (stamp.isdigit() and int(stamp) % _HOUR_SECONDS == 0)) + and not (stamp.isdigit() and int(stamp) % 3600 == 0)) def _source_change(table: Table, recorded: Any) -> Dict[str, Any]: @@ -340,13 +336,8 @@ def load_from_connector(cls, connector: Connector, *, path: str, root, manifest = read(connector, path) graph_key = manifest['graph_key'] if graph is not None and graph.id != graph_key: - warnings.warn( - f"Snapshot at {path} was taken of graph {graph_key!r}, not of " - f"the given graph {graph.id!r}; tables, edges, schemas or " - "connectors differ from the exported graph.", - UserWarning, - stacklevel=2, - ) + warnings.warn(f"Snapshot at {path} was taken of another graph.", + UserWarning, stacklevel=2) target = _store().path( DiskgraphSnapshotHandle( graph_key=graph_key, @@ -367,12 +358,9 @@ def load_from_connector(cls, connector: Connector, *, path: str, ] if changed: warnings.warn( - f"Snapshot at {path} predates the current data of " - f"source table(s) {', '.join(changed)}; see " - "`snapshot.source_changes(graph)`.", - UserWarning, - stacklevel=2, - ) + 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]]: @@ -381,18 +369,12 @@ def source_changes(self, graph: 'Graph') -> Dict[str, Dict[str, Any]]: try: manifest = self.manifest recorded = ({} if not manifest or manifest.get('reconstructed') else - manifest['parents']['source_tables']) + dict(manifest['parents']['source_tables'])) except Exception: recorded = {} - if not isinstance(recorded, dict): - recorded = {} - try: - tables = graph._named_tables - except Exception: - return {} return { alias: _source_change(table, recorded.get(alias)) - for alias, table in tables.items() + for alias, table in graph._named_tables.items() } @property diff --git a/packages/kumo-engine/test/test_snapshot_export.py b/packages/kumo-engine/test/test_snapshot_export.py index 62822cd26..9a08c1e02 100644 --- a/packages/kumo-engine/test/test_snapshot_export.py +++ b/packages/kumo-engine/test/test_snapshot_export.py @@ -16,7 +16,6 @@ import pandas as pd import pytest from diskgraph import format_versions -from kumoapi.data_source import DataSourceType from pyarrow.fs import LocalFileSystem, SubTreeFileSystem import kumo.engine as kumo @@ -37,50 +36,34 @@ def use(name): 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', 'snapshot-manifest.json') + and not str(file.relative_to(path)).startswith('train_tables/') } -def _exported_files(snapshot): - return { - name: data - for name, data in _index_files(snapshot.index_path).items() - if not name.startswith('train_tables/') - } - - -@pytest.fixture -def exports(tmp_path): - return kumo.FileConnector(root_dir=str(tmp_path / 'exports')) - - def _export(snapshot, exports, tmp_path): location = snapshot.export(exports, 'one') return location, tmp_path / 'exports' / location -def _edit_manifest(exported, **changes): +def _edit_manifest(exported, edit): path = exported / 'kumo_snapshot.json' - path.write_text(json.dumps({**json.loads(path.read_text()), **changes})) - - -@pytest.fixture -def workers(monkeypatch): - @contextmanager - def four(found, worker_ceiling=None): - yield 4 - - monkeypatch.setattr('kumo_core.sources.object_store.copies', four) + manifest = json.loads(path.read_text()) + edit(manifest) + path.write_text(json.dumps(manifest)) -def _no_ingest(monkeypatch): - monkeypatch.setattr( - 'kumo_core.diskgraph.ingest.execute_diskgraph_ingest', - Mock(side_effect=AssertionError('ingest must not run'))) +def _load(exports, location, **kwargs): + return GraphSnapshot.load_from_connector(exports, path=location, **kwargs) def test_export_load_round_trip_relocates_and_is_reused( @@ -93,80 +76,57 @@ def test_export_load_round_trip_relocates_and_is_reused( assert location == f'one/{graph_id}/{data_version}' assert snapshot.export(exports, 'one') == location exported = tmp_path / 'exports' / location - assert _index_files(exported / 'index') == _exported_files(snapshot) manifest = json.loads((exported / 'kumo_snapshot.json').read_text()) - assert manifest['snapshot_id'] == snapshot.id assert manifest['snapshot'] == snapshot.manifest shutil.rmtree(tmp_path / 'cpu') scratch('gpu') - loaded = GraphSnapshot.load_from_connector(exports, path=location, - graph=graph) - assert loaded.id == snapshot.id - assert loaded.graph_config == snapshot.graph_config + 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(exported / 'index') - marker = json.loads((loaded.index_path / '_SUCCESS').read_text()) - assert marker['index_path'] == str(loaded.index_path) recorded = loaded.manifest - assert recorded.pop('loaded') == { - 'formats': format_versions(), - 'environment': manifest['environment'], - } + assert recorded.pop('loaded')['formats'] == format_versions() assert recorded == manifest['snapshot'] reexported = tmp_path / 'exports' / loaded.export(exports, 'two') assert json.loads((reexported / 'kumo_snapshot.json').read_text())[ 'snapshot'] == manifest['snapshot'] - _no_ingest(monkeypatch) + 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 GraphSnapshot.load_from_connector(exports, path=location) == loaded + assert _load(exports, location) == loaded assert graph.snapshot() == loaded -@pytest.mark.parametrize('damage', [False, True]) -def test_in_place_index_reuse_is_unchanged(scratch, graph, monkeypatch, - damage): - built = graph.snapshot() - marker = (built.index_path / '_SUCCESS').read_bytes() - if damage: - stats = built.index_path / 'col_stats.json' - stats.write_bytes(stats.read_bytes()[:-1]) - (built.index_path / 'extra.bin').write_bytes(b'x') - _no_ingest(monkeypatch) +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() == built - assert (built.index_path / '_SUCCESS').read_bytes() == marker + 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_moved_index_is_rebuilt_as_before(tmp_path, scratch, graph): - built = graph.snapshot() - shutil.move(tmp_path / 'cpu', tmp_path / 'moved') - scratch('moved') - rebuilt = graph.snapshot() - assert rebuilt.id == built.id - marker = json.loads((rebuilt.index_path / '_SUCCESS').read_text()) - assert marker['index_path'] == str(rebuilt.index_path) +def test_interrupted_export_is_not_loadable_and_resumes( + tmp_path, scratch, graph, exports, monkeypatch): + @contextmanager + def four(found, worker_ceiling=None): + yield 4 -def test_load_replaces_a_damaged_local_copy( - tmp_path, scratch, graph, exports): + monkeypatch.setattr('kumo_core.sources.object_store.copies', four) snapshot = graph.snapshot() - location, _ = _export(snapshot, exports, tmp_path) - scratch('gpu') - loaded = GraphSnapshot.load_from_connector(exports, path=location) - stats = loaded.index_path / 'col_stats.json' - stats.write_bytes(stats.read_bytes()[:-1]) - with pytest.warns(UserWarning, match='Reloading the damaged'): - again = GraphSnapshot.load_from_connector(exports, path=location) - assert _index_files(again.index_path) == _exported_files(snapshot) - - -def _interrupt_export(snapshot, exports, monkeypatch): - stream = snapshot_io._stream - finished = threading.Event() + stream, finished = snapshot_io._stream, threading.Event() def interrupted(source, size, stop, write=None): if write is not None and source.name == 'diskgraph_config.json': @@ -177,42 +137,25 @@ def interrupted(source, size, stop, write=None): finished.set() return crcs - monkeypatch.setattr(snapshot_io, '_stream', interrupted) - monkeypatch.setattr(snapshot_io, '_ATTEMPTS', 1) - with pytest.raises(OSError, match='connection reset'): - snapshot.export(exports, 'one') - monkeypatch.setattr(snapshot_io, '_stream', stream) - monkeypatch.setattr(snapshot_io, '_ATTEMPTS', 3) - - -@pytest.mark.parametrize('change', [False, True]) -def test_interrupted_export_is_not_loadable_and_resumes( - tmp_path, scratch, graph, exports, monkeypatch, workers, change): - snapshot = graph.snapshot() - _interrupt_export(snapshot, exports, monkeypatch) + with monkeypatch.context() as patch: + patch.setattr(snapshot_io, '_stream', interrupted) + patch.setattr(snapshot_io, '_ATTEMPTS', 1) + with pytest.raises(OSError, match='connection reset'): + snapshot.export(exports, 'one') graph_id, data_version = snapshot.id.split('@') location = f'one/{graph_id}/{data_version}' with pytest.raises(FileNotFoundError, match='kumo_snapshot.json'): - GraphSnapshot.load_from_connector(exports, path=location) - if change: - stats = snapshot.index_path / 'col_stats.json' - data = stats.read_bytes() - stats.write_bytes(data[:-1] + bytes([data[-1] ^ 1])) - index_manifest = snapshot.index_path / '_MANIFEST.json' - index_manifest.write_bytes(index_manifest.read_bytes() + b' ') - + _load(exports, location) assert snapshot.export(exports, 'one') == location scratch('gpu') - loaded = GraphSnapshot.load_from_connector(exports, path=location) - assert (loaded.index_path / 'col_stats.json').read_bytes() == ( - snapshot.index_path / 'col_stats.json').read_bytes() + assert _index_files(_load(exports, location).index_path) == _index_files( + snapshot.index_path) -def test_concurrent_identical_exports_both_succeed(tmp_path, scratch, graph, - exports, monkeypatch): +def test_concurrent_exports_of_one_snapshot_both_succeed( + tmp_path, scratch, graph, exports, monkeypatch): snapshot = graph.snapshot() - parallel = snapshot_io._parallel - raced = [] + parallel, raced = snapshot_io._parallel, [] def racing(*args): if not raced: @@ -220,32 +163,24 @@ def racing(*args): raced[0] = snapshot.export(exports, 'one') return parallel(*args) - graph_id, data_version = snapshot.id.split('@') - location = f'one/{graph_id}/{data_version}' monkeypatch.setattr(snapshot_io, '_parallel', racing) - assert snapshot.export(exports, 'one') == location - assert raced == [location] - scratch('gpu') - loaded = GraphSnapshot.load_from_connector(exports, path=location) - assert _index_files(loaded.index_path) == _exported_files(snapshot) + assert raced == [snapshot.export(exports, 'one')] -def test_export_committed_by_a_late_identical_export_succeeds( - tmp_path, scratch, graph, exports, monkeypatch): - snapshot = graph.snapshot() - write = snapshot_io._write_json - raced = [] +def test_failed_upload_is_retried_with_a_fresh_client( + scratch, graph, exports, monkeypatch): + open_output, clients = snapshot_io._open_output, [] - def racing(filesystem, path, value): - if path.endswith('kumo_snapshot.json') and not raced: - raced.append(None) - raced[0] = snapshot.export(exports, 'one') - return write(filesystem, path, value) + def expiring(client, path): + if path.endswith('col_stats.json'): + clients.append(client) + if len(clients) == 1: + raise OSError('ExpiredToken') + return open_output(client, path) - monkeypatch.setattr(snapshot_io, '_write_json', racing) - location = snapshot.export(exports, 'one') - assert raced == [location] - assert not (tmp_path / 'exports' / location / '_EXPORTING').exists() + monkeypatch.setattr(snapshot_io, '_open_output', expiring) + graph.snapshot().export(exports, 'one') + assert len(clients) == 2 and clients[0] is not clients[1] def test_rerunning_export_repairs_a_committed_export( @@ -258,9 +193,6 @@ def test_rerunning_export_repairs_a_committed_export( (exported / 'index/index_metadata.json').unlink() assert snapshot.export(exports, 'one') == location assert stats.read_bytes() == original - scratch('gpu') - loaded = GraphSnapshot.load_from_connector(exports, path=location) - assert _index_files(loaded.index_path) == _exported_files(snapshot) def test_failed_load_leaves_nothing_and_restarts_cleanly( @@ -272,102 +204,69 @@ def test_failed_load_leaves_nothing_and_restarts_cleanly( stats.write_bytes(original[:-1] + bytes([original[-1] ^ 1])) scratch('gpu') - with pytest.raises(ValueError, match='checksum.*corrupt. Delete one/'): - GraphSnapshot.load_from_connector(exports, path=location) + 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() - (graph_dir / f'{snapshot.index_path.name}.loading/stale').write_text('x') - loaded = GraphSnapshot.load_from_connector(exports, path=location) - assert _index_files(loaded.index_path) == _exported_files(snapshot) + 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_needs_the_full_size_while_an_old_copy_remains( +def test_load_refuses_before_downloading_when_disk_is_too_small( tmp_path, scratch, graph, exports, monkeypatch): - snapshot = graph.snapshot() - location, exported = _export(snapshot, exports, tmp_path) - total = sum(entry['size'] for entry in json.loads( - (exported / 'kumo_snapshot.json').read_text())['files']) + location, _ = _export(graph.snapshot(), exports, tmp_path) scratch('gpu') - target = (tmp_path / 'gpu/snapshots' / snapshot.index_path.parent.name / - snapshot.index_path.name) - target.mkdir(parents=True) - (target / 'stale').write_bytes(b'x' * total) - monkeypatch.setattr( - snapshot_io.os, 'statvfs', - lambda path: SimpleNamespace(f_frsize=1, f_bavail=total - 1)) - with pytest.raises(OSError, match=f'needs {total} bytes'): - GraphSnapshot.load_from_connector(exports, path=location) - assert (target / 'stale').is_file() - - -@pytest.mark.parametrize('field, value', [ - ('data_version', '../../victim'), - ('files', [{ - 'path': '../../victim/x', - 'size': 1, - 'crc32': [0] - }]), + 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('edit', [ + lambda manifest: manifest.update(data_version='../../victim'), + lambda manifest: manifest['files'][0].update(path='../../victim/x'), + lambda manifest: manifest.pop('files'), + lambda manifest: manifest['files'].append(manifest['files'][0]), ]) -def test_load_rejects_a_hostile_manifest(tmp_path, scratch, graph, exports, - field, value): - snapshot = graph.snapshot() - location, exported = _export(snapshot, exports, tmp_path) +def test_load_rejects_a_hostile_or_malformed_manifest(tmp_path, scratch, graph, + exports, edit): + location, exported = _export(graph.snapshot(), exports, tmp_path) victim = tmp_path / 'victim' victim.mkdir() (victim / 'keep').write_text('x') - _edit_manifest(exported, **{field: value}) + _edit_manifest(exported, edit) scratch('gpu/a/b') with pytest.raises(ValueError): - GraphSnapshot.load_from_connector(exports, path=location) + _load(exports, location) assert (victim / 'keep').read_text() == 'x' -@pytest.mark.parametrize('edit', [ - lambda manifest: manifest.pop('files'), - lambda manifest: manifest.update(snapshot=[1]), - lambda manifest: manifest['files'][0].pop('crc32'), - lambda manifest: manifest['files'].append(manifest['files'][0]), -]) -def test_load_rejects_a_malformed_manifest_with_value_error( - tmp_path, scratch, graph, exports, edit): +def test_load_warns_instead_of_refusing_newer_or_foreign_exports( + tmp_path, scratch, graph, exports): snapshot = graph.snapshot() location, exported = _export(snapshot, exports, tmp_path) - path = exported / 'kumo_snapshot.json' - manifest = json.loads(path.read_text()) - edit(manifest) - path.write_text(json.dumps(manifest)) - scratch('gpu') - with pytest.raises(ValueError, match='corrupt. Delete one/'): - GraphSnapshot.load_from_connector(exports, path=location) - path.write_text('[]') - with pytest.raises(ValueError, match='not a JSON object'): - GraphSnapshot.load_from_connector(exports, path=location) + def newer(manifest): + manifest.update(format=2, added_later={'x': 1}) + for entry in manifest['files']: + entry['added_later'] = 1 -def test_load_records_differences_instead_of_refusing( - tmp_path, scratch, graph, exports, monkeypatch): - exporter_formats = {**format_versions(), 'layout': -1} - monkeypatch.setattr('diskgraph.format_versions', lambda: exporter_formats) - snapshot = graph.snapshot() - monkeypatch.undo() - location, exported = _export(snapshot, exports, tmp_path) - _edit_manifest(exported, format=2) + _edit_manifest(exported, newer) other = kumo.Graph(tables={'users': graph.tables['users']}) scratch('gpu') with pytest.warns(UserWarning) as caught: - loaded = GraphSnapshot.load_from_connector(exports, path=location, - graph=other) + loaded = _load(exports, location, graph=other) messages = ' '.join(str(warning.message) for warning in caught) - assert 'layout format 2' in messages and 'was taken of graph' in messages + assert 'layout format 2' in messages and 'another graph' in messages assert loaded.id == snapshot.id - assert loaded.manifest['formats'] == exporter_formats - assert loaded.manifest['loaded']['formats'] == format_versions() def test_export_refuses_a_path_holding_other_data(tmp_path, scratch, graph, @@ -381,10 +280,9 @@ def test_export_refuses_a_path_holding_other_data(tmp_path, scratch, graph, snapshot.export(exports, 'occupied') _, exported = _export(snapshot, exports, tmp_path) - _edit_manifest(exported, snapshot_id='other') + _edit_manifest(exported, lambda m: m.update(snapshot_id='other')) with pytest.raises(ValueError, match="snapshot 'other', not"): snapshot.export(exports, 'one') - assert (exported / 'index/col_stats.json').is_file() def test_rebuilt_snapshot_reexports_to_the_existing_export( @@ -392,96 +290,38 @@ def test_rebuilt_snapshot_reexports_to_the_existing_export( snapshot = graph.snapshot() location, exported = _export(snapshot, exports, tmp_path) committed = (exported / 'kumo_snapshot.json').read_bytes() - scratch('rebuilt/on/another/machine') + scratch('rebuilt') rebuilt = graph.snapshot() assert rebuilt.id == snapshot.id - assert (rebuilt.index_path / '_MANIFEST.json').read_bytes() != ( - snapshot.index_path / '_MANIFEST.json').read_bytes() - assert rebuilt.export(exports, 'one') == location assert (exported / 'kumo_snapshot.json').read_bytes() == committed - scratch('gpu') - loaded = GraphSnapshot.load_from_connector(exports, path=location) - assert _index_files(loaded.index_path) == _exported_files(snapshot) -@pytest.mark.parametrize('kind, reason', [ - ('SNOWFLAKE', 'Snowflake stages'), - ('DATABRICKS', 'Unity Catalog Volume'), - ('BIGQUERY', 'BigQuery'), -]) -def test_warehouse_connectors_are_not_supported(scratch, graph, kind, reason): - connector = type(f'{kind.title()}Connector', (), { - 'source_type': DataSourceType[kind], - })() - with pytest.raises(NotImplementedError, match=reason): +def test_warehouse_connectors_are_not_supported(scratch, graph): + connector = SimpleNamespace(sink_backend=object()) + with pytest.raises(NotImplementedError, match='FileConnector'): graph.snapshot().export(connector, 'x') - with pytest.raises(NotImplementedError, match=reason): - GraphSnapshot.load_from_connector(connector, path='x') - - -def test_interrupted_load_joins_workers_before_cleaning_up( - tmp_path, scratch, graph, exports, monkeypatch, workers): - snapshot = graph.snapshot() - location, _ = _export(snapshot, exports, tmp_path) - download = snapshot_io._download - running = threading.Event() - - def interrupted(root, staging, client, item, stop): - if item[0] == 'col_stats.json': - running.wait(10) - raise KeyboardInterrupt - running.set() - stop.wait(1) - download(root, staging, client, item, stop) - - monkeypatch.setattr(snapshot_io, '_download', interrupted) - scratch('gpu') - with pytest.raises(KeyboardInterrupt): - GraphSnapshot.load_from_connector(exports, path=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'] + 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): - snapshot = graph.snapshot() - location, exported = _export(snapshot, exports, tmp_path) + location, exported = _export(graph.snapshot(), exports, tmp_path) metadata = exported / 'index/index_metadata.json' garbage = b'{' * metadata.stat().st_size metadata.write_bytes(garbage) - manifest = json.loads((exported / 'kumo_snapshot.json').read_text()) - for entry in manifest['files']: - if entry['path'] == 'index_metadata.json': - entry['crc32'] = [zlib.crc32(garbage)] - manifest['snapshot']['formats'] = {'layout': -1} - manifest['snapshot']['environment']['kumo_engine'] = '0.0.1' - manifest['format'] = 2 - (exported / 'kumo_snapshot.json').write_text(json.dumps(manifest)) - scratch('gpu') - with pytest.raises(RuntimeError, match=r"formats \{'layout': -1\} by " - r'kumo-engine 0\.0\.1; installed are formats .*' - r'export layout format 2; this engine reads 1\)\. ' - 'The on-disk formats differ: rebuild'): - GraphSnapshot.load_from_connector(exports, path=location) + def unreadable(manifest): + for entry in manifest['files']: + if entry['path'] == 'index_metadata.json': + entry['crc32'] = [zlib.crc32(garbage)] + manifest['snapshot']['formats'] = {'layout': -1} - -def test_load_ignores_unknown_keys_and_a_missing_snapshot_manifest( - tmp_path, scratch, graph, exports): - snapshot = graph.snapshot() - location, exported = _export(snapshot, exports, tmp_path) - path = exported / 'kumo_snapshot.json' - manifest = json.loads(path.read_text()) - manifest['added_later'] = {'x': 1} - manifest['snapshot'] = None - for entry in manifest['files']: - entry['added_later'] = 1 - path.write_text(json.dumps(manifest)) + _edit_manifest(exported, unreadable) scratch('gpu') - loaded = GraphSnapshot.load_from_connector(exports, path=location) - assert loaded.id == snapshot.id - assert loaded.manifest is None + with pytest.raises(RuntimeError, match=r"formats \{'layout': -1\}"): + _load(exports, location) @pytest.mark.parametrize('scheme', ['s3', 'gs']) @@ -489,58 +329,21 @@ def test_object_store_round_trip(tmp_path, scratch, graph, monkeypatch, scheme): storage = tmp_path / 'storage' storage.mkdir() - bucket = SubTreeFileSystem(str(storage), LocalFileSystem()) - seen = [] - - def s3(bucket_name, *, options=None, **kwargs): - seen.append(options) - return bucket - monkeypatch.setattr('kumo_s3.filesystem', s3) - - class Gcs(SubTreeFileSystem): - def __init__(self): + class Bucket(SubTreeFileSystem): + def __init__(self, *args, **kwargs): super().__init__(str(storage), LocalFileSystem()) - monkeypatch.setattr('kumo_gcs.backend.fs.GcsFileSystem', Gcs) - + 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(connector, f'{scheme}://bucket/snapshots/') graph_id, data_version = snapshot.id.split('@') assert location == f'{scheme}://bucket/snapshots/{graph_id}/{data_version}' - assert (storage / 'bucket/snapshots' / graph_id / data_version / - 'kumo_snapshot.json').is_file() - scratch('gpu') - loaded = GraphSnapshot.load_from_connector(connector, path=location, - graph=graph) - assert loaded.id == snapshot.id - assert _index_files(loaded.index_path) == _exported_files(snapshot) - if scheme == 's3': - assert seen and all( - option['allow_delayed_open'] and not option['background_writes'] - for option in seen) - - -def test_transfer_renews_an_expired_client(tmp_path, scratch, graph, exports, - monkeypatch): - snapshot = graph.snapshot() - location, _ = _export(snapshot, exports, tmp_path) - download = snapshot_io._download - failed = [] - - def expiring(root, staging, client, item, stop): - if item[0] == 'col_stats.json': - if not failed: - failed.append(client) - raise OSError('AWS Error ACCESS_DENIED: ExpiredToken') - assert client is not failed[0] - return download(root, staging, client, item, stop) - - monkeypatch.setattr(snapshot_io, '_download', expiring) scratch('gpu') - GraphSnapshot.load_from_connector(exports, path=location) - assert failed + loaded = _load(connector, location, graph=graph) + assert _index_files(loaded.index_path) == _index_files(snapshot.index_path) def _set_mtime(path, seconds): @@ -561,87 +364,46 @@ def dated(tmp_path, graph): return data -def _forbid_staging(monkeypatch): +def test_source_changes_reads_stamps_without_staging( + scratch, graph, dated, monkeypatch): 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() 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) - - -def test_source_changes_reads_stamps_without_staging( - scratch, graph, dated, monkeypatch): - snapshot = graph.snapshot() - _forbid_staging(monkeypatch) - unchanged = {'status': 'unchanged', 'recorded': '1700000123', - 'current': '1700000123'} - assert snapshot.source_changes(graph) == {'users': unchanged, - 'events': unchanged} - _rewrite_events(dated) - assert snapshot.source_changes(graph) == { - 'users': unchanged, - 'events': {'status': 'changed', 'recorded': '1700000123', - 'current': '1700009123'}, - } - - -def test_source_changes_unknown_for_hour_stamps_and_read_errors( - scratch, graph, dated, monkeypatch): - from kumo.engine.connector.file import _FileSourceBackend - - _set_mtime(dated / 'users.parquet', 1_700_002_800) - snapshot = graph.snapshot() - stamp = _FileSourceBackend.get_data_version - - def failing(self, connector, table_name, project_id=None): - if table_name == 'events.parquet': - raise OSError('listing failed') - return stamp(self, connector, table_name) - - monkeypatch.setattr(_FileSourceBackend, 'get_data_version', failing) assert snapshot.source_changes(graph) == { 'users': {'status': 'unknown', 'recorded': '1700002800', 'current': '1700002800'}, - 'events': {'status': 'unknown', 'recorded': '1700000123', - 'current': None}, + '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} -@pytest.mark.parametrize('edit', ['absent', 'reconstructed', 'missing', - 'malformed']) +@pytest.mark.parametrize('edit', ['absent', 'reconstructed']) def test_source_changes_unknown_without_a_recorded_version( scratch, graph, dated, edit): snapshot = graph.snapshot() path = snapshot.index_path / 'snapshot-manifest.json' - manifest = json.loads(path.read_text()) if edit == 'absent': path.unlink() - elif edit == 'reconstructed': - path.write_text(json.dumps({**manifest, 'reconstructed': True})) - elif edit == 'missing': - del manifest['parents']['source_tables']['users'] - path.write_text(json.dumps(manifest)) else: - path.write_text(json.dumps({**manifest, 'parents': []})) - changes = snapshot.source_changes(graph) - assert changes['users'] == {'status': 'unknown', 'recorded': None, - 'current': '1700000123'} - assert changes['events']['status'] == ( - 'unchanged' if edit == 'missing' else 'unknown') - - -def test_source_changes_unknown_for_a_moved_table(scratch, graph, dated): - snapshot = graph.snapshot() - path = snapshot.index_path / 'snapshot-manifest.json' - manifest = json.loads(path.read_text()) - manifest['parents']['source_tables']['users']['location'] += '.old' - path.write_text(json.dumps(manifest)) - changes = snapshot.source_changes(graph) - assert changes['users']['status'] == 'unknown' - assert changes['events']['status'] == 'unchanged' + path.write_text( + json.dumps({**json.loads(path.read_text()), 'reconstructed': True})) + assert {change['status'] for change in snapshot.source_changes( + graph).values()} == {'unknown'} def test_load_warns_only_with_graph_and_only_for_changed_tables( @@ -650,10 +412,10 @@ def test_load_warns_only_with_graph_and_only_for_changed_tables( scratch('gpu') with warnings.catch_warnings(): warnings.simplefilter('error') - GraphSnapshot.load_from_connector(exports, path=location, graph=graph) + _load(exports, location, graph=graph) _rewrite_events(dated) - GraphSnapshot.load_from_connector(exports, path=location) + _load(exports, location) with pytest.warns(UserWarning) as caught: - GraphSnapshot.load_from_connector(exports, path=location, graph=graph) + _load(exports, location, graph=graph) [message] = [str(warning.message) for warning in caught] assert 'source table(s) events;' in message diff --git a/packages/kumo-engine/test/test_snapshot_export_live.py b/packages/kumo-engine/test/test_snapshot_export_live.py index 9628cedbf..3c69b548a 100644 --- a/packages/kumo-engine/test/test_snapshot_export_live.py +++ b/packages/kumo-engine/test/test_snapshot_export_live.py @@ -1,15 +1,12 @@ # SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. # All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Snapshot export and load through real S3 and GCS; each store is skipped -unless its writable root is set.""" import contextlib import os import uuid import pytest -from pyarrow.fs import FileSelector import kumo.engine as kumo from kumo.engine._scratch import get_scratch_dir @@ -25,9 +22,7 @@ def test_snapshot_round_trip_through_object_store(tmp_path, graph, scheme): root = ROOTS[scheme] if not root: - pytest.skip(f'LIVE SNAPSHOT {scheme.upper()} NOT RUN: set ' - f"KUMO_SNAPSHOT_LIVE_{'S3' if scheme == 's3' else 'GCS'}" - '_ROOT to a writable prefix') + 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() @@ -44,13 +39,6 @@ def test_snapshot_round_trip_through_object_store(tmp_path, graph, scheme): filesystem.delete_dir_contents(prefix, missing_dir_ok=True) with contextlib.suppress(OSError): filesystem.delete_dir(prefix) - assert not filesystem.get_file_info( - FileSelector(prefix, allow_not_found=True)) assert loaded.id == snapshot.id - for file in snapshot.index_path.rglob('*'): - relative = file.relative_to(snapshot.index_path) - if (file.is_file() - and file.name not in ('_SUCCESS', 'snapshot-manifest.json') - and relative.parts[0] != 'train_tables'): - assert (loaded.index_path / relative).read_bytes() == \ - file.read_bytes() + assert (loaded.index_path / 'col_stats.json').read_bytes() == ( + snapshot.index_path / 'col_stats.json').read_bytes() diff --git a/packages/kumo-s3/src/kumo_s3/backend.py b/packages/kumo-s3/src/kumo_s3/backend.py index d34effa5c..78e9c5390 100644 --- a/packages/kumo-s3/src/kumo_s3/backend.py +++ b/packages/kumo-s3/src/kumo_s3/backend.py @@ -109,8 +109,7 @@ def _filesystem(self, connector: Any, bucket: str, ) -> S3FileSystem: r"""A ready filesystem is taken as it is: it may address a service that is not AWS at all, and probing its bucket name against AWS would - answer for somebody else's bucket. ``options`` are extra - ``S3FileSystem`` keywords. + answer for somebody else's bucket. """ session, allow_anonymous, connection = self._resolved_session( From 5f8239a366edcddebb14845ece674c7c4962979c Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Wed, 21 Oct 2026 18:43:55 +0000 Subject: [PATCH 57/80] Move snapshot export checksums into DiskGraph and reuse its manifest Export and load computed per-range CRC32s in Python over every index byte, kept a second file list and commit protocol beside DiskGraph's _MANIFEST.json/_SUCCESS contract, and re-implemented helpers kumo-core already has (parallel runner, cleanup, directory fsync, snapshot verification, source location). DiskGraph gains snapshot_range_checksums next to verify_snapshot_integrity (shared manifest reader, rayon, 1 MiB buffers), exposed through the pybind module. Export computes it alongside the upload; load downloads, then runs verify_snapshot_integrity and the checksum pass on the staged copy. The export now lists files only in index/_MANIFEST.json and commits with index/_SUCCESS written last; kumo_snapshot.json keeps the snapshot metadata plus range_bytes and a crc32 map. Transfers run on kumo-core's run_in_parallel; a load pre-allocates files with posix_fallocate and removes a damaged copy before staging; _loaded uses LocalSnapshotStore.verify + rebased; source_changes reuses source_table and checks tables in parallel; GraphSnapshot gains graph_key and data_version. Verified: cargo test -p diskgraph-ingest snapshot_mirror (21 passed), pybind test_exception_exports (2), test_snapshot_export + test_snapshot_manifest (39), kumo-s3 (51 passed, 1 skipped). On a 13 GB index the Rust checksum pass takes 0.55 s versus 1.98 s for 64 Python threads with zlib; local export and load times are unchanged within noise (load is disk-bound) and load peak RSS is +65 MB. 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 ++++++- .../python/diskgraph/__init__.py | 2 + .../python/tests/test_exception_exports.py | 19 + .../diskgraph/diskgraph-pybind/src/lib.rs | 16 + .../third_party/artifact-manifest.json | 9 +- .../src/kumo_core/sources/object_store.py | 4 +- .../kumo-engine/kumo/engine/_snapshot_io.py | 379 +++++++++--------- packages/kumo-engine/kumo/engine/graph.py | 84 ++-- .../kumo-engine/test/test_snapshot_export.py | 59 +-- 13 files changed, 456 insertions(+), 258 deletions(-) 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/python/diskgraph/__init__.py b/packages/diskgraph/diskgraph-pybind/python/diskgraph/__init__.py index 548081716..5b9f5d0b6 100644 --- a/packages/diskgraph/diskgraph-pybind/python/diskgraph/__init__.py +++ b/packages/diskgraph/diskgraph-pybind/python/diskgraph/__init__.py @@ -40,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, @@ -141,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 aa4eca8e9..d4b29806e 100644 --- a/packages/diskgraph/diskgraph-pybind/python/tests/test_exception_exports.py +++ b/packages/diskgraph/diskgraph-pybind/python/tests/test_exception_exports.py @@ -16,3 +16,22 @@ def test_public_corruption_error_catches_snapshot_verifier(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 e02bda485..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)) @@ -1608,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-core/src/kumo_core/sources/object_store.py b/packages/kumo-core/src/kumo_core/sources/object_store.py index bf1058d07..701675ab3 100644 --- a/packages/kumo-core/src/kumo_core/sources/object_store.py +++ b/packages/kumo-core/src/kumo_core/sources/object_store.py @@ -68,6 +68,8 @@ from .types import ReadMode, StagedTable STREAMED_ONLY = staged_first(False) +#: The version of a source with no modification time. +UNDATED_VERSION = "-1" logger = logging.getLogger(__name__) @@ -102,7 +104,7 @@ 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 (SourceVersion(epoch_second(latest), "timestamp") - if latest is not None else SourceVersion("-1")) + if latest is not None else SourceVersion(UNDATED_VERSION)) # A maximum mtime misses changes to older files and added/deleted parts. metadata = [ (entry.path, entry.size, getattr(entry, "mtime_ns", None), diff --git a/packages/kumo-engine/kumo/engine/_snapshot_io.py b/packages/kumo-engine/kumo/engine/_snapshot_io.py index c1e3f5fde..b09db4f37 100644 --- a/packages/kumo-engine/kumo/engine/_snapshot_io.py +++ b/packages/kumo-engine/kumo/engine/_snapshot_io.py @@ -7,12 +7,10 @@ import logging import os import re -import shutil import threading import time import warnings -import zlib -from concurrent.futures import FIRST_EXCEPTION, ThreadPoolExecutor, wait +from concurrent.futures import ThreadPoolExecutor from functools import partial from pathlib import Path from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple @@ -21,10 +19,10 @@ _FORMAT = 1 _MANIFEST = 'kumo_snapshot.json' -_EXPORTING = '_EXPORTING' +_INDEX_MANIFEST = '_MANIFEST.json' +_SUCCESS = '_SUCCESS' _RANGE_BYTES = 64 << 20 _CHUNK = 8 << 20 -_INDEX_MANIFEST = '_MANIFEST.json' _ATTEMPTS = 3 _KEY = re.compile('[0-9a-f]{64}') _PROGRESS_SECONDS = 30 @@ -46,10 +44,18 @@ def _location(connector: Any, path: str) -> Tuple[FileSystem, str]: return sink.filesystem(connector, path, parallel_uploads=True) +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 == f'{graph_key}@{data_version[:16]}'): + and snapshot_id == _snapshot_id(graph_key, data_version)): raise ValueError(f'{snapshot_id!r} is not a snapshot id.') @@ -65,42 +71,50 @@ def renew(self, stale: FileSystem) -> None: self.current = self._make() -def _parallel(connector: Any, path: str, items: Sequence[Any], +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() - total = sum(sizes) - with copies(items) as workers, ThreadPoolExecutor(workers) as pool: + 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), total, workers) - started = time.monotonic() - futures = { - pool.submit(_retried, client, work, item, stop): size - for item, size in zip(items, sizes) - } - try: - pending, moved = set(futures), 0 - while pending: - done, pending = wait(pending, timeout=_PROGRESS_SECONDS, - return_when=FIRST_EXCEPTION) - for future in done: - if future.exception() is not None: - raise future.exception() - moved += futures[future] - elapsed = time.monotonic() - started - logger.info('Snapshot %s of %s: %d of %d bytes, %.2f GB/s', - action, path, moved, total, - moved / max(elapsed, 1e-9) / 1e9) - return [future.result() for future in futures] - except BaseException: - stop.set() - for future in futures: - future.cancel() - raise + 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, @@ -130,59 +144,45 @@ def _stopped(stop: threading.Event) -> None: raise WorkCancelled('abandoned after another transfer failed') -def _stream(source: Path, size: int, stop: threading.Event, - write: Optional[Callable[[bytes], Any]] = None) -> List[int]: - crcs: List[int] = [] - with open(source, 'rb', buffering=0) as reader: - 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} changed while it was exported') - if offset % _RANGE_BYTES == 0: - crcs.append(0) - crcs[-1] = zlib.crc32(chunk, crcs[-1]) - if write is not None: - write(chunk) - return crcs or [0] - - -def _read_json(filesystem: FileSystem, path: str) -> Optional[dict]: +def _read_json(filesystem: FileSystem, path: str) -> Optional[Any]: try: - with filesystem.open_input_stream(path) as reader: - return json.loads(reader.readall()) + return json.loads(_files(filesystem).read_text(path)) except FileNotFoundError: return None def _open_output(filesystem: FileSystem, path: str) -> Any: - if filesystem.type_name not in ('s3', 'gcs'): - filesystem.create_dir(path.rsplit('/', 1)[0], recursive=True) + _files(filesystem).makedirs(path.rsplit('/', 1)[0]) return filesystem.open_output_stream(path) -def _write_json(filesystem: FileSystem, path: str, value: dict) -> None: +def _write(filesystem: FileSystem, path: str, payload: bytes) -> None: with _open_output(filesystem, path) as writer: - writer.write(json.dumps(value).encode()) + writer.write(payload) + + +def _fingerprint(index_manifest: bytes) -> str: + return hashlib.sha256(index_manifest).hexdigest() def export(snapshot: Any, connector: Any, path: str) -> str: + from diskgraph import snapshot_range_checksums + from kumo.engine._provenance import environment source = Path(snapshot.index_path) - graph_key, data_version = source.parent.name, source.name + graph_key, data_version = snapshot.graph_key, snapshot.data_version _check_key(graph_key, data_version, snapshot.id) location = f"{path.rstrip('/')}/{graph_key}/{data_version[:16]}" index_manifest = (source / _INDEX_MANIFEST).read_bytes() - files = sorted([(entry['rel_path'], int(entry['size_bytes'])) - for entry in json.loads(index_manifest)['files']] + - [(_INDEX_MANIFEST, len(index_manifest))], - key=lambda file: -file[1]) - fingerprint = hashlib.sha256(index_manifest).hexdigest() - identity = {'snapshot_id': snapshot.id, 'index_manifest': fingerprint} + sizes = { + entry['rel_path']: int(entry['size_bytes']) + for entry in json.loads(index_manifest)['files'] + } + fingerprint = _fingerprint(index_manifest) filesystem, root = _location(connector, location) - def committed() -> Optional[dict]: + def recorded() -> Optional[dict]: existing = _read_json(filesystem, f'{root}/{_MANIFEST}') found = (existing.get('snapshot_id') if isinstance(existing, dict) else None) @@ -191,12 +191,11 @@ def committed() -> Optional[dict]: f'{snapshot.id!r}; remove it and re-export.') return existing - def finished() -> bool: - existing = committed() - if existing and existing.get('index_manifest') != fingerprint: - raise RuntimeError(f'Another build of {snapshot.id} was exported ' - f'to {location} meanwhile; re-export.') - return existing is not None + def build() -> Optional[str]: + return (recorded() or {}).get('index_manifest') + + def committed() -> bool: + return _files(filesystem).exists(f'{root}/index/{_SUCCESS}') def present() -> Dict[str, int]: return { @@ -207,67 +206,72 @@ def present() -> Dict[str, int]: if info.type is FileType.File } - def upload(client: FileSystem, file: Tuple[str, int], - stop: threading.Event) -> List[int]: - name, size = file - if uploaded.get(name) == size: - return _stream(source / name, size, stop) - with _open_output(client, f'{root}/index/{name}') as writer: - return _stream(source / name, size, stop, writer.write) - - existing = committed() - if existing is not None and existing.get('index_manifest') != fingerprint: - logger.info('Keeping the existing export of %s at %s', snapshot.id, - location) - return location - if existing is not None: + def upload(client: FileSystem, name: str, stop: threading.Event) -> None: + size = len(index_manifest) if name == _INDEX_MANIFEST else sizes[name] + with open(source / name, 'rb', buffering=0) as reader, _open_output( + client, f'{root}/index/{name}') 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) + + def send(names: List[str], action: str) -> None: + names.sort(key=lambda name: -sizes.get(name, 0)) + _transfer(connector, location, names, + [sizes.get(name, len(index_manifest)) for name in names], + upload, action) + + if committed(): + if build() != fingerprint: + logger.info('Keeping the existing export of %s at %s', + snapshot.id, location) + return location uploaded = present() - stale = [file for file in files if uploaded.get(file[0]) != file[1]] - if stale: - _parallel(connector, location, stale, - [size for _, size in stale], upload, 'repair') + send([ + name for name, size in { + **sizes, _INDEX_MANIFEST: len(index_manifest) + }.items() if uploaded.get(name) != size + ], 'repair') return location - marker = _read_json(filesystem, f'{root}/{_EXPORTING}') - if not isinstance(marker, (dict, type(None))): - marker = {} - if marker is None and filesystem.get_file_info( + existing = recorded() + if existing is None and filesystem.get_file_info( FileSelector(root, allow_not_found=True)): raise ValueError(f'{location} is not empty; remove it and re-export.') - if marker is not None and marker.get('snapshot_id') != snapshot.id: - raise ValueError(f'{location} holds a partial export of another ' - 'snapshot; remove it and re-export.') - resume = (marker or {}).get('index_manifest') == fingerprint - _write_json(filesystem, f'{root}/{_EXPORTING}', identity) - uploaded = present() if resume else {} - crcs = _parallel(connector, location, files, [size for _, size in files], - upload, 'export') + resume = (existing or {}).get('index_manifest') == fingerprint + recorded_snapshot = snapshot.manifest + metadata = { + 'format': _FORMAT, + 'snapshot_id': snapshot.id, + 'graph_key': graph_key, + 'data_version': data_version, + 'index_manifest': fingerprint, + 'graph_config': snapshot.graph_config, + 'snapshot': None if recorded_snapshot is None else + {k: v + for k, v in recorded_snapshot.items() if k != 'loaded'}, + 'environment': environment(), + 'range_bytes': _RANGE_BYTES, + } + _write(filesystem, f'{root}/{_MANIFEST}', json.dumps(metadata).encode()) + with ThreadPoolExecutor(1) as pool: + checksums = pool.submit(snapshot_range_checksums, str(source), + _RANGE_BYTES) + uploaded = present() if resume else {} + send([ + name for name, size in sizes.items() if uploaded.get(name) != size + ], 'export') + metadata['crc32'] = checksums.result() filesystem, root = _location(connector, location) - if finished(): - return location - if (_read_json(filesystem, f'{root}/{_EXPORTING}') != identity - and not finished()): - raise RuntimeError(f'Another export to {location} started ' - 'meanwhile; re-export.') - recorded = snapshot.manifest - _write_json( - filesystem, f'{root}/{_MANIFEST}', { - 'format': _FORMAT, - 'snapshot_id': snapshot.id, - 'graph_key': graph_key, - 'data_version': data_version, - 'index_manifest': fingerprint, - 'graph_config': snapshot.graph_config, - 'snapshot': None if recorded is None else - {k: v for k, v in recorded.items() if k != 'loaded'}, - 'environment': environment(), - 'range_bytes': _RANGE_BYTES, - 'files': [{'path': name, 'size': size, 'crc32': file_crcs} - for (name, size), file_crcs in zip(files, crcs)], - }) - try: - filesystem.delete_file(f'{root}/{_EXPORTING}') - except FileNotFoundError: - pass + if build() != fingerprint: + raise RuntimeError(f'Another build of {snapshot.id} was exported to ' + f'{location} meanwhile; re-export.') + if not committed(): + _write(filesystem, f'{root}/{_MANIFEST}', json.dumps(metadata).encode()) + _write(filesystem, f'{root}/index/{_INDEX_MANIFEST}', index_manifest) + _write(filesystem, f'{root}/index/{_SUCCESS}', fingerprint.encode()) return location @@ -279,9 +283,10 @@ def _corrupt(path: str, detail: str) -> ValueError: def read(connector: Any, path: str) -> Tuple[str, dict]: filesystem, root = _location(connector, path) manifest = _read_json(filesystem, f'{root}/{_MANIFEST}') - if manifest is None: - raise FileNotFoundError(f'No complete snapshot export at {path} ' - f'({_MANIFEST} is missing); re-run export.') + if manifest is None or not _files(filesystem).exists( + f'{root}/index/{_SUCCESS}'): + 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: @@ -291,31 +296,18 @@ def read(connector: Any, path: str) -> Tuple[str, dict]: UserWarning, stacklevel=3) _check_key(manifest.get('graph_key'), manifest.get('data_version'), manifest.get('snapshot_id')) - span, files = manifest.get('range_bytes'), manifest.get('files') - if (type(span) is not int or span <= 0 or not isinstance(files, list) + span, crcs = manifest.get('range_bytes'), manifest.get('crc32') + 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(manifest.get('snapshot'), (dict, type(None)))): raise _corrupt(path, 'malformed manifest') - seen = set() - for entry in files: - name, size, crcs = ((entry.get('path'), entry.get('size'), - entry.get('crc32')) - if isinstance(entry, dict) else (None, None, None)) - if (not isinstance(name, str) or name.startswith('/') or '\\' in name - or any(part in ('', '.', '..') for part in name.split('/')) - or name in seen or type(size) is not int or size < 0 - or not isinstance(crcs, list) - or any(type(crc) is not int for crc in crcs) - or len(crcs) != max(-(-size // span), 1)): - raise _corrupt(path, f'bad file entry {name!r}') - seen.add(name) return root, manifest def _download(root: str, staging: Path, client: FileSystem, - item: Tuple[str, int, int, int], stop: threading.Event) -> None: - name, offset, length, expected = item - crc = 0 + 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) @@ -326,60 +318,96 @@ def _download(root: str, staging: Path, client: FileSystem, chunk = reader.read_at(want, start) if len(chunk) != want: raise ValueError(f'{name} is truncated') - crc = zlib.crc32(chunk, crc) os.pwrite(fd, chunk, start) getattr(os, 'fdatasync', os.fsync)(fd) finally: os.close(fd) - if crc != expected: - raise ValueError(f'{name} failed its checksum at offset {offset}') + + +def _index_files(path: str, root: str, manifest: dict, + connector: Any) -> Tuple[bytes, List[Tuple[str, int]]]: + 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() + 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') + return index_manifest, files def load_into(connector: Any, path: str, root: str, manifest: dict, target: Path) -> None: - from diskgraph import format_versions, verify_snapshot_integrity + from diskgraph import ( + DiskgraphCorruptionError, + format_versions, + snapshot_range_checksums, + verify_snapshot_integrity, + ) + from kumo_core.artifacts import FilesystemArtifactStore from kumo_core.diskgraph.snapshot import ( DiskgraphIngestResult, write_diskgraph_success_marker, ) + from kumo_core.sources.transfer import discard_destination from kumo.engine._manifest import write_manifest from kumo.engine._provenance import environment from kumo.engine.graph import SNAPSHOT_MANIFEST_FILE span = manifest['range_bytes'] - files = [(entry['path'], entry['size'], entry['crc32']) - for entry in manifest['files']] + index_manifest, files = _index_files(path, root, manifest, connector) staging = target.with_name(f'{target.name}.loading') - shutil.rmtree(staging, ignore_errors=True) + discard_destination(str(target)) + discard_destination(str(staging)) try: staging.mkdir(parents=True) - # Any old target stays until the rename. - needed = sum(size for _, size, _ in files) + 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: + 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: - os.ftruncate(fd, size) + if size: + os.posix_fallocate(fd, 0, size) finally: os.close(fd) - ranges = sorted( - ((name, offset, min(span, size - offset), crcs[index]) - for name, size, crcs in files - for index, offset in enumerate(range(0, size, span))), - key=lambda item: -item[2]) + 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: - _parallel(connector, path, ranges, [item[2] for item in ranges], + _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)) - verify_snapshot_integrity(str(staging)) + 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) if manifest.get('snapshot') is not None: write_manifest( @@ -390,13 +418,12 @@ def load_into(connector: Any, path: str, root: str, manifest: dict, } }) for directory in sorted({(staging / name).parent - for name, _, _ in files} | {staging}): - _fsync(directory) - shutil.rmtree(target, ignore_errors=True) + for name, _ in files} | {staging}): + FilesystemArtifactStore._fsync_directory(directory) os.rename(staging, target) - _fsync(target.parent) + FilesystemArtifactStore._fsync_directory(target.parent) except BaseException: - shutil.rmtree(staging, ignore_errors=True) + discard_destination(str(staging)) raise @@ -411,11 +438,3 @@ def _check_readable(index: Path, manifest: dict) -> None: f'This DiskGraph cannot open the snapshot: built with formats ' f'{built}, installed are {format_versions()}. Rebuild it with ' 'this kumo-engine or load it with the exporting one.') from error - - -def _fsync(path: Path) -> None: - fd = os.open(path, os.O_RDONLY) - try: - os.fsync(fd) - finally: - os.close(fd) diff --git a/packages/kumo-engine/kumo/engine/graph.py b/packages/kumo-engine/kumo/engine/graph.py index 7d0835559..88427fe77 100644 --- a/packages/kumo-engine/kumo/engine/graph.py +++ b/packages/kumo-engine/kumo/engine/graph.py @@ -60,10 +60,10 @@ """ import json -import os +import threading import warnings from dataclasses import dataclass, field, replace -from functools import cached_property +from functools import cached_property, partial from pathlib import Path from typing import Any, Dict, Mapping, Optional, Sequence @@ -80,7 +80,7 @@ stamp, write_manifest, ) -from kumo.engine._provenance import source_location, source_table +from kumo.engine._provenance import source_table from kumo.engine._scratch import ( fingerprint, get_scratch_dir, @@ -223,26 +223,31 @@ def _lock(graph_key: str) -> Any: return FileLock(f'{persist_dir}.lock') -def _loaded(path: str) -> Optional[Any]: - from diskgraph import DiskgraphCorruptionError, verify_snapshot_integrity - from kumo_core.diskgraph.snapshot import check_diskgraph_output +def _loaded(store: Any, handle: Any) -> Optional[Any]: + from diskgraph import DiskgraphCorruptionError, read_success_marker + from kumo_core.diskgraph.snapshot import rebased - marker = check_diskgraph_output(path) - if marker is None or (os.path.realpath(marker.index_path) - != os.path.realpath(path)): + path = store.path(handle) + if read_success_marker(path) is None: return None try: - verify_snapshot_integrity(path) - except DiskgraphCorruptionError as error: + 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 - return marker + + +def _snapshot_id(graph_key: str, data_version: str) -> str: + return f'{graph_key}@{data_version[:16]}' def _comparable(stamp: Any) -> bool: - # Views and unversioned sources are stamped on the hour; -1 is undated. - return (isinstance(stamp, str) and stamp != '-1' + from kumo_core.sources.object_store import UNDATED_VERSION + + # Views and unversioned sources are stamped on the hour. + return (isinstance(stamp, str) and stamp != UNDATED_VERSION and not (stamp.isdigit() and int(stamp) % 3600 == 0)) @@ -255,14 +260,13 @@ def _source_change(table: Table, recorded: Any) -> Dict[str, Any]: source = table.definition.source_table if not isinstance(source, CustomSQLSourceTable): connector = table.connector - kind = connector.source_type - same_source = (recorded.get('connector'), - recorded.get('location')) == ( - str(getattr(kind, 'value', kind)), - source_location(connector, table.table_name)) stamp = connector.source_backend.get_data_version( connector, table.table_name) - current = None if stamp is None else str(stamp) + now = source_table(connector, table.table_name, None, stamp) + current = now['version'] + same_source = (recorded.get('connector'), + recorded.get('location')) == (now['connector'], + now['location']) except Exception: current = None if not (same_source and _comparable(before) and _comparable(current)): @@ -307,6 +311,17 @@ class GraphSnapshot: 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 Path(self.index_path).parent.name + + @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 + def export(self, connector: Connector, path: str) -> str: r"""Copy the index to `//` and return that location; re-running resumes. @@ -338,15 +353,15 @@ def load_from_connector(cls, connector: Connector, *, path: str, if graph is not None and graph.id != graph_key: warnings.warn(f"Snapshot at {path} was taken of another graph.", UserWarning, stacklevel=2) - target = _store().path( - DiskgraphSnapshotHandle( - graph_key=graph_key, - graph_data_version=manifest['data_version'])) + store = _store() + handle = DiskgraphSnapshotHandle( + graph_key=graph_key, graph_data_version=manifest['data_version']) with _lock(graph_key): - result = _loaded(target) + result = _loaded(store, handle) if result is None: - load_into(connector, path, root, manifest, Path(target)) - result = _loaded(target) + load_into(connector, path, root, manifest, + Path(store.path(handle))) + result = _loaded(store, handle) snapshot = cls(id=manifest['snapshot_id'], index_path=Path(result.index_path), graph_config=manifest['graph_config']) @@ -372,10 +387,15 @@ def source_changes(self, graph: 'Graph') -> Dict[str, Dict[str, Any]]: dict(manifest['parents']['source_tables'])) except Exception: recorded = {} - return { - alias: _source_change(table, recorded.get(alias)) - for alias, table in graph._named_tables.items() - } + from kumo_core.sources.parallel import run_in_parallel + + tables = graph._named_tables + changes = run_in_parallel( + [ + partial(_source_change, table, recorded.get(alias)) + 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: @@ -713,7 +733,7 @@ def _snapshot_no_lock(self, graph_key: str, persist_dir: Path, ), ) - snapshot_id = f'{graph_key}@{data_version[:16]}' + snapshot_id = _snapshot_id(graph_key, data_version) _write_snapshot_manifest(result, snapshot_id, source_tables, record_sources, fresh=existing is None) diff --git a/packages/kumo-engine/test/test_snapshot_export.py b/packages/kumo-engine/test/test_snapshot_export.py index 9a08c1e02..be75eeeaa 100644 --- a/packages/kumo-engine/test/test_snapshot_export.py +++ b/packages/kumo-engine/test/test_snapshot_export.py @@ -2,6 +2,7 @@ # All rights reserved. # SPDX-License-Identifier: Apache-2.0 +import hashlib import json import os import shutil @@ -62,6 +63,16 @@ def _edit_manifest(exported, edit): path.write_text(json.dumps(manifest)) +def _edit_index_manifest(exported, edit): + path = 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) @@ -126,25 +137,24 @@ def four(found, worker_ceiling=None): monkeypatch.setattr('kumo_core.sources.object_store.copies', four) snapshot = graph.snapshot() - stream, finished = snapshot_io._stream, threading.Event() + open_output, finished = snapshot_io._open_output, threading.Event() - def interrupted(source, size, stop, write=None): - if write is not None and source.name == 'diskgraph_config.json': + def interrupted(client, path): + if path.endswith('diskgraph_config.json'): finished.wait(10) raise OSError('connection reset') - crcs = stream(source, size, stop, write) - if write is not None and source.name == 'col_stats.json': + if path.endswith('col_stats.json'): finished.set() - return crcs + return open_output(client, path) with monkeypatch.context() as patch: - patch.setattr(snapshot_io, '_stream', interrupted) + patch.setattr(snapshot_io, '_open_output', interrupted) patch.setattr(snapshot_io, '_ATTEMPTS', 1) with pytest.raises(OSError, match='connection reset'): snapshot.export(exports, 'one') graph_id, data_version = snapshot.id.split('@') location = f'one/{graph_id}/{data_version}' - with pytest.raises(FileNotFoundError, match='kumo_snapshot.json'): + with pytest.raises(FileNotFoundError, match='No complete snapshot'): _load(exports, location) assert snapshot.export(exports, 'one') == location scratch('gpu') @@ -155,15 +165,15 @@ def interrupted(source, size, stop, write=None): def test_concurrent_exports_of_one_snapshot_both_succeed( tmp_path, scratch, graph, exports, monkeypatch): snapshot = graph.snapshot() - parallel, raced = snapshot_io._parallel, [] + transfer, raced = snapshot_io._transfer, [] def racing(*args): if not raced: raced.append(None) raced[0] = snapshot.export(exports, 'one') - return parallel(*args) + return transfer(*args) - monkeypatch.setattr(snapshot_io, '_parallel', racing) + monkeypatch.setattr(snapshot_io, '_transfer', racing) assert raced == [snapshot.export(exports, 'one')] @@ -230,19 +240,20 @@ def test_load_refuses_before_downloading_when_disk_is_too_small( _load(exports, location) -@pytest.mark.parametrize('edit', [ - lambda manifest: manifest.update(data_version='../../victim'), - lambda manifest: manifest['files'][0].update(path='../../victim/x'), - lambda manifest: manifest.pop('files'), - lambda manifest: manifest['files'].append(manifest['files'][0]), +@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, edit): + exports, index, edit): location, exported = _export(graph.snapshot(), exports, tmp_path) victim = tmp_path / 'victim' victim.mkdir() (victim / 'keep').write_text('x') - _edit_manifest(exported, edit) + (_edit_index_manifest if index else _edit_manifest)(exported, edit) scratch('gpu/a/b') with pytest.raises(ValueError): _load(exports, location) @@ -254,12 +265,8 @@ def test_load_warns_instead_of_refusing_newer_or_foreign_exports( snapshot = graph.snapshot() location, exported = _export(snapshot, exports, tmp_path) - def newer(manifest): - manifest.update(format=2, added_later={'x': 1}) - for entry in manifest['files']: - entry['added_later'] = 1 - - _edit_manifest(exported, newer) + _edit_manifest(exported, + lambda m: m.update(format=2, added_later={'x': 1})) other = kumo.Graph(tables={'users': graph.tables['users']}) scratch('gpu') with pytest.warns(UserWarning) as caught: @@ -313,9 +320,7 @@ def test_load_explains_an_index_the_installed_diskgraph_cannot_open( metadata.write_bytes(garbage) def unreadable(manifest): - for entry in manifest['files']: - if entry['path'] == 'index_metadata.json': - entry['crc32'] = [zlib.crc32(garbage)] + manifest['crc32']['index_metadata.json'] = [zlib.crc32(garbage)] manifest['snapshot']['formats'] = {'layout': -1} _edit_manifest(exported, unreadable) From 66fb6d29c3ba4492e3791ecc79302bb24842af4f Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Wed, 21 Oct 2026 18:45:54 +0000 Subject: [PATCH 58/80] Check loaded snapshots with what DiskGraph already stamps The rebased snapshot manifest no longer records format versions and DiskGraph no longer exposes format_versions(). A loaded snapshot's manifest now records only the loading environment; an index this DiskGraph cannot open surfaces SamplerEngine's own error with the exporting kumo-engine version, and directory fsyncs use kumo-core's fsync_directory. Verified: test_snapshot_export + test_snapshot_manifest (42 passed). Co-Authored-By: Claude Opus 5.5 --- .../kumo-engine/kumo/engine/_snapshot_io.py | 20 +++++++++---------- .../kumo-engine/test/test_snapshot_export.py | 7 +++---- 2 files changed, 12 insertions(+), 15 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/_snapshot_io.py b/packages/kumo-engine/kumo/engine/_snapshot_io.py index b09db4f37..7982bb51d 100644 --- a/packages/kumo-engine/kumo/engine/_snapshot_io.py +++ b/packages/kumo-engine/kumo/engine/_snapshot_io.py @@ -348,15 +348,14 @@ def load_into(connector: Any, path: str, root: str, manifest: dict, target: Path) -> None: from diskgraph import ( DiskgraphCorruptionError, - format_versions, snapshot_range_checksums, verify_snapshot_integrity, ) - from kumo_core.artifacts import FilesystemArtifactStore 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._manifest import write_manifest @@ -413,28 +412,27 @@ def load_into(connector: Any, path: str, root: str, manifest: dict, write_manifest( staging / SNAPSHOT_MANIFEST_FILE, { **manifest['snapshot'], 'loaded': { - 'formats': format_versions(), - 'environment': environment(), + 'environment': environment() } }) for directory in sorted({(staging / name).parent for name, _ in files} | {staging}): - FilesystemArtifactStore._fsync_directory(directory) + fsync_directory(directory) os.rename(staging, target) - FilesystemArtifactStore._fsync_directory(target.parent) + fsync_directory(target.parent) except BaseException: discard_destination(str(staging)) raise def _check_readable(index: Path, manifest: dict) -> None: - from diskgraph import SamplerEngine, format_versions + from diskgraph import SamplerEngine try: SamplerEngine(str(index), num_threads=1) except Exception as error: - built = (manifest.get('snapshot') or {}).get('formats') + built = (manifest.get('environment') or {}).get('kumo_engine') raise RuntimeError( - f'This DiskGraph cannot open the snapshot: built with formats ' - f'{built}, installed are {format_versions()}. Rebuild it with ' - 'this kumo-engine or load it with the exporting one.') from error + 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/test/test_snapshot_export.py b/packages/kumo-engine/test/test_snapshot_export.py index be75eeeaa..446dd4b77 100644 --- a/packages/kumo-engine/test/test_snapshot_export.py +++ b/packages/kumo-engine/test/test_snapshot_export.py @@ -16,7 +16,6 @@ import pandas as pd import pytest -from diskgraph import format_versions from pyarrow.fs import LocalFileSystem, SubTreeFileSystem import kumo.engine as kumo @@ -98,7 +97,7 @@ def test_export_load_round_trip_relocates_and_is_reused( assert loaded.index_path.is_relative_to(tmp_path / 'gpu') assert _index_files(loaded.index_path) == _index_files(exported / 'index') recorded = loaded.manifest - assert recorded.pop('loaded')['formats'] == format_versions() + assert set(recorded.pop('loaded')) == {'environment'} assert recorded == manifest['snapshot'] reexported = tmp_path / 'exports' / loaded.export(exports, 'two') assert json.loads((reexported / 'kumo_snapshot.json').read_text())[ @@ -321,11 +320,11 @@ def test_load_explains_an_index_the_installed_diskgraph_cannot_open( def unreadable(manifest): manifest['crc32']['index_metadata.json'] = [zlib.crc32(garbage)] - manifest['snapshot']['formats'] = {'layout': -1} + manifest['environment']['kumo_engine'] = '0.0.1' _edit_manifest(exported, unreadable) scratch('gpu') - with pytest.raises(RuntimeError, match=r"formats \{'layout': -1\}"): + with pytest.raises(RuntimeError, match='exported by kumo-engine 0.0.1'): _load(exports, location) From 6bc6d6cad8e11dcd62a672c0d32d26eab01fac14 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Wed, 21 Oct 2026 18:47:07 +0000 Subject: [PATCH 59/80] Derive GraphSnapshot.graph_key from the snapshot id A snapshot built or loaded by this package sits at / in the store, but a GraphSnapshot constructed elsewhere need not, so read the graph id from the id itself. Export still checks that the id and the index directory agree. Verified: test_snapshot_export (23 passed). Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/graph.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/kumo-engine/kumo/engine/graph.py b/packages/kumo-engine/kumo/engine/graph.py index 88427fe77..2bb43bc29 100644 --- a/packages/kumo-engine/kumo/engine/graph.py +++ b/packages/kumo-engine/kumo/engine/graph.py @@ -314,7 +314,7 @@ class GraphSnapshot: @property def graph_key(self) -> str: """The `Graph.id` this snapshot was built from.""" - return Path(self.index_path).parent.name + return self.id.partition('@')[0] @property def data_version(self) -> str: From 51f5e8c321c87e43c396db09155545b49932b4a1 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Wed, 21 Oct 2026 21:14:51 +0000 Subject: [PATCH 60/80] Make snapshot export atomic and pass the snapshot manifest through Export resumed by file name and size and reused an existing export by a hash of the file listing, so an interrupted export resumed against a different build could commit mismatched files. An unknown export format was read as format 1, a newer snapshot manifest was dropped on export, and source_changes guessed hour buckets from round timestamps. Export now removes kumo_snapshot.json and the rest of the location first, uploads every file, and writes kumo_snapshot.json with the range CRCs last; there is no resume, reuse or index/_SUCCESS. snapshot-manifest.json is copied byte for byte (a reconstructed one is written when the snapshot has none) and restored unchanged on load. Loads reject unknown formats. source_changes compares only exact or timestamp versions, using the snapshot's record_locations choice. Verified with the snapshot export, manifest and kumo-s3 tests. Co-Authored-By: Claude Opus 5.5 --- .../src/kumo_core/sources/object_store.py | 4 +- .../kumo-engine/kumo/engine/_snapshot_io.py | 165 +++++++----------- packages/kumo-engine/kumo/engine/graph.py | 27 +-- .../kumo-engine/test/test_snapshot_export.py | 108 ++++++------ 4 files changed, 139 insertions(+), 165 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 701675ab3..bf1058d07 100644 --- a/packages/kumo-core/src/kumo_core/sources/object_store.py +++ b/packages/kumo-core/src/kumo_core/sources/object_store.py @@ -68,8 +68,6 @@ from .types import ReadMode, StagedTable STREAMED_ONLY = staged_first(False) -#: The version of a source with no modification time. -UNDATED_VERSION = "-1" logger = logging.getLogger(__name__) @@ -104,7 +102,7 @@ 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 (SourceVersion(epoch_second(latest), "timestamp") - if latest is not None else SourceVersion(UNDATED_VERSION)) + 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), diff --git a/packages/kumo-engine/kumo/engine/_snapshot_io.py b/packages/kumo-engine/kumo/engine/_snapshot_io.py index 7982bb51d..4f0b7ac97 100644 --- a/packages/kumo-engine/kumo/engine/_snapshot_io.py +++ b/packages/kumo-engine/kumo/engine/_snapshot_io.py @@ -9,18 +9,16 @@ import re import threading import time -import warnings from concurrent.futures import ThreadPoolExecutor from functools import partial from pathlib import Path -from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple +from typing import Any, Callable, List, Optional, Sequence, Tuple -from pyarrow.fs import FileSelector, FileSystem, FileType +from pyarrow.fs import FileSystem _FORMAT = 1 _MANIFEST = 'kumo_snapshot.json' _INDEX_MANIFEST = '_MANIFEST.json' -_SUCCESS = '_SUCCESS' _RANGE_BYTES = 64 << 20 _CHUNK = 8 << 20 _ATTEMPTS = 3 @@ -165,10 +163,26 @@ def _fingerprint(index_manifest: bytes) -> str: return hashlib.sha256(index_manifest).hexdigest() +def _snapshot_manifest(snapshot: Any) -> bytes: + from kumo.engine._manifest import new_manifest + from kumo.engine.graph import SNAPSHOT_MANIFEST_FILE + + try: + return (Path(snapshot.index_path) / SNAPSHOT_MANIFEST_FILE).read_bytes() + except OSError: + manifest = new_manifest('graph_snapshot', snapshot.id, + {'source_tables': {}}, created_at=None, + environment=None, scheduler=None, + reconstructed=True) + return (json.dumps(manifest, indent=2, sort_keys=True) + + '\n').encode() + + 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 @@ -179,35 +193,14 @@ def export(snapshot: Any, connector: Any, path: str) -> str: entry['rel_path']: int(entry['size_bytes']) for entry in json.loads(index_manifest)['files'] } - fingerprint = _fingerprint(index_manifest) + snapshot_manifest = _snapshot_manifest(snapshot) filesystem, root = _location(connector, location) - - def recorded() -> Optional[dict]: - existing = _read_json(filesystem, f'{root}/{_MANIFEST}') - found = (existing.get('snapshot_id') - if isinstance(existing, dict) else None) - if existing is not None and found != snapshot.id: - raise ValueError(f'{location} holds snapshot {found!r}, not ' - f'{snapshot.id!r}; remove it and re-export.') - return existing - - def build() -> Optional[str]: - return (recorded() or {}).get('index_manifest') - - def committed() -> bool: - return _files(filesystem).exists(f'{root}/index/{_SUCCESS}') - - def present() -> Dict[str, int]: - return { - info.path[len(root) + len('/index/'):]: info.size - for info in filesystem.get_file_info( - FileSelector(f'{root}/index', recursive=True, - allow_not_found=True)) - if info.type is FileType.File - } + if _files(filesystem).exists(f'{root}/{_MANIFEST}'): + filesystem.delete_file(f'{root}/{_MANIFEST}') + filesystem.delete_dir_contents(root, missing_dir_ok=True) def upload(client: FileSystem, name: str, stop: threading.Event) -> None: - size = len(index_manifest) if name == _INDEX_MANIFEST else sizes[name] + size = sizes[name] with open(source / name, 'rb', buffering=0) as reader, _open_output( client, f'{root}/index/{name}') as writer: for offset in range(0, size, _CHUNK): @@ -218,60 +211,29 @@ def upload(client: FileSystem, name: str, stop: threading.Event) -> None: 'exported') writer.write(chunk) - def send(names: List[str], action: str) -> None: - names.sort(key=lambda name: -sizes.get(name, 0)) - _transfer(connector, location, names, - [sizes.get(name, len(index_manifest)) for name in names], - upload, action) - - if committed(): - if build() != fingerprint: - logger.info('Keeping the existing export of %s at %s', - snapshot.id, location) - return location - uploaded = present() - send([ - name for name, size in { - **sizes, _INDEX_MANIFEST: len(index_manifest) - }.items() if uploaded.get(name) != size - ], 'repair') - return location - existing = recorded() - if existing is None and filesystem.get_file_info( - FileSelector(root, allow_not_found=True)): - raise ValueError(f'{location} is not empty; remove it and re-export.') - resume = (existing or {}).get('index_manifest') == fingerprint - recorded_snapshot = snapshot.manifest - metadata = { - 'format': _FORMAT, - 'snapshot_id': snapshot.id, - 'graph_key': graph_key, - 'data_version': data_version, - 'index_manifest': fingerprint, - 'graph_config': snapshot.graph_config, - 'snapshot': None if recorded_snapshot is None else - {k: v - for k, v in recorded_snapshot.items() if k != 'loaded'}, - 'environment': environment(), - 'range_bytes': _RANGE_BYTES, - } - _write(filesystem, f'{root}/{_MANIFEST}', json.dumps(metadata).encode()) + names = sorted(sizes, key=lambda name: -sizes[name]) with ThreadPoolExecutor(1) as pool: checksums = pool.submit(snapshot_range_checksums, str(source), _RANGE_BYTES) - uploaded = present() if resume else {} - send([ - name for name, size in sizes.items() if uploaded.get(name) != size - ], 'export') - metadata['crc32'] = checksums.result() + _transfer(connector, location, names, [sizes[name] for name in names], + upload, 'export') + crc32 = checksums.result() filesystem, root = _location(connector, location) - if build() != fingerprint: - raise RuntimeError(f'Another build of {snapshot.id} was exported to ' - f'{location} meanwhile; re-export.') - if not committed(): - _write(filesystem, f'{root}/{_MANIFEST}', json.dumps(metadata).encode()) - _write(filesystem, f'{root}/index/{_INDEX_MANIFEST}', index_manifest) - _write(filesystem, f'{root}/index/{_SUCCESS}', fingerprint.encode()) + _write(filesystem, f'{root}/index/{_INDEX_MANIFEST}', index_manifest) + _write(filesystem, f'{root}/{SNAPSHOT_MANIFEST_FILE}', snapshot_manifest) + _write( + filesystem, f'{root}/{_MANIFEST}', + json.dumps({ + 'format': _FORMAT, + 'snapshot_id': snapshot.id, + 'graph_key': graph_key, + 'data_version': data_version, + 'index_manifest': _fingerprint(index_manifest), + 'graph_config': snapshot.graph_config, + 'environment': environment(), + 'range_bytes': _RANGE_BYTES, + 'crc32': crc32, + }).encode()) return location @@ -283,24 +245,22 @@ def _corrupt(path: str, detail: str) -> ValueError: def read(connector: Any, path: str) -> Tuple[str, dict]: filesystem, root = _location(connector, path) manifest = _read_json(filesystem, f'{root}/{_MANIFEST}') - if manifest is None or not _files(filesystem).exists( - f'{root}/index/{_SUCCESS}'): + 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: - warnings.warn( - f"Snapshot export at {path} has layout format " - f"{manifest.get('format')!r}; reading it as format {_FORMAT}.", - UserWarning, stacklevel=3) + 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') 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(manifest.get('snapshot'), (dict, type(None)))): + or not isinstance(manifest.get('graph_config'), dict)): raise _corrupt(path, 'malformed manifest') return root, manifest @@ -324,10 +284,13 @@ def _download(root: str, staging: Path, client: FileSystem, os.close(fd) -def _index_files(path: str, root: str, manifest: dict, - connector: Any) -> Tuple[bytes, List[Tuple[str, int]]]: +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: @@ -341,7 +304,13 @@ def _index_files(path: str, root: str, manifest: dict, 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') - return index_manifest, files + 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, @@ -358,12 +327,11 @@ 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._manifest import write_manifest - from kumo.engine._provenance import environment from kumo.engine.graph import SNAPSHOT_MANIFEST_FILE span = manifest['range_bytes'] - index_manifest, files = _index_files(path, root, manifest, connector) + 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)) @@ -408,13 +376,8 @@ def load_into(connector: Any, path: str, root: str, manifest: dict, raise _corrupt( path, f'{name} failed its checksum at offset {offset}') _check_readable(staging, manifest) - if manifest.get('snapshot') is not None: - write_manifest( - staging / SNAPSHOT_MANIFEST_FILE, { - **manifest['snapshot'], 'loaded': { - 'environment': environment() - } - }) + 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) diff --git a/packages/kumo-engine/kumo/engine/graph.py b/packages/kumo-engine/kumo/engine/graph.py index 2bb43bc29..ea45f5c26 100644 --- a/packages/kumo-engine/kumo/engine/graph.py +++ b/packages/kumo-engine/kumo/engine/graph.py @@ -243,33 +243,33 @@ def _snapshot_id(graph_key: str, data_version: str) -> str: return f'{graph_key}@{data_version[:16]}' -def _comparable(stamp: Any) -> bool: - from kumo_core.sources.object_store import UNDATED_VERSION +_COMPARABLE = ('exact', 'timestamp') - # Views and unversioned sources are stamped on the hour. - return (isinstance(stamp, str) and stamp != UNDATED_VERSION - and not (stamp.isdigit() and int(stamp) % 3600 == 0)) - -def _source_change(table: Table, recorded: Any) -> Dict[str, Any]: +def _source_change(table: Table, recorded: Any, + record_locations: bool) -> Dict[str, Any]: from kumoapi.source_table import CustomSQLSourceTable recorded = recorded if isinstance(recorded, dict) else {} before, current, same_source = recorded.get('version'), None, False + comparable = recorded.get('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) + now = source_table(connector, table.table_name, None, stamp, + record_locations=record_locations) current = now['version'] + comparable = comparable and now['version_kind'] in _COMPARABLE same_source = (recorded.get('connector'), recorded.get('location')) == (now['connector'], now['location']) except Exception: current = None - if not (same_source and _comparable(before) and _comparable(current)): + 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' @@ -324,7 +324,8 @@ def data_version(self) -> str: def export(self, connector: Connector, path: str) -> str: r"""Copy the index to `//` and - return that location; re-running resumes. + return that location. Any earlier export there is replaced; it is not + loadable until the new one completes. Args: connector: an S3, GCS or file connector (on Databricks or @@ -385,14 +386,16 @@ def source_changes(self, graph: 'Graph') -> Dict[str, Dict[str, Any]]: manifest = self.manifest recorded = ({} if not manifest or manifest.get('reconstructed') else dict(manifest['parents']['source_tables'])) + record_locations = manifest.get('record_locations', True) except Exception: - recorded = {} + recorded, record_locations = {}, 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)) + partial(_source_change, table, recorded.get(alias), + record_locations) for alias, table in tables.items() ], workers=max(1, min(len(tables), 16)), stop=threading.Event()) return dict(zip(tables, changes)) diff --git a/packages/kumo-engine/test/test_snapshot_export.py b/packages/kumo-engine/test/test_snapshot_export.py index 446dd4b77..f75e2f178 100644 --- a/packages/kumo-engine/test/test_snapshot_export.py +++ b/packages/kumo-engine/test/test_snapshot_export.py @@ -86,8 +86,8 @@ def test_export_load_round_trip_relocates_and_is_reused( assert location == f'one/{graph_id}/{data_version}' assert snapshot.export(exports, 'one') == location exported = tmp_path / 'exports' / location - manifest = json.loads((exported / 'kumo_snapshot.json').read_text()) - assert manifest['snapshot'] == snapshot.manifest + recorded = (snapshot.index_path / 'snapshot-manifest.json').read_bytes() + assert (exported / 'snapshot-manifest.json').read_bytes() == recorded shutil.rmtree(tmp_path / 'cpu') scratch('gpu') @@ -96,12 +96,10 @@ def test_export_load_round_trip_relocates_and_is_reused( snapshot.graph_config) assert loaded.index_path.is_relative_to(tmp_path / 'gpu') assert _index_files(loaded.index_path) == _index_files(exported / 'index') - recorded = loaded.manifest - assert set(recorded.pop('loaded')) == {'environment'} - assert recorded == manifest['snapshot'] + assert (loaded.index_path / + 'snapshot-manifest.json').read_bytes() == recorded reexported = tmp_path / 'exports' / loaded.export(exports, 'two') - assert json.loads((reexported / 'kumo_snapshot.json').read_text())[ - 'snapshot'] == manifest['snapshot'] + assert (reexported / 'snapshot-manifest.json').read_bytes() == recorded monkeypatch.setattr( 'kumo_core.diskgraph.ingest.execute_diskgraph_ingest', @@ -128,7 +126,7 @@ def test_damaged_index_is_reused_by_snapshot_but_reloaded_by_load( assert _index_files(loaded.index_path) == expected -def test_interrupted_export_is_not_loadable_and_resumes( +def test_interrupted_reexport_is_not_loadable_and_reruns_from_scratch( tmp_path, scratch, graph, exports, monkeypatch): @contextmanager def four(found, worker_ceiling=None): @@ -136,6 +134,8 @@ def four(found, worker_ceiling=None): monkeypatch.setattr('kumo_core.sources.object_store.copies', four) snapshot = graph.snapshot() + location, exported = _export(snapshot, exports, tmp_path) + (exported / 'index/stray.bin').write_bytes(b'x') open_output, finished = snapshot_io._open_output, threading.Event() def interrupted(client, path): @@ -151,11 +151,10 @@ def interrupted(client, path): patch.setattr(snapshot_io, '_ATTEMPTS', 1) with pytest.raises(OSError, match='connection reset'): snapshot.export(exports, 'one') - graph_id, data_version = snapshot.id.split('@') - location = f'one/{graph_id}/{data_version}' with pytest.raises(FileNotFoundError, match='No complete snapshot'): _load(exports, location) assert snapshot.export(exports, 'one') == location + assert not (exported / 'index/stray.bin').exists() scratch('gpu') assert _index_files(_load(exports, location).index_path) == _index_files( snapshot.index_path) @@ -174,6 +173,7 @@ def racing(*args): monkeypatch.setattr(snapshot_io, '_transfer', racing) assert raced == [snapshot.export(exports, 'one')] + assert _load(exports, raced[0]).id == snapshot.id def test_failed_upload_is_retried_with_a_fresh_client( @@ -192,16 +192,23 @@ def expiring(client, path): assert len(clients) == 2 and clients[0] is not clients[1] -def test_rerunning_export_repairs_a_committed_export( +def test_reexport_replaces_whatever_the_location_holds( tmp_path, scratch, graph, exports): snapshot = graph.snapshot() + graph_id, data_version = snapshot.id.split('@') + occupied = tmp_path / 'exports/one' / graph_id / data_version + occupied.mkdir(parents=True) + (occupied / 'other.parquet').write_bytes(b'x') location, exported = _export(snapshot, exports, tmp_path) + assert not (exported / 'other.parquet').exists() stats = exported / 'index/col_stats.json' original = stats.read_bytes() stats.write_bytes(original[:3]) (exported / 'index/index_metadata.json').unlink() + _edit_manifest(exported, lambda m: m.update(snapshot_id='other')) assert snapshot.export(exports, 'one') == location assert stats.read_bytes() == original + assert _load(exports, location).id == snapshot.id def test_failed_load_leaves_nothing_and_restarts_cleanly( @@ -259,48 +266,46 @@ def test_load_rejects_a_hostile_or_malformed_manifest(tmp_path, scratch, graph, assert (victim / 'keep').read_text() == 'x' -def test_load_warns_instead_of_refusing_newer_or_foreign_exports( +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, added_later={'x': 1})) + _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']}) - scratch('gpu') - with pytest.warns(UserWarning) as caught: + with pytest.warns(UserWarning, match='another graph'): loaded = _load(exports, location, graph=other) - messages = ' '.join(str(warning.message) for warning in caught) - assert 'layout format 2' in messages and 'another graph' in messages assert loaded.id == snapshot.id -def test_export_refuses_a_path_holding_other_data(tmp_path, scratch, graph, - exports): - snapshot = graph.snapshot() - graph_id, data_version = snapshot.id.split('@') - occupied = tmp_path / 'exports/occupied' / graph_id / data_version - occupied.mkdir(parents=True) - (occupied / 'other.parquet').write_bytes(b'x') - with pytest.raises(ValueError, match='not empty'): - snapshot.export(exports, 'occupied') - - _, exported = _export(snapshot, exports, tmp_path) - _edit_manifest(exported, lambda m: m.update(snapshot_id='other')) - with pytest.raises(ValueError, match="snapshot 'other', not"): - snapshot.export(exports, 'one') - - -def test_rebuilt_snapshot_reexports_to_the_existing_export( - tmp_path, scratch, graph, exports): +@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 / 'snapshot-manifest.json' + if recorded == 'newer': + path.write_bytes(b'{"schema_version": 99, "future": [1]}') + else: + path.unlink() location, exported = _export(snapshot, exports, tmp_path) - committed = (exported / 'kumo_snapshot.json').read_bytes() - scratch('rebuilt') - rebuilt = graph.snapshot() - assert rebuilt.id == snapshot.id - assert rebuilt.export(exports, 'one') == location - assert (exported / 'kumo_snapshot.json').read_bytes() == committed + scratch('gpu') + with warnings.catch_warnings(): + warnings.simplefilter('ignore') + loaded = _load(exports, location) + shipped = (exported / 'snapshot-manifest.json').read_bytes() + assert (loaded.index_path / 'snapshot-manifest.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): @@ -368,20 +373,21 @@ def dated(tmp_path, graph): return data +@pytest.mark.parametrize('record_locations', [True, False]) def test_source_changes_reads_stamps_without_staging( - scratch, graph, dated, monkeypatch): + scratch, graph, dated, monkeypatch, record_locations): 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() + snapshot = graph.snapshot(record_locations=record_locations) 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': 'unknown', 'recorded': '1700002800', + 'users': {'status': 'unchanged', 'recorded': '1700002800', 'current': '1700002800'}, 'events': {'status': 'unchanged', 'recorded': '1700000123', 'current': '1700000123'}, @@ -396,16 +402,20 @@ def test_source_changes_reads_stamps_without_staging( 'status': 'unknown', 'recorded': '1700000123', 'current': None} -@pytest.mark.parametrize('edit', ['absent', 'reconstructed']) -def test_source_changes_unknown_without_a_recorded_version( +@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 / 'snapshot-manifest.json' + manifest = json.loads(path.read_text()) if edit == 'absent': path.unlink() + elif edit == 'reconstructed': + path.write_text(json.dumps({**manifest, 'reconstructed': True})) else: - path.write_text( - json.dumps({**json.loads(path.read_text()), 'reconstructed': True})) + 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'} From e910bc2773d240933ababbb2409650e7237780a4 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Wed, 21 Oct 2026 22:25:43 +0000 Subject: [PATCH 61/80] Drop the S3 upload tuning from snapshot export Export passed allow_delayed_open and background_writes to S3 (with a pyarrow capability probe) and rebuilt every anonymous S3 source read with an explicit retry strategy. Neither was measured to help export, and the anonymous change touched every S3 read, so both go: ObjectStoreSinkBackend.filesystem() returns the sink's own filesystem. Verified with kumo-s3 tests and test_snapshot_export.py. Co-Authored-By: Claude Opus 5.5 --- .../src/kumo_core/object_store_sink.py | 5 +- packages/kumo-s3/src/kumo_s3/backend.py | 29 ++++----- packages/kumo-s3/src/kumo_s3/sink.py | 21 ------- packages/kumo-s3/tests/test_sink.py | 60 ------------------- 4 files changed, 12 insertions(+), 103 deletions(-) delete mode 100644 packages/kumo-s3/tests/test_sink.py 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 e4a42be99..93eed57ce 100644 --- a/packages/kumo-core/src/kumo_core/object_store_sink.py +++ b/packages/kumo-core/src/kumo_core/object_store_sink.py @@ -100,10 +100,9 @@ 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, *, - parallel_uploads: bool = False) -> Tuple[FileSystem, str]: + def filesystem(self, connector: SourceConnector, + path: str) -> Tuple[FileSystem, str]: """The authenticated filesystem and resolved ``path``.""" - del parallel_uploads filesystem, resolved = self._destination(connector, path) return filesystem, resolved.rstrip('/') diff --git a/packages/kumo-s3/src/kumo_s3/backend.py b/packages/kumo-s3/src/kumo_s3/backend.py index 78e9c5390..dfd31df9d 100644 --- a/packages/kumo-s3/src/kumo_s3/backend.py +++ b/packages/kumo-s3/src/kumo_s3/backend.py @@ -1,7 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. # All rights reserved. # SPDX-License-Identifier: Apache-2.0 -from typing import Any, Mapping, Optional, Union +from typing import Any, Optional, Union import kumo_s3 from kumo_core.sources.errors import ReconnectRequiredError @@ -13,11 +13,7 @@ ) from kumo_core.sources.principals import IdentityRef, Principal from kumo_core.sources.types import ReadMode -from kumo_s3.session import ( - anonymous_filesystem, - arrow_connection_options, - retry_strategy, -) +from kumo_s3.session import anonymous_filesystem, arrow_connection_options from kumoapi.data_source import ( DataSourceType, S3AccessKeyAuthentication, @@ -104,9 +100,7 @@ def _source_access(self, connector: Any, bucket: str, if connection else None) def _filesystem(self, connector: Any, bucket: str, - region: Optional[str], - options: Optional[Mapping[str, Any]] = None - ) -> S3FileSystem: + region: Optional[str]) -> S3FileSystem: r"""A ready filesystem is taken as it is: it may address a service that is not AWS at all, and probing its bucket name against AWS would answer for somebody else's bucket. @@ -116,12 +110,6 @@ def _filesystem(self, connector: Any, bucket: str, connector, bucket) if isinstance(session, S3FileSystem): return session - arrow_options = { - **(arrow_connection_options(connection.endpoint_url, - connection.addressing_style) - if connection else {}), - **(options or {}) - } with kumo_s3.translated(f'Signing for {bucket}', table=bucket): _require_credentials(session) # Resolve through the package so applications/tests replacing its @@ -131,7 +119,9 @@ def _filesystem(self, connector: Any, bucket: str, default_region=(connection.region or region) if connection else region, endpoint_url=connection.endpoint_url if connection else None, - options=arrow_options or None) + options=arrow_connection_options(connection.endpoint_url, + connection.addressing_style) + if connection else None) if resolved is not None: return resolved if not allow_anonymous: @@ -151,11 +141,12 @@ def _filesystem(self, connector: Any, bucket: str, if connection is not None: return anonymous_filesystem( bucket, region=connection.region or region - or session.region_name, options=arrow_options) + or session.region_name, + options=arrow_connection_options(connection.endpoint_url, + connection.addressing_style)) from pyarrow.fs import FileSystem filesystem, _ = FileSystem.from_uri(f's3://{bucket}') - return S3FileSystem(region=filesystem.region, - retry_strategy=retry_strategy(), **(options or {})) + return filesystem class S3SourceBackend(S3Sessions, ObjectStoreSourceBackend): diff --git a/packages/kumo-s3/src/kumo_s3/sink.py b/packages/kumo-s3/src/kumo_s3/sink.py index 557551b64..5e2a5a2d8 100644 --- a/packages/kumo-s3/src/kumo_s3/sink.py +++ b/packages/kumo-s3/src/kumo_s3/sink.py @@ -2,9 +2,6 @@ # All rights reserved. # SPDX-License-Identifier: Apache-2.0 -from functools import lru_cache -from typing import Dict - from kumo_core.object_store_sink import ObjectStoreSinkBackend from kumoapi.data_source import DataSourceType @@ -17,23 +14,5 @@ class S3SinkBackend(S3Sessions, ObjectStoreSinkBackend): """ source_type = DataSourceType.S3 - def filesystem(self, connector, path, *, parallel_uploads=False): - bucket, resolved, region = self._resolved(connector, path) - options = _parallel_upload_options() if parallel_uploads else None - return (self._filesystem(connector, bucket, region, options), - resolved.rstrip('/')) - def __repr__(self) -> str: return f'{type(self).__name__}(identity={self._identity!r})' - - -@lru_cache(maxsize=None) -def _parallel_upload_options() -> Dict[str, bool]: - from pyarrow.fs import S3FileSystem - - try: - S3FileSystem(anonymous=True, region='us-east-1', - allow_delayed_open=True) - except TypeError: - return {'background_writes': False} - return {'allow_delayed_open': True, 'background_writes': False} diff --git a/packages/kumo-s3/tests/test_sink.py b/packages/kumo-s3/tests/test_sink.py deleted file mode 100644 index c143321df..000000000 --- a/packages/kumo-s3/tests/test_sink.py +++ /dev/null @@ -1,60 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. -# All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -from __future__ import annotations - -from types import SimpleNamespace - -import kumo_s3 -import pytest -from fakes import FakeSession -from kumo_s3 import backend, sink -from kumo_s3.sink import S3SinkBackend -from kumoapi.data_source import ( - S3AnyURIConnectorTarget, - S3ConnectorResourceConfig, -) - - -@pytest.mark.parametrize('delayed_open', [True, False]) -@pytest.mark.parametrize('credentials', ['default', None]) -def test_parallel_uploads_reach_every_filesystem_construction( - monkeypatch, credentials, delayed_open): - known = { - 'access_key', 'secret_key', 'session_token', 'region', 'anonymous', - 'retry_strategy', 'background_writes', 'endpoint_override', 'scheme', - 'force_virtual_addressing' - } | ({'allow_delayed_open'} if delayed_open else set()) - built = [] - - class Record: - def __init__(self, **kwargs): - if set(kwargs) - known: - raise TypeError(f'unexpected {set(kwargs) - known}') - built.append(kwargs) - - monkeypatch.setattr('pyarrow.fs.S3FileSystem', Record) - monkeypatch.setattr(backend, 'S3FileSystem', Record) - monkeypatch.setattr(kumo_s3, 'ambient_session', - lambda: FakeSession(credentials=credentials)) - monkeypatch.setattr('kumo_s3.session.bucket_region', lambda bucket: 'x') - monkeypatch.setattr( - 'pyarrow.fs.FileSystem', - SimpleNamespace(from_uri=lambda uri: - (SimpleNamespace(region='us-west-2'), ''))) - sink._parallel_upload_options.cache_clear() - connector = SimpleNamespace( - id='connector', identity=None, identity_id=None, - config=S3ConnectorResourceConfig(target=S3AnyURIConnectorTarget())) - - filesystem, path = S3SinkBackend().filesystem( - connector, 's3://bucket/prefix/', parallel_uploads=True) - sink._parallel_upload_options.cache_clear() - - assert isinstance(filesystem, Record) - assert path == 'bucket/prefix' - assert built[-1].get('allow_delayed_open') is (True if delayed_open else - None) - assert built[-1]['background_writes'] is False - assert 'retry_strategy' in built[-1] From b1924ef79f124b587cd265d0c3afb94e5a1c532f Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Wed, 21 Oct 2026 22:25:43 +0000 Subject: [PATCH 62/80] Repair files a concurrent export wiped before committing Two exports to one location could interleave so that one committed kumo_snapshot.json while the other had just wiped the index; if the second export then died, the location claimed a complete export whose files were gone. Export now deletes the marker again after its wipe, and after committing lists the index: files another export removed are uploaded again and the marker rewritten, so a committed marker always covers complete files with their checksums. Re-export still rewrites everything. Also follows #50's record_sources rename. Verified with test_snapshot_export.py, including a wipe injected just before the commit (fails without the repair). Co-Authored-By: Claude Opus 5.5 --- .../kumo-engine/kumo/engine/_snapshot_io.py | 85 +++++++++++++------ packages/kumo-engine/kumo/engine/graph.py | 10 +-- .../kumo-engine/test/test_snapshot_export.py | 25 +++++- 3 files changed, 88 insertions(+), 32 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/_snapshot_io.py b/packages/kumo-engine/kumo/engine/_snapshot_io.py index 4f0b7ac97..5f473157d 100644 --- a/packages/kumo-engine/kumo/engine/_snapshot_io.py +++ b/packages/kumo-engine/kumo/engine/_snapshot_io.py @@ -12,7 +12,7 @@ from concurrent.futures import ThreadPoolExecutor from functools import partial from pathlib import Path -from typing import Any, Callable, List, Optional, Sequence, Tuple +from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple from pyarrow.fs import FileSystem @@ -39,7 +39,7 @@ def _location(connector: Any, path: str) -> Tuple[FileSystem, str]: raise NotImplementedError( f'{type(connector).__name__} cannot hold a snapshot export; use ' 'an S3Connector, GCSConnector or a FileConnector on a volume.') - return sink.filesystem(connector, path, parallel_uploads=True) + return sink.filesystem(connector, path) def _files(filesystem: FileSystem) -> Any: @@ -195,9 +195,10 @@ def export(snapshot: Any, connector: Any, path: str) -> str: } snapshot_manifest = _snapshot_manifest(snapshot) filesystem, root = _location(connector, location) - if _files(filesystem).exists(f'{root}/{_MANIFEST}'): - filesystem.delete_file(f'{root}/{_MANIFEST}') + _delete_marker(filesystem, root) filesystem.delete_dir_contents(root, missing_dir_ok=True) + # An export that committed while this one wiped no longer holds its files. + _delete_marker(filesystem, root) def upload(client: FileSystem, name: str, stop: threading.Event) -> None: size = sizes[name] @@ -211,30 +212,66 @@ def upload(client: FileSystem, name: str, stop: threading.Event) -> None: 'exported') writer.write(chunk) - names = sorted(sizes, key=lambda name: -sizes[name]) + def send(names: List[str], action: str) -> None: + names = sorted(names, key=lambda name: -sizes[name]) + _transfer(connector, location, names, [sizes[name] for name in names], + upload, action) + 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') + send(list(sizes), 'export') crc32 = checksums.result() - filesystem, root = _location(connector, location) - _write(filesystem, f'{root}/index/{_INDEX_MANIFEST}', index_manifest) - _write(filesystem, f'{root}/{SNAPSHOT_MANIFEST_FILE}', snapshot_manifest) - _write( - filesystem, f'{root}/{_MANIFEST}', - json.dumps({ - 'format': _FORMAT, - 'snapshot_id': snapshot.id, - 'graph_key': graph_key, - 'data_version': data_version, - 'index_manifest': _fingerprint(index_manifest), - 'graph_config': snapshot.graph_config, - 'environment': environment(), - 'range_bytes': _RANGE_BYTES, - 'crc32': crc32, - }).encode()) - return location + marker = json.dumps({ + 'format': _FORMAT, + 'snapshot_id': snapshot.id, + 'graph_key': graph_key, + 'data_version': data_version, + 'index_manifest': _fingerprint(index_manifest), + 'graph_config': snapshot.graph_config, + 'environment': environment(), + 'range_bytes': _RANGE_BYTES, + 'crc32': crc32, + }).encode() + for _ in range(_ATTEMPTS): + filesystem, root = _location(connector, location) + _write(filesystem, f'{root}/index/{_INDEX_MANIFEST}', index_manifest) + _write(filesystem, f'{root}/{SNAPSHOT_MANIFEST_FILE}', + snapshot_manifest) + _write(filesystem, f'{root}/{_MANIFEST}', marker) + if not _files(filesystem).exists(f'{root}/{_MANIFEST}'): + logger.warning('A later export to %s took it over.', location) + return location + missing = _missing(filesystem, root, + {**sizes, _INDEX_MANIFEST: len(index_manifest)}) + if not missing: + return location + logger.warning('Another export to %s removed %d files; uploading ' + 'them again.', location, len(missing)) + send([name for name in missing if name in sizes], 'repair') + raise OSError(f'Another export to {location} keeps replacing it; run ' + 'GraphSnapshot.export again once it finishes.') + + +def _delete_marker(filesystem: FileSystem, root: str) -> None: + try: + filesystem.delete_file(f'{root}/{_MANIFEST}') + except FileNotFoundError: + pass + + +def _missing(filesystem: FileSystem, root: str, + sizes: Dict[str, int]) -> List[str]: + from pyarrow.fs import FileSelector, FileType + + prefix = f'{root}/index/' + found = { + info.path[len(prefix):]: info.size + for info in filesystem.get_file_info( + FileSelector(prefix, recursive=True, allow_not_found=True)) + if info.type is FileType.File + } + return [name for name, size in sizes.items() if found.get(name) != size] def _corrupt(path: str, detail: str) -> ValueError: diff --git a/packages/kumo-engine/kumo/engine/graph.py b/packages/kumo-engine/kumo/engine/graph.py index ea45f5c26..aa413107c 100644 --- a/packages/kumo-engine/kumo/engine/graph.py +++ b/packages/kumo-engine/kumo/engine/graph.py @@ -247,7 +247,7 @@ def _snapshot_id(graph_key: str, data_version: str) -> str: def _source_change(table: Table, recorded: Any, - record_locations: bool) -> Dict[str, Any]: + record_sources: bool) -> Dict[str, Any]: from kumoapi.source_table import CustomSQLSourceTable recorded = recorded if isinstance(recorded, dict) else {} @@ -260,7 +260,7 @@ def _source_change(table: Table, recorded: Any, stamp = connector.source_backend.get_data_version( connector, table.table_name) now = source_table(connector, table.table_name, None, stamp, - record_locations=record_locations) + record_sources=record_sources) current = now['version'] comparable = comparable and now['version_kind'] in _COMPARABLE same_source = (recorded.get('connector'), @@ -386,16 +386,16 @@ def source_changes(self, graph: 'Graph') -> Dict[str, Dict[str, Any]]: manifest = self.manifest recorded = ({} if not manifest or manifest.get('reconstructed') else dict(manifest['parents']['source_tables'])) - record_locations = manifest.get('record_locations', True) + record_sources = manifest.get('record_sources', True) except Exception: - recorded, record_locations = {}, True + 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_locations) + record_sources) for alias, table in tables.items() ], workers=max(1, min(len(tables), 16)), stop=threading.Event()) return dict(zip(tables, changes)) diff --git a/packages/kumo-engine/test/test_snapshot_export.py b/packages/kumo-engine/test/test_snapshot_export.py index f75e2f178..1c5e4811a 100644 --- a/packages/kumo-engine/test/test_snapshot_export.py +++ b/packages/kumo-engine/test/test_snapshot_export.py @@ -176,6 +176,25 @@ def racing(*args): assert _load(exports, raced[0]).id == snapshot.id +def test_export_repairs_files_a_dying_concurrent_export_wiped( + tmp_path, scratch, graph, exports, monkeypatch): + snapshot = graph.snapshot() + write, wiped = snapshot_io._write, [] + + def wipe_before_commit(filesystem, path, payload): + if path.endswith('/index/_MANIFEST.json') and not wiped: + wiped.append(path) + shutil.rmtree(Path(path).parent) + return write(filesystem, path, payload) + + monkeypatch.setattr(snapshot_io, '_write', wipe_before_commit) + location, _ = _export(snapshot, exports, tmp_path) + assert wiped + scratch('gpu') + assert _index_files(_load(exports, location).index_path) == _index_files( + snapshot.index_path) + + def test_failed_upload_is_retried_with_a_fresh_client( scratch, graph, exports, monkeypatch): open_output, clients = snapshot_io._open_output, [] @@ -373,15 +392,15 @@ def dated(tmp_path, graph): return data -@pytest.mark.parametrize('record_locations', [True, False]) +@pytest.mark.parametrize('record_sources', [True, False]) def test_source_changes_reads_stamps_without_staging( - scratch, graph, dated, monkeypatch, record_locations): + 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_locations=record_locations) + 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) From 6839eea29abe9c319e1c2ab0e2c974f9675bae1b Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Wed, 21 Oct 2026 22:32:09 +0000 Subject: [PATCH 63/80] Compare directory listings in source_changes #50 types a Parquet or CSV directory's version as "listing" (a fingerprint of its parts' paths, sizes and mtimes). source_changes treats it as comparable, like a file's mtime, instead of reporting such tables as unknown. Verified with test_snapshot_export.py (new directory case). Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/graph.py | 2 +- packages/kumo-engine/test/test_snapshot_export.py | 15 +++++++++++++++ 2 files changed, 16 insertions(+), 1 deletion(-) diff --git a/packages/kumo-engine/kumo/engine/graph.py b/packages/kumo-engine/kumo/engine/graph.py index aa413107c..33f668d27 100644 --- a/packages/kumo-engine/kumo/engine/graph.py +++ b/packages/kumo-engine/kumo/engine/graph.py @@ -243,7 +243,7 @@ def _snapshot_id(graph_key: str, data_version: str) -> str: return f'{graph_key}@{data_version[:16]}' -_COMPARABLE = ('exact', 'timestamp') +_COMPARABLE = ('exact', 'timestamp', 'listing') def _source_change(table: Table, recorded: Any, diff --git a/packages/kumo-engine/test/test_snapshot_export.py b/packages/kumo-engine/test/test_snapshot_export.py index 1c5e4811a..40e484382 100644 --- a/packages/kumo-engine/test/test_snapshot_export.py +++ b/packages/kumo-engine/test/test_snapshot_export.py @@ -421,6 +421,21 @@ def test_source_changes_reads_stamps_without_staging( '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('edit', ['absent', 'reconstructed', 'hour_bucket']) def test_source_changes_unknown_without_a_comparable_version( scratch, graph, dated, edit): From 1e2b14c05becbaf95ab08ce6d56706390a7a68e6 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Wed, 21 Oct 2026 22:33:57 +0000 Subject: [PATCH 64/80] Drop the connector-image allowlist entry for the removed sink test Only the kumo-s3 sink test, removed with the S3 upload tuning, needed kumo_core/object_store_sink.py in the public connector image. Co-Authored-By: Claude Opus 5.5 --- tests/public/Dockerfile.dockerignore | 2 -- 1 file changed, 2 deletions(-) diff --git a/tests/public/Dockerfile.dockerignore b/tests/public/Dockerfile.dockerignore index 03a5e01fb..4bedfe32e 100644 --- a/tests/public/Dockerfile.dockerignore +++ b/tests/public/Dockerfile.dockerignore @@ -12,8 +12,6 @@ !packages/kumo-core/src/kumo_core/csv_rules.py !packages/kumo-core/src/kumo_core/custom_sql.py !packages/kumo-core/src/kumo_core/artifact_staging.py -# The S3 sink tests import the shared object-store sink. -!packages/kumo-core/src/kumo_core/object_store_sink.py # Model registries validate bundles through the portable artifact contract. !packages/kumo-core/src/kumo_core/artifacts.py !packages/kumo-core/src/kumo_core/config.py From cc48acd1bde3e1a3a7f8c6fcef7a02aaf41258c5 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Thu, 22 Oct 2026 05:34:51 +0000 Subject: [PATCH 65/80] Raise ConcurrentExportError when another export keeps replacing the destination Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/_snapshot_io.py | 9 +++++++-- packages/kumo-engine/test/test_snapshot_export.py | 15 +++++++++++++++ 2 files changed, 22 insertions(+), 2 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/_snapshot_io.py b/packages/kumo-engine/kumo/engine/_snapshot_io.py index 5f473157d..e2c2cfe0a 100644 --- a/packages/kumo-engine/kumo/engine/_snapshot_io.py +++ b/packages/kumo-engine/kumo/engine/_snapshot_io.py @@ -28,6 +28,10 @@ logger = logging.getLogger(__name__) +class ConcurrentExportError(TimeoutError): + pass + + def _location(connector: Any, path: str) -> Tuple[FileSystem, str]: from kumo_core.object_store_sink import ObjectStoreSinkBackend @@ -249,8 +253,9 @@ def send(names: List[str], action: str) -> None: logger.warning('Another export to %s removed %d files; uploading ' 'them again.', location, len(missing)) send([name for name in missing if name in sizes], 'repair') - raise OSError(f'Another export to {location} keeps replacing it; run ' - 'GraphSnapshot.export again once it finishes.') + raise ConcurrentExportError( + f'Another export to {location} keeps replacing it; run ' + 'GraphSnapshot.export again once it finishes.') def _delete_marker(filesystem: FileSystem, root: str) -> None: diff --git a/packages/kumo-engine/test/test_snapshot_export.py b/packages/kumo-engine/test/test_snapshot_export.py index 40e484382..b4d19c5f1 100644 --- a/packages/kumo-engine/test/test_snapshot_export.py +++ b/packages/kumo-engine/test/test_snapshot_export.py @@ -195,6 +195,21 @@ def wipe_before_commit(filesystem, path, payload): snapshot.index_path) +def test_export_gives_up_when_another_export_keeps_wiping_it( + scratch, graph, exports, monkeypatch): + write = snapshot_io._write + + def wipe_before_commit(filesystem, path, payload): + if path.endswith('/index/_MANIFEST.json'): + shutil.rmtree(Path(path).parent) + return write(filesystem, path, payload) + + monkeypatch.setattr(snapshot_io, '_write', wipe_before_commit) + with pytest.raises(snapshot_io.ConcurrentExportError, + match='keeps replacing it'): + graph.snapshot().export(exports, 'one') + + def test_failed_upload_is_retried_with_a_fresh_client( scratch, graph, exports, monkeypatch): open_output, clients = snapshot_io._open_output, [] From 318c18b18c3b9753365b6227842df09b500495fd Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Thu, 22 Oct 2026 18:49:29 +0000 Subject: [PATCH 66/80] Read snapshot manifests as dataclasses in export and source_changes The snapshot manifest is now a SnapshotManifest, so source_changes reads its source tables by attribute, the reconstructed manifest an export ships is built from the dataclass, and the export manifest stores the environment as its dict. The exported files are unchanged. Verified with test_snapshot_export.py. Co-Authored-By: Claude Opus 5.5 --- .../kumo-engine/kumo/engine/_snapshot_io.py | 17 +++++++----- packages/kumo-engine/kumo/engine/graph.py | 26 ++++++++++--------- .../kumo-engine/test/test_snapshot_export.py | 5 ++-- 3 files changed, 26 insertions(+), 22 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/_snapshot_io.py b/packages/kumo-engine/kumo/engine/_snapshot_io.py index e2c2cfe0a..33823d18d 100644 --- a/packages/kumo-engine/kumo/engine/_snapshot_io.py +++ b/packages/kumo-engine/kumo/engine/_snapshot_io.py @@ -168,17 +168,20 @@ def _fingerprint(index_manifest: bytes) -> str: def _snapshot_manifest(snapshot: Any) -> bytes: - from kumo.engine._manifest import new_manifest + 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 = new_manifest('graph_snapshot', snapshot.id, - {'source_tables': {}}, created_at=None, - environment=None, scheduler=None, - reconstructed=True) - return (json.dumps(manifest, indent=2, sort_keys=True) + + manifest = SnapshotManifest( + id=snapshot.id, parents=SnapshotParents(source_tables={}), + **RECONSTRUCTED) + return (json.dumps(manifest.to_dict(), indent=2, sort_keys=True) + '\n').encode() @@ -233,7 +236,7 @@ def send(names: List[str], action: str) -> None: 'data_version': data_version, 'index_manifest': _fingerprint(index_manifest), 'graph_config': snapshot.graph_config, - 'environment': environment(), + 'environment': environment().to_dict(), 'range_bytes': _RANGE_BYTES, 'crc32': crc32, }).encode() diff --git a/packages/kumo-engine/kumo/engine/graph.py b/packages/kumo-engine/kumo/engine/graph.py index 33f668d27..0717706f5 100644 --- a/packages/kumo-engine/kumo/engine/graph.py +++ b/packages/kumo-engine/kumo/engine/graph.py @@ -246,13 +246,14 @@ def _snapshot_id(graph_key: str, data_version: str) -> str: _COMPARABLE = ('exact', 'timestamp', 'listing') -def _source_change(table: Table, recorded: Any, +def _source_change(table: Table, recorded: Optional[SourceTableEntry], record_sources: bool) -> Dict[str, Any]: from kumoapi.source_table import CustomSQLSourceTable - recorded = recorded if isinstance(recorded, dict) else {} - before, current, same_source = recorded.get('version'), None, False - comparable = recorded.get('version_kind') in _COMPARABLE + 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): @@ -261,11 +262,11 @@ def _source_change(table: Table, recorded: Any, 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.get('connector'), - recorded.get('location')) == (now['connector'], - now['location']) + current = now.version + comparable = comparable and now.version_kind in _COMPARABLE + same_source = recorded is not None and ( + recorded.connector, recorded.location) == (now.connector, + now.location) except Exception: current = None if not (same_source and comparable and before is not None @@ -384,9 +385,10 @@ def source_changes(self, graph: 'Graph') -> Dict[str, Dict[str, Any]]: `unchanged`, `changed` or `unknown`. Never raises.""" try: manifest = self.manifest - recorded = ({} if not manifest or manifest.get('reconstructed') else - dict(manifest['parents']['source_tables'])) - record_sources = manifest.get('record_sources', True) + 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 diff --git a/packages/kumo-engine/test/test_snapshot_export.py b/packages/kumo-engine/test/test_snapshot_export.py index b4d19c5f1..81c367199 100644 --- a/packages/kumo-engine/test/test_snapshot_export.py +++ b/packages/kumo-engine/test/test_snapshot_export.py @@ -337,9 +337,8 @@ def test_snapshot_manifest_travels_unchanged_or_is_reconstructed( 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': {}} + assert (manifest.id, manifest.reconstructed) == (snapshot.id, True) + assert manifest.parents.source_tables == {} def test_warehouse_connectors_are_not_supported(scratch, graph): From 46465ffdd173f304980fea7394480552d388a9a2 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sat, 24 Oct 2026 17:47:06 +0000 Subject: [PATCH 67/80] Export snapshots through a config to any file store; store in MLOps scope Snapshot export took (connector, path), gated on object-store sinks, and was the only place a snapshot reached durable storage. - GraphSnapshot.export(SnapshotOutputConfig(connector, path)), in the style of the other output configs; any sink that holds files at the path (FileStoreSink, #112) can take it, and every upload tells a chunked store its size so files up to one part stay plain. - Graph.snapshot(mlops=): with MLOps on, the snapshot is exported under /snapshots/ unless a complete export of the same format and id is there; failures warn. - ConcurrentExportError is a RuntimeError (no retry role). GraphSnapshot.graph_key/data_version are private. The README section moves to docs/mlops.md. Verified: test_snapshot_export.py (30: config form, a filesystem-only chunked sink round trip, stored once in a block, existing export, retry and repair cases), snapshot manifest and training-table suites. Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/docs/mlops.md | 21 +++ packages/kumo-engine/kumo/engine/__init__.py | 2 + .../kumo-engine/kumo/engine/_snapshot_io.py | 61 +++++--- packages/kumo-engine/kumo/engine/graph.py | 39 +++-- .../kumo/engine/trainer/__init__.py | 8 +- .../kumo/engine/trainer/output_config.py | 9 ++ .../kumo-engine/test/test_snapshot_export.py | 141 ++++++++++++++---- .../test/test_snapshot_export_live.py | 3 +- 8 files changed, 224 insertions(+), 60 deletions(-) diff --git a/packages/kumo-engine/docs/mlops.md b/packages/kumo-engine/docs/mlops.md index 5c93f0f1f..8ebc092c0 100644 --- a/packages/kumo-engine/docs/mlops.md +++ b/packages/kumo-engine/docs/mlops.md @@ -46,3 +46,24 @@ Every fitted bundle carries `_provenance.json` (`ModelManifest`): 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` and `Trainer.manifest` describe the latest fit or the loaded bundle. With MLOps on, each bundle is stored under `/models//`. A complete copy that is already there is kept; a partial one is replaced. + +### 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. +- **Transfer:** files move in parallel with a CRC32 per range. An interrupted export resumes, and a load either completes or leaves nothing. +- **Reuse:** loading again reuses the local copy, unless that copy fails DiskGraph's integrity check; then it is downloaded again. +- **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. diff --git a/packages/kumo-engine/kumo/engine/__init__.py b/packages/kumo-engine/kumo/engine/__init__.py index e2d3c1500..db8c0743e 100644 --- a/packages/kumo-engine/kumo/engine/__init__.py +++ b/packages/kumo-engine/kumo/engine/__init__.py @@ -57,6 +57,7 @@ ModelOutputConfig, PredictionOutputConfig, PredictionResult, + SnapshotOutputConfig, Trainer, TrainingResult, TrainingTableOutputConfig, @@ -103,6 +104,7 @@ "ModelOutputConfig", "PredictionOutputConfig", "TrainingTableOutputConfig", + "SnapshotOutputConfig", ] ensure_nofile() diff --git a/packages/kumo-engine/kumo/engine/_snapshot_io.py b/packages/kumo-engine/kumo/engine/_snapshot_io.py index 33823d18d..9fc228d83 100644 --- a/packages/kumo-engine/kumo/engine/_snapshot_io.py +++ b/packages/kumo-engine/kumo/engine/_snapshot_io.py @@ -28,22 +28,20 @@ logger = logging.getLogger(__name__) -class ConcurrentExportError(TimeoutError): - pass +class ConcurrentExportError(RuntimeError): + r"""Another export to the same location kept replacing this one.""" def _location(connector: Any, path: str) -> Tuple[FileSystem, str]: - from kumo_core.object_store_sink import ObjectStoreSinkBackend + from kumo.engine.connector import file_store - try: - sink = connector.sink_backend - except Exception: - sink = None - if not isinstance(sink, ObjectStoreSinkBackend): + store = file_store(connector, path) + if store is None: raise NotImplementedError( - f'{type(connector).__name__} cannot hold a snapshot export; use ' - 'an S3Connector, GCSConnector or a FileConnector on a volume.') - return sink.filesystem(connector, path) + 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: @@ -153,13 +151,15 @@ def _read_json(filesystem: FileSystem, path: str) -> Optional[Any]: return None -def _open_output(filesystem: FileSystem, path: str) -> Any: +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 filesystem.open_output_stream(path) + return open_sized(filesystem, path, size) def _write(filesystem: FileSystem, path: str, payload: bytes) -> None: - with _open_output(filesystem, path) as writer: + with _open_output(filesystem, path, len(payload)) as writer: writer.write(payload) @@ -185,6 +185,32 @@ def _snapshot_manifest(snapshot: Any) -> bytes: '\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)) + marker = _read_json(filesystem, f'{root}/{_MANIFEST}') + 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 @@ -192,9 +218,8 @@ def export(snapshot: Any, connector: Any, path: str) -> str: from kumo.engine.graph import SNAPSHOT_MANIFEST_FILE source = Path(snapshot.index_path) - graph_key, data_version = snapshot.graph_key, snapshot.data_version - _check_key(graph_key, data_version, snapshot.id) - location = f"{path.rstrip('/')}/{graph_key}/{data_version[:16]}" + graph_key, data_version = snapshot._graph_key, snapshot._data_version + location = export_location(snapshot, path) index_manifest = (source / _INDEX_MANIFEST).read_bytes() sizes = { entry['rel_path']: int(entry['size_bytes']) @@ -210,7 +235,7 @@ def export(snapshot: Any, connector: Any, path: str) -> str: 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'{root}/index/{name}') as writer: + client, f'{root}/index/{name}', size) as writer: for offset in range(0, size, _CHUNK): _stopped(stop) chunk = reader.read(min(_CHUNK, size - offset)) diff --git a/packages/kumo-engine/kumo/engine/graph.py b/packages/kumo-engine/kumo/engine/graph.py index 0717706f5..3c182605a 100644 --- a/packages/kumo-engine/kumo/engine/graph.py +++ b/packages/kumo-engine/kumo/engine/graph.py @@ -65,7 +65,7 @@ from dataclasses import dataclass, field, replace from functools import cached_property, partial from pathlib import Path -from typing import Any, Dict, Mapping, Optional, Sequence +from typing import TYPE_CHECKING, Any, Dict, Mapping, Optional, Sequence from kumo_core.artifacts import PROVENANCE_FILE from kumo_core.pquery import NativeGraph, normalize_graph @@ -88,6 +88,9 @@ ) 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 @@ -313,28 +316,26 @@ class GraphSnapshot: record_sources: bool | None = field(default=None, compare=False) @property - def graph_key(self) -> str: + def _graph_key(self) -> str: """The `Graph.id` this snapshot was built from.""" return self.id.partition('@')[0] @property - def data_version(self) -> str: + 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 - def export(self, connector: Connector, path: str) -> str: + def export(self, output_config: 'SnapshotOutputConfig') -> str: r"""Copy the index to `//` and return that location. Any earlier export there is replaced; it is not - loadable until the new one completes. - - Args: - connector: an S3, GCS or file connector (on Databricks or - Snowflake, a `FileConnector` on a mounted volume). + loadable until the new one completes. 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, connector, path) + return export(self, output_config.connector, output_config.path) @classmethod def load_from_connector(cls, connector: Connector, *, path: str, @@ -581,7 +582,8 @@ def validate(self) -> None: """ self.native_graph - def snapshot(self, *, record_sources: bool = True) -> 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` @@ -592,6 +594,9 @@ def snapshot(self, *, record_sources: bool = True) -> GraphSnapshot: 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). @@ -620,9 +625,17 @@ def snapshot(self, *, record_sources: bool = True) -> 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. + from kumo.engine._mlops import resolve + + settings = resolve(mlops) with _lock(graph_key): - return self._snapshot_no_lock(graph_key, persist_dir, - record_sources) + 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: 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/output_config.py b/packages/kumo-engine/kumo/engine/trainer/output_config.py index 43da33eac..ddc9a9ef3 100644 --- a/packages/kumo-engine/kumo/engine/trainer/output_config.py +++ b/packages/kumo-engine/kumo/engine/trainer/output_config.py @@ -82,6 +82,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/test/test_snapshot_export.py b/packages/kumo-engine/test/test_snapshot_export.py index 81c367199..98f5ba1c4 100644 --- a/packages/kumo-engine/test/test_snapshot_export.py +++ b/packages/kumo-engine/test/test_snapshot_export.py @@ -20,7 +20,7 @@ import kumo.engine as kumo import kumo.engine._snapshot_io as snapshot_io -from kumo.engine import GraphSnapshot +from kumo.engine import GraphSnapshot, SnapshotOutputConfig from kumo.engine._scratch import get_scratch_dir @@ -45,13 +45,13 @@ 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', 'snapshot-manifest.json') + and file.name not in ('_SUCCESS', '_provenance.json') and not str(file.relative_to(path)).startswith('train_tables/') } def _export(snapshot, exports, tmp_path): - location = snapshot.export(exports, 'one') + location = snapshot.export(SnapshotOutputConfig(exports, 'one')) return location, tmp_path / 'exports' / location @@ -81,13 +81,13 @@ 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(exports, 'one/') + location = snapshot.export(SnapshotOutputConfig(exports, 'one/')) graph_id, data_version = snapshot.id.split('@') assert location == f'one/{graph_id}/{data_version}' - assert snapshot.export(exports, 'one') == location + assert snapshot.export(SnapshotOutputConfig(exports, 'one')) == location exported = tmp_path / 'exports' / location - recorded = (snapshot.index_path / 'snapshot-manifest.json').read_bytes() - assert (exported / 'snapshot-manifest.json').read_bytes() == recorded + recorded = (snapshot.index_path / '_provenance.json').read_bytes() + assert (exported / '_provenance.json').read_bytes() == recorded shutil.rmtree(tmp_path / 'cpu') scratch('gpu') @@ -97,9 +97,10 @@ def test_export_load_round_trip_relocates_and_is_reused( assert loaded.index_path.is_relative_to(tmp_path / 'gpu') assert _index_files(loaded.index_path) == _index_files(exported / 'index') assert (loaded.index_path / - 'snapshot-manifest.json').read_bytes() == recorded - reexported = tmp_path / 'exports' / loaded.export(exports, 'two') - assert (reexported / 'snapshot-manifest.json').read_bytes() == recorded + '_provenance.json').read_bytes() == recorded + reexported = tmp_path / 'exports' / loaded.export( + SnapshotOutputConfig(exports, 'two')) + assert (reexported / '_provenance.json').read_bytes() == recorded monkeypatch.setattr( 'kumo_core.diskgraph.ingest.execute_diskgraph_ingest', @@ -138,22 +139,22 @@ def four(found, worker_ceiling=None): (exported / 'index/stray.bin').write_bytes(b'x') open_output, finished = snapshot_io._open_output, threading.Event() - def interrupted(client, path): + 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) + return open_output(client, path, size) 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'): - snapshot.export(exports, 'one') + snapshot.export(SnapshotOutputConfig(exports, 'one')) with pytest.raises(FileNotFoundError, match='No complete snapshot'): _load(exports, location) - assert snapshot.export(exports, 'one') == location + assert snapshot.export(SnapshotOutputConfig(exports, 'one')) == location assert not (exported / 'index/stray.bin').exists() scratch('gpu') assert _index_files(_load(exports, location).index_path) == _index_files( @@ -168,11 +169,11 @@ def test_concurrent_exports_of_one_snapshot_both_succeed( def racing(*args): if not raced: raced.append(None) - raced[0] = snapshot.export(exports, 'one') + raced[0] = snapshot.export(SnapshotOutputConfig(exports, 'one')) return transfer(*args) monkeypatch.setattr(snapshot_io, '_transfer', racing) - assert raced == [snapshot.export(exports, 'one')] + assert raced == [snapshot.export(SnapshotOutputConfig(exports, 'one'))] assert _load(exports, raced[0]).id == snapshot.id @@ -207,22 +208,22 @@ def wipe_before_commit(filesystem, path, payload): monkeypatch.setattr(snapshot_io, '_write', wipe_before_commit) with pytest.raises(snapshot_io.ConcurrentExportError, match='keeps replacing it'): - graph.snapshot().export(exports, 'one') + graph.snapshot().export(SnapshotOutputConfig(exports, 'one')) 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): + 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) + return open_output(client, path, size) monkeypatch.setattr(snapshot_io, '_open_output', expiring) - graph.snapshot().export(exports, 'one') + graph.snapshot().export(SnapshotOutputConfig(exports, 'one')) assert len(clients) == 2 and clients[0] is not clients[1] @@ -240,7 +241,7 @@ def test_reexport_replaces_whatever_the_location_holds( stats.write_bytes(original[:3]) (exported / 'index/index_metadata.json').unlink() _edit_manifest(exported, lambda m: m.update(snapshot_id='other')) - assert snapshot.export(exports, 'one') == location + assert snapshot.export(SnapshotOutputConfig(exports, 'one')) == location assert stats.read_bytes() == original assert _load(exports, location).id == snapshot.id @@ -321,7 +322,7 @@ def test_load_rejects_unknown_formats_and_accepts_unknown_fields( def test_snapshot_manifest_travels_unchanged_or_is_reconstructed( tmp_path, scratch, graph, exports, recorded): snapshot = graph.snapshot() - path = snapshot.index_path / 'snapshot-manifest.json' + path = snapshot.index_path / '_provenance.json' if recorded == 'newer': path.write_bytes(b'{"schema_version": 99, "future": [1]}') else: @@ -331,8 +332,8 @@ def test_snapshot_manifest_travels_unchanged_or_is_reconstructed( with warnings.catch_warnings(): warnings.simplefilter('ignore') loaded = _load(exports, location) - shipped = (exported / 'snapshot-manifest.json').read_bytes() - assert (loaded.index_path / 'snapshot-manifest.json').read_bytes() == shipped + shipped = (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: @@ -344,7 +345,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='FileConnector'): - graph.snapshot().export(connector, 'x') + graph.snapshot().export(SnapshotOutputConfig(connector, 'x')) with pytest.raises(NotImplementedError, match='FileConnector'): _load(connector, 'x') @@ -380,7 +381,8 @@ 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(connector, f'{scheme}://bucket/snapshots/') + 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') @@ -454,7 +456,7 @@ def test_source_changes_compares_directory_listings(tmp_path, scratch): def test_source_changes_unknown_without_a_comparable_version( scratch, graph, dated, edit): snapshot = graph.snapshot() - path = snapshot.index_path / 'snapshot-manifest.json' + path = snapshot.index_path / '_provenance.json' manifest = json.loads(path.read_text()) if edit == 'absent': path.unlink() @@ -481,3 +483,88 @@ def test_load_warns_only_with_graph_and_only_for_changed_tables( _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 + '/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 (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 index 3c69b548a..c28b19712 100644 --- a/packages/kumo-engine/test/test_snapshot_export_live.py +++ b/packages/kumo-engine/test/test_snapshot_export_live.py @@ -9,6 +9,7 @@ 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 @@ -29,7 +30,7 @@ def test_snapshot_round_trip_through_object_store(tmp_path, graph, scheme): try: kumo.configure(scratch_dir=tmp_path / 'cpu') snapshot = graph.snapshot() - location = snapshot.export(connector, path) + location = snapshot.export(SnapshotOutputConfig(connector, path)) kumo.configure(scratch_dir=tmp_path / 'gpu') loaded = kumo.GraphSnapshot.load_from_connector( connector, path=location, graph=graph) From 68152b02ef716eee473a629a9a7fc1e8a5784bed Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Tue, 20 Oct 2026 20:05:48 +0000 Subject: [PATCH 68/80] Add PredictiveQuery.from_snapshot for graph-free scoring A job that loads a model bundle and a stored snapshot has no Python Graph, but PredictiveQuery required one, so it could not compile a prediction table. Prediction-table compilation only needs the snapshot's graph config. PredictiveQuery.from_snapshot(snapshot, query) binds the query to a GraphSnapshot: parsing uses the snapshot's graph_config and the snapshot match check uses the graph id embedded in snapshot.id, so a snapshot of a different graph is still refused. PredictiveQuery(graph, query) is unchanged. Verified: a snapshot-bound query generates prediction tables (generated and supplied rows) that Trainer.predict scores identically to the Graph-bound path; existing predictive-query and custom-source tests pass. Co-Authored-By: Claude Opus 5.5 --- .../kumo/engine/pquery/predictive_query.py | 83 ++++++++++++------- .../test/test_query_from_snapshot.py | 54 ++++++++++++ 2 files changed, 107 insertions(+), 30 deletions(-) create mode 100644 packages/kumo-engine/test/test_query_from_snapshot.py diff --git a/packages/kumo-engine/kumo/engine/pquery/predictive_query.py b/packages/kumo-engine/kumo/engine/pquery/predictive_query.py index 5cbf4771c..a967774f1 100644 --- a/packages/kumo-engine/kumo/engine/pquery/predictive_query.py +++ b/packages/kumo-engine/kumo/engine/pquery/predictive_query.py @@ -6,15 +6,19 @@ `PredictiveQuery(graph, query)` binds a `Graph` at construction, matching the cloud SDK's `kumo.PredictiveQuery(graph=graph, query=QUERY)` -- not a -`GraphSnapshot`. Parsing (`kumo_core.pquery.parse_and_validate_query` against -`self.graph.native_graph` -- `Graph`'s own `cached_property`, needing no -materialized index) happens lazily, on first access to the `_parsed` -`cached_property`, and only once: `validate()`/`get_task_type()` both read -that same property rather than re-parsing on every call. Confirmed against -the existing local backend's own `_LocalPQueryBackend.validate` -(`kumoai/backend/local/pquery.py`), which parses against -`to_portable(definition)` -- the schema-level conversion, not an ingested -snapshot -- that this doesn't need a `GraphSnapshot` either. +`GraphSnapshot`. `PredictiveQuery.from_snapshot(snapshot, query)` binds a +snapshot instead, for callers that load a snapshot without its `Graph`; the +query then parses against `snapshot.graph_config`, the same portable config +the source `Graph` produces. + +Parsing (`kumo_core.pquery.parse_and_validate_query` against +`self._native_graph`, needing no materialized index) happens lazily, on first +access to the `_parsed` `cached_property`, and only once: +`validate()`/`get_task_type()` both read that same property rather than +re-parsing on every call. Confirmed against the existing local backend's own +`_LocalPQueryBackend.validate` (`kumoai/backend/local/pquery.py`), which parses +against `to_portable(definition)` -- the schema-level conversion, not an +ingested snapshot -- that this doesn't need a `GraphSnapshot` either. Construction itself stays I/O-free, consistent with `Table`/`Graph`: nothing parses until `_parsed` is actually touched (by `validate()`, `get_task_type()`, @@ -25,8 +29,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,7 +61,9 @@ ) from kumo_core.model_plan_conversion import public_model_plan_from_internal from kumo_core.pquery import ( + NativeGraph, ValidatedNativeQuery, + normalize_graph, parse_and_validate_query, query_is_temporal, resolve_query_split, @@ -97,15 +103,33 @@ @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.""" + return cls(graph=snapshot, query=query) + + @cached_property + def _native_graph(self) -> NativeGraph: + if isinstance(self.graph, GraphSnapshot): + return normalize_graph(self.graph.graph_config) + return self.graph.native_graph + + @property + def _graph_id(self) -> str: + if isinstance(self.graph, GraphSnapshot): + return self.graph.id.partition('@')[0] + return self.graph.id + @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._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 @@ -126,22 +150,21 @@ def _parsed(self) -> ValidatedNativeQuery: `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. + For a `Graph`, reuses its cached `native_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) 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 + # 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 snapshot.id != self.graph.id and not snapshot.id.startswith( + f'{self._graph_id}@'): 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 {self._graph_id!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 +196,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 +218,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 @@ -223,7 +246,7 @@ def generate_training_table( 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 @@ -231,7 +254,7 @@ 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. @@ -330,7 +353,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. @@ -342,7 +365,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/test/test_query_from_snapshot.py b/packages/kumo-engine/test/test_query_from_snapshot.py new file mode 100644 index 000000000..8c2d4ee23 --- /dev/null +++ b/packages/kumo-engine/test/test_query_from_snapshot.py @@ -0,0 +1,54 @@ +# 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 +from kumoapi.task import TaskType + +import kumo.engine as kumo + + +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_snapshot_bound_query_predicts_like_graph_bound(fitted, tmp_path): + trainer, graph_table, snapshot = fitted + query = kumo.PredictiveQuery.from_snapshot(snapshot, graph_table.query) + expected = _scores(trainer, graph_table, snapshot) + + assert query.get_task_type() == TaskType.BINARY_CLASSIFICATION + table = query.generate_prediction_table( + snapshot, plan=trainer.training_table_generation_plan) + assert table.timestamps == graph_table.timestamps + pd.testing.assert_frame_equal(_scores(trainer, table, snapshot), expected) + + anchor = pd.Timestamp(graph_table.timestamps[0]) + pd.DataFrame({ + 'ENTITY1': range(50), + 'TIMESTAMP': [anchor] * 50, + }).to_parquet(tmp_path / 'rows.parquet') + supplied = kumo.PredictionTable.from_source_table( + kumo.FileConnector(root_dir=str(tmp_path))['rows.parquet'], pq=query, + snapshot=snapshot) + pd.testing.assert_frame_equal(_scores(trainer, supplied, snapshot), + expected) + + +def test_snapshot_bound_query_refuses_other_graph(fitted): + _, graph_table, snapshot = fitted + query = kumo.PredictiveQuery.from_snapshot(snapshot, graph_table.query) + newer = replace(snapshot, id=snapshot.id.split('@')[0] + '@newer') + other = replace(snapshot, id='other@' + snapshot.id.split('@')[1]) + + query._require_matching_snapshot(newer) + plain = replace(snapshot, id='plain') + kumo.PredictiveQuery.from_snapshot(plain, '')._require_matching_snapshot( + plain) + with pytest.raises(ValueError, match='was not materialized'): + query.generate_prediction_table(other) From 5256d2c881580c3e8e82fe0f7ebca193a0767af4 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Wed, 21 Oct 2026 08:08:23 +0000 Subject: [PATCH 69/80] Expose Trainer.training_query and keep the graph-bound snapshot check A job scoring a loaded model had no public way to get the query it was trained on, so binding it with PredictiveQuery.from_snapshot meant copying the text out of training-spec.json. Trainer.training_query now returns it. The snapshot check accepted a snapshot whose id equals a Graph's bare id; equality now applies only to a snapshot-bound query, so a Graph-bound query keeps the original prefix check. Verified with test_query_from_snapshot.py. Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/pquery/predictive_query.py | 6 ++++-- packages/kumo-engine/kumo/engine/trainer/trainer.py | 6 ++++++ packages/kumo-engine/test/test_query_from_snapshot.py | 3 ++- 3 files changed, 12 insertions(+), 3 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/pquery/predictive_query.py b/packages/kumo-engine/kumo/engine/pquery/predictive_query.py index a967774f1..f7cd0f5c8 100644 --- a/packages/kumo-engine/kumo/engine/pquery/predictive_query.py +++ b/packages/kumo-engine/kumo/engine/pquery/predictive_query.py @@ -159,8 +159,10 @@ 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 snapshot.id != self.graph.id and not snapshot.id.startswith( - f'{self._graph_id}@'): + if isinstance(self.graph, GraphSnapshot) and (snapshot.id + == self.graph.id): + return + if not snapshot.id.startswith(f'{self._graph_id}@'): raise ValueError( f"snapshot {snapshot.id!r} was not materialized from " f"this query's graph (id {self._graph_id!r}) -- pass a " diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index 1acc5e100..c3d6d04d6 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -344,6 +344,12 @@ def training_table_generation_plan( """Effective table-generation settings saved by the last fit/load.""" return self._training_table_generation_plan + @property + def training_query(self) -> Optional[str]: + """Predictive query of the last fit/load, for + :meth:`PredictiveQuery.from_snapshot`.""" + return self._training_query + @property def scratch_dir(self) -> Path: return trainer_scratch_dir(self._id) diff --git a/packages/kumo-engine/test/test_query_from_snapshot.py b/packages/kumo-engine/test/test_query_from_snapshot.py index 8c2d4ee23..1e90285ea 100644 --- a/packages/kumo-engine/test/test_query_from_snapshot.py +++ b/packages/kumo-engine/test/test_query_from_snapshot.py @@ -19,7 +19,8 @@ def _scores(trainer, table, snapshot): def test_snapshot_bound_query_predicts_like_graph_bound(fitted, tmp_path): trainer, graph_table, snapshot = fitted - query = kumo.PredictiveQuery.from_snapshot(snapshot, graph_table.query) + assert trainer.training_query == graph_table.query + query = kumo.PredictiveQuery.from_snapshot(snapshot, trainer.training_query) expected = _scores(trainer, graph_table, snapshot) assert query.get_task_type() == TaskType.BINARY_CLASSIFICATION From 81fb21c075407a967448120a35407845b1b5939c Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Wed, 21 Oct 2026 08:19:53 +0000 Subject: [PATCH 70/80] Trim comments and implementation-detail tests The from_snapshot change carried docstring prose that restated the code, and its graph-check test called a private helper. Shorten the module note and training_query docstring, drop the restating _parsed paragraph, and exercise the snapshot check through the public suggest_model_plan instead of _require_matching_snapshot. test_query_from_snapshot.py: 2 passed. Co-Authored-By: Claude Opus 5.5 --- .../kumo/engine/pquery/predictive_query.py | 7 +------ packages/kumo-engine/kumo/engine/trainer/trainer.py | 3 +-- .../kumo-engine/test/test_query_from_snapshot.py | 12 ++++++------ 3 files changed, 8 insertions(+), 14 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/pquery/predictive_query.py b/packages/kumo-engine/kumo/engine/pquery/predictive_query.py index f7cd0f5c8..300fa1410 100644 --- a/packages/kumo-engine/kumo/engine/pquery/predictive_query.py +++ b/packages/kumo-engine/kumo/engine/pquery/predictive_query.py @@ -7,9 +7,7 @@ `PredictiveQuery(graph, query)` binds a `Graph` at construction, matching the cloud SDK's `kumo.PredictiveQuery(graph=graph, query=QUERY)` -- not a `GraphSnapshot`. `PredictiveQuery.from_snapshot(snapshot, query)` binds a -snapshot instead, for callers that load a snapshot without its `Graph`; the -query then parses against `snapshot.graph_config`, the same portable config -the source `Graph` produces. +snapshot instead and parses against `snapshot.graph_config`. Parsing (`kumo_core.pquery.parse_and_validate_query` against `self._native_graph`, needing no materialized index) happens lazily, on first @@ -149,9 +147,6 @@ 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`. - - For a `Graph`, reuses its cached `native_graph`, shared by - `Graph.validate()` and every `PredictiveQuery` bound to it. """ return parse_and_validate_query(self.query, self._native_graph) diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index c3d6d04d6..4396b8731 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -346,8 +346,7 @@ def training_table_generation_plan( @property def training_query(self) -> Optional[str]: - """Predictive query of the last fit/load, for - :meth:`PredictiveQuery.from_snapshot`.""" + """Predictive query of the last fit/load.""" return self._training_query @property diff --git a/packages/kumo-engine/test/test_query_from_snapshot.py b/packages/kumo-engine/test/test_query_from_snapshot.py index 1e90285ea..b106248aa 100644 --- a/packages/kumo-engine/test/test_query_from_snapshot.py +++ b/packages/kumo-engine/test/test_query_from_snapshot.py @@ -41,15 +41,15 @@ def test_snapshot_bound_query_predicts_like_graph_bound(fitted, tmp_path): expected) -def test_snapshot_bound_query_refuses_other_graph(fitted): +def test_snapshot_bound_query_checks_snapshot_graph(fitted): _, graph_table, snapshot = fitted query = kumo.PredictiveQuery.from_snapshot(snapshot, graph_table.query) newer = replace(snapshot, id=snapshot.id.split('@')[0] + '@newer') other = replace(snapshot, id='other@' + snapshot.id.split('@')[1]) - - query._require_matching_snapshot(newer) plain = replace(snapshot, id='plain') - kumo.PredictiveQuery.from_snapshot(plain, '')._require_matching_snapshot( - plain) + + query.suggest_model_plan(newer) + kumo.PredictiveQuery.from_snapshot(plain, graph_table.query) \ + .suggest_model_plan(plain) with pytest.raises(ValueError, match='was not materialized'): - query.generate_prediction_table(other) + query.suggest_model_plan(other) From 4ffea50945c1348fdc463294aee854f0d8cc5a8c Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Wed, 21 Oct 2026 18:47:20 +0000 Subject: [PATCH 71/80] Read the native graph and graph id from GraphSnapshot PredictiveQuery kept its own normalized-graph cache and parsed the graph id out of snapshot ids. GraphSnapshot now offers native_graph (cached, like Graph.native_graph), and the snapshot check compares graph_key. Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/graph.py | 6 +++ .../kumo/engine/pquery/predictive_query.py | 38 ++++++------------- 2 files changed, 18 insertions(+), 26 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/graph.py b/packages/kumo-engine/kumo/engine/graph.py index 3c182605a..4174ca94d 100644 --- a/packages/kumo-engine/kumo/engine/graph.py +++ b/packages/kumo-engine/kumo/engine/graph.py @@ -326,6 +326,12 @@ def _data_version(self) -> str: 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. Any earlier export there is replaced; it is not diff --git a/packages/kumo-engine/kumo/engine/pquery/predictive_query.py b/packages/kumo-engine/kumo/engine/pquery/predictive_query.py index 300fa1410..0eee72f01 100644 --- a/packages/kumo-engine/kumo/engine/pquery/predictive_query.py +++ b/packages/kumo-engine/kumo/engine/pquery/predictive_query.py @@ -10,7 +10,7 @@ snapshot instead and parses against `snapshot.graph_config`. Parsing (`kumo_core.pquery.parse_and_validate_query` against -`self._native_graph`, needing no materialized index) happens lazily, on first +`self.graph.native_graph`, needing no materialized index) happens lazily, on first access to the `_parsed` `cached_property`, and only once: `validate()`/`get_task_type()` both read that same property rather than re-parsing on every call. Confirmed against the existing local backend's own @@ -59,9 +59,7 @@ ) from kumo_core.model_plan_conversion import public_model_plan_from_internal from kumo_core.pquery import ( - NativeGraph, ValidatedNativeQuery, - normalize_graph, parse_and_validate_query, query_is_temporal, resolve_query_split, @@ -112,21 +110,9 @@ def from_snapshot(cls, snapshot: GraphSnapshot, r"""Bind `query` to `snapshot`'s graph without its `Graph` object.""" return cls(graph=snapshot, query=query) - @cached_property - def _native_graph(self) -> NativeGraph: - if isinstance(self.graph, GraphSnapshot): - return normalize_graph(self.graph.graph_config) - return self.graph.native_graph - - @property - def _graph_id(self) -> str: - if isinstance(self.graph, GraphSnapshot): - return self.graph.id.partition('@')[0] - return self.graph.id - @cached_property def _parsed(self) -> ValidatedNativeQuery: - r"""Parses and validates `self.query` against `self._native_graph`, + 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` @@ -148,19 +134,19 @@ def _parsed(self) -> ValidatedNativeQuery: `ValidatedNativeQuery` are all plain `str`; `get_task_type` below wraps `task_type` into the shared, public `kumoapi.task.TaskType`. """ - return parse_and_validate_query(self.query, self._native_graph) + return parse_and_validate_query(self.query, 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 isinstance(self.graph, GraphSnapshot) and (snapshot.id - == self.graph.id): - return - 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"this query's graph (id {self._graph_id!r}) -- pass a " + f"this query's graph (id {graph_key!r}) -- pass a " "snapshot of that graph.") def validate(self, verbose: bool = True) -> ValidationResponse: @@ -193,7 +179,7 @@ def get_task_type(self) -> TaskType: def _model_plan_validation_facts(self, run_mode: RunMode) -> dict[str, Any]: parsed = self._parsed - graph = self._native_graph + graph = self.graph.native_graph entity_tables = [parsed.entity_table] if parsed.rhs_node_type is not None: entity_tables.append(parsed.rhs_node_type) From 76da3dc9010bb98a86efc05c5020fa1427883b4f Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Wed, 21 Oct 2026 21:18:29 +0000 Subject: [PATCH 72/80] Add Trainer.predict_snapshot for one-call graph-free scoring A graph-free scoring job had to bind the saved query to the snapshot, pass the saved generation plan to generate_prediction_table and then call predict; each caller (MLflow model, job scripts) repeated that. Trainer.predict_snapshot(snapshot, *, entities=None, prediction_time=None, ...) does it in one call with predict's options. The docs now say the query's graph check only ties a snapshot to the query's graph, not to a trained model. Verified with the query-from-snapshot and predictive-query tests. Co-Authored-By: Claude Opus 5.5 --- .../kumo/engine/pquery/predictive_query.py | 8 +++- .../kumo/engine/trainer/trainer.py | 42 +++++++++++++++++++ .../test/test_query_from_snapshot.py | 20 +++++++++ 3 files changed, 69 insertions(+), 1 deletion(-) diff --git a/packages/kumo-engine/kumo/engine/pquery/predictive_query.py b/packages/kumo-engine/kumo/engine/pquery/predictive_query.py index 0eee72f01..d36621e05 100644 --- a/packages/kumo-engine/kumo/engine/pquery/predictive_query.py +++ b/packages/kumo-engine/kumo/engine/pquery/predictive_query.py @@ -107,7 +107,13 @@ class PredictiveQuery: @classmethod def from_snapshot(cls, snapshot: GraphSnapshot, query: str) -> 'PredictiveQuery': - r"""Bind `query` to `snapshot`'s graph without its `Graph` object.""" + 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/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index 4396b8731..16a508a85 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -55,6 +55,7 @@ from kumo.engine._manifest import ModelManifest from kumo.engine._scratch import new_run_id, trainer_scratch_dir from kumo.engine.connector.base import Connector +from kumo.engine.connector.source_table import SourceTable from kumo.engine.graph import GraphSnapshot from kumo.engine.pquery.prediction_table import PredictionTable from kumo.engine.pquery.training_table import TrainingTable @@ -685,6 +686,47 @@ def _finish_fit( return training_result + def predict_snapshot( + self, + snapshot: GraphSnapshot, + *, + 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, + ) -> PredictionResult: + r"""Score ``snapshot`` with the query and table-generation plan saved + by the last fit/load. + + Without ``entities``, scores the 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. Other arguments are as in + :meth:`predict`. + """ + from kumo.engine.pquery.predictive_query import PredictiveQuery + + if self._training_query is None: + raise RuntimeError( + "Trainer.predict_snapshot() 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.") + query = PredictiveQuery.from_snapshot(snapshot, self._training_query) + plan = self._training_table_generation_plan + table = (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)) + 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) + @_without_torch_scatter_hint def predict( self, diff --git a/packages/kumo-engine/test/test_query_from_snapshot.py b/packages/kumo-engine/test/test_query_from_snapshot.py index b106248aa..65629f9b4 100644 --- a/packages/kumo-engine/test/test_query_from_snapshot.py +++ b/packages/kumo-engine/test/test_query_from_snapshot.py @@ -41,6 +41,26 @@ def test_snapshot_bound_query_predicts_like_graph_bound(fitted, tmp_path): expected) +def test_predict_snapshot_scores_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( + 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) + with pytest.raises(ValueError, match='not both'): + scores(entities=rows, prediction_time=graph_table.timestamps[0]) + + def test_snapshot_bound_query_checks_snapshot_graph(fitted): _, graph_table, snapshot = fitted query = kumo.PredictiveQuery.from_snapshot(snapshot, graph_table.query) From aab9d02ab40e25533b552cfd4caf2a51561ec8fe Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Wed, 21 Oct 2026 22:36:39 +0000 Subject: [PATCH 73/80] Test graph-free scoring of an exported snapshot end to end Exports the fitted snapshot, loads it into a fresh scratch directory without a Graph, reloads the model from its bundle and scores with predict_snapshot and with PredictiveQuery.from_snapshot; both match the graph-bound scores. Co-Authored-By: Claude Opus 5.5 --- .../test/test_query_from_snapshot.py | 27 +++++++++++++++++++ 1 file changed, 27 insertions(+) diff --git a/packages/kumo-engine/test/test_query_from_snapshot.py b/packages/kumo-engine/test/test_query_from_snapshot.py index 65629f9b4..87866165c 100644 --- a/packages/kumo-engine/test/test_query_from_snapshot.py +++ b/packages/kumo-engine/test/test_query_from_snapshot.py @@ -9,6 +9,7 @@ from kumoapi.task import TaskType import kumo.engine as kumo +from kumo.engine._scratch import get_scratch_dir def _scores(trainer, table, snapshot): @@ -61,6 +62,32 @@ def scores(**options): scores(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) + exports = kumo.FileConnector(root_dir=str(tmp_path / 'exports')) + location = snapshot.export(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.predict_snapshot( + loaded, binary_classification_threshold=0.5).frame + pd.testing.assert_frame_equal( + scored.sort_values('ENTITY').reset_index(drop=True), expected) + query = kumo.PredictiveQuery.from_snapshot(loaded, model.training_query) + table = query.generate_prediction_table( + loaded, plan=model.training_table_generation_plan) + pd.testing.assert_frame_equal(_scores(model, table, loaded), expected) + finally: + kumo.configure(scratch_dir=previous) + + def test_snapshot_bound_query_checks_snapshot_graph(fitted): _, graph_table, snapshot = fitted query = kumo.PredictiveQuery.from_snapshot(snapshot, graph_table.query) From 32e9cdda3cb6b6c24ded1497c2d3312c3f7f613b Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Thu, 22 Oct 2026 01:14:39 +0000 Subject: [PATCH 74/80] Warn when predict_snapshot scores another graph's snapshot predict_snapshot bound the saved query to whatever snapshot it was given, so scoring a snapshot of a different graph than the model was trained on went through silently. It now warns (and still scores) when the snapshot's graph_key differs from the graph key of the snapshot recorded at fit/load. Verified with test_query_from_snapshot.py. Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/trainer/trainer.py | 9 ++++++++- packages/kumo-engine/test/test_query_from_snapshot.py | 6 ++++++ 2 files changed, 14 insertions(+), 1 deletion(-) diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index 16a508a85..6a5ba60b3 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -704,7 +704,8 @@ def predict_snapshot( ``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. Other arguments are as in - :meth:`predict`. + :meth:`predict`. Warns when ``snapshot`` was taken of a different + graph than the one the model was trained on. """ from kumo.engine.pquery.predictive_query import PredictiveQuery @@ -715,6 +716,12 @@ def predict_snapshot( 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: + warnings.warn( + f"Snapshot {snapshot.id} was taken of graph " + f"{snapshot.graph_key}, but this model was trained on a " + f"snapshot of graph {trained}.", UserWarning, stacklevel=2) query = PredictiveQuery.from_snapshot(snapshot, self._training_query) plan = self._training_table_generation_plan table = (query.generate_prediction_table( diff --git a/packages/kumo-engine/test/test_query_from_snapshot.py b/packages/kumo-engine/test/test_query_from_snapshot.py index 87866165c..da611aa38 100644 --- a/packages/kumo-engine/test/test_query_from_snapshot.py +++ b/packages/kumo-engine/test/test_query_from_snapshot.py @@ -60,6 +60,12 @@ def scores(**options): pd.testing.assert_frame_equal(scores(entities=rows), 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]) + with pytest.warns(UserWarning, match='trained on a snapshot of graph'): + scored = trainer.predict_snapshot( + 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( From c5b4bc8fdf3676fb11a57f4086a9f919c2172724 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sat, 24 Oct 2026 18:03:16 +0000 Subject: [PATCH 75/80] One predict path: predict(snapshot, ...); results predict too predict_snapshot sat next to predict, so scoring had two entry points with overlapping arguments. - Trainer.predict(snapshot, *, table=, entities=, prediction_time=, ...) scores a snapshot with the trained query and generation plan (entities or a generated table instead when given) and warns when the snapshot is of another graph. main's predict(prediction_table, snapshot=...) still works, positionally or by keyword. Argument checks run first, as before. - TrainingResult.predict(...) scores with the model that produced or loaded it (TrainingResult._trainer, set by fit and load). - predict_snapshot is gone (never released); Trainer.training_query and GraphSnapshot.native_graph are private. Docs move to docs/mlops.md. Verified: test_query_from_snapshot.py (snapshot, entities, table and legacy forms, other-graph warning, a fresh scratch scoring an exported snapshot through a loaded result), prediction workers, output columns, completion, preflight, registry, warm start, embedding and file connector suites. Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/docs/mlops.md | 17 +++ packages/kumo-engine/kumo/engine/graph.py | 2 +- .../kumo/engine/pquery/predictive_query.py | 15 +- .../kumo/engine/trainer/trainer.py | 138 ++++++++++-------- .../test/test_query_from_snapshot.py | 21 ++- 5 files changed, 122 insertions(+), 71 deletions(-) diff --git a/packages/kumo-engine/docs/mlops.md b/packages/kumo-engine/docs/mlops.md index 8ebc092c0..bc2b84697 100644 --- a/packages/kumo-engine/docs/mlops.md +++ b/packages/kumo-engine/docs/mlops.md @@ -67,3 +67,20 @@ snapshot = GraphSnapshot.load_from_connector(S3Connector(), path=location) - **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; `table=` scores a prediction table generated beforehand. +- It warns when the snapshot was taken of a different graph than the one the model was trained on. +- `predict(prediction_table, snapshot=...)` keeps working as before. +- `fit` results score too: `result.predict(snapshot, ...)`. +- `PredictiveQuery.from_snapshot(snapshot, query)` binds a query to a snapshot without its `Graph`. Its graph check only means the snapshot belongs to the query's graph, not that it is compatible with a trained model. diff --git a/packages/kumo-engine/kumo/engine/graph.py b/packages/kumo-engine/kumo/engine/graph.py index 4174ca94d..e82f13cc4 100644 --- a/packages/kumo-engine/kumo/engine/graph.py +++ b/packages/kumo-engine/kumo/engine/graph.py @@ -327,7 +327,7 @@ def _data_version(self) -> str: return Path(self.index_path).name @cached_property - def native_graph(self) -> NativeGraph: + def _native_graph(self) -> NativeGraph: """`graph_config` in `kumo_core`'s normalized form, as `Graph.native_graph`.""" return normalize_graph(self.graph_config) diff --git a/packages/kumo-engine/kumo/engine/pquery/predictive_query.py b/packages/kumo-engine/kumo/engine/pquery/predictive_query.py index d36621e05..7ac4a9d8d 100644 --- a/packages/kumo-engine/kumo/engine/pquery/predictive_query.py +++ b/packages/kumo-engine/kumo/engine/pquery/predictive_query.py @@ -59,6 +59,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, @@ -140,16 +141,22 @@ def _parsed(self) -> ValidatedNativeQuery: `ValidatedNativeQuery` are all plain `str`; `get_task_type` below wraps `task_type` into the shared, public `kumoapi.task.TaskType`. """ - 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: if isinstance(self.graph, GraphSnapshot): if snapshot.id == self.graph.id: return - graph_key = self.graph.graph_key + graph_key = self.graph._graph_key else: graph_key = self.graph.id - if snapshot.graph_key != graph_key: + if snapshot._graph_key != graph_key: raise ValueError( f"snapshot {snapshot.id!r} was not materialized from " f"this query's graph (id {graph_key!r}) -- pass a " @@ -185,7 +192,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) diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index 6a5ba60b3..82612fe8f 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -10,7 +10,7 @@ import logging 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 +21,7 @@ Mapping, Optional, TypeVar, + Union, cast, ) @@ -94,6 +95,8 @@ class TrainingResult: 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]: @@ -107,6 +110,13 @@ def manifest(self) -> Optional[ModelManifest]: return read_provenance(self.bundle_dir) + def predict(self, *args: Any, **kwargs: Any) -> 'PredictionResult': + """Score with this model, as :meth:`Trainer.predict`.""" + if self._trainer is None: + raise RuntimeError('This result has no model to predict with; ' + 'use Trainer.load on its bundle_dir.') + return self._trainer.predict(*args, **kwargs) + def holdout_df(self) -> pd.DataFrame: r"""Read the holdout predictions as a pandas DataFrame. @@ -345,11 +355,6 @@ def training_table_generation_plan( """Effective table-generation settings saved by the last fit/load.""" return self._training_table_generation_plan - @property - def training_query(self) -> Optional[str]: - """Predictive query of the last fit/load.""" - return self._training_query - @property def scratch_dir(self) -> Path: return trainer_scratch_dir(self._id) @@ -679,6 +684,7 @@ def _finish_fit( 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._training_query = predictive_query @@ -686,56 +692,98 @@ def _finish_fit( return training_result - def predict_snapshot( + def predict( self, - snapshot: GraphSnapshot, + source: Union[GraphSnapshot, PredictionTable, None] = None, + /, *, + snapshot: Optional[GraphSnapshot] = None, + table: Optional[PredictionTable] = 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"""Score ``snapshot`` with the query and table-generation plan saved - by the last fit/load. - - Without ``entities``, scores the 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. Other arguments are as in - :meth:`predict`. Warns when ``snapshot`` was taken of a different - graph than the one the model was trained on. + 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, and ``table`` a + prediction table generated beforehand. It warns when the snapshot + was taken of another graph than the one the model was trained on. + ``predict(prediction_table, snapshot=...)`` is accepted as before. + + The trainer must contain fitted state from :meth:`fit`, :meth:`load`, + 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 + scorers over disjoint partitions, capped by the prediction population. + CUDA execution requires one visible device per concurrent scorer. + + ``num_classes_to_return`` must be a positive integer when supplied. + 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. """ + from kumo_core.prediction_options import ( + validate_num_classes_to_return, + validate_prediction_workers, + ) + + validate_prediction_workers(num_workers) + validate_num_classes_to_return(num_classes_to_return) + if source is not None and not isinstance(source, GraphSnapshot): + source, prediction_table = None, source + snapshot = snapshot if snapshot is not None else source + if snapshot is None: + raise TypeError('predict() needs a snapshot') + table = table if table is not None else prediction_table + if table is None: + table = 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_snapshot() requires a prior Trainer.fit() or " + "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: + if trained and snapshot._graph_key != trained: warnings.warn( f"Snapshot {snapshot.id} was taken of graph " - f"{snapshot.graph_key}, but this model was trained on a " - f"snapshot of graph {trained}.", UserWarning, stacklevel=2) + f"{snapshot._graph_key}, but this model was trained on a " + f"snapshot of graph {trained}.", UserWarning, stacklevel=3) query = PredictiveQuery.from_snapshot(snapshot, self._training_query) plan = self._training_table_generation_plan - table = (query.generate_prediction_table( + 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)) - 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) + if entities is None else PredictionTable.from_source_table( + entities, pq=query, snapshot=snapshot, plan=plan)) @_without_torch_scatter_hint - def predict( + def _predict( self, prediction_table: PredictionTable, *, @@ -745,33 +793,6 @@ def predict( num_workers: int = 1, binary_classification_threshold: Optional[float] = None, ) -> PredictionResult: - r"""Run batch prediction and return durable file-backed results. - - 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. - - ``num_workers`` controls scorer processes, not DataLoader workers. - The default stays in-process; values above one spawn independent - scorers over disjoint partitions, capped by the prediction population. - CUDA execution requires one visible device per concurrent scorer. - - ``num_classes_to_return`` must be a positive integer when supplied. - 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. - """ - from kumo_core.prediction_options import ( - validate_num_classes_to_return, - validate_prediction_workers, - ) - - validate_prediction_workers(num_workers) - validate_num_classes_to_return(num_classes_to_return) if self._training_result is None: raise RuntimeError( "Trainer.predict() requires a prior Trainer.fit() call on " @@ -1015,5 +1036,6 @@ def _load_bundle( trainer._training_table_generation_plan = from_json_dict( spec.training_table_generation, TrainingTableGenerationPlan) trainer._training_result = TrainingResult(metrics, paths, bundle_dir, - bundle_id=bundle_id) + bundle_id=bundle_id, + _trainer=trainer) return trainer diff --git a/packages/kumo-engine/test/test_query_from_snapshot.py b/packages/kumo-engine/test/test_query_from_snapshot.py index da611aa38..488d09d0d 100644 --- a/packages/kumo-engine/test/test_query_from_snapshot.py +++ b/packages/kumo-engine/test/test_query_from_snapshot.py @@ -20,8 +20,8 @@ def _scores(trainer, table, snapshot): def test_snapshot_bound_query_predicts_like_graph_bound(fitted, tmp_path): trainer, graph_table, snapshot = fitted - assert trainer.training_query == graph_table.query - query = kumo.PredictiveQuery.from_snapshot(snapshot, trainer.training_query) + assert trainer._training_query == graph_table.query + query = kumo.PredictiveQuery.from_snapshot(snapshot, trainer._training_query) expected = _scores(trainer, graph_table, snapshot) assert query.get_task_type() == TaskType.BINARY_CLASSIFICATION @@ -42,12 +42,12 @@ def test_snapshot_bound_query_predicts_like_graph_bound(fitted, tmp_path): expected) -def test_predict_snapshot_scores_with_the_saved_query(fitted, tmp_path): +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( + return trainer.predict( snapshot, binary_classification_threshold=0.5, **options).frame.sort_values('ENTITY').reset_index(drop=True) @@ -58,11 +58,16 @@ def scores(**options): }).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(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]) with pytest.warns(UserWarning, match='trained on a snapshot of graph'): - scored = trainer.predict_snapshot( + scored = trainer.predict( other, binary_classification_threshold=0.5).frame pd.testing.assert_frame_equal( scored.sort_values('ENTITY').reset_index(drop=True), expected) @@ -73,7 +78,7 @@ def test_exported_snapshot_scores_in_a_fresh_scratch_without_a_graph( trainer, graph_table, snapshot = fitted expected = _scores(trainer, graph_table, snapshot) exports = kumo.FileConnector(root_dir=str(tmp_path / 'exports')) - location = snapshot.export(exports, 'snapshots') + location = snapshot.export(kumo.SnapshotOutputConfig(exports, 'snapshots')) bundle = trainer._training_result.bundle_dir previous = get_scratch_dir() kumo.configure(scratch_dir=tmp_path / 'fresh') @@ -82,11 +87,11 @@ 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.predict_snapshot( + 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) - query = kumo.PredictiveQuery.from_snapshot(loaded, model.training_query) + query = kumo.PredictiveQuery.from_snapshot(loaded, model._training_query) table = query.generate_prediction_table( loaded, plan=model.training_table_generation_plan) pd.testing.assert_frame_equal(_scores(model, table, loaded), expected) From 4fd38b9f7e9b7cb9334dfabb72584bbab811ae2d Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sat, 24 Oct 2026 19:10:01 +0000 Subject: [PATCH 76/80] Hide the torch-scatter hint across all of predict() 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 82612fe8f..6e903e98d 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -692,6 +692,7 @@ def _finish_fit( return training_result + @_without_torch_scatter_hint def predict( self, source: Union[GraphSnapshot, PredictionTable, None] = None, @@ -782,7 +783,6 @@ def _prediction_table(self, snapshot: GraphSnapshot, if entities is None else PredictionTable.from_source_table( entities, pq=query, snapshot=snapshot, plan=plan)) - @_without_torch_scatter_hint def _predict( self, prediction_table: PredictionTable, From 803cd82e2cdfa251e2a7cfa39f297243b6742816 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sat, 24 Oct 2026 19:40:50 +0000 Subject: [PATCH 77/80] Compare reloaded-model scores within float tolerance Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/test/test_query_from_snapshot.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/packages/kumo-engine/test/test_query_from_snapshot.py b/packages/kumo-engine/test/test_query_from_snapshot.py index 488d09d0d..70cd9e66b 100644 --- a/packages/kumo-engine/test/test_query_from_snapshot.py +++ b/packages/kumo-engine/test/test_query_from_snapshot.py @@ -90,11 +90,13 @@ def test_exported_snapshot_scores_in_a_fresh_scratch_without_a_graph( 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) + scored.sort_values('ENTITY').reset_index(drop=True), expected, + atol=1e-5) query = kumo.PredictiveQuery.from_snapshot(loaded, model._training_query) table = query.generate_prediction_table( loaded, plan=model.training_table_generation_plan) - pd.testing.assert_frame_equal(_scores(model, table, loaded), expected) + pd.testing.assert_frame_equal(_scores(model, table, loaded), expected, + atol=1e-5) finally: kumo.configure(scratch_dir=previous) From 7cdc1effbf520756b6119e0a7e017dbf60d9b7be Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sat, 24 Oct 2026 18:05:09 +0000 Subject: [PATCH 78/80] Report fit progress events to callbacks fit(callbacks=[...]) receives each kumo.engine.tracking.Event: fit_started, trial_started, epoch_ended, trial_ended and fit_ended, with status and a short error on ended events. A callback that raises is logged and disabled; it never fails the fit. fit resolves its MLOps settings once on the caller's thread and passes them down. Co-Authored-By: Claude Opus 5.5 --- .../src/kumo_core/automl/auto_trainer.py | 37 ++++- .../src/kumo_core/automl/optuna_searcher.py | 1 + .../src/kumo_core/automl/searcher.py | 15 +- .../callbacks/train_progress_logger.py | 25 ++- .../src/kumo_core/instrumentation.py | 10 ++ .../src/kumo_core/training/runtime.py | 4 +- .../kumo-core/tests/test_searched_params.py | 38 +++++ packages/kumo-engine/docs/mlops.md | 16 ++ packages/kumo-engine/kumo/engine/_events.py | 86 ++++++++++ packages/kumo-engine/kumo/engine/tracking.py | 10 ++ .../kumo/engine/trainer/_runtime.py | 6 +- .../kumo/engine/trainer/trainer.py | 68 ++++++-- packages/kumo-engine/test/test_fit_events.py | 154 ++++++++++++++++++ 13 files changed, 447 insertions(+), 23 deletions(-) create mode 100644 packages/kumo-core/tests/test_searched_params.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/test/test_fit_events.py 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 5dd818966..3abe36170 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 @@ -458,6 +459,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. @@ -713,7 +715,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 @@ -847,7 +849,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() @@ -872,12 +874,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( @@ -885,7 +892,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 " @@ -902,9 +911,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) @@ -957,6 +978,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/searcher.py b/packages/kumo-core/src/kumo_core/automl/searcher.py index f8810e7b6..9116bf389 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,17 +88,25 @@ 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 def __repr__(self) -> str: return f"{self.__class__.__name__}()" + 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/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/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_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-engine/docs/mlops.md b/packages/kumo-engine/docs/mlops.md index bc2b84697..f8f21b41d 100644 --- a/packages/kumo-engine/docs/mlops.md +++ b/packages/kumo-engine/docs/mlops.md @@ -84,3 +84,19 @@ result = trainer.predict(snapshot, binary_classification_threshold=0.5) - `predict(prediction_table, snapshot=...)` keeps working as before. - `fit` results score too: `result.predict(snapshot, ...)`. - `PredictiveQuery.from_snapshot(snapshot, query)` binds a query to a snapshot without its `Graph`. Its graph check only means the snapshot belongs to the query's graph, not that it is compatible with a trained model. + +## 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()))]) +``` 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/tracking.py b/packages/kumo-engine/kumo/engine/tracking.py new file mode 100644 index 000000000..ef29cd7f5 --- /dev/null +++ b/packages/kumo-engine/kumo/engine/tracking.py @@ -0,0 +1,10 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +r"""Progress events and experiment-tracking types for `Trainer.fit`.""" + +from kumo.engine._events import Event + +__all__ = [ + 'Event', +] 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/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index 6e903e98d..085f2d8fb 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -20,6 +20,7 @@ Literal, Mapping, Optional, + Sequence, TypeVar, Union, cast, @@ -27,6 +28,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 @@ -53,6 +55,7 @@ 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._scratch import new_run_id, trainer_scratch_dir from kumo.engine.connector.base import Connector @@ -384,6 +387,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`. @@ -414,10 +418,60 @@ def fit( 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. + 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 and not callable( + getattr(export_config.connector.sink_backend, 'publish_model', + None)): + raise NotImplementedError( + f'{type(export_config.connector).__name__} cannot publish ' + f'models') + from kumo.engine._mlops import resolve + + settings = resolve(mlops) + run_id = f'{self._id}-fit-{new_run_id()}' + events = EventEmitter(run_id, 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) + except BaseException as error: + events.emit('fit_ended', params={ + 'status': 'failed', + 'error': format_error(error), + }) + raise + events.emit('fit_ended', metrics={ + **result.metrics, 'seconds': time.monotonic() - fit_start + }, params={'status': 'succeeded'}) + return result + + def _fit( + self, + train_table: TrainingTable, + *, + snapshot: GraphSnapshot, + export_config: Optional[ModelOutputConfig], + warm_start: Optional[Self], + run_id: str, + events: EventEmitter, + settings: Optional[Any], + ) -> TrainingResult: from kumo_core.automl.runtime import ( inherit_warm_start_model_plan, validate_warm_start_task_type, @@ -437,20 +491,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(export_config.connector.sink_backend, '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) @@ -569,7 +613,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(), ) 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 == [] From d97b285e315734f839ebcade4439d21535da9296 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sat, 24 Oct 2026 18:11:15 +0000 Subject: [PATCH 79/80] Record Trainer.fit in MLflow through MLOps tracking kumo.MLOps(tracking=MlflowTracking(...)) records each fit in MLflow: one run per fit and one child run per search-trial attempt, with params, per-epoch metrics, the model's _provenance.json and the source tables as native MLflow inputs read from the training table's embedded parents. Runs carry pointer tags (kumo.kind, kumo.run_id, kumo.bundle_id); the model manifest records the tracked run in tracking.run_id. Tracking runs on a background thread and never fails a fit. Verified with test_tracking_contract.py, test_mlflow_tracking.py, test_optional_imports.py and test_bundle_provenance.py. Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/docs/mlops.md | 48 ++ packages/kumo-engine/kumo/engine/_manifest.py | 8 + packages/kumo-engine/kumo/engine/_tracking.py | 385 +++++++++++ .../kumo/engine/integrations/__init__.py | 3 + .../kumo/engine/integrations/mlflow.py | 279 ++++++++ packages/kumo-engine/kumo/engine/tracking.py | 5 +- .../kumo/engine/trainer/_provenance.py | 4 + .../kumo/engine/trainer/trainer.py | 27 +- packages/kumo-engine/pyproject.toml | 5 + .../kumo-engine/test/test_mlflow_tracking.py | 316 +++++++++ .../kumo-engine/test/test_optional_imports.py | 9 + .../test/test_tracking_contract.py | 357 ++++++++++ setup.py | 2 + uv.lock | 616 ++++++++++++++---- 14 files changed, 1928 insertions(+), 136 deletions(-) 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.py create mode 100644 packages/kumo-engine/test/test_mlflow_tracking.py create mode 100644 packages/kumo-engine/test/test_tracking_contract.py diff --git a/packages/kumo-engine/docs/mlops.md b/packages/kumo-engine/docs/mlops.md index f8f21b41d..2c1c7f7dd 100644 --- a/packages/kumo-engine/docs/mlops.md +++ b/packages/kumo-engine/docs/mlops.md @@ -100,3 +100,51 @@ 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, one run per hyperparameter-search trial attempt. 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) +``` + +- Tracking errors are logged as warnings and never fail the fit. 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. +- MLflow runs go to `MLFLOW_TRACKING_URI`: on Databricks this is automatic; for SageMaker managed MLflow set it to the tracking server ARN with the `sagemaker-mlflow` plugin installed. With nothing configured MLflow writes a local store in the working directory. +- Per-epoch metrics use step = epoch. After the fit, `tracking.run_id` is the fit run's MLflow run id. +- 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`. + +Runs carry pointers, not lineage: + +| tag | on | value | +|---|---|---| +| `kumo.kind` | every run | `fit` or `trial` | +| `kumo.run_id` | every run | the kumo fit 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 | 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`. + +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`, and the best-effort hooks `log_model`, +`log_prediction_run` and `log_monitoring`. Params are the resolved model plan +(fit) or searched hyperparameters and `attempt` (trial). `RunRecord.inputs` +maps each source table to its `connector`, `location`, `version` and +`version_kind`. A tracker's `tags=` are added as given, except `kumo.*` names +and param names. diff --git a/packages/kumo-engine/kumo/engine/_manifest.py b/packages/kumo-engine/kumo/engine/_manifest.py index 81265178f..974f361f6 100644 --- a/packages/kumo-engine/kumo/engine/_manifest.py +++ b/packages/kumo-engine/kumo/engine/_manifest.py @@ -238,6 +238,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, @@ -248,6 +253,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 write_manifest(path: Path, manifest: Manifest) -> None: diff --git a/packages/kumo-engine/kumo/engine/_tracking.py b/packages/kumo-engine/kumo/engine/_tracking.py new file mode 100644 index 000000000..c79aa608b --- /dev/null +++ b/packages/kumo-engine/kumo/engine/_tracking.py @@ -0,0 +1,385 @@ +# 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.engine._events import Event, _finite_floats + +if TYPE_CHECKING: + 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 + + +@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: + logger.warning('Could not read the fit\'s source tables', + exc_info=True) + 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); and the best-effort hooks + ``log_model(result, snapshot)``, ``log_prediction_run(record)`` and + ``log_monitoring(record)``. + """ + 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 TrackingSession: + def __init__(self, tracking: Optional[Tracking]) -> None: + self._tracker = tracking if _rank() == 0 else None + 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 + + @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: + logger.warning('Experiment tracking failed; the fit continues', + exc_info=True) + 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 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) + except Exception: + logger.warning('Could not build the tracking record', + exc_info=True) + record = None + + def finish_tracker() -> None: + tracker, self._tracker = self._tracker, None + if tracker is not None: + tracker.finish(record, error) + + 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: + logger.warning('Experiment tracking failed', exc_info=True) + + def _call(self, call: Callable[[Tracking], Any]) -> None: + tracker = self._tracker + if tracker is None: + return + try: + call(tracker) + except Exception: + logger.warning( + 'Experiment tracking failed; it is off for the rest of this ' + 'run', exc_info=True) + self._tracker = None + _abandon(tracker) + + +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/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.py b/packages/kumo-engine/kumo/engine/integrations/mlflow.py new file mode 100644 index 000000000..e799aa68f --- /dev/null +++ b/packages/kumo-engine/kumo/engine/integrations/mlflow.py @@ -0,0 +1,279 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +r"""MLflow tracking for `Trainer.fit` (``pip install 'kumo-engine[mlflow]'``). +""" + +import hashlib +import json +import logging +import os +import posixpath +import re +import time +import warnings +from contextlib import ExitStack, contextmanager +from typing import Any, Callable, Dict, Iterator, Mapping, Optional, Set, Tuple + +from kumo.engine._events import Event +from kumo.engine._tracking import DROPPED_PARAMS, RunRecord +from kumo.engine.connector._extras import import_provider + +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', +} +_PATH_CONNECTORS = ('LOCAL_FILE', 'S3', 'GCS', 'ADLS') + + +class MlflowTracking: + r"""Records `Trainer.fit` in MLflow: one run per fit (nested under the + active run, if any) 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 the current experiment. + run_name: name of the fit run. `None` lets MLflow pick one. + 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. + + ``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, + ) -> None: + self._mlflow = import_provider('mlflow', 'mlflow') + self.experiment_name = experiment_name + self.run_name = run_name + self.log_system_metrics = log_system_metrics + self.tags = dict(tags or {}) + 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._guard = ExitStack() + + def prepare(self) -> None: + active = self._mlflow.active_run() + self._outer = active.info.run_id if active is not None else None + + def start(self, record: RunRecord) -> None: + self._reset() + self.run_id = None + self._guard.enter_context(_limited_http()) + 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: + logger.warning('Could not suspend MLflow autologging', + exc_info=True) + mlflow = self._mlflow + self._client = mlflow.MlflowClient() + experiment_id = None + if self.experiment_name is not None: + experiment = self._client.get_experiment_by_name( + self.experiment_name) + experiment_id = (experiment.experiment_id + if experiment else self._client.create_experiment( + self.experiment_name)) + run = mlflow.start_run( + experiment_id=experiment_id, + run_name=self.run_name, + parent_run_id=self._outer, + tags=record.tags, + log_system_metrics=self.log_system_metrics, + ) + self._parent = self.run_id = run.info.run_id + self._experiment_id = run.info.experiment_id + self._log_params(self._parent, record.params) + self._log_artifacts(record.artifacts) + try: + self._log_inputs(record.inputs) + except Exception: + logger.warning('Could not log the fit\'s source tables to MLflow', + exc_info=True) + + def on_trial(self, record: RunRecord, event: Event) -> None: + key = (event.trial or 0, int(event.params['attempt'])) + if event.name == 'trial_started': + 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' + self._client.set_terminated(run_id, + 'FINISHED' if succeeded else 'FAILED') + + def finish(self, record: Optional[RunRecord], + error: Optional[BaseException]) -> None: + status = 'FAILED' + try: + if record is not None and self._parent is not None: + self._log_metrics(self._parent, record.metrics) + self._log_tags(self._parent, record.tags) + self._log_artifacts(record.artifacts) + if error is None: + status = 'FINISHED' + finally: + _quietly(self._mlflow.flush_async_logging) + for run_id in self._children.values(): + _quietly(self._client.set_terminated, run_id, 'FAILED') + if self._parent is not None: + active = self._mlflow.active_run() + if active is not None and active.info.run_id == self._parent: + _quietly(self._mlflow.end_run, status) + else: + _quietly(self._client.set_terminated, self._parent, status) + self._guard.close() + self._reset() + + def _log_artifacts(self, artifacts: Mapping[str, str]) -> None: + for name, text in artifacts.items(): + if name not in self._artifacts: + self._client.log_text(self._parent, text, name) + self._artifacts.add(name) + + def _log_inputs(self, + inputs: Mapping[str, Mapping[str, Optional[str]]]) -> None: + meta = import_provider('mlflow.data.meta_dataset', 'mlflow') + entities = self._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: + self._client.log_inputs(self._parent, datasets=datasets) + + def _log_tags(self, run_id: str, tags: Mapping[str, str]) -> None: + entities = self._mlflow.entities + 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 [] + 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) -> None: + if run_id is None or not metrics: + return + entities = self._mlflow.entities + timestamp = int(time.time() * 1000) + self._client.log_batch( + run_id, metrics=[ + entities.Metric(key, float(value), timestamp, step) + for key, value in metrics.items() + ]) + + +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() -> Iterator[None]: + unset = [name for name in _HTTP_LIMITS if name not in os.environ] + os.environ.update({name: _HTTP_LIMITS[name] for name in unset}) + try: + yield + finally: + for name in unset: + os.environ.pop(name, None) diff --git a/packages/kumo-engine/kumo/engine/tracking.py b/packages/kumo-engine/kumo/engine/tracking.py index ef29cd7f5..9f0f217f4 100644 --- a/packages/kumo-engine/kumo/engine/tracking.py +++ b/packages/kumo-engine/kumo/engine/tracking.py @@ -1,10 +1,13 @@ # SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. # All rights reserved. # SPDX-License-Identifier: Apache-2.0 -r"""Progress events and experiment-tracking types for `Trainer.fit`.""" +r"""API for writing an experiment-tracking adapter for ``fit(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/_provenance.py b/packages/kumo-engine/kumo/engine/trainer/_provenance.py index 4429186dd..28e23acf4 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) diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index 085f2d8fb..6efd488c4 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -58,6 +58,7 @@ from kumo.engine._events import Event, EventEmitter, query_sha256 from kumo.engine._manifest import ModelManifest from kumo.engine._scratch import new_run_id, trainer_scratch_dir +from kumo.engine._tracking import FitContext, TrackingSession, fit_inputs from kumo.engine.connector.base import Connector from kumo.engine.connector.source_table import SourceTable from kumo.engine.graph import GraphSnapshot @@ -417,7 +418,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. + 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. @@ -435,7 +437,11 @@ def fit( settings = resolve(mlops) run_id = f'{self._id}-fit-{new_run_id()}' - events = EventEmitter(run_id, callbacks) + session = TrackingSession( + None if settings is None else settings.tracking) + events = EventEmitter( + run_id, + [*callbacks, session.on_event] if session.enabled else callbacks) fit_start = time.monotonic() try: started = { @@ -449,16 +455,19 @@ 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) + 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'}) + session.finish(result, None) return result def _fit( @@ -471,12 +480,13 @@ def _fit( 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, @@ -601,6 +611,12 @@ def _fit( tt_id=ingest_tt_id, ) del task_dataset + 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, @@ -664,7 +680,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, ) diff --git a/packages/kumo-engine/pyproject.toml b/packages/kumo-engine/pyproject.toml index ce96f58f7..25dc85914 100644 --- a/packages/kumo-engine/pyproject.toml +++ b/packages/kumo-engine/pyproject.toml @@ -68,6 +68,11 @@ text=[ gcs=[ "kumo-gcs>=0.1.0.dev1", ] +# nvidia-ml-py adds GPU system metrics. +mlflow=[ + "mlflow>=3.13,<4", + "nvidia-ml-py", +] # 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/test_mlflow_tracking.py b/packages/kumo-engine/test/test_mlflow_tracking.py new file mode 100644 index 000000000..1d4f639d5 --- /dev/null +++ b/packages/kumo-engine/test/test_mlflow_tracking.py @@ -0,0 +1,316 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import json +import os +import socket +import threading +import time +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 + 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_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) == {} diff --git a/packages/kumo-engine/test/test_optional_imports.py b/packages/kumo-engine/test/test_optional_imports.py index c70368fed..f9ee60309 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,14 @@ 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") print("ok") ''') 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..2061c267b --- /dev/null +++ b/packages/kumo-engine/test/test_tracking_contract.py @@ -0,0 +1,357 @@ +# 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(params=['mlflow']) +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/setup.py b/setup.py index a8446e82b..48b4680d0 100644 --- a/setup.py +++ b/setup.py @@ -293,6 +293,8 @@ def parse_project_version(project_file_path: Path) -> str: "respx==0.22.0", "pytorch-frame==0.3.0", "pytorch-lightning==2.6.6", + # kumo-engine's `mlflow` extra; its wheel is tested with --no-deps. + "mlflow>=3.13,<4", ] dev_requires = test_requires + [ diff --git a/uv.lock b/uv.lock index 6c8b4649c..32d82ea28 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" }, @@ -3343,6 +3487,10 @@ databricks = [ gcs = [ { name = "kumo-gcs" }, ] +mlflow = [ + { name = "mlflow" }, + { name = "nvidia-ml-py" }, +] s3 = [ { name = "kumo-s3" }, ] @@ -3356,11 +3504,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,7 +3530,9 @@ 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", marker = "extra == 'mlflow'", specifier = ">=3.13,<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" }, @@ -3391,7 +3541,7 @@ requires-dist = [ { name = "torch", marker = "extra == 'test'", specifier = ">=2.9,<2.13" }, { name = "typing-extensions", specifier = ">=4.5.0" }, ] -provides-extras = ["test", "snowflake", "databricks", "bigquery", "s3", "text", "gcs"] +provides-extras = ["test", "snowflake", "databricks", "bigquery", "s3", "text", "gcs", "mlflow"] [[package]] name = "kumo-gcs" @@ -3772,8 +3922,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" }, @@ -3860,6 +4010,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" @@ -4087,7 +4320,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 = [ @@ -4285,7 +4518,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" }, @@ -4428,7 +4661,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" }, @@ -4466,7 +4699,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" }, @@ -4578,9 +4811,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" }, @@ -4618,7 +4851,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" }, @@ -4636,6 +4869,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" @@ -5089,7 +5331,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 = [ @@ -5256,6 +5498,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" @@ -5769,7 +6023,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 = [ @@ -5802,7 +6056,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" }, @@ -5819,7 +6073,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 = [ @@ -5922,6 +6176,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" @@ -5953,11 +6216,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" } @@ -5974,11 +6237,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" }, @@ -5997,6 +6260,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" @@ -6213,8 +6498,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" } @@ -6277,7 +6562,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 = [ @@ -6371,7 +6656,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 = [ @@ -6482,6 +6767,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" @@ -6491,6 +6793,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" @@ -6588,8 +6899,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" }, @@ -6646,7 +6957,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" }, @@ -6800,7 +7111,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" } @@ -6867,7 +7178,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 = [ @@ -7073,14 +7384,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" }, @@ -7157,14 +7468,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" }, @@ -7207,21 +7518,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 = [ @@ -7256,21 +7567,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" }, @@ -7321,11 +7632,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 = [ @@ -7337,7 +7648,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 = [ @@ -7405,7 +7716,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 = [ @@ -7422,7 +7733,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 = [ @@ -7490,7 +7801,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 = [ @@ -7534,6 +7845,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" @@ -7578,6 +7898,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 ef5ed4b91e9f9cd162b187f359fe70ee210686a2 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sat, 24 Oct 2026 18:13:56 +0000 Subject: [PATCH 80/80] Record Trainer.fit in Snowflake Experiment Tracking kumo.MLOps(tracking=SnowflakeTracking(session, experiment_name=...)) records each fit as one Snowflake run: trial attempts as prefixed params and metrics, tags as params, the model's _provenance.json as an artifact, Snowflake's per-run limits counted in kumo.dropped_*. The user's Snowpark session is not altered and an existing run is never resumed. snowflake-ml-python is capped below 3 for the internals read. Verified with test_snowflake_tracking.py and test_tracking_contract.py. Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/docs/mlops.md | 22 +- .../kumo/engine/integrations/snowflake.py | 251 ++++++++++++++++++ .../test/_snowflake_experiment_fake.py | 192 ++++++++++++++ .../kumo-engine/test/test_optional_imports.py | 8 + .../test/test_snowflake_tracking.py | 216 +++++++++++++++ .../test/test_tracking_contract.py | 45 +++- packages/kumo-snowflake/pyproject.toml | 8 +- uv.lock | 2 +- 8 files changed, 738 insertions(+), 6 deletions(-) create mode 100644 packages/kumo-engine/kumo/engine/integrations/snowflake.py create mode 100644 packages/kumo-engine/test/_snowflake_experiment_fake.py create mode 100644 packages/kumo-engine/test/test_snowflake_tracking.py diff --git a/packages/kumo-engine/docs/mlops.md b/packages/kumo-engine/docs/mlops.md index 2c1c7f7dd..a71d07f71 100644 --- a/packages/kumo-engine/docs/mlops.md +++ b/packages/kumo-engine/docs/mlops.md @@ -104,7 +104,8 @@ trainer.fit(training_table, snapshot=snapshot, ## Experiment tracking `kumo.MLOps(tracking=...)` records each fit in an experiment tracker: one run -per fit, one run per hyperparameter-search trial attempt. Tracking is on only +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 @@ -123,6 +124,25 @@ with kumo.mlops(cfg): - Per-epoch metrics use step = epoch. After the fit, `tracking.run_id` is the fit run's MLflow run id. - 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 | 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..4d78d5b56 --- /dev/null +++ b/packages/kumo-engine/kumo/engine/integrations/snowflake.py @@ -0,0 +1,251 @@ +# 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.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. + 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, + *, + 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') + self.session = 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 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: + logger.warning('Could not end the Snowflake run %s', status, + 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/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_optional_imports.py b/packages/kumo-engine/test/test_optional_imports.py index f9ee60309..c18d0310d 100644 --- a/packages/kumo-engine/test/test_optional_imports.py +++ b/packages/kumo-engine/test/test_optional_imports.py @@ -123,6 +123,14 @@ def find_spec(self, name, path=None, target=None): 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_snowflake_tracking.py b/packages/kumo-engine/test/test_snowflake_tracking.py new file mode 100644 index 000000000..ebf705d57 --- /dev/null +++ b/packages/kumo-engine/test/test_snowflake_tracking.py @@ -0,0 +1,216 @@ +# 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=None, **kwargs): + from kumo.engine.integrations.snowflake import SnowflakeTracking + return SnowflakeTracking(session or FakeSession(), 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' diff --git a/packages/kumo-engine/test/test_tracking_contract.py b/packages/kumo-engine/test/test_tracking_contract.py index 2061c267b..2efe3cb35 100644 --- a/packages/kumo-engine/test/test_tracking_contract.py +++ b/packages/kumo-engine/test/test_tracking_contract.py @@ -242,7 +242,50 @@ def history(run_id, key): mlflow.set_tracking_uri(previous) -@pytest.fixture(params=['mlflow']) +@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') 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/uv.lock b/uv.lock index 32d82ea28..4868d05f5 100644 --- a/uv.lock +++ b/uv.lock @@ -3687,7 +3687,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"]