From 231a309b1bf704b619d993d945fdf253e7cd7bb3 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Tue, 20 Oct 2026 20:30:39 +0000 Subject: [PATCH 01/19] 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 02/19] 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 03/19] 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 04/19] 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 05/19] 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 06/19] 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 07/19] 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 08/19] 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 09/19] 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 10/19] 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 11/19] 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 12/19] 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 13/19] 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 14/19] 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 15/19] 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 16/19] 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 17/19] 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 18/19] 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 19/19] 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)