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 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/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 new file mode 100644 index 000000000..9fc228d83 --- /dev/null +++ b/packages/kumo-engine/kumo/engine/_snapshot_io.py @@ -0,0 +1,471 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import hashlib +import json +import logging +import os +import re +import threading +import time +from concurrent.futures import ThreadPoolExecutor +from functools import partial +from pathlib import Path +from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple + +from pyarrow.fs import FileSystem + +_FORMAT = 1 +_MANIFEST = 'kumo_snapshot.json' +_INDEX_MANIFEST = '_MANIFEST.json' +_RANGE_BYTES = 64 << 20 +_CHUNK = 8 << 20 +_ATTEMPTS = 3 +_KEY = re.compile('[0-9a-f]{64}') +_PROGRESS_SECONDS = 30 + +logger = logging.getLogger(__name__) + + +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.engine.connector import file_store + + store = file_store(connector, path) + if store is None: + raise NotImplementedError( + f'{type(connector).__name__} cannot hold a snapshot export: export ' + 'needs a file destination (artifact root), such as an S3Connector, ' + 'GCSConnector or a FileConnector on a volume.') + return store + + +def _files(filesystem: FileSystem) -> Any: + from kumo_core.object_store_sink import _ArrowExportFiles + + return _ArrowExportFiles(filesystem, None) + + +def _check_key(graph_key: Any, data_version: Any, snapshot_id: Any) -> None: + from kumo.engine.graph import _snapshot_id + + if not (isinstance(graph_key, str) and isinstance(data_version, str) + and _KEY.fullmatch(graph_key) and _KEY.fullmatch(data_version) + and snapshot_id == _snapshot_id(graph_key, data_version)): + raise ValueError(f'{snapshot_id!r} is not a snapshot id.') + + +class _Client: + def __init__(self, connector: Any, path: str) -> None: + self._make = lambda: _location(connector, path)[0] + self.current = self._make() + self._lock = threading.Lock() + + def renew(self, stale: FileSystem) -> None: + with self._lock: + if self.current is stale: + self.current = self._make() + + +class _Progress: + def __init__(self, action: str, path: str, total: int) -> None: + self.action, self.path, self.total = action, path, total + self.moved = 0 + self.started = self.logged = time.monotonic() + self._lock = threading.Lock() + + def add(self, size: int) -> None: + with self._lock: + self.moved += size + now = time.monotonic() + if (now - self.logged < _PROGRESS_SECONDS + and self.moved < self.total): + return + self.logged = now + logger.info('Snapshot %s of %s: %d of %d bytes, %.2f GB/s', + self.action, self.path, self.moved, self.total, + self.moved / max(now - self.started, 1e-9) / 1e9) + + +def _transfer(connector: Any, path: str, items: Sequence[Any], + sizes: Sequence[int], work: Callable[..., Any], + action: str) -> List[Any]: + from kumo_core.sources.object_store import copies + from kumo_core.sources.parallel import run_in_parallel + + if not items: + return [] + stop = threading.Event() + with copies(items) as workers: + # After `copies`: Arrow's IO pool size sets the connection pool. + client = _Client(connector, path) + logger.info('Snapshot %s of %s: %d files, %d bytes, %d workers', + action, path, len(items), sum(sizes), workers) + progress = _Progress(action, path, sum(sizes)) + + def run(item: Any, size: int) -> Any: + result = _retried(client, work, item, stop) + progress.add(size) + return result + + return run_in_parallel( + [partial(run, item, size) for item, size in zip(items, sizes)], + workers=workers, stop=stop) + + +def _retried(client: _Client, work: Callable[..., Any], item: Any, + stop: threading.Event) -> Any: + attempt = 0 + while True: + _stopped(stop) + current = client.current + try: + return work(current, item, stop) + except (FileNotFoundError, PermissionError): + raise + except OSError as error: + if attempt + 1 == _ATTEMPTS: + raise + logger.warning('Retrying a snapshot transfer after: %s', error) + # Arrow keeps the credentials a client was built with. + client.renew(current) + stop.wait(2**attempt) + attempt += 1 + + +def _stopped(stop: threading.Event) -> None: + from kumo_core.sources.parallel import WorkCancelled + + if stop.is_set(): + raise WorkCancelled('abandoned after another transfer failed') + + +def _read_json(filesystem: FileSystem, path: str) -> Optional[Any]: + try: + return json.loads(_files(filesystem).read_text(path)) + except FileNotFoundError: + return None + + +def _open_output(filesystem: FileSystem, path: str, size: int) -> Any: + from kumo_core.object_store_sink import open_sized + + _files(filesystem).makedirs(path.rsplit('/', 1)[0]) + return open_sized(filesystem, path, size) + + +def _write(filesystem: FileSystem, path: str, payload: bytes) -> None: + with _open_output(filesystem, path, len(payload)) as writer: + writer.write(payload) + + +def _fingerprint(index_manifest: bytes) -> str: + return hashlib.sha256(index_manifest).hexdigest() + + +def _snapshot_manifest(snapshot: Any) -> bytes: + from kumo.engine._manifest import ( + RECONSTRUCTED, + SnapshotManifest, + SnapshotParents, + ) + from kumo.engine.graph import SNAPSHOT_MANIFEST_FILE + + try: + return (Path(snapshot.index_path) / SNAPSHOT_MANIFEST_FILE).read_bytes() + except OSError: + manifest = SnapshotManifest( + id=snapshot.id, parents=SnapshotParents(source_tables={}), + **RECONSTRUCTED) + return (json.dumps(manifest.to_dict(), indent=2, sort_keys=True) + + '\n').encode() + + +def export_location(snapshot: Any, path: str) -> str: + graph_key, data_version = snapshot._graph_key, snapshot._data_version + _check_key(graph_key, data_version, snapshot.id) + return f"{path.rstrip('/')}/{graph_key}/{data_version[:16]}" + + +def store(snapshot: Any, settings: Any) -> None: + r"""Export ``snapshot`` under ``/snapshots/`` unless a complete + export of it is there; failures warn.""" + import warnings + + try: + connector = settings.root.connector + path = settings.root.location('snapshots') + filesystem, root = _location(connector, + export_location(snapshot, path)) + 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 + + from kumo.engine._provenance import environment + from kumo.engine.graph import SNAPSHOT_MANIFEST_FILE + + source = Path(snapshot.index_path) + graph_key, data_version = snapshot._graph_key, snapshot._data_version + location = export_location(snapshot, path) + index_manifest = (source / _INDEX_MANIFEST).read_bytes() + sizes = { + entry['rel_path']: int(entry['size_bytes']) + for entry in json.loads(index_manifest)['files'] + } + snapshot_manifest = _snapshot_manifest(snapshot) + filesystem, root = _location(connector, location) + _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] + with open(source / name, 'rb', buffering=0) as reader, _open_output( + client, f'{root}/index/{name}', size) as writer: + for offset in range(0, size, _CHUNK): + _stopped(stop) + chunk = reader.read(min(_CHUNK, size - offset)) + if len(chunk) != min(_CHUNK, size - offset): + raise OSError(f'{source / name} changed while it was ' + 'exported') + writer.write(chunk) + + 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) + send(list(sizes), 'export') + crc32 = checksums.result() + 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().to_dict(), + '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 ConcurrentExportError( + 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: + 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}; ' + 're-run export.') + if not isinstance(manifest, dict): + raise _corrupt(path, 'manifest is not a JSON object') + if manifest.get('format') != _FORMAT: + raise ValueError( + f"Snapshot export at {path} has format " + f"{manifest.get('format')!r}; this kumo-engine reads format " + f"{_FORMAT}. Load it with the kumo-engine that exported it.") + _check_key(manifest.get('graph_key'), manifest.get('data_version'), + manifest.get('snapshot_id')) + span, crcs = manifest.get('range_bytes'), manifest.get('crc32') + 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)): + raise _corrupt(path, 'malformed manifest') + return root, manifest + + +def _download(root: str, staging: Path, client: FileSystem, + item: Tuple[str, int, int], stop: threading.Event) -> None: + name, offset, length = item + # Each open is a HEAD: pyarrow cannot open a file from a known size. + with client.open_input_file(f'{root}/index/{name}') as reader: + fd = os.open(staging / name, os.O_WRONLY) + try: + for start in range(offset, offset + length, _CHUNK): + _stopped(stop) + want = min(_CHUNK, offset + length - start) + chunk = reader.read_at(want, start) + if len(chunk) != want: + raise ValueError(f'{name} is truncated') + os.pwrite(fd, chunk, start) + getattr(os, 'fdatasync', os.fsync)(fd) + finally: + os.close(fd) + + +def _index_files( + path: str, root: str, manifest: dict, connector: Any +) -> Tuple[bytes, List[Tuple[str, int]], Optional[bytes]]: + from kumo_core.sources.layout import confined + + from kumo.engine.graph import SNAPSHOT_MANIFEST_FILE + + filesystem, _ = _location(connector, path) + with filesystem.open_input_stream( + f'{root}/index/{_INDEX_MANIFEST}') as reader: + index_manifest = reader.readall() + if _fingerprint(index_manifest) != manifest['index_manifest']: + raise _corrupt(path, f'{_INDEX_MANIFEST} is not the exported one') + try: + files = [(confined(entry['rel_path']), int(entry['size_bytes'])) + for entry in json.loads(index_manifest)['files']] + except Exception as error: + raise _corrupt(path, f'bad {_INDEX_MANIFEST}: {error}') from error + if len({name for name, _ in files}) != len(files): + raise _corrupt(path, f'{_INDEX_MANIFEST} lists a file twice') + try: + with filesystem.open_input_stream( + f'{root}/{SNAPSHOT_MANIFEST_FILE}') as reader: + snapshot_manifest = reader.readall() + except FileNotFoundError: + snapshot_manifest = None + return index_manifest, files, snapshot_manifest + + +def load_into(connector: Any, path: str, root: str, manifest: dict, + target: Path) -> None: + from diskgraph import ( + DiskgraphCorruptionError, + snapshot_range_checksums, + verify_snapshot_integrity, + ) + from kumo_core.diskgraph.snapshot import ( + DiskgraphIngestResult, + write_diskgraph_success_marker, + ) + from kumo_core.io.filesystem import fsync_directory + from kumo_core.sources.transfer import discard_destination + + from kumo.engine.graph import SNAPSHOT_MANIFEST_FILE + + span = manifest['range_bytes'] + index_manifest, files, snapshot_manifest = _index_files( + path, root, manifest, connector) + staging = target.with_name(f'{target.name}.loading') + discard_destination(str(target)) + discard_destination(str(staging)) + try: + staging.mkdir(parents=True) + needed = sum(size for _, size in files) + disk = os.statvfs(staging) + if disk.f_bavail * disk.f_frsize < needed: + raise OSError(f'Loading this snapshot needs {needed} free bytes ' + f'in {staging}.') + for name, size in files: + (staging / name).parent.mkdir(parents=True, exist_ok=True) + fd = os.open(staging / name, os.O_WRONLY | os.O_CREAT | os.O_EXCL) + try: + if size: + os.posix_fallocate(fd, 0, size) + finally: + os.close(fd) + ranges = sorted(((name, offset, min(span, size - offset)) + for name, size in files + for offset in range(0, size, span)), + key=lambda item: -item[2]) + try: + _transfer(connector, path, ranges, [item[2] for item in ranges], + partial(_download, root, staging), 'load') + except ValueError as error: + raise _corrupt(path, str(error)) from error + (staging / _INDEX_MANIFEST).write_bytes(index_manifest) + write_diskgraph_success_marker(str(staging), + DiskgraphIngestResult.for_index(target)) + try: + verify_snapshot_integrity(str(staging)) + loaded = snapshot_range_checksums(str(staging), span) + except DiskgraphCorruptionError as error: + raise _corrupt(path, str(error)) from error + for name, crcs in sorted(loaded.items()): + expected = manifest['crc32'].get(name) + if crcs != expected: + offset = next((index * span for index, crc in enumerate(crcs) + if not expected or index >= len(expected) + or crc != expected[index]), 0) + raise _corrupt( + path, f'{name} failed its checksum at offset {offset}') + _check_readable(staging, manifest) + if snapshot_manifest is not None: + (staging / SNAPSHOT_MANIFEST_FILE).write_bytes(snapshot_manifest) + for directory in sorted({(staging / name).parent + for name, _ in files} | {staging}): + fsync_directory(directory) + os.rename(staging, target) + fsync_directory(target.parent) + except BaseException: + discard_destination(str(staging)) + raise + + +def _check_readable(index: Path, manifest: dict) -> None: + from diskgraph import SamplerEngine + + try: + SamplerEngine(str(index), num_threads=1) + except Exception as error: + built = (manifest.get('environment') or {}).get('kumo_engine') + raise RuntimeError( + f'This DiskGraph cannot open the snapshot exported by kumo-engine ' + f'{built} ({error}). Rebuild it with this kumo-engine or load it ' + 'with the exporting one.') from error diff --git a/packages/kumo-engine/kumo/engine/graph.py b/packages/kumo-engine/kumo/engine/graph.py index 6b25908ee..3c182605a 100644 --- a/packages/kumo-engine/kumo/engine/graph.py +++ b/packages/kumo-engine/kumo/engine/graph.py @@ -60,11 +60,12 @@ """ import json +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, 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 @@ -86,6 +87,10 @@ snapshot_scratch_dir, ) from kumo.engine._storage import network_storage +from kumo.engine.connector.base import Connector + +if TYPE_CHECKING: + from kumo.engine.trainer.output_config import SnapshotOutputConfig from kumo.engine.table import Table SNAPSHOT_MANIFEST_FILE = PROVENANCE_FILE @@ -203,6 +208,78 @@ 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(store: Any, handle: Any) -> Optional[Any]: + from diskgraph import DiskgraphCorruptionError, read_success_marker + from kumo_core.diskgraph.snapshot import rebased + + path = store.path(handle) + if read_success_marker(path) is None: + return None + try: + store.verify(handle) + return rebased(store, handle) + except (DiskgraphCorruptionError, ValueError) as error: + warnings.warn(f"Reloading the damaged graph index at {path}: {error}", + UserWarning, stacklevel=3) + return None + + +def _snapshot_id(graph_key: str, data_version: str) -> str: + return f'{graph_key}@{data_version[:16]}' + + +_COMPARABLE = ('exact', 'timestamp', 'listing') + + +def _source_change(table: Table, recorded: Optional[SourceTableEntry], + record_sources: bool) -> Dict[str, Any]: + from kumoapi.source_table import CustomSQLSourceTable + + before = None if recorded is None else recorded.version + current, same_source = None, False + comparable = (recorded is not None + and recorded.version_kind in _COMPARABLE) + try: + source = table.definition.source_table + if not isinstance(source, CustomSQLSourceTable): + connector = table.connector + stamp = connector.source_backend.get_data_version( + connector, table.table_name) + now = source_table(connector, table.table_name, None, stamp, + record_sources=record_sources) + current = now.version + comparable = comparable and now.version_kind in _COMPARABLE + same_source = recorded is not None and ( + recorded.connector, recorded.location) == (now.connector, + now.location) + except Exception: + current = None + if not (same_source and comparable and before is not None + and current is not None): + status = 'unknown' + else: + status = 'unchanged' if before == current else 'changed' + return {'status': status, 'recorded': before, 'current': current} + + @dataclass(frozen=True) class GraphSnapshot: r"""The materialized artifact produced by `Graph.snapshot()`: an on-disk @@ -238,6 +315,94 @@ 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 self.id.partition('@')[0] + + @property + def _data_version(self) -> str: + """The full fingerprint of the source versions the index was built + from; `id` carries its first 16 characters.""" + return Path(self.index_path).name + + 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. The connector must hold files + at the path: an S3, GCS or file connector, or a volume or stage + whose sink offers a filesystem. + """ + from kumo.engine._snapshot_io import export + + return export(self, output_config.connector, output_config.path) + + @classmethod + def load_from_connector(cls, connector: Connector, *, path: str, + graph: Optional['Graph'] = None) -> 'GraphSnapshot': + r"""Copy the snapshot exported at `path` into the local store. + + Args: + path: the location `export` returned. + graph: if given, warn when its source tables changed since the + snapshot was built (see `source_changes`). + """ + from kumo_core.diskgraph.snapshot import DiskgraphSnapshotHandle + + from kumo.engine._snapshot_io import 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 another graph.", + UserWarning, stacklevel=2) + store = _store() + handle = DiskgraphSnapshotHandle( + graph_key=graph_key, graph_data_version=manifest['data_version']) + with _lock(graph_key): + result = _loaded(store, handle) + if result is None: + 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']) + if graph is not None: + changed = [ + alias + for alias, change in snapshot.source_changes(graph).items() + if change['status'] == 'changed' + ] + if changed: + warnings.warn( + f"Snapshot at {path} predates source table(s) " + f"{', '.join(changed)}; see `snapshot.source_changes`.", + UserWarning, stacklevel=2) + return snapshot + + def source_changes(self, graph: 'Graph') -> Dict[str, Dict[str, Any]]: + r"""Each of `graph`'s tables' source status since this snapshot: + `unchanged`, `changed` or `unknown`. Never raises.""" + try: + manifest = self.manifest + recorded = ({} if manifest is None or manifest.reconstructed else + dict(manifest.parents.source_tables)) + record_sources = (manifest is None + or manifest.record_sources is not False) + except Exception: + recorded, record_sources = {}, True + from kumo_core.sources.parallel import run_in_parallel + + tables = graph._named_tables + changes = run_in_parallel( + [ + partial(_source_change, table, recorded.get(alias), + record_sources) + for alias, table in tables.items() + ], workers=max(1, min(len(tables), 16)), stop=threading.Event()) + return dict(zip(tables, changes)) + @property def manifest(self) -> SnapshotManifest | None: """The snapshot manifest stored next to the index, if any.""" @@ -417,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` @@ -428,12 +594,13 @@ 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). """ - from filelock import FileLock - self.validate() for table in self._named_tables.values(): table.connector.validate_scratch_dir(get_scratch_dir()) @@ -441,7 +608,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,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. - with FileLock(f'{persist_dir}.lock'): - return self._snapshot_no_lock(graph_key, persist_dir, - record_sources) + from kumo.engine._mlops import resolve + + settings = resolve(mlops) + with _lock(graph_key): + snapshot = self._snapshot_no_lock(graph_key, persist_dir, + record_sources) + if settings is not None and settings.wants('snapshot'): + from kumo.engine._snapshot_io import store + + store(snapshot, settings) + return snapshot def _snapshot_no_lock(self, graph_key: str, persist_dir: Path, record_sources: bool = True) -> GraphSnapshot: @@ -478,9 +652,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 +662,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] = {} @@ -581,7 +751,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/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/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..98f5ba1c4 --- /dev/null +++ b/packages/kumo-engine/test/test_snapshot_export.py @@ -0,0 +1,570 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import hashlib +import json +import os +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 pandas as pd +import pytest +from pyarrow.fs import LocalFileSystem, SubTreeFileSystem + +import kumo.engine as kumo +import kumo.engine._snapshot_io as snapshot_io +from kumo.engine import GraphSnapshot, SnapshotOutputConfig +from kumo.engine._scratch import get_scratch_dir + + +@pytest.fixture +def scratch(tmp_path): + previous = get_scratch_dir() + + def use(name): + kumo.configure(scratch_dir=tmp_path / name) + + use('cpu') + yield use + kumo.configure(scratch_dir=previous) + + +@pytest.fixture +def exports(tmp_path): + return kumo.FileConnector(root_dir=str(tmp_path / 'exports')) + + +def _index_files(path): + return { + str(file.relative_to(path)): file.read_bytes() + for file in Path(path).rglob('*') if file.is_file() + and file.name not in ('_SUCCESS', '_provenance.json') + and not str(file.relative_to(path)).startswith('train_tables/') + } + + +def _export(snapshot, exports, tmp_path): + location = snapshot.export(SnapshotOutputConfig(exports, 'one')) + return location, tmp_path / 'exports' / location + + +def _edit_manifest(exported, edit): + path = exported / 'kumo_snapshot.json' + manifest = json.loads(path.read_text()) + edit(manifest) + path.write_text(json.dumps(manifest)) + + +def _edit_index_manifest(exported, edit): + path = exported / 'index/_MANIFEST.json' + manifest = json.loads(path.read_text()) + edit(manifest) + path.write_text(json.dumps(manifest)) + _edit_manifest( + exported, lambda m: m.update( + index_manifest=hashlib.sha256(path.read_bytes()).hexdigest())) + + +def _load(exports, location, **kwargs): + return GraphSnapshot.load_from_connector(exports, path=location, **kwargs) + + +def test_export_load_round_trip_relocates_and_is_reused( + tmp_path, scratch, graph, exports, monkeypatch): + monkeypatch.setattr(snapshot_io, '_CHUNK', 1024) + monkeypatch.setattr(snapshot_io, '_RANGE_BYTES', 4096) + snapshot = graph.snapshot() + location = snapshot.export(SnapshotOutputConfig(exports, 'one/')) + graph_id, data_version = snapshot.id.split('@') + assert location == f'one/{graph_id}/{data_version}' + assert snapshot.export(SnapshotOutputConfig(exports, 'one')) == location + exported = tmp_path / 'exports' / location + recorded = (snapshot.index_path / '_provenance.json').read_bytes() + assert (exported / '_provenance.json').read_bytes() == recorded + + shutil.rmtree(tmp_path / 'cpu') + scratch('gpu') + loaded = _load(exports, location, graph=graph) + assert (loaded.id, loaded.graph_config) == (snapshot.id, + snapshot.graph_config) + assert loaded.index_path.is_relative_to(tmp_path / 'gpu') + assert _index_files(loaded.index_path) == _index_files(exported / 'index') + assert (loaded.index_path / + '_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', + Mock(side_effect=AssertionError('ingest must not run'))) + monkeypatch.setattr(snapshot_io, 'load_into', + Mock(side_effect=AssertionError('already loaded'))) + assert _load(exports, location) == loaded + assert graph.snapshot() == loaded + + +def test_damaged_index_is_reused_by_snapshot_but_reloaded_by_load( + tmp_path, scratch, graph, exports, monkeypatch): + snapshot = graph.snapshot() + location, _ = _export(snapshot, exports, tmp_path) + expected = _index_files(snapshot.index_path) + stats = snapshot.index_path / 'col_stats.json' + stats.write_bytes(stats.read_bytes()[:-1]) + with warnings.catch_warnings(): + warnings.simplefilter('error') + assert graph.snapshot() == snapshot + + with pytest.warns(UserWarning, match='Reloading the damaged'): + loaded = _load(exports, location) + assert _index_files(loaded.index_path) == expected + + +def test_interrupted_reexport_is_not_loadable_and_reruns_from_scratch( + tmp_path, scratch, graph, exports, monkeypatch): + @contextmanager + def four(found, worker_ceiling=None): + yield 4 + + monkeypatch.setattr('kumo_core.sources.object_store.copies', four) + snapshot = graph.snapshot() + location, exported = _export(snapshot, exports, tmp_path) + (exported / 'index/stray.bin').write_bytes(b'x') + open_output, finished = snapshot_io._open_output, threading.Event() + + def interrupted(client, path, size): + if path.endswith('diskgraph_config.json'): + finished.wait(10) + raise OSError('connection reset') + if path.endswith('col_stats.json'): + finished.set() + return open_output(client, path, size) + + 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(SnapshotOutputConfig(exports, 'one')) + with pytest.raises(FileNotFoundError, match='No complete snapshot'): + _load(exports, 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( + snapshot.index_path) + + +def test_concurrent_exports_of_one_snapshot_both_succeed( + tmp_path, scratch, graph, exports, monkeypatch): + snapshot = graph.snapshot() + transfer, raced = snapshot_io._transfer, [] + + def racing(*args): + if not raced: + raced.append(None) + raced[0] = snapshot.export(SnapshotOutputConfig(exports, 'one')) + return transfer(*args) + + monkeypatch.setattr(snapshot_io, '_transfer', racing) + assert raced == [snapshot.export(SnapshotOutputConfig(exports, 'one'))] + assert _load(exports, raced[0]).id == snapshot.id + + +def 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_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(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, size): + if path.endswith('col_stats.json'): + clients.append(client) + if len(clients) == 1: + raise OSError('ExpiredToken') + return open_output(client, path, size) + + monkeypatch.setattr(snapshot_io, '_open_output', expiring) + graph.snapshot().export(SnapshotOutputConfig(exports, 'one')) + assert len(clients) == 2 and clients[0] is not clients[1] + + +def test_reexport_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(SnapshotOutputConfig(exports, 'one')) == location + assert stats.read_bytes() == original + assert _load(exports, location).id == snapshot.id + + +def test_failed_load_leaves_nothing_and_restarts_cleanly( + tmp_path, scratch, graph, exports): + snapshot = graph.snapshot() + location, exported = _export(snapshot, exports, tmp_path) + stats = exported / 'index/col_stats.json' + original = stats.read_bytes() + stats.write_bytes(original[:-1] + bytes([original[-1] ^ 1])) + + scratch('gpu') + with pytest.raises(ValueError, match='corrupt.*checksum'): + _load(exports, location) + graph_dir = tmp_path / 'gpu/snapshots' / snapshot.index_path.parent.name + assert not [path for path in graph_dir.iterdir() if path.suffix != '.lock'] + + stats.write_bytes(original) + (graph_dir / f'{snapshot.index_path.name}.loading').mkdir() + loaded = _load(exports, location) + assert _index_files(loaded.index_path) == _index_files(snapshot.index_path) + assert [ + path.name for path in graph_dir.iterdir() if path.suffix != '.lock' + ] == [snapshot.index_path.name] + + +def test_load_refuses_before_downloading_when_disk_is_too_small( + tmp_path, scratch, graph, exports, monkeypatch): + location, _ = _export(graph.snapshot(), exports, tmp_path) + scratch('gpu') + monkeypatch.setattr(snapshot_io.os, 'statvfs', + lambda path: SimpleNamespace(f_frsize=1, f_bavail=0)) + monkeypatch.setattr(snapshot_io, '_download', + Mock(side_effect=AssertionError('downloaded'))) + with pytest.raises(OSError, match='free bytes'): + _load(exports, location) + + +@pytest.mark.parametrize('index, edit', [ + (False, lambda manifest: manifest.update(data_version='../../victim')), + (False, lambda manifest: manifest.pop('crc32')), + (True, lambda manifest: manifest['files'][0].update( + rel_path='../../victim/x')), + (True, lambda manifest: manifest['files'].append(manifest['files'][0])), +]) +def test_load_rejects_a_hostile_or_malformed_manifest(tmp_path, scratch, graph, + exports, index, edit): + location, exported = _export(graph.snapshot(), exports, tmp_path) + victim = tmp_path / 'victim' + victim.mkdir() + (victim / 'keep').write_text('x') + (_edit_index_manifest if index else _edit_manifest)(exported, edit) + scratch('gpu/a/b') + with pytest.raises(ValueError): + _load(exports, location) + assert (victim / 'keep').read_text() == 'x' + + +def test_load_rejects_unknown_formats_and_accepts_unknown_fields( + tmp_path, scratch, graph, exports): + snapshot = graph.snapshot() + location, exported = _export(snapshot, exports, tmp_path) + scratch('gpu') + + _edit_manifest(exported, lambda m: m.update(format=2)) + with pytest.raises(ValueError, match='format 2'): + _load(exports, location) + + _edit_manifest(exported, lambda m: m.update(format=1, added_later={'x': 1})) + other = kumo.Graph(tables={'users': graph.tables['users']}) + with pytest.warns(UserWarning, match='another graph'): + loaded = _load(exports, location, graph=other) + assert loaded.id == snapshot.id + + +@pytest.mark.parametrize('recorded', ['newer', 'missing']) +def test_snapshot_manifest_travels_unchanged_or_is_reconstructed( + tmp_path, scratch, graph, exports, recorded): + snapshot = graph.snapshot() + path = snapshot.index_path / '_provenance.json' + if recorded == 'newer': + path.write_bytes(b'{"schema_version": 99, "future": [1]}') + else: + path.unlink() + location, exported = _export(snapshot, exports, tmp_path) + scratch('gpu') + with warnings.catch_warnings(): + warnings.simplefilter('ignore') + loaded = _load(exports, location) + shipped = (exported / '_provenance.json').read_bytes() + assert (loaded.index_path / '_provenance.json').read_bytes() == shipped + if recorded == 'newer': + assert shipped == b'{"schema_version": 99, "future": [1]}' + else: + manifest = loaded.manifest + assert (manifest.id, manifest.reconstructed) == (snapshot.id, True) + assert manifest.parents.source_tables == {} + + +def test_warehouse_connectors_are_not_supported(scratch, graph): + connector = SimpleNamespace(sink_backend=object()) + with pytest.raises(NotImplementedError, match='FileConnector'): + graph.snapshot().export(SnapshotOutputConfig(connector, 'x')) + with pytest.raises(NotImplementedError, match='FileConnector'): + _load(connector, 'x') + + +def test_load_explains_an_index_the_installed_diskgraph_cannot_open( + tmp_path, scratch, graph, exports): + location, exported = _export(graph.snapshot(), exports, tmp_path) + metadata = exported / 'index/index_metadata.json' + garbage = b'{' * metadata.stat().st_size + metadata.write_bytes(garbage) + + def unreadable(manifest): + manifest['crc32']['index_metadata.json'] = [zlib.crc32(garbage)] + manifest['environment']['kumo_engine'] = '0.0.1' + + _edit_manifest(exported, unreadable) + scratch('gpu') + with pytest.raises(RuntimeError, match='exported by kumo-engine 0.0.1'): + _load(exports, location) + + +@pytest.mark.parametrize('scheme', ['s3', 'gs']) +def test_object_store_round_trip(tmp_path, scratch, graph, monkeypatch, + scheme): + storage = tmp_path / 'storage' + storage.mkdir() + + class Bucket(SubTreeFileSystem): + def __init__(self, *args, **kwargs): + super().__init__(str(storage), LocalFileSystem()) + + monkeypatch.setattr('kumo_s3.filesystem', Bucket) + monkeypatch.setattr('kumo_gcs.backend.fs.GcsFileSystem', Bucket) + snapshot = graph.snapshot() + connector = kumo.S3Connector() if scheme == 's3' else kumo.GCSConnector() + location = snapshot.export( + SnapshotOutputConfig(connector, f'{scheme}://bucket/snapshots/')) + graph_id, data_version = snapshot.id.split('@') + assert location == f'{scheme}://bucket/snapshots/{graph_id}/{data_version}' + scratch('gpu') + loaded = _load(connector, location, graph=graph) + assert _index_files(loaded.index_path) == _index_files(snapshot.index_path) + + +def _set_mtime(path, seconds): + os.utime(path, (seconds, seconds)) + + +def _rewrite_events(data): + frame = pd.read_parquet(data / 'events.parquet') + frame.iloc[::-1].to_parquet(data / 'events.parquet') + _set_mtime(data / 'events.parquet', 1_700_009_123) + + +@pytest.fixture +def dated(tmp_path, graph): + data = tmp_path / 'data' + for name in ('users', 'events'): + _set_mtime(data / f'{name}.parquet', 1_700_000_123) + return data + + +@pytest.mark.parametrize('record_sources', [True, False]) +def test_source_changes_reads_stamps_without_staging( + scratch, graph, dated, monkeypatch, record_sources): + from kumo_core.sources.object_store import ObjectStoreSourceReader + + from kumo.engine.connector.file import _FileSourceBackend + + _set_mtime(dated / 'users.parquet', 1_700_002_800) + snapshot = graph.snapshot(record_sources=record_sources) + fail = Mock(side_effect=AssertionError('must not stage')) + monkeypatch.setattr('kumo_core.sources.prepare.read_table', fail) + monkeypatch.setattr(ObjectStoreSourceReader, 'prepare', fail) + monkeypatch.setattr(_FileSourceBackend, 'prepare_ingest', fail) + assert snapshot.source_changes(graph) == { + 'users': {'status': 'unchanged', 'recorded': '1700002800', + 'current': '1700002800'}, + 'events': {'status': 'unchanged', 'recorded': '1700000123', + 'current': '1700000123'}, + } + _rewrite_events(dated) + assert snapshot.source_changes(graph)['events'] == { + 'status': 'changed', 'recorded': '1700000123', 'current': '1700009123'} + + monkeypatch.setattr(_FileSourceBackend, 'get_data_version', + Mock(side_effect=OSError('listing failed'))) + assert snapshot.source_changes(graph)['events'] == { + 'status': 'unknown', 'recorded': '1700000123', 'current': None} + + +def test_source_changes_compares_directory_listings(tmp_path, scratch): + parts = tmp_path / 'dir' / 'users' + parts.mkdir(parents=True) + pd.DataFrame({'ID': range(5)}).to_parquet(parts / 'part-0.parquet') + connector = kumo.FileConnector(root_dir=str(tmp_path / 'dir')) + graph = kumo.Graph(tables={ + 'users': kumo.Table.from_source_table(connector['users'], + primary_key='ID') + }, edges=[]) + snapshot = graph.snapshot() + assert snapshot.source_changes(graph)['users']['status'] == 'unchanged' + pd.DataFrame({'ID': range(5, 9)}).to_parquet(parts / 'part-1.parquet') + assert snapshot.source_changes(graph)['users']['status'] == 'changed' + + +@pytest.mark.parametrize('edit', ['absent', 'reconstructed', 'hour_bucket']) +def test_source_changes_unknown_without_a_comparable_version( + scratch, graph, dated, edit): + snapshot = graph.snapshot() + path = snapshot.index_path / '_provenance.json' + manifest = json.loads(path.read_text()) + if edit == 'absent': + path.unlink() + elif edit == 'reconstructed': + path.write_text(json.dumps({**manifest, 'reconstructed': True})) + else: + for entry in manifest['parents']['source_tables'].values(): + entry['version_kind'] = 'hour_bucket' + path.write_text(json.dumps(manifest)) + assert {change['status'] for change in snapshot.source_changes( + graph).values()} == {'unknown'} + + +def test_load_warns_only_with_graph_and_only_for_changed_tables( + tmp_path, scratch, graph, dated, exports): + location, _ = _export(graph.snapshot(), exports, tmp_path) + scratch('gpu') + with warnings.catch_warnings(): + warnings.simplefilter('error') + _load(exports, location, graph=graph) + _rewrite_events(dated) + _load(exports, location) + with pytest.warns(UserWarning) as caught: + _load(exports, location, graph=graph) + [message] = [str(warning.message) for warning in caught] + assert 'source table(s) events;' in message + + +class _Bytes: + def __init__(self, data): + self.data = data + + def read_at(self, offset, length): + return self.data[offset:offset + length] + + def close(self): + pass + + +def test_a_sink_with_only_a_filesystem_holds_exports(tmp_path, scratch, + graph): + from kumo_core.chunked_fs import ChunkedFileHandler + from pyarrow.fs import FileSelector, FileType, PyFileSystem + + class Store(ChunkedFileHandler): + part_bytes = 1000 + + def __init__(self): + self.objects = {} + + def get_type_name(self): + return 'memory' + + def _size(self, path): + data = self.objects.get(path) + return None if data is None else len(data) + + def _list(self, directory): + return {p: len(d) for p, d in self.objects.items() + if p.startswith(directory + '/')} + + def _put(self, local, path): + self.objects[path] = Path(local).read_bytes() + + def _part(self, path, size): + return _Bytes(self.objects[path]) + + def _remove(self, path, *, tree): + for name in list(self.objects): + if name.startswith(path + '/') if tree else name == path: + del self.objects[name] + + store = Store() + sink = SimpleNamespace( + holds_files=lambda connector, path: path.startswith('@'), + filesystem=lambda connector, path: (PyFileSystem(store), + path.rstrip('/'))) + connector = SimpleNamespace(sink_backend=sink) + snapshot = graph.snapshot() + location = snapshot.export(SnapshotOutputConfig(connector, '@stage/s')) + assert any('.chunks/' in name for name in store.objects) + scratch('gpu') + loaded = _load(connector, location) + assert loaded.id == snapshot.id + filesystem, root = sink.filesystem(connector, location + '/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 new file mode 100644 index 000000000..c28b19712 --- /dev/null +++ b/packages/kumo-engine/test/test_snapshot_export_live.py @@ -0,0 +1,45 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import contextlib +import os +import uuid + +import pytest + +import kumo.engine as kumo +from kumo.engine import SnapshotOutputConfig +from kumo.engine._scratch import get_scratch_dir +from kumo.engine._snapshot_io import _location + +ROOTS = { + 's3': os.environ.get('KUMO_SNAPSHOT_LIVE_S3_ROOT'), + 'gs': os.environ.get('KUMO_SNAPSHOT_LIVE_GCS_ROOT'), +} + + +@pytest.mark.parametrize('scheme', ['s3', 'gs']) +def test_snapshot_round_trip_through_object_store(tmp_path, graph, scheme): + root = ROOTS[scheme] + if not root: + pytest.skip('set KUMO_SNAPSHOT_LIVE_S3_ROOT / _GCS_ROOT to run') + connector = kumo.S3Connector() if scheme == 's3' else kumo.GCSConnector() + path = f"{root.rstrip('/')}/engine-snapshot-{uuid.uuid4().hex[:12]}" + previous = get_scratch_dir() + try: + kumo.configure(scratch_dir=tmp_path / 'cpu') + snapshot = graph.snapshot() + location = snapshot.export(SnapshotOutputConfig(connector, path)) + kumo.configure(scratch_dir=tmp_path / 'gpu') + loaded = kumo.GraphSnapshot.load_from_connector( + connector, path=location, graph=graph) + finally: + kumo.configure(scratch_dir=previous) + filesystem, prefix = _location(connector, path) + filesystem.delete_dir_contents(prefix, missing_dir_ok=True) + with contextlib.suppress(OSError): + filesystem.delete_dir(prefix) + assert loaded.id == snapshot.id + assert (loaded.index_path / 'col_stats.json').read_bytes() == ( + snapshot.index_path / 'col_stats.json').read_bytes()