diff --git a/packages/kumo-api/kumoapi/json_serde.py b/packages/kumo-api/kumoapi/json_serde.py index 8f3356019..0a62c949f 100644 --- a/packages/kumo-api/kumoapi/json_serde.py +++ b/packages/kumo-api/kumoapi/json_serde.py @@ -9,7 +9,7 @@ from decimal import Decimal from enum import Enum from pathlib import PurePath -from typing import Any, Dict, Type, TypeVar, Union +from typing import Any, Dict, Type, TypeVar, Union, cast from uuid import UUID from pydantic import BaseModel, SecretStr @@ -139,5 +139,5 @@ def from_json(obj: Any, cls: Union[Type[T], object]) -> T: if isinstance(obj, str): obj = json.loads(obj) from pydantic import TypeAdapter - adapter = TypeAdapter(cls) + adapter = TypeAdapter(cast(Any, cls)) return adapter.validate_python(obj) diff --git a/packages/kumo-api/kumoapi/plugin.py b/packages/kumo-api/kumoapi/plugin.py index 2bb9dd998..b649e302b 100644 --- a/packages/kumo-api/kumoapi/plugin.py +++ b/packages/kumo-api/kumoapi/plugin.py @@ -29,6 +29,7 @@ Mapping, Optional, Protocol, + Tuple, TypeVar, runtime_checkable, ) @@ -43,6 +44,7 @@ __all__ = [ 'ArtifactLocation', 'ArtifactSink', + 'FileStoreSink', 'IngestResult', 'ModelFetchingSink', 'ModelListingSink', @@ -604,6 +606,22 @@ def write_files( ... +@runtime_checkable +class FileStoreSink(Protocol): + r"""A sink that also holds files: snapshot exports, model bundles and + export manifests go through its filesystem.""" + def holds_files(self, connector: SourceConnector, path: str) -> bool: + r"""Whether ``path`` is a file location of this sink rather than a + table name; decided from the path alone, without I/O.""" + ... + + def filesystem(self, connector: SourceConnector, + path: str) -> Tuple[Any, str]: + r"""A ``pyarrow.fs.FileSystem`` and ``path`` resolved on it, for a + path :meth:`holds_files` accepts.""" + ... + + @runtime_checkable class ModelPublishingSink(ArtifactSink, Protocol): r"""A sink that can also receive a trained model bundle. diff --git a/packages/kumo-core/src/kumo_core/artifacts.py b/packages/kumo-core/src/kumo_core/artifacts.py index b6e7a49c4..7a4226974 100644 --- a/packages/kumo-core/src/kumo_core/artifacts.py +++ b/packages/kumo-core/src/kumo_core/artifacts.py @@ -19,6 +19,10 @@ from .io.json import load_json, save_json from .manifest import RunManifest, RunPaths +#: Every engine artifact's manifest file; underscore-prefixed so Parquet +#: readers skip it. +PROVENANCE_FILE = '_provenance.json' + @runtime_checkable class ArtifactStore(Protocol): diff --git a/packages/kumo-core/src/kumo_core/chunked_fs.py b/packages/kumo-core/src/kumo_core/chunked_fs.py new file mode 100644 index 000000000..8972badd8 --- /dev/null +++ b/packages/kumo-core/src/kumo_core/chunked_fs.py @@ -0,0 +1,441 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +r"""A pyarrow filesystem over stores that only put, list, read and remove. + +A file larger than :attr:`ChunkedFileHandler.part_bytes` is stored as +``.chunks/00000 ... NNNNN`` followed by an empty ``_SUCCESS``; anything +smaller is one plain object. Readers see one logical file either way. A write +removes the other form; if a crash leaves both, the plain object wins. A writer +spools one part at a time on local disk, under :func:`spool_root`. A handler +that can stream an upload implements ``_stream``: a writer told the size +(:data:`SIZE_KEY` in the stream metadata) streams a file of up to +``part_bytes`` as one plain object and a larger one as parts, without local +disk; without a size, up to ``plain_bytes`` is spooled and stays plain. +""" + +import io +import os +import posixpath +import shutil +import tempfile +from typing import ( + BinaryIO, + Callable, + Dict, + List, + Optional, + Protocol, + Tuple, + Union, +) + +import pyarrow as pa +from pyarrow.fs import FileInfo, FileSystemHandler, FileType + +PART_BYTES = 256 << 20 +CHUNKS = '.chunks' +DONE = '_SUCCESS' +#: Stream metadata key carrying the size a written file will have. +SIZE_KEY = 'kumo-size' +_spool_dir: Union[None, str, Callable[[], str]] = None + + +def set_spool_dir(path: Union[None, str, Callable[[], str]]) -> None: + r"""Where writers spool parts unless a handler sets ``spool_dir``: a + path, a callable read at each write, or ``None`` (tempfile's).""" + global _spool_dir + _spool_dir = path + + +class Part(Protocol): + def read_at(self, offset: int, length: int) -> bytes: + ... + + def close(self) -> None: + ... + + +def spool_root(handler: 'ChunkedFileHandler') -> Optional[str]: + r"""Where ``handler`` stages local copies: its ``spool_dir``, else the + :func:`set_spool_dir` directory (created), else None (tempfile's).""" + root = handler.spool_dir or (_spool_dir() + if callable(_spool_dir) else _spool_dir) + if root is not None: + os.makedirs(root, exist_ok=True) + return root + + +class ChunkedFileHandler(FileSystemHandler): + r"""Subclasses supply the five store primitives below, and optionally + ``_stream``.""" + part_bytes = PART_BYTES + spool_dir: Optional[str] = None + #: Largest plain file when ``_stream`` is implemented (<= part_bytes). + plain_bytes = PART_BYTES + + def _stream(self, path: str) -> BinaryIO: + r"""Override to stream parts: a writable whose ``close`` commits the + object at ``path``.""" + raise NotImplementedError + + def _size(self, path: str) -> Optional[int]: + r"""Size of exactly the object at ``path``, or ``None``.""" + raise NotImplementedError + + def _list(self, directory: str) -> Dict[str, int]: + r"""Every object below ``directory/``, by full path.""" + raise NotImplementedError + + def _put(self, local: str, path: str) -> None: + r"""Uploads ``local``, whose basename matches ``path``'s.""" + raise NotImplementedError + + def _part(self, path: str, size: int) -> Part: + raise NotImplementedError + + def _remove(self, path: str, *, tree: bool) -> None: + r"""Removes the object at ``path``, or everything below it.""" + raise NotImplementedError + + def normalize_path(self, path): + return path.rstrip('/') + + def _layout(self, path: str) -> Optional[List[Tuple[str, int]]]: + cache = self.__dict__.setdefault('_layouts', {}) + layout = cache.get(path) + if layout is not None: + return layout + size = self._size(path) + layout = ([(path, size)] if size is not None else _logical_parts( + self._list(path + CHUNKS), path + CHUNKS)) + if layout is not None: + cache[path] = layout + return layout + + def _forget(self, path: str) -> None: + r"""Drops cached layouts at or below ``path``; writes and deletes + through this handler call it.""" + cache = self.__dict__.get('_layouts', {}) + path = path.rstrip('/') + for name in list(cache): + if name == path or name.startswith(path + '/'): + cache.pop(name, None) + + def get_file_info(self, paths): + infos = [] + for path in paths: + layout = self._layout(path) + if layout is not None: + infos.append( + FileInfo(path, FileType.File, + size=sum(size for _, size in layout))) + elif self._list(path): + infos.append(FileInfo(path, FileType.Directory)) + else: + infos.append(FileInfo(path, FileType.NotFound)) + return infos + + def get_file_info_selector(self, selector): + base = selector.base_dir.rstrip('/') + files = _logical_files(self._list(base)) + if not files and not selector.allow_not_found: + raise FileNotFoundError(f'{base} does not exist.') + infos: Dict[str, FileInfo] = {} + for path, size in files.items(): + parts = path[len(base) + 1:].split('/') + depths = len(parts) if selector.recursive else min(len(parts), 2) + for depth in range(1, depths): + directory = f"{base}/{'/'.join(parts[:depth])}" + infos[directory] = FileInfo(directory, FileType.Directory) + if selector.recursive or len(parts) == 1: + infos[path] = FileInfo(path, FileType.File, size=size) + return [infos[path] for path in sorted(infos)] + + def create_dir(self, path, recursive): + pass + + def delete_dir(self, path): + path = path.rstrip('/') + self._forget(path) + self._remove(path, tree=True) + + def delete_dir_contents(self, path, missing_dir_ok=False): + path = path.rstrip('/') + self._forget(path) + if not missing_dir_ok and not self._list(path): + raise FileNotFoundError(f'{path} does not exist.') + self._remove(path, tree=True) + + def delete_root_dir_contents(self): + raise NotImplementedError('Refusing to delete the whole store.') + + def delete_file(self, path): + self._forget(path) + plain = self._size(path) is not None + chunked = bool(self._list(path + CHUNKS)) + if not plain and not chunked: + raise FileNotFoundError(f'{path} does not exist.') + if plain: + self._remove(path, tree=False) + if chunked: + self._remove(path + CHUNKS, tree=True) + + def move(self, src, dest): + raise NotImplementedError(f'{self.get_type_name()} cannot rename.') + + def copy_file(self, src, dest): + raise NotImplementedError(f'{self.get_type_name()} cannot copy.') + + def open_append_stream(self, path, metadata): + raise NotImplementedError(f'{self.get_type_name()} cannot append.') + + def open_input_stream(self, path): + return self.open_input_file(path) + + def open_input_file(self, path): + layout = self._layout(path) + if layout is None: + raise FileNotFoundError(f'{path} does not exist.') + return pa.PythonFile(_Reader(self, layout), mode='r') + + def open_output_stream(self, path, metadata): + self._forget(path) + size = (metadata or {}).get(SIZE_KEY) + return pa.PythonFile( + _Writer(self, path, None if size is None else int(size)), + mode='w') + + +def _logical_parts(objects: Dict[str, int], + directory: str) -> Optional[List[Tuple[str, int]]]: + if f'{directory}/{DONE}' not in objects: + return None + return sorted((path, size) for path, size in objects.items() + if posixpath.basename(path) != DONE) + + +def _logical_files(objects: Dict[str, int]) -> Dict[str, int]: + r"""Collapses committed ``.chunks`` directories into their files.""" + files: Dict[str, int] = {} + chunked: Dict[str, Dict[str, int]] = {} + for path, size in objects.items(): + directory = posixpath.dirname(path) + if directory.endswith(CHUNKS): + chunked.setdefault(directory, {})[path] = size + else: + files[path] = size + for directory, parts in chunked.items(): + name = directory[:-len(CHUNKS)] + layout = _logical_parts(parts, directory) + if layout is not None and name not in files: + files[name] = sum(size for _, size in layout) + return files + + +class _Reader(io.RawIOBase): + def __init__(self, handler: ChunkedFileHandler, + layout: List[Tuple[str, int]]) -> None: + self._handler = handler + self._layout: List[Tuple[str, int, int]] = [] + start = 0 + for path, size in layout: + self._layout.append((path, start, size)) + start += size + self._size = start + self._open: Dict[str, Part] = {} + self._position = 0 + + def readable(self): + return True + + def seekable(self): + return True + + def tell(self): + return self._position + + def seek(self, offset, whence=io.SEEK_SET): + base = { + io.SEEK_SET: 0, + io.SEEK_CUR: self._position, + io.SEEK_END: self._size + }[whence] + self._position = max(base + offset, 0) + return self._position + + def _pieces(self, end: int): + for path, start, length in self._layout: + if start + length <= self._position or start >= end: + continue + offset = max(self._position, start) - start + want = min(end, start + length) - start - offset + part = self._open.get(path) + if part is None: + part = self._open[path] = self._handler._part(path, length) + chunk = part.read_at(offset, want) + if len(chunk) != want: + raise OSError(f'{path} returned {len(chunk)} of {want} bytes') + yield chunk + + def read(self, size=-1): + end = self._size if size is None or size < 0 else min( + self._position + size, self._size) + if end <= self._position: + return b'' + pieces = self._pieces(end) + first = next(pieces) + if len(first) == end - self._position: + self._position = end + return first + out = bytearray(end - self._position) + out[:len(first)] = first + filled = len(first) + for chunk in pieces: + out[filled:filled + len(chunk)] = chunk + filled += len(chunk) + self._position = end + return out + + def readall(self): + return self.read() + + def readinto(self, buffer): + view = memoryview(buffer).cast('B') + end = min(self._position + len(view), self._size) + filled = 0 + for chunk in self._pieces(end): + view[filled:filled + len(chunk)] = chunk + filled += len(chunk) + self._position += filled + return filled + + def close(self): + if not self.closed: + for part in self._open.values(): + part.close() + super().close() + + +class _Writer(io.RawIOBase): + def __init__(self, handler: ChunkedFileHandler, path: str, + size: Optional[int] = None) -> None: + self._handler = handler + self._path = path + self._dir = tempfile.mkdtemp(prefix='kumo-chunked-', + dir=spool_root(handler)) + self._spool = open(os.path.join(self._dir, 'spool'), 'wb') + self._filled = 0 + self._parts = 0 + self._written = 0 + self._part: Optional[BinaryIO] = None + self._streams = type(handler)._stream is not ChunkedFileHandler._stream + self._limit = (min(handler.plain_bytes, handler.part_bytes) + if self._streams else handler.part_bytes) + self._plain: Optional[BinaryIO] = None + if self._streams and size is not None: + if size <= handler.part_bytes: + self._plain = handler._stream(path) + else: + self._clear() + self._part = self._open_part() + + def writable(self): + return True + + def tell(self): + return self._written + + def write(self, data): + view = memoryview(data).cast('B') + total = len(view) + if self._plain is not None: + self._plain.write(view) + self._written += total + return total + while view: + if self._part is not None: + if self._filled == self._handler.part_bytes: + self._next_stream() + take = min(len(view), self._handler.part_bytes - self._filled) + self._part.write(view[:take]) + else: + if self._filled == self._limit: + if self._streams: + self._start_stream() + continue + self._flush_part() + take = min(len(view), self._limit - self._filled) + self._spool.write(view[:take]) + self._filled += take + view = view[take:] + self._written += total + return total + + def _upload(self, target: str) -> None: + self._spool.close() + local = os.path.join(self._dir, posixpath.basename(target)) + os.replace(os.path.join(self._dir, 'spool'), local) + try: + self._handler._put(local, target) + finally: + os.remove(local) + self._spool = open(os.path.join(self._dir, 'spool'), 'wb') + self._filled = 0 + + def _clear(self) -> None: + self._handler._remove(self._path + CHUNKS, tree=True) + if self._handler._size(self._path) is not None: + self._handler._remove(self._path, tree=False) + + def _flush_part(self) -> None: + if not self._parts: + self._clear() + self._upload(f'{self._path}{CHUNKS}/{self._parts:05d}') + self._parts += 1 + + def _open_part(self) -> BinaryIO: + return self._handler._stream( + f'{self._path}{CHUNKS}/{self._parts:05d}') + + def _start_stream(self) -> None: + self._clear() + self._spool.close() + spooled = os.path.join(self._dir, 'spool') + self._part = self._open_part() + with open(spooled, 'rb') as reader: + shutil.copyfileobj(reader, self._part, 8 << 20) + self._spool = open(spooled, 'wb') + + def _next_stream(self) -> None: + assert self._part is not None + self._part.close() + self._parts += 1 + self._part = self._open_part() + self._filled = 0 + + def close(self): + if self.closed: + return + try: + if self._plain is not None: + self._plain.close() + self._plain = None + if self._handler._list(self._path + CHUNKS): + self._handler._remove(self._path + CHUNKS, tree=True) + elif self._part is not None: + self._part.close() + self._part = None + self._upload(f'{self._path}{CHUNKS}/{DONE}') + elif self._parts: + self._flush_part() + self._upload(f'{self._path}{CHUNKS}/{DONE}') + else: + self._upload(self._path) + if self._handler._list(self._path + CHUNKS): + self._handler._remove(self._path + CHUNKS, tree=True) + finally: + self._handler._forget(self._path) + self._spool.close() + shutil.rmtree(self._dir, ignore_errors=True) + super().close() diff --git a/packages/kumo-core/src/kumo_core/object_store_sink.py b/packages/kumo-core/src/kumo_core/object_store_sink.py index ec34f89f9..3951ea5b3 100644 --- a/packages/kumo-core/src/kumo_core/object_store_sink.py +++ b/packages/kumo-core/src/kumo_core/object_store_sink.py @@ -2,11 +2,13 @@ # All rights reserved. # SPDX-License-Identifier: Apache-2.0 +import os import posixpath import shutil from abc import ABC, abstractmethod +from dataclasses import replace from pathlib import Path -from typing import Mapping, Optional, Tuple +from typing import Any, Mapping, Optional, Tuple import pandas as pd from kumoapi.jobs import WriteMode @@ -21,6 +23,69 @@ _OBJECT_STORES = frozenset({'s3', 'gcs', 'abfs'}) +def _chunked(filesystem: FileSystem) -> bool: + from kumo_core.chunked_fs import ChunkedFileHandler + from pyarrow.fs import PyFileSystem + + return isinstance(filesystem, PyFileSystem) and isinstance( + filesystem.handler, ChunkedFileHandler) + + +def open_sized(filesystem: FileSystem, path: str, size: int) -> Any: + """``open_output_stream(path)``, telling a chunked filesystem the size + the file will have.""" + if _chunked(filesystem): + from kumo_core.chunked_fs import SIZE_KEY + + return filesystem.open_output_stream( + path, metadata={SIZE_KEY: str(size)}) + return filesystem.open_output_stream(path) + + +def upload_bundle(filesystem: FileSystem, destination: str, + bundle_dir: str) -> None: + """Copy every file of ``bundle_dir`` under ``destination``, the + provenance file last.""" + from kumo_core.artifacts import PROVENANCE_FILE + + root = Path(bundle_dir) + for source in sorted((item for item in root.rglob('*') if item.is_file()), + key=lambda item: + (item.name == PROVENANCE_FILE, item)): + target = f'{destination}/{source.relative_to(root).as_posix()}' + _prepared(filesystem, posixpath.dirname(target)) + with open(source, 'rb') as reader, open_sized( + filesystem, target, source.stat().st_size) as writer: + shutil.copyfileobj(reader, writer, _CHUNK) + + +def download_bundle(filesystem: FileSystem, source: str, + local_dir: str) -> bool: + """Replace ``local_dir`` with the files under ``source``; False if there + are none.""" + from kumo_core.sources.layout import confined + from pyarrow.fs import FileSelector, FileType + + destination = Path(local_dir) + if destination.exists(): + shutil.rmtree(destination) + destination.mkdir(parents=True) + found = False + selector = FileSelector(source, recursive=True, allow_not_found=True) + for info in filesystem.get_file_info(selector): + if info.type is not FileType.File: + continue + found = True + # Remote keys are untrusted; each relative path must stay inside + # the local bundle, even when it contains parent components. + target = destination / confined(info.path[len(source) + 1:]) + target.parent.mkdir(parents=True, exist_ok=True) + with filesystem.open_input_stream(info.path) as reader, \ + open(target, 'wb') as writer: + shutil.copyfileobj(reader, writer, _CHUNK) + return found + + class ObjectStoreSinkBackend(ABC): scheme: str provider: str @@ -31,6 +96,16 @@ def _destination(self, connector: SourceConnector, """Return the authenticated filesystem and validated bucket/key path.""" raise NotImplementedError + def holds_files(self, connector: SourceConnector, path: str) -> bool: + """Every path of an object store is a file location.""" + return True + + def filesystem(self, connector: SourceConnector, + path: str) -> Tuple[FileSystem, str]: + """The authenticated filesystem and resolved ``path``.""" + filesystem, resolved = self._destination(connector, path) + return filesystem, resolved.rstrip('/') + def write_table( self, connector: SourceConnector, @@ -98,14 +173,7 @@ def publish_model( filesystem, destination = self._destination(connector, path) destination = destination.rstrip('/') - root = Path(bundle_dir) - for source in sorted(item for item in root.rglob('*') - if item.is_file()): - target = f'{destination}/{source.relative_to(root).as_posix()}' - _prepared(filesystem, posixpath.dirname(target)) - with open(source, 'rb') as reader, \ - filesystem.open_output_stream(target) as writer: - shutil.copyfileobj(reader, writer, _CHUNK) + upload_bundle(filesystem, destination, bundle_dir) return PublishedModel( name=posixpath.basename(destination), version=destination, @@ -124,9 +192,6 @@ def fetch_model( ) -> PublishedModel: """Replace the local directory with the published bundle.""" del project_id - from kumo_core.sources.layout import confined - from pyarrow.fs import FileSelector, FileType - if location is not None: raise NotImplementedError( f"The local backend does not import a model by name from " @@ -141,25 +206,7 @@ def fetch_model( filesystem, source = self._destination(connector, path) source = source.rstrip('/') - destination = Path(local_dir) - if destination.exists(): - shutil.rmtree(destination) - destination.mkdir(parents=True) - - found = False - selector = FileSelector(source, recursive=True, allow_not_found=True) - for info in filesystem.get_file_info(selector): - if info.type is not FileType.File: - continue - found = True - # Remote keys are untrusted; each relative path must stay inside - # the local bundle, even when it contains parent components. - target = destination / confined(info.path[len(source) + 1:]) - target.parent.mkdir(parents=True, exist_ok=True) - with filesystem.open_input_stream(info.path) as reader, \ - open(target, 'wb') as writer: - shutil.copyfileobj(reader, writer, _CHUNK) - if not found: + if not download_bundle(filesystem, source, local_dir): raise ValueError( f"Nothing is published at {self.scheme}://{source}. Check the " f"path the model was exported to.") @@ -232,8 +279,21 @@ def _write_one(filesystem: FileSystem, frame: pd.DataFrame, path: str, with filesystem.open_output_stream(path) as writer: csv.write_csv(table, writer) return - pq.write_table(pa.Table.from_pandas(frame, preserve_index=False), path, - filesystem=filesystem) + table = pa.Table.from_pandas(frame, preserve_index=False) + if not _chunked(filesystem): + pq.write_table(table, path, filesystem=filesystem) + return + import tempfile + + from kumo_core.chunked_fs import spool_root + # Spool so the size is known and a file up to one part stays plain. + with tempfile.TemporaryDirectory( + dir=spool_root(filesystem.handler)) as spool: + local = os.path.join(spool, 'frame.parquet') + pq.write_table(table, local) + with open(local, 'rb') as reader, open_sized( + filesystem, path, os.path.getsize(local)) as writer: + shutil.copyfileobj(reader, writer, _CHUNK) def _prepared(filesystem: FileSystem, directory: str) -> None: @@ -309,9 +369,9 @@ def copy_file(self, source, destination, overwrite=True): # CSV conversion already applied this destination's metadata. columns = self.extra_columns if source.endswith('.parquet') else None with prepared_artifact_file(source, columns) as prepared: - with open(prepared, - 'rb') as reader, self.filesystem.open_output_stream( - destination) as writer: + with open(prepared, 'rb') as reader, open_sized( + self.filesystem, destination, + os.path.getsize(prepared)) as writer: shutil.copyfileobj(reader, writer, _CHUNK) def copy_files_threaded(self, pairs, overwrite, max_workers): @@ -329,3 +389,35 @@ def convert_csv(self, source, destination): ) with prepared_artifact_file(source, self.extra_columns) as prepared: parquet_to_csv(prepared, destination) + + +class FilesystemSink(ObjectStoreSinkBackend): + """The object-store file writers over any sink's ``filesystem()``, for + the file paths of a sink that also writes tables: paths in, bare + resolved paths out.""" + scheme = 'filesystem' + + def __init__(self, sink: Any) -> None: + self._sink = sink + self.provider = type(sink).__name__ + + def _destination(self, connector, coordinate): + filesystem, path = self._sink.filesystem(connector, coordinate) + return filesystem, path + ('/' if coordinate.endswith('/') else '') + + def _bare(self, uri: str) -> str: + return uri.removeprefix(f'{self.scheme}://') + + def write_table(self, *args, **kwargs): + return self._bare(super().write_table(*args, **kwargs)) + + def write_files(self, *args, **kwargs): + return self._bare(super().write_files(*args, **kwargs)) + + def publish_model(self, *args, **kwargs): + published = super().publish_model(*args, **kwargs) + return replace(published, uri=self._bare(published.uri)) + + def fetch_model(self, *args, **kwargs): + fetched = super().fetch_model(*args, **kwargs) + return replace(fetched, uri=self._bare(fetched.uri)) diff --git a/packages/kumo-core/tests/test_chunked_fs.py b/packages/kumo-core/tests/test_chunked_fs.py new file mode 100644 index 000000000..f93b19e60 --- /dev/null +++ b/packages/kumo-core/tests/test_chunked_fs.py @@ -0,0 +1,320 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import io +import os + +import pytest +from kumo_core import chunked_fs +from kumo_core.chunked_fs import ChunkedFileHandler +from pyarrow.fs import FileSelector, FileType, PyFileSystem + +PART = 16 + + +class _Bytes: + def __init__(self, data: bytes) -> None: + self.data = data + + def read_at(self, offset, length): + return self.data[offset:offset + length] + + def close(self): + pass + + +class MemoryHandler(ChunkedFileHandler): + part_bytes = PART + + def __init__(self): + self.objects = {} + self.spooled = [] + + def get_type_name(self): + return 'memory' + + def _size(self, path): + data = self.objects.get(path) + return None if data is None else len(data) + + def _list(self, directory): + return { + path: len(data) + for path, data in self.objects.items() + if path.startswith(directory + '/') + } + + def _put(self, local, path): + assert os.path.basename(local) == os.path.basename(path) + self.spooled.append(os.path.dirname(os.path.dirname(local))) + with open(local, 'rb') as reader: + self.objects[path] = reader.read() + + def _part(self, path, size): + return _Bytes(self.objects[path]) + + def _remove(self, path, *, tree): + for name in list(self.objects): + if name.startswith(path + '/') if tree else name == path: + del self.objects[name] + + +@pytest.fixture +def store(): + handler = MemoryHandler() + return handler, PyFileSystem(handler) + + +def _write(fs, path, data): + with fs.open_output_stream(path) as writer: + for start in range(0, len(data), 5): + writer.write(data[start:start + 5]) + + +def _read(fs, path): + with fs.open_input_stream(path) as reader: + return reader.readall() + + +@pytest.mark.parametrize('size', [0, 3, PART]) +def test_files_up_to_one_part_are_one_plain_object(store, size): + handler, fs = store + data = bytes(range(size)) + _write(fs, 'root/a.bin', data) + assert handler.objects == {'root/a.bin': data} + assert fs.get_file_info('root/a.bin').size == size + assert _read(fs, 'root/a.bin') == data + + +@pytest.mark.parametrize('size', [2 * PART + 5, 3 * PART]) +def test_larger_files_are_parts_committed_last_and_read_back_exactly( + store, size): + handler, fs = store + data = os.urandom(size) + _write(fs, 'root/a.bin', data) + assert sorted(handler.objects) == [ + 'root/a.bin.chunks/00000', 'root/a.bin.chunks/00001', + 'root/a.bin.chunks/00002', 'root/a.bin.chunks/_SUCCESS' + ] + info = fs.get_file_info('root/a.bin') + assert (info.type, info.size) == (FileType.File, len(data)) + assert _read(fs, 'root/a.bin') == data + with fs.open_input_file('root/a.bin') as reader: + assert reader.size() == len(data) + assert reader.read_at(10, PART - 4) == data[PART - 4:PART + 6] + assert reader.read_at(0, 2 * PART) == b'' + assert reader.read_at(PART, PART) == data[PART:2 * PART] + assert reader.read_at(100, PART + 1) == data[PART + 1:] + + +def test_an_uncommitted_chunked_write_is_not_a_file(store): + handler, fs = store + handler.objects['root/a.bin.chunks/00000'] = b'x' * PART + assert fs.get_file_info('root/a.bin').type is not FileType.File + with pytest.raises(FileNotFoundError): + fs.open_input_file('root/a.bin') + + +@pytest.mark.parametrize('first,second', [(3, 2 * PART), (2 * PART, 3)]) +def test_a_write_removes_the_other_form(store, first, second): + handler, fs = store + _write(fs, 'root/a.bin', b'1' * first) + _write(fs, 'root/a.bin', b'2' * second) + assert _read(fs, 'root/a.bin') == b'2' * second + plain = 'root/a.bin' in handler.objects + assert plain == (second <= PART) + assert any('.chunks/' in name for name in handler.objects) != plain + + +def test_plain_wins_when_a_crash_left_both_forms(store): + handler, fs = store + _write(fs, 'root/a.bin', b'c' * (2 * PART)) + handler.objects['root/a.bin'] = b'plain' + assert _read(fs, 'root/a.bin') == b'plain' + infos = fs.get_file_info(FileSelector('root', recursive=True)) + assert [(i.path, i.size) for i in infos] == [('root/a.bin', 5)] + + +def test_listing_shows_logical_files_and_directories(store): + _, fs = store + _write(fs, 'root/top.bin', b'1') + _write(fs, 'root/dir/big.bin', b'2' * (PART + 1)) + _write(fs, 'root/dir/sub/small.bin', b'3') + recursive = fs.get_file_info(FileSelector('root', recursive=True)) + assert [(i.path, i.type, i.size) for i in recursive] == [ + ('root/dir', FileType.Directory, None), + ('root/dir/big.bin', FileType.File, PART + 1), + ('root/dir/sub', FileType.Directory, None), + ('root/dir/sub/small.bin', FileType.File, 1), + ('root/top.bin', FileType.File, 1), + ] + shallow = fs.get_file_info(FileSelector('root')) + assert [i.path for i in shallow] == ['root/dir', 'root/top.bin'] + assert fs.get_file_info('root/dir').type is FileType.Directory + assert fs.get_file_info('root/none').type is FileType.NotFound + assert fs.get_file_info(FileSelector('nothing', allow_not_found=True)) == [] + with pytest.raises(FileNotFoundError): + fs.get_file_info(FileSelector('nothing')) + + +def test_deletes(store): + handler, fs = store + _write(fs, 'root/a.bin', b'a' * (PART + 1)) + _write(fs, 'root/b.bin', b'b') + fs.delete_file('root/a.bin') + assert sorted(handler.objects) == ['root/b.bin'] + with pytest.raises(FileNotFoundError): + fs.delete_file('root/a.bin') + fs.delete_dir_contents('root') + assert handler.objects == {} + fs.delete_dir_contents('root', missing_dir_ok=True) + with pytest.raises(FileNotFoundError): + fs.delete_dir_contents('root') + with pytest.raises(NotImplementedError): + fs.move('root/a', 'root/b') + + +def test_parts_spool_under_the_spool_dir_and_are_removed( + store, tmp_path, monkeypatch): + handler, fs = store + monkeypatch.setattr(chunked_fs, '_spool_dir', None) + chunked_fs.set_spool_dir(str(tmp_path / 'scratch')) + _write(fs, 'root/a.bin', b'a' * (2 * PART + 1)) + handler.spool_dir = str(tmp_path / 'own') + _write(fs, 'root/b.bin', b'b') + assert handler.spooled == [str(tmp_path / 'scratch')] * 4 + [ + str(tmp_path / 'own') + ] + assert not os.listdir(tmp_path / 'scratch') + assert not os.listdir(tmp_path / 'own') + + +class _Upload(io.BytesIO): + def __init__(self, handler, path): + super().__init__() + self.handler, self.path = handler, path + + def close(self): + self.handler.objects[self.path] = self.getvalue() + self.handler.streamed.append(self.path) + super().close() + + +class StreamingHandler(MemoryHandler): + plain_bytes = 6 + + def __init__(self): + super().__init__() + self.streamed = [] + + def _stream(self, path): + return _Upload(self, path) + + +@pytest.mark.parametrize('size', [0, 6, 7, PART, 2 * PART + 5, 3 * PART]) +def test_a_streaming_handler_streams_parts_and_spools_only_small_files(size): + handler = StreamingHandler() + fs = PyFileSystem(handler) + _write(fs, 'root/a.bin', b'x' * (2 * PART)) + data = os.urandom(size) + _write(fs, 'root/a.bin', data) + assert _read(fs, 'root/a.bin') == data + if size <= 6: + assert handler.objects == {'root/a.bin': data} + return + parts = -(-size // PART) + assert sorted(handler.objects) == [ + f'root/a.bin.chunks/{i:05d}' for i in range(parts) + ] + ['root/a.bin.chunks/_SUCCESS'] + assert handler.streamed[-parts:] == [ + f'root/a.bin.chunks/{i:05d}' for i in range(parts) + ] + assert len(handler.spooled) == 2 + with fs.open_input_file('root/a.bin') as reader: + assert reader.read_at(4, PART - 2) == data[PART - 2:PART + 2] + + +def test_opening_a_file_again_reuses_its_layout(store): + handler, fs = store + _write(fs, 'root/a.bin', b'a' * (2 * PART)) + probes = [] + size, listing = handler._size, handler._list + handler._size = lambda path: probes.append(path) or size(path) + handler._list = lambda path: probes.append(path) or listing(path) + for _ in range(3): + assert _read(fs, 'root/a.bin') == b'a' * (2 * PART) + assert len(probes) == 2 + _write(fs, 'root/a.bin', b'new') + assert _read(fs, 'root/a.bin') == b'new' + fs.delete_dir_contents('root/') + with pytest.raises(FileNotFoundError): + fs.open_input_file('root/a.bin') + + +def test_layout_cache_survives_concurrent_opens_and_writes(store): + import threading + + handler, fs = store + for i in range(20): + _write(fs, f'root/{i}.bin', b'x' * (PART + 1)) + errors = [] + + def work(i): + try: + for _ in range(50): + _read(fs, f'root/{i}.bin') + _write(fs, f'root/{i}.bin', b'x' * (PART + 1)) + except Exception as error: + errors.append(error) + + threads = [threading.Thread(target=work, args=(i, )) for i in range(20)] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + assert errors == [] + + +@pytest.mark.parametrize('size', [7, PART, PART + 1, 3 * PART]) +def test_a_known_size_streams_without_spooling(size): + from kumo_core.object_store_sink import open_sized + + handler = StreamingHandler() + fs = PyFileSystem(handler) + _write(fs, 'root/a.bin', b'x' * (2 * PART)) + handler.spooled.clear() + data = os.urandom(size) + with open_sized(fs, 'root/a.bin', size) as writer: + for start in range(0, size, 5): + writer.write(data[start:start + 5]) + assert _read(fs, 'root/a.bin') == data + plain = 'root/a.bin' in handler.objects + assert plain == (size <= PART) + assert any('.chunks/' in name for name in handler.objects) != plain + # Only the empty _SUCCESS marker of a chunked file is put from disk. + assert len(handler.spooled) == (0 if plain else 1) + + +def test_the_spool_dir_can_be_read_at_each_write(store, tmp_path, + monkeypatch): + handler, fs = store + root = [tmp_path / 'one'] + monkeypatch.setattr(chunked_fs, '_spool_dir', None) + chunked_fs.set_spool_dir(lambda: str(root[0])) + _write(fs, 'root/a.bin', b'a') + root[0] = tmp_path / 'two' + _write(fs, 'root/b.bin', b'b') + assert handler.spooled == [str(tmp_path / 'one'), str(tmp_path / 'two')] + + +def test_reads_spanning_parts_fill_the_callers_buffer(store): + _, fs = store + data = os.urandom(3 * PART + 3) + _write(fs, 'root/a.bin', data) + with fs.open_input_file('root/a.bin') as reader: + buffer = bytearray(2 * PART) + reader.seek(PART - 3) + assert reader.readinto(buffer) == 2 * PART + assert bytes(buffer) == data[PART - 3:3 * PART - 3] + assert reader.read_at(4, 2) == data[2:6] diff --git a/packages/kumo-engine/kumo/engine/_scratch.py b/packages/kumo-engine/kumo/engine/_scratch.py index f3642e9f5..7381e5cb5 100644 --- a/packages/kumo-engine/kumo/engine/_scratch.py +++ b/packages/kumo-engine/kumo/engine/_scratch.py @@ -68,6 +68,19 @@ def get_scratch_dir() -> Path: return _root +def _spool_dir() -> str: + return str(get_scratch_dir() / "spool") + + +def _spool_parts() -> None: + from kumo_core.chunked_fs import set_spool_dir + + set_spool_dir(_spool_dir) + + +_spool_parts() + + def new_run_id() -> str: r"""A fresh, random id for one run -- a :class:`~kumo.engine.trainer.Trainer` instance, a `Trainer.predict` call, a `PredictiveQuery. diff --git a/packages/kumo-engine/kumo/engine/connector/__init__.py b/packages/kumo-engine/kumo/engine/connector/__init__.py index d328cf3ee..64c869506 100644 --- a/packages/kumo-engine/kumo/engine/connector/__init__.py +++ b/packages/kumo-engine/kumo/engine/connector/__init__.py @@ -17,6 +17,8 @@ `_response`) and the `_bind_connector` round-trip. See `kumo.engine.connector.base`. """ +from typing import Any, Optional, Tuple + from kumo.engine.connector.base import Connector from kumo.engine.connector.bigquery import BigQueryConnector from kumo.engine.connector.databricks import DatabricksConnector @@ -34,6 +36,31 @@ from kumo.engine.connector.snowflake import SnowflakeConnector from kumo.engine.connector.source_table import SourceTable, SQLSourceTable + +def _file_sink(connector: Any, path: str) -> Any: + from kumoapi.plugin import FileStoreSink + + try: + sink = connector.sink_backend + except Exception: + return None + return sink if isinstance(sink, FileStoreSink) and sink.holds_files( + connector, path) else None + + +def is_file_destination(connector: Any, path: str) -> bool: + r"""Whether ``path`` is a file location of ``connector``'s sink + (`kumoapi.plugin.FileStoreSink`) rather than a table; no I/O.""" + return _file_sink(connector, path) is not None + + +def file_store(connector: Any, path: str) -> Optional[Tuple[Any, str]]: + r"""``connector``'s ``(pyarrow filesystem, resolved path)`` for a file + destination; None for a table.""" + sink = _file_sink(connector, path) + return None if sink is None else sink.filesystem(connector, path) + + __all__ = [ "Connector", "FileConnector", diff --git a/packages/kumo-engine/kumo/engine/trainer/output_config.py b/packages/kumo-engine/kumo/engine/trainer/output_config.py index 4a325126d..43da33eac 100644 --- a/packages/kumo-engine/kumo/engine/trainer/output_config.py +++ b/packages/kumo-engine/kumo/engine/trainer/output_config.py @@ -11,10 +11,9 @@ from kumoapi.plugin import ModelRegistryLocation, SourceConnector from kumoapi.prediction_outputs import resolve_prediction_output_paths +from kumo.engine.connector import is_file_destination from kumo.engine.connector.base import Connector -_FILE_SOURCES = frozenset({'S3', 'GCS', 'ADLS', 'LOCAL_FILE'}) - @dataclass(frozen=True) class PredictionOutputConfig: @@ -36,12 +35,15 @@ class PredictionOutputConfig: def __post_init__(self) -> None: if self.output_paths is not None: - if self.connector.source_type not in _FILE_SOURCES: - raise ValueError( - 'output_paths requires a file or object-store connector') paths = resolve_prediction_output_paths( self.output_paths, output_types=self.artifact_type, output_table_name=self.table_name) + if not all( + is_file_destination(self.connector, path) + for path in paths.values()): + raise ValueError( + 'output_paths requires a file or object-store connector ' + '(or a volume or stage path it holds)') sink = self.connector.sink_backend if isinstance(sink, ObjectStoreSessions): # Reject aliases; retain caller paths. @@ -75,8 +77,8 @@ class TrainingTableOutputConfig: write_mode: WriteMode = WriteMode.OVERWRITE def __post_init__(self) -> None: - if (self.connector.source_type in _FILE_SOURCES - and not self.table_name.endswith('/')): + if (not self.table_name.endswith('/') + and is_file_destination(self.connector, self.table_name)): object.__setattr__(self, 'table_name', self.table_name + '/') diff --git a/packages/kumo-engine/test/test_filestore_sink.py b/packages/kumo-engine/test/test_filestore_sink.py new file mode 100644 index 000000000..21a59362d --- /dev/null +++ b/packages/kumo-engine/test/test_filestore_sink.py @@ -0,0 +1,155 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import os +from pathlib import Path +from types import SimpleNamespace + +import pandas as pd +import pytest +from kumo_core.chunked_fs import ChunkedFileHandler +from kumo_core.object_store_sink import FilesystemSink +from kumoapi.jobs import PredictionArtifactType +from kumoapi.plugin import FileStoreSink +from pyarrow.fs import PyFileSystem + +import kumo.engine as kumo +from kumo.engine._scratch import get_scratch_dir +from kumo.engine.connector import file_store, is_file_destination +from kumo.engine.trainer.output_config import ( + PredictionOutputConfig, + TrainingTableOutputConfig, +) + + +class _Bytes: + def __init__(self, data): + self.data = data + + def read_at(self, offset, length): + return self.data[offset:offset + length] + + def close(self): + pass + + +class MemoryStore(ChunkedFileHandler): + part_bytes = 4096 + + def __init__(self): + self.objects, self.spooled = {}, set() + + def get_type_name(self): + return 'memory' + + def _size(self, path): + data = self.objects.get(path) + return None if data is None else len(data) + + def _list(self, directory): + return { + path: len(data) + for path, data in self.objects.items() + if path.startswith(directory + '/') + } + + def _put(self, local, path): + self.spooled.add(Path(local).parent.parent) + self.objects[path] = Path(local).read_bytes() + + def _part(self, path, size): + return _Bytes(self.objects[path]) + + def _remove(self, path, *, tree): + for name in list(self.objects): + if name.startswith(path + '/') if tree else name == path: + del self.objects[name] + + +class StageSink: + r"""Not an object store: holds '@stage/...' paths through + ``filesystem``; anything else is a table.""" + def __init__(self): + self.store = MemoryStore() + + def holds_files(self, connector, path): + return path.startswith('@') + + def filesystem(self, connector, path): + assert self.holds_files(connector, path) + return PyFileSystem(self.store), path.rstrip('/') + + def write_files(self, *args, **kwargs): + return FilesystemSink(self).write_files(*args, **kwargs) + + +@pytest.fixture +def connector(tmp_path): + previous = get_scratch_dir() + kumo.configure(scratch_dir=tmp_path / 'scratch') + yield SimpleNamespace(source_type='SNOWFLAKE', sink_backend=StageSink()) + kumo.configure(scratch_dir=previous) + + +def test_sinks_with_a_filesystem_hold_files(tmp_path, connector): + local = kumo.FileConnector(root_dir=str(tmp_path)) + assert isinstance(connector.sink_backend, FileStoreSink) + assert isinstance(local.sink_backend, FileStoreSink) + assert is_file_destination(local, 'any/path') + assert is_file_destination(connector, '@stage/p') + assert not is_file_destination(connector, 'DB.SCHEMA.T') + assert file_store(connector, 'DB.SCHEMA.T') is None + assert not is_file_destination(SimpleNamespace(), 'x') + + +def test_output_configs_take_a_sinks_file_paths(connector): + assert TrainingTableOutputConfig(connector, '@stage/tt').table_name == ( + '@stage/tt/') + assert TrainingTableOutputConfig(connector, 'TT').table_name == 'TT' + PredictionOutputConfig(connector, + output_paths={'predictions': '@stage/p/'}) + with pytest.raises(ValueError, match='output_paths requires'): + PredictionOutputConfig(connector, output_paths={'predictions': 'T'}) + PredictionOutputConfig(connector, PredictionArtifactType.PREDICTIONS, + 'PREDS') + + +def test_bundles_round_trip_through_a_filesystem_only_sink( + tmp_path, connector): + bundle = tmp_path / 'bundle' + (bundle / 'sub').mkdir(parents=True) + (bundle / '_provenance.json').write_text('{}') + (bundle / 'sub' / 'weights.bin').write_bytes(os.urandom(10000)) + files = FilesystemSink(connector.sink_backend) + published = files.publish_model(connector, str(bundle), + path='@stage/models/m') + assert published.uri == '@stage/models/m' + store = connector.sink_backend.store + assert list(store.objects)[-1] == '@stage/models/m/_provenance.json' + assert '@stage/models/m/sub/weights.bin.chunks/_SUCCESS' in store.objects + files.fetch_model(connector, str(tmp_path / 'fetched'), + path='@stage/models/m') + for name in ('_provenance.json', 'sub/weights.bin'): + assert (tmp_path / 'fetched' / name).read_bytes() == ( + bundle / name).read_bytes() + + +def test_parquet_writes_stay_plain_up_to_one_part(tmp_path, connector): + files = FilesystemSink(connector.sink_backend) + frame = pd.DataFrame({'x': range(10)}) + files.write_table(connector, '@stage/small.parquet', frame) + files.write_table(connector, '@stage/big.parquet', + pd.DataFrame({'x': [os.urandom(50).hex() for _ in range(100)]})) + store = connector.sink_backend.store + assert '@stage/small.parquet' in store.objects + assert '@stage/big.parquet.chunks/_SUCCESS' in store.objects + assert pd.read_parquet( + '@stage/big.parquet', + filesystem=PyFileSystem(store)).shape == (100, 1) + + part = tmp_path / 'part-0.parquet' + frame.to_parquet(part) + assert files.write_files(connector, '@stage/out/', [str(part)]) + assert pd.read_parquet('@stage/out/part-0.parquet', + filesystem=PyFileSystem(store)).equals(frame) diff --git a/packages/kumo-engine/test/test_prediction_completion.py b/packages/kumo-engine/test/test_prediction_completion.py index a1ad9a53f..526d87a3e 100644 --- a/packages/kumo-engine/test/test_prediction_completion.py +++ b/packages/kumo-engine/test/test_prediction_completion.py @@ -261,6 +261,7 @@ def test_export_missing_artifact_does_not_publish_partial_outputs( writes = [] connector = SimpleNamespace( source_type='LOCAL_FILE', sink_backend=SimpleNamespace( + holds_files=lambda *a: True, filesystem=lambda *a: None, write_files=lambda *a, **kw: writes.append(a))) config = PredictionOutputConfig( connector, output_paths={