From d293a3e334d619677454394d003db429becb0da8 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Tue, 27 Oct 2026 09:50:58 +0000 Subject: [PATCH 01/25] Hold MLOps artifacts on Unity Catalog volumes through the Files API Co-Authored-By: Claude Opus 5.5 --- packages/kumo-databricks/README.md | 30 +++ .../src/kumo_databricks/registry.py | 3 +- .../src/kumo_databricks/sink.py | 90 ++++++- .../src/kumo_databricks/volume.py | 240 ++++++++++++++++++ packages/kumo-databricks/tests/fakes.py | 81 +++++- packages/kumo-databricks/tests/test_volume.py | 180 +++++++++++++ tests/public/Dockerfile.dockerignore | 4 + 7 files changed, 622 insertions(+), 6 deletions(-) create mode 100644 packages/kumo-databricks/src/kumo_databricks/volume.py create mode 100644 packages/kumo-databricks/tests/test_volume.py diff --git a/packages/kumo-databricks/README.md b/packages/kumo-databricks/README.md index 6769362fc..865ae85a1 100644 --- a/packages/kumo-databricks/README.md +++ b/packages/kumo-databricks/README.md @@ -67,6 +67,36 @@ events = Table.from_source_table( snapshot = Graph(tables=[events]).snapshot() ``` +## Artifacts on a volume + +A `/Volumes////...` path is a file destination of a +`DatabricksConnector` from any machine, not only on Databricks compute: the +sink reads and writes it through the Files API with the connector's identity. + +```python +from kumo.engine import ModelOutputConfig, PredictionOutputConfig, Trainer + +root = "/Volumes/main/kumo/artifacts" +trainer.fit(training_table, snapshot=snapshot, export_config=ModelOutputConfig( + databricks, path=f"{root}/models/churn")) +trainer = Trainer.load_from_connector(databricks, path=f"{root}/models/churn") +trainer.predict(prediction_table, snapshot=snapshot, output_config=PredictionOutputConfig( + databricks, output_paths={"predictions": f"{root}/scores/"})) +``` + +For an MLOps artifact root on a volume, see `packages/kumo-engine/docs/mlops.md`. + +A file whose size is known when it is written (bundle files, exported +Parquet parts) is stored as one object up to 4 GiB, so Spark +and other tools read it directly. A larger one streams into +`.chunks/00000`, `00001`, ... (4 GiB each, under the Files API's 5 GiB +upload cap) with an empty `_SUCCESS` written last, and Kumo reads the parts +back as one file. A file written without a known size stays one object up to +256 MiB, staged on local disk under the engine scratch directory, and streams +into parts above that. The identity needs `USE CATALOG`, `USE SCHEMA`, and +`READ VOLUME` (loads) and `WRITE VOLUME` (exports) on the volume; no +warehouse is used. + ## Model registry The sink publishes Kumo bundles directly to Unity Catalog using the connector's diff --git a/packages/kumo-databricks/src/kumo_databricks/registry.py b/packages/kumo-databricks/src/kumo_databricks/registry.py index d2d88de09..1f7d2e352 100644 --- a/packages/kumo-databricks/src/kumo_databricks/registry.py +++ b/packages/kumo-databricks/src/kumo_databricks/registry.py @@ -34,7 +34,8 @@ def _address(connector: SourceConnector, path: Optional[str], location: Optional[ModelRegistryLocation]) -> str: if path is not None: raise ValueError('Databricks models take a ModelRegistryLocation, ' - 'rather than a path.') + 'or a /Volumes////... ' + 'path.') if location is None: raise ValueError('Pass a ModelRegistryLocation naming the model.') config = _target_config(connector) diff --git a/packages/kumo-databricks/src/kumo_databricks/sink.py b/packages/kumo-databricks/src/kumo_databricks/sink.py index b258c70a5..56d7717fb 100644 --- a/packages/kumo-databricks/src/kumo_databricks/sink.py +++ b/packages/kumo-databricks/src/kumo_databricks/sink.py @@ -22,16 +22,23 @@ import os import tempfile from pathlib import Path -from typing import Optional, Sequence +from typing import Mapping, Optional, Sequence import pandas as pd from databricks.sql.client import Connection +from kumo_core.object_store_sink import FilesystemSink from kumo_databricks.registry import DatabricksModelRegistryMixin from kumo_databricks.session import _target_config +from kumo_databricks.volume import is_volume_path, volume_filesystem from kumo_databricks.write import RunVolume from kumoapi.data_source import DataSourceType from kumoapi.jobs import WriteMode -from kumoapi.plugin import SourceConnector +from kumoapi.plugin import ( + ModelRegistryLocation, + PublishedModel, + SourceConnector, +) +from pyarrow.fs import FileSystem class DatabricksSinkBackend(DatabricksModelRegistryMixin): @@ -125,6 +132,9 @@ def write_table( names a schema other than the connector's, or if the staging volume is neither reachable nor creatable. """ + if self.holds_files(connector, table_name): + return FilesystemSink(self).write_table( + connector, table_name, frame, write_mode, export_id=export_id) del project_id, export_id from kumo_databricks import output, sql, write @@ -152,8 +162,13 @@ def write_table( def write_files(self, connector, table_name, files, write_mode=WriteMode.OVERWRITE, project_id=None, *, - export_id=None, extra_columns=None): - """Pass selected Parquet directly to the existing hosted COPY owner.""" + export_id=None, extra_columns=None, **options): + """Copy Parquet to a volume path, or COPY it into a table.""" + if self.holds_files(connector, table_name): + return FilesystemSink(self).write_files( + connector, table_name, files, write_mode, + export_id=export_id, extra_columns=extra_columns, + **options) del project_id, export_id from functools import partial @@ -187,6 +202,73 @@ def write_files(self, connector, table_name, files, prepare_file=prepare) return sql.qualified_name(config.catalog, schema, table) + def publish_model( + self, + connector: SourceConnector, + bundle_dir: str, + *, + path: Optional[str] = None, + location: Optional[ModelRegistryLocation] = None, + metadata: Mapping[str, str] = {}, + metrics: Mapping[str, float] = {}, + project_id: Optional[str] = None, + ) -> PublishedModel: + """Register in Unity Catalog, or copy to a volume ``path``.""" + if location is None and path and self.holds_files(connector, path): + return FilesystemSink(self).publish_model( + connector, bundle_dir, path=path, metadata=metadata) + return super().publish_model(connector, bundle_dir, path=path, + location=location, metadata=metadata, + metrics=metrics, project_id=project_id) + + def fetch_model( + self, + connector: SourceConnector, + local_dir: str, + *, + path: Optional[str] = None, + location: Optional[ModelRegistryLocation] = None, + project_id: Optional[str] = None, + ) -> PublishedModel: + """Download from Unity Catalog, or from a volume ``path``.""" + if location is None and path and self.holds_files(connector, path): + return FilesystemSink(self).fetch_model(connector, local_dir, + path=path) + return super().fetch_model(connector, local_dir, path=path, + location=location, project_id=project_id) + + def preflight_write_files(self, connector: SourceConnector, + table_name: str, + write_mode: Optional[WriteMode] = None) -> None: + """Refuse an occupied volume directory before scoring, unless + ``write_mode`` is ``OVERWRITE``; tables are not checked.""" + from kumo_core.artifact_export import ArtifactDestinationConflict + from pyarrow.fs import FileSelector + + if (write_mode is WriteMode.OVERWRITE or not table_name.endswith('/') + or not self.holds_files(connector, table_name)): + return + filesystem, path = self.filesystem(connector, table_name) + if filesystem.get_file_info(FileSelector(path, allow_not_found=True)): + raise ArtifactDestinationConflict( + f"Output directory '{path}/' already contains data from " + f"another export. Choose an empty destination path, or " + f"replace it with write_mode=WriteMode.OVERWRITE.") + + def holds_files(self, connector: SourceConnector, path: str) -> bool: + """``/Volumes////...``; anything else is a + table.""" + return is_volume_path(path) + + def filesystem(self, connector: SourceConnector, + path: str) -> tuple[FileSystem, str]: + """A volume path as a filesystem, as the connector's identity.""" + if not is_volume_path(path): + raise ValueError( + f"{path!r} is not a Unity Catalog volume path; Databricks " + f"artifacts go under /Volumes////.") + return volume_filesystem(self._client(connector)), path.rstrip('/') + def __repr__(self) -> str: return f'DatabricksSinkBackend(identity={self._identity!r})' diff --git a/packages/kumo-databricks/src/kumo_databricks/volume.py b/packages/kumo-databricks/src/kumo_databricks/volume.py new file mode 100644 index 000000000..3a7b39e61 --- /dev/null +++ b/packages/kumo-databricks/src/kumo_databricks/volume.py @@ -0,0 +1,240 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Unity Catalog volumes as a pyarrow filesystem, over the Files API.""" + +from __future__ import annotations + +import io +import posixpath +import queue +import threading +from contextlib import contextmanager +from typing import BinaryIO, Dict, Iterator, Optional +from urllib.parse import quote + +from databricks.sdk import WorkspaceClient +from databricks.sdk.errors import DatabricksError, NotFound, PermissionDenied +from kumo_core.chunked_fs import ChunkedFileHandler +from pyarrow.fs import FileSystem, PyFileSystem + +_FILES = '/api/2.0/fs/files' +_READ = 8 << 20 + + +def is_volume_path(path: str) -> bool: + parts = path.rstrip('/').split('/') + return (path.startswith('/Volumes/') and len(parts) >= 5 + and all(parts[1:]) and posixpath.normpath(path) == path.rstrip('/')) + + +def volume_filesystem(client: WorkspaceClient) -> FileSystem: + return PyFileSystem(VolumeHandler(client)) + + +@contextmanager +def _translated(path: str) -> Iterator[None]: + """Nothing is chained, and only SDK errors keep their text: transport + errors from presigned uploads carry the signed URL.""" + from kumo_core.sources.diagnostics import safe_summary + + try: + yield + except NotFound as error: + raise FileNotFoundError(f'{path}: {error}') from None + except PermissionDenied as error: + raise PermissionError(f'{path}: {error}') from None + except DatabricksError as error: + raise OSError(f'{path}: {error}') from None + except Exception as error: + raise OSError(f'{path}: {safe_summary(error)}') from None + + +class VolumeHandler(ChunkedFileHandler): + """``/Volumes////...`` through ``client.files``. + + Parts stay under the Files API's 5 GiB single-upload cap. + """ + part_bytes = 4 << 30 + plain_bytes = 256 << 20 + + def __init__(self, client: WorkspaceClient) -> None: + self._client = client + + def __eq__(self, other: object) -> bool: + return (isinstance(other, VolumeHandler) + and other._client is self._client) + + def __ne__(self, other: object) -> bool: + return not self == other + + def get_type_name(self) -> str: + return 'databricks-volume' + + def _size(self, path: str) -> Optional[int]: + path = path.rstrip('/') + with _translated(path): + try: + return int(self._client.files.get_metadata(path).content_length) + except NotFound: + return None + + def _list(self, directory: str) -> Dict[str, int]: + directory = directory.rstrip('/') + with _translated(directory): + try: + entries = list( + self._client.files.list_directory_contents(directory)) + except NotFound: + return {} + found: Dict[str, int] = {} + for entry in entries: + path = entry.path.rstrip('/') + if entry.is_directory: + found.update(self._list(path)) + else: + found[path] = int(entry.file_size) + return found + + def _put(self, local: str, path: str) -> None: + with _translated(path), open(local, 'rb') as contents: + self._client.files.upload(path, contents, overwrite=True) + + def _stream(self, path: str) -> BinaryIO: + return _Upload(self._client, path) + + def _part(self, path: str, size: int) -> _Part: + return _Part(self._client, path) + + def _remove(self, path: str, *, tree: bool) -> None: + from kumo_databricks.staging import clear_volume + + path = path.rstrip('/') + with _translated(path): + if tree: + clear_volume(self._client, path) + return + try: + self._client.files.delete(path) + except NotFound: + pass + + +class _Part: + """Serves consecutive reads from one open-ended ranged download.""" + def __init__(self, client: WorkspaceClient, path: str) -> None: + self._client, self._path = client, path + self._stream = None + self._at = 0 + + def read_at(self, offset: int, length: int) -> bytes: + if length <= 0: + return b'' + if self._stream is None or offset != self._at: + self._open(offset) + with _translated(self._path): + pieces, got = [], 0 + while got < length: + piece = self._stream.read(min(_READ, length - got)) + if not piece: + break + pieces.append(piece) + got += len(piece) + self._at += got + return pieces[0] if len(pieces) == 1 else b''.join(pieces) + + def _open(self, offset: int) -> None: + self.close() + with _translated(self._path): + response = self._client.api_client.do( + 'GET', f'{_FILES}{quote(self._path)}', raw=True, + response_headers=['content-range'], headers={ + 'Accept': 'application/octet-stream', + 'Range': f'bytes={offset}-' + }) + self._stream = response['contents'] + served = response.get('content-range') or '' + if not served.startswith(f'bytes {offset}-'): + self.close() + raise OSError(f'{self._path}: asked for bytes from {offset}, ' + f'got {served or "the whole file"}') + self._stream.set_chunk_size(_READ) + self._at = offset + + def close(self) -> None: + if self._stream is not None: + self._stream.close() + self._stream = None + + +class _Feed(io.RawIOBase): + def __init__(self) -> None: + self.pieces: queue.Queue[Optional[bytes]] = queue.Queue(maxsize=4) + self._pending = b'' + self._done = False + + def readable(self) -> bool: + return True + + def readinto(self, buffer) -> int: + while not self._pending and not self._done: + piece = self.pieces.get() + if piece is None: + self._done = True + else: + self._pending = piece + size = min(len(buffer), len(self._pending)) + buffer[:size] = self._pending[:size] + self._pending = self._pending[size:] + return size + + +class _Upload(io.RawIOBase): + """One ``files.upload`` fed from writes; ``close`` commits it.""" + def __init__(self, client: WorkspaceClient, path: str) -> None: + self._path = path + self._feed = _Feed() + self._error: Optional[BaseException] = None + source = io.BufferedReader(self._feed, 8 << 20) + + def upload() -> None: + try: + with _translated(path): + client.files.upload(path, source, overwrite=True) + except BaseException as error: + self._error = error + + self._thread = threading.Thread(target=upload, daemon=True) + self._thread.start() + + def writable(self) -> bool: + return True + + def write(self, data) -> int: + piece = bytes(data) + self._put(piece) + return len(piece) + + def _put(self, piece: Optional[bytes]) -> None: + while self._thread.is_alive(): + try: + self._feed.pieces.put(piece, timeout=1) + return + except queue.Full: + continue + self._thread.join() + if self._error is not None: + raise self._error + if piece is not None: + raise OSError(f'{self._path}: the upload stopped reading') + + def close(self) -> None: + if self.closed: + return + try: + self._put(None) + self._thread.join() + if self._error is not None: + raise self._error + finally: + super().close() diff --git a/packages/kumo-databricks/tests/fakes.py b/packages/kumo-databricks/tests/fakes.py index a3f941b0a..b629e404d 100644 --- a/packages/kumo-databricks/tests/fakes.py +++ b/packages/kumo-databricks/tests/fakes.py @@ -8,9 +8,10 @@ import io from dataclasses import dataclass from typing import Any, Callable, Mapping +from urllib.parse import unquote import pyarrow as pa -from databricks.sdk.errors import NotFound +from databricks.sdk.errors import BadRequest, NotFound from databricks.sql.exc import DataError from sources.recording import RecordingConnection, RecordingCursor @@ -197,7 +198,85 @@ def delete_directory(self, path: str) -> None: return None +class FakeVolume(FakeFiles): + """A volume with real directories, uploads and ranged downloads.""" + def __init__(self, objects: Mapping[str, bytes] | None = None) -> None: + super().__init__(objects) + self.directories: set[str] = set() + self.ignore_range = False + + def _is_directory(self, path: str) -> bool: + return path in self.directories or any( + name.startswith(path + "/") for name in self.objects) + + def list_directory_contents(self, path: str) -> list[FakeEntry]: + path = path.rstrip("/") + if not self._is_directory(path): + raise NotFound(f"No such directory: {path}") + children: dict[str, FakeEntry] = {} + for name in sorted(self.objects) + sorted(self.directories): + if not name.startswith(path + "/"): + continue + child = name[len(path) + 1:].split("/")[0] + full = f"{path}/{child}" + children[full] = (FakeEntry(full, len(self.objects[full])) + if full in self.objects else FakeEntry( + full + "/", 0, is_directory=True)) + return list(children.values()) + + def upload(self, path: str, contents: Any, *, overwrite: bool) -> None: + assert overwrite + self.objects[path] = contents.read() + + def get_metadata(self, path: str) -> Any: + if path.endswith("/"): + raise BadRequest("Bad Request") + if path not in self.objects: + raise NotFound(f"No such file: {path}") + return FakeMetadata(len(self.objects[path])) + + def get_directory_metadata(self, path: str) -> None: + if not self._is_directory(path): + raise NotFound(f"No such directory: {path}") + + def create_directory(self, path: str) -> None: + self.directories.add(path) + + def delete_directory(self, path: str) -> None: + self.directories.discard(path.rstrip("/")) + + def do(self, method: str, path: str, *, raw: bool, + response_headers: list[str], headers: Mapping[str, str]) -> dict: + assert method == "GET" and raw + remote = unquote(path.removeprefix("/api/2.0/fs/files")) + if remote not in self.objects: + raise NotFound(f"No such file: {remote}") + if self.ignore_range: + return {"contents": FakeStreamingResponse(self.objects[remote])} + first, last = headers["Range"].removeprefix("bytes=").split("-") + size = len(self.objects[remote]) + last = last or str(size - 1) + stream = FakeStreamingResponse( + self.objects[remote][int(first):int(last) + 1]) + self.streams.append(stream) + return { + "contents": stream, + "content-range": f"bytes {first}-{last}/{size}" + } + + +@dataclass +class FakeMetadata: + content_length: int + + +class FakeStreamingResponse(TrackedStream): + def set_chunk_size(self, size: int) -> None: + pass + + class FakeWorkspaceClient: """Exposes only the files API the reader uses.""" def __init__(self, files: FakeFiles) -> None: self.files = files + self.api_client = files diff --git a/packages/kumo-databricks/tests/test_volume.py b/packages/kumo-databricks/tests/test_volume.py new file mode 100644 index 000000000..db60e65fd --- /dev/null +++ b/packages/kumo-databricks/tests/test_volume.py @@ -0,0 +1,180 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from types import SimpleNamespace + +import pyarrow as pa +import pyarrow.parquet as pq +import pytest +from databricks.sdk.errors import PermissionDenied +from fakes import FakeVolume, FakeWorkspaceClient +from kumo_databricks.sink import DatabricksSinkBackend +from kumo_databricks.volume import VolumeHandler +from pyarrow.fs import FileType, PyFileSystem + +ROOT = '/Volumes/main/analytics/artifacts' +BODY = bytes(range(256)) * 10 + + +def filesystem(objects=None, part_bytes=None): + volume = FakeVolume(objects) + handler = VolumeHandler(FakeWorkspaceClient(volume)) + if part_bytes is not None: + handler.part_bytes = part_bytes + return PyFileSystem(handler), volume + + +def write(fs, path, body, piece=700): + with fs.open_output_stream(path) as writer: + for start in range(0, len(body), piece): + writer.write(body[start:start + piece]) + + +def test_reads_cross_part_boundaries(): + fs, _ = filesystem(part_bytes=1000) + write(fs, f'{ROOT}/run/big.bin', BODY) + with fs.open_input_file(f'{ROOT}/run/big.bin') as reader: + assert reader.read_at(300, 900) == BODY[900:1200] + assert reader.read_at(10, 1200) == BODY[1200:1210] + + +def test_sized_writes_stay_plain_up_to_a_part(): + from kumo_core.object_store_sink import open_sized + + fs, volume = filesystem(part_bytes=1000) + fs.handler.plain_bytes = 100 + for name, body in (('plain', BODY[:1000]), ('parts', BODY[:1500])): + with open_sized(fs, f'{ROOT}/{name}', len(body)) as writer: + writer.write(body) + with fs.open_input_stream(f'{ROOT}/{name}') as reader: + assert reader.readall() == body + assert volume.objects[f'{ROOT}/plain'] == BODY[:1000] + assert f'{ROOT}/parts' not in volume.objects + assert f'{ROOT}/parts.chunks/_SUCCESS' in volume.objects + + +def test_a_failed_upload_raises_without_the_signed_url(): + from kumo_core.object_store_sink import open_sized + + fs, volume = filesystem() + signed = 'https://bucket.example/part?X-Amz-Signature=secret' + + def drop(path, contents, *, overwrite): + contents.read(100) + raise ConnectionError(f'Max retries exceeded with url: {signed}') + + volume.upload = drop + with pytest.raises(OSError) as raised: + with open_sized(fs, f'{ROOT}/x', len(BODY)) as writer: + for start in range(0, len(BODY), 64): + writer.write(BODY[start:start + 64]) + assert 'Signature' not in str(raised.value) + assert raised.value.__cause__ is None + + +def test_a_download_that_ignores_the_range_is_refused(): + fs, volume = filesystem({f'{ROOT}/a': BODY}) + volume.ignore_range = True + with fs.open_input_file(f'{ROOT}/a') as reader: + with pytest.raises(OSError, match='asked for bytes from 100'): + reader.read_at(10, 100) + + +def test_directories_and_missing_paths(): + fs, volume = filesystem({ + f'{ROOT}/s/kumo_snapshot.json': b'{}', + f'{ROOT}/s/index/sub/b': b'de', + }) + kinds = [info.type for info in fs.get_file_info( + [f'{ROOT}/s/kumo_snapshot.json', f'{ROOT}/s/index/', f'{ROOT}/s/no'])] + assert kinds == [FileType.File, FileType.Directory, FileType.NotFound] + with pytest.raises(FileNotFoundError): + fs.open_input_stream(f'{ROOT}/s/no') + fs.delete_dir_contents(f'{ROOT}/s/') + assert volume.objects == {} + + +def test_refused_upload_raises_permission_error(): + fs, volume = filesystem() + + def refuse(path, contents, *, overwrite): + raise PermissionDenied('no WRITE VOLUME') + + volume.upload = refuse + with pytest.raises(PermissionError): + write(fs, f'{ROOT}/x', b'data') + + +def test_sink_filesystem_only_takes_volume_paths(): + sink = DatabricksSinkBackend() + volume = FakeVolume() + sink._client = lambda connector: FakeWorkspaceClient(volume) + connector = SimpleNamespace() + fs, path = sink.filesystem(connector, f'{ROOT}/exports/') + assert path == f'{ROOT}/exports' + write(fs, f'{path}/x', b'1') + assert volume.objects == {f'{path}/x': b'1'} + assert sink.holds_files(connector, f'{ROOT}/exports/') + for bad in ('s3://bucket/x', '/Volumes/main/analytics', 'users_scores', + f'{ROOT}/../other/x', 'dbfs:/Volumes/main/analytics/v'): + assert not sink.holds_files(connector, bad) + with pytest.raises(ValueError, match='volume path'): + sink.filesystem(connector, bad) + + +def test_bundle_publish_and_fetch_through_a_volume_path(tmp_path): + sink = DatabricksSinkBackend() + volume = FakeVolume() + sink._client = lambda connector: FakeWorkspaceClient(volume) + bundle = tmp_path / 'bundle' + (bundle / 'model').mkdir(parents=True) + (bundle / 'training-spec.json').write_text('{}') + (bundle / 'model' / 'weights.pt').write_bytes(BODY) + published = sink.publish_model(SimpleNamespace(), str(bundle), + path=f'{ROOT}/models/b1') + assert published.uri == f'{ROOT}/models/b1' + assert volume.objects[f'{ROOT}/models/b1/model/weights.pt'] == BODY + fetched = tmp_path / 'fetched' + sink.fetch_model(SimpleNamespace(), str(fetched), path=f'{ROOT}/models/b1') + assert (fetched / 'model' / 'weights.pt').read_bytes() == BODY + with pytest.raises(ValueError, match='Nothing is published'): + sink.fetch_model(SimpleNamespace(), str(fetched), + path=f'{ROOT}/models/none') + + +def test_predictions_and_frames_written_to_a_volume_path(tmp_path): + sink = DatabricksSinkBackend() + volume = FakeVolume() + sink._client = lambda connector: FakeWorkspaceClient(volume) + part = tmp_path / 'part-0.parquet' + pq.write_table(pa.table({'users': [1, 2], 'score': [0.1, 0.9]}), part) + written = sink.write_files(SimpleNamespace(), f'{ROOT}/predictions/', + [str(part)], export_id='run-1') + assert written == f'{ROOT}/predictions/' + names = {path.rsplit('/', 1)[-1] for path in volume.objects} + assert {'part-0.parquet', '_SUCCESS'} <= names + frame = pa.table({'users': [3]}).to_pandas() + assert sink.write_table(SimpleNamespace(), f'{ROOT}/one', + frame) == f'{ROOT}/one.parquet' + + +def test_an_occupied_volume_directory_is_refused_before_scoring(tmp_path): + from kumo_core.artifact_export import ArtifactDestinationConflict + from kumoapi.jobs import WriteMode + + sink = DatabricksSinkBackend() + volume = FakeVolume() + sink._client = lambda connector: FakeWorkspaceClient(volume) + target = f'{ROOT}/predictions/' + sink.preflight_write_files(SimpleNamespace(), target) + part = tmp_path / 'part-0.parquet' + pq.write_table(pa.table({'users': [1]}), part) + sink.write_files(SimpleNamespace(), target, [str(part)], export_id='r1') + + with pytest.raises(ArtifactDestinationConflict, match='/predictions/'): + sink.preflight_write_files(SimpleNamespace(), target) + sink.preflight_write_files(SimpleNamespace(), target, + write_mode=WriteMode.OVERWRITE) + sink.preflight_write_files(SimpleNamespace(), f'{ROOT}/other/') + sink.preflight_write_files(SimpleNamespace(), 'main.ml.scores') diff --git a/tests/public/Dockerfile.dockerignore b/tests/public/Dockerfile.dockerignore index c6d21c61a..00eca5c99 100644 --- a/tests/public/Dockerfile.dockerignore +++ b/tests/public/Dockerfile.dockerignore @@ -14,6 +14,10 @@ !packages/kumo-core/src/kumo_core/artifact_staging.py # Model registries validate bundles through the portable artifact contract. !packages/kumo-core/src/kumo_core/artifacts.py +# Volume and stage sinks hold files through the shared chunked filesystem. +!packages/kumo-core/src/kumo_core/chunked_fs.py +!packages/kumo-core/src/kumo_core/object_store_sink.py +!packages/kumo-core/src/kumo_core/artifact_export.py !packages/kumo-core/src/kumo_core/config.py !packages/kumo-core/src/kumo_core/manifest.py !packages/kumo-core/src/kumo_core/io/__init__.py From 35418404858ca5ad24052bb3f1ac56aaa870330d Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Fri, 30 Oct 2026 23:42:25 +0000 Subject: [PATCH 02/25] Expect volume chunks under a generation directory Co-Authored-By: Claude Opus 5.5 --- packages/kumo-databricks/README.md | 6 +++--- packages/kumo-databricks/tests/test_volume.py | 3 ++- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/packages/kumo-databricks/README.md b/packages/kumo-databricks/README.md index 865ae85a1..03f99568d 100644 --- a/packages/kumo-databricks/README.md +++ b/packages/kumo-databricks/README.md @@ -89,9 +89,9 @@ For an MLOps artifact root on a volume, see `packages/kumo-engine/docs/mlops.md` A file whose size is known when it is written (bundle files, exported Parquet parts) is stored as one object up to 4 GiB, so Spark and other tools read it directly. A larger one streams into -`.chunks/00000`, `00001`, ... (4 GiB each, under the Files API's 5 GiB -upload cap) with an empty `_SUCCESS` written last, and Kumo reads the parts -back as one file. A file written without a known size stays one object up to +`.chunks//00000`, `00001`, ... (4 GiB each, under the Files +API's 5 GiB upload cap) with an empty `_SUCCESS` written last, and Kumo reads +the newest committed generation back as one file. A file written without a known size stays one object up to 256 MiB, staged on local disk under the engine scratch directory, and streams into parts above that. The identity needs `USE CATALOG`, `USE SCHEMA`, and `READ VOLUME` (loads) and `WRITE VOLUME` (exports) on the volume; no diff --git a/packages/kumo-databricks/tests/test_volume.py b/packages/kumo-databricks/tests/test_volume.py index db60e65fd..25cb8fc05 100644 --- a/packages/kumo-databricks/tests/test_volume.py +++ b/packages/kumo-databricks/tests/test_volume.py @@ -51,7 +51,8 @@ def test_sized_writes_stay_plain_up_to_a_part(): assert reader.readall() == body assert volume.objects[f'{ROOT}/plain'] == BODY[:1000] assert f'{ROOT}/parts' not in volume.objects - assert f'{ROOT}/parts.chunks/_SUCCESS' in volume.objects + assert any(path.startswith(f'{ROOT}/parts.chunks/') + and path.endswith('/_SUCCESS') for path in volume.objects) def test_a_failed_upload_raises_without_the_signed_url(): From 434814130dc34385d5b96930d0f683c5c2de0c7f Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Fri, 30 Oct 2026 23:40:03 +0000 Subject: [PATCH 03/25] Refuse an occupied volume directory whatever the write mode, through FilesystemSink Co-Authored-By: Claude Opus 5.5 --- .../kumo-databricks/src/kumo_databricks/sink.py | 17 +++-------------- packages/kumo-databricks/tests/test_volume.py | 8 ++++---- 2 files changed, 7 insertions(+), 18 deletions(-) diff --git a/packages/kumo-databricks/src/kumo_databricks/sink.py b/packages/kumo-databricks/src/kumo_databricks/sink.py index 56d7717fb..bdbdf7554 100644 --- a/packages/kumo-databricks/src/kumo_databricks/sink.py +++ b/packages/kumo-databricks/src/kumo_databricks/sink.py @@ -240,20 +240,9 @@ def fetch_model( def preflight_write_files(self, connector: SourceConnector, table_name: str, write_mode: Optional[WriteMode] = None) -> None: - """Refuse an occupied volume directory before scoring, unless - ``write_mode`` is ``OVERWRITE``; tables are not checked.""" - from kumo_core.artifact_export import ArtifactDestinationConflict - from pyarrow.fs import FileSelector - - if (write_mode is WriteMode.OVERWRITE or not table_name.endswith('/') - or not self.holds_files(connector, table_name)): - return - filesystem, path = self.filesystem(connector, table_name) - if filesystem.get_file_info(FileSelector(path, allow_not_found=True)): - raise ArtifactDestinationConflict( - f"Output directory '{path}/' already contains data from " - f"another export. Choose an empty destination path, or " - f"replace it with write_mode=WriteMode.OVERWRITE.") + """Refuse an occupied volume directory; tables are not checked.""" + if self.holds_files(connector, table_name): + FilesystemSink(self).preflight_write_files(connector, table_name) def holds_files(self, connector: SourceConnector, path: str) -> bool: """``/Volumes////...``; anything else is a diff --git a/packages/kumo-databricks/tests/test_volume.py b/packages/kumo-databricks/tests/test_volume.py index 25cb8fc05..15cef7d32 100644 --- a/packages/kumo-databricks/tests/test_volume.py +++ b/packages/kumo-databricks/tests/test_volume.py @@ -173,9 +173,9 @@ def test_an_occupied_volume_directory_is_refused_before_scoring(tmp_path): pq.write_table(pa.table({'users': [1]}), part) sink.write_files(SimpleNamespace(), target, [str(part)], export_id='r1') - with pytest.raises(ArtifactDestinationConflict, match='/predictions/'): - sink.preflight_write_files(SimpleNamespace(), target) - sink.preflight_write_files(SimpleNamespace(), target, - write_mode=WriteMode.OVERWRITE) + for mode in (None, WriteMode.OVERWRITE): + with pytest.raises(ArtifactDestinationConflict, match='/predictions/'): + sink.preflight_write_files(SimpleNamespace(), target, + write_mode=mode) sink.preflight_write_files(SimpleNamespace(), f'{ROOT}/other/') sink.preflight_write_files(SimpleNamespace(), 'main.ml.scores') From b5394e003965548670407e8ba0833b7c2731fe46 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Fri, 30 Oct 2026 23:40:03 +0000 Subject: [PATCH 04/25] Drop the unused write_files options passthrough Co-Authored-By: Claude Opus 5.5 --- packages/kumo-databricks/src/kumo_databricks/sink.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/packages/kumo-databricks/src/kumo_databricks/sink.py b/packages/kumo-databricks/src/kumo_databricks/sink.py index bdbdf7554..8cfd36711 100644 --- a/packages/kumo-databricks/src/kumo_databricks/sink.py +++ b/packages/kumo-databricks/src/kumo_databricks/sink.py @@ -162,13 +162,12 @@ def write_table( def write_files(self, connector, table_name, files, write_mode=WriteMode.OVERWRITE, project_id=None, *, - export_id=None, extra_columns=None, **options): + export_id=None, extra_columns=None): """Copy Parquet to a volume path, or COPY it into a table.""" if self.holds_files(connector, table_name): return FilesystemSink(self).write_files( connector, table_name, files, write_mode, - export_id=export_id, extra_columns=extra_columns, - **options) + export_id=export_id, extra_columns=extra_columns) del project_id, export_id from functools import partial From 36911ae19c606ddb454f4ea1aaa4b48223c8f605 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Fri, 30 Oct 2026 23:40:17 +0000 Subject: [PATCH 05/25] Shorten the volume sink docstrings Co-Authored-By: Claude Opus 5.5 --- packages/kumo-databricks/src/kumo_databricks/sink.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/packages/kumo-databricks/src/kumo_databricks/sink.py b/packages/kumo-databricks/src/kumo_databricks/sink.py index 8cfd36711..4b21d4d55 100644 --- a/packages/kumo-databricks/src/kumo_databricks/sink.py +++ b/packages/kumo-databricks/src/kumo_databricks/sink.py @@ -244,14 +244,13 @@ def preflight_write_files(self, connector: SourceConnector, FilesystemSink(self).preflight_write_files(connector, table_name) def holds_files(self, connector: SourceConnector, path: str) -> bool: - """``/Volumes////...``; anything else is a - table.""" + """Volume paths hold files; anything else is a table.""" return is_volume_path(path) def filesystem(self, connector: SourceConnector, path: str) -> tuple[FileSystem, str]: - """A volume path as a filesystem, as the connector's identity.""" - if not is_volume_path(path): + """A volume path's filesystem, as the connector's identity.""" + if not self.holds_files(connector, path): raise ValueError( f"{path!r} is not a Unity Catalog volume path; Databricks " f"artifacts go under /Volumes////.") From 2325f8d00c1e5b538eaed1cbe92d4bb2212e65c3 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Fri, 30 Oct 2026 23:40:24 +0000 Subject: [PATCH 06/25] Tidy the volume handler: one-line docs, top-level imports, no custom equality Co-Authored-By: Claude Opus 5.5 --- .../src/kumo_databricks/volume.py | 32 ++++++------------- 1 file changed, 9 insertions(+), 23 deletions(-) diff --git a/packages/kumo-databricks/src/kumo_databricks/volume.py b/packages/kumo-databricks/src/kumo_databricks/volume.py index 3a7b39e61..8d65c0b30 100644 --- a/packages/kumo-databricks/src/kumo_databricks/volume.py +++ b/packages/kumo-databricks/src/kumo_databricks/volume.py @@ -16,10 +16,12 @@ from databricks.sdk import WorkspaceClient from databricks.sdk.errors import DatabricksError, NotFound, PermissionDenied from kumo_core.chunked_fs import ChunkedFileHandler +from kumo_core.sources.diagnostics import safe_summary +from kumo_databricks.staging import clear_volume from pyarrow.fs import FileSystem, PyFileSystem _FILES = '/api/2.0/fs/files' -_READ = 8 << 20 +_BUFFER = 8 << 20 def is_volume_path(path: str) -> bool: @@ -34,10 +36,7 @@ def volume_filesystem(client: WorkspaceClient) -> FileSystem: @contextmanager def _translated(path: str) -> Iterator[None]: - """Nothing is chained, and only SDK errors keep their text: transport - errors from presigned uploads carry the signed URL.""" - from kumo_core.sources.diagnostics import safe_summary - + """OS errors, unchained; non-SDK text may hold a signed URL.""" try: yield except NotFound as error: @@ -51,23 +50,12 @@ def _translated(path: str) -> Iterator[None]: class VolumeHandler(ChunkedFileHandler): - """``/Volumes////...`` through ``client.files``. - - Parts stay under the Files API's 5 GiB single-upload cap. - """ - part_bytes = 4 << 30 + part_bytes = 4 << 30 # under the Files API's 5 GiB upload cap plain_bytes = 256 << 20 def __init__(self, client: WorkspaceClient) -> None: self._client = client - def __eq__(self, other: object) -> bool: - return (isinstance(other, VolumeHandler) - and other._client is self._client) - - def __ne__(self, other: object) -> bool: - return not self == other - def get_type_name(self) -> str: return 'databricks-volume' @@ -107,8 +95,6 @@ def _part(self, path: str, size: int) -> _Part: return _Part(self._client, path) def _remove(self, path: str, *, tree: bool) -> None: - from kumo_databricks.staging import clear_volume - path = path.rstrip('/') with _translated(path): if tree: @@ -135,13 +121,13 @@ def read_at(self, offset: int, length: int) -> bytes: with _translated(self._path): pieces, got = [], 0 while got < length: - piece = self._stream.read(min(_READ, length - got)) + piece = self._stream.read(min(_BUFFER, length - got)) if not piece: break pieces.append(piece) got += len(piece) self._at += got - return pieces[0] if len(pieces) == 1 else b''.join(pieces) + return b''.join(pieces) def _open(self, offset: int) -> None: self.close() @@ -158,7 +144,7 @@ def _open(self, offset: int) -> None: self.close() raise OSError(f'{self._path}: asked for bytes from {offset}, ' f'got {served or "the whole file"}') - self._stream.set_chunk_size(_READ) + self._stream.set_chunk_size(_BUFFER) self._at = offset def close(self) -> None: @@ -195,7 +181,7 @@ def __init__(self, client: WorkspaceClient, path: str) -> None: self._path = path self._feed = _Feed() self._error: Optional[BaseException] = None - source = io.BufferedReader(self._feed, 8 << 20) + source = io.BufferedReader(self._feed, _BUFFER) def upload() -> None: try: From 9241e50098ab25c92b978d2b2e8e8e70e96f3ac8 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Fri, 30 Oct 2026 23:40:29 +0000 Subject: [PATCH 07/25] Show the chunk generation directory in the volume README Co-Authored-By: Claude Opus 5.5 --- packages/kumo-databricks/README.md | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/packages/kumo-databricks/README.md b/packages/kumo-databricks/README.md index 03f99568d..2c0a1a53f 100644 --- a/packages/kumo-databricks/README.md +++ b/packages/kumo-databricks/README.md @@ -91,7 +91,8 @@ Parquet parts) is stored as one object up to 4 GiB, so Spark and other tools read it directly. A larger one streams into `.chunks//00000`, `00001`, ... (4 GiB each, under the Files API's 5 GiB upload cap) with an empty `_SUCCESS` written last, and Kumo reads -the newest committed generation back as one file. A file written without a known size stays one object up to +the newest committed generation back as one file. A file written without a +known size stays one object up to 256 MiB, staged on local disk under the engine scratch directory, and streams into parts above that. The identity needs `USE CATALOG`, `USE SCHEMA`, and `READ VOLUME` (loads) and `WRITE VOLUME` (exports) on the volume; no From a64b33d4580ba1882e4a3db77c0655776f1a0d9e Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sat, 31 Oct 2026 00:03:23 +0000 Subject: [PATCH 08/25] Cancel a volume upload whose write failed instead of leaving it open Co-Authored-By: Claude Opus 5.5 --- .../src/kumo_databricks/volume.py | 21 +++++++++++++++++-- packages/kumo-databricks/tests/test_volume.py | 1 + 2 files changed, 20 insertions(+), 2 deletions(-) diff --git a/packages/kumo-databricks/src/kumo_databricks/volume.py b/packages/kumo-databricks/src/kumo_databricks/volume.py index 8d65c0b30..ae7def3fe 100644 --- a/packages/kumo-databricks/src/kumo_databricks/volume.py +++ b/packages/kumo-databricks/src/kumo_databricks/volume.py @@ -10,7 +10,7 @@ import queue import threading from contextlib import contextmanager -from typing import BinaryIO, Dict, Iterator, Optional +from typing import Any, BinaryIO, Dict, Iterator, Optional from urllib.parse import quote from databricks.sdk import WorkspaceClient @@ -153,9 +153,12 @@ def close(self) -> None: self._stream = None +_CANCEL = object() + + class _Feed(io.RawIOBase): def __init__(self) -> None: - self.pieces: queue.Queue[Optional[bytes]] = queue.Queue(maxsize=4) + self.pieces: queue.Queue[Any] = queue.Queue(maxsize=4) self._pending = b'' self._done = False @@ -165,6 +168,8 @@ def readable(self) -> bool: def readinto(self, buffer) -> int: while not self._pending and not self._done: piece = self.pieces.get() + if piece is _CANCEL: + raise OSError('upload cancelled') if piece is None: self._done = True else: @@ -214,6 +219,18 @@ def _put(self, piece: Optional[bytes]) -> None: if piece is not None: raise OSError(f'{self._path}: the upload stopped reading') + def abort(self) -> None: + if self.closed: + return + while self._thread.is_alive(): + try: + self._feed.pieces.put(_CANCEL, timeout=1) + break + except queue.Full: + continue + self._thread.join() + super().close() + def close(self) -> None: if self.closed: return diff --git a/packages/kumo-databricks/tests/test_volume.py b/packages/kumo-databricks/tests/test_volume.py index 15cef7d32..c13650ffc 100644 --- a/packages/kumo-databricks/tests/test_volume.py +++ b/packages/kumo-databricks/tests/test_volume.py @@ -179,3 +179,4 @@ def test_an_occupied_volume_directory_is_refused_before_scoring(tmp_path): write_mode=mode) sink.preflight_write_files(SimpleNamespace(), f'{ROOT}/other/') sink.preflight_write_files(SimpleNamespace(), 'main.ml.scores') + From 3b246839374510517090ab00d509e0f53779b144 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sat, 31 Oct 2026 00:03:35 +0000 Subject: [PATCH 09/25] Drop a trailing blank line flake8 flags in the volume tests Co-Authored-By: Claude Opus 5.5 --- packages/kumo-databricks/tests/test_volume.py | 1 - 1 file changed, 1 deletion(-) diff --git a/packages/kumo-databricks/tests/test_volume.py b/packages/kumo-databricks/tests/test_volume.py index c13650ffc..15cef7d32 100644 --- a/packages/kumo-databricks/tests/test_volume.py +++ b/packages/kumo-databricks/tests/test_volume.py @@ -179,4 +179,3 @@ def test_an_occupied_volume_directory_is_refused_before_scoring(tmp_path): write_mode=mode) sink.preflight_write_files(SimpleNamespace(), f'{ROOT}/other/') sink.preflight_write_files(SimpleNamespace(), 'main.ml.scores') - From 197e661959164c1d0511e988275942570d744d9e Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sat, 31 Oct 2026 01:08:20 +0000 Subject: [PATCH 10/25] Pass the write mode to the shared preflight for volume directories Co-Authored-By: Claude Opus 5.5 --- packages/kumo-databricks/src/kumo_databricks/sink.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/packages/kumo-databricks/src/kumo_databricks/sink.py b/packages/kumo-databricks/src/kumo_databricks/sink.py index 4b21d4d55..47734fddf 100644 --- a/packages/kumo-databricks/src/kumo_databricks/sink.py +++ b/packages/kumo-databricks/src/kumo_databricks/sink.py @@ -241,7 +241,8 @@ def preflight_write_files(self, connector: SourceConnector, write_mode: Optional[WriteMode] = None) -> None: """Refuse an occupied volume directory; tables are not checked.""" if self.holds_files(connector, table_name): - FilesystemSink(self).preflight_write_files(connector, table_name) + FilesystemSink(self).preflight_write_files( + connector, table_name, write_mode) def holds_files(self, connector: SourceConnector, path: str) -> bool: """Volume paths hold files; anything else is a table.""" From 6b00e0acea5d4de587811f14f05ced2209c6c1d2 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sat, 31 Oct 2026 21:46:57 +0000 Subject: [PATCH 11/25] Route volume model publish, fetch and preflight through FileStoreRouting Co-Authored-By: Claude Opus 5.5 --- .../src/kumo_databricks/sink.py | 55 ++----------------- 1 file changed, 4 insertions(+), 51 deletions(-) diff --git a/packages/kumo-databricks/src/kumo_databricks/sink.py b/packages/kumo-databricks/src/kumo_databricks/sink.py index 47734fddf..fd0a7a353 100644 --- a/packages/kumo-databricks/src/kumo_databricks/sink.py +++ b/packages/kumo-databricks/src/kumo_databricks/sink.py @@ -22,26 +22,22 @@ import os import tempfile from pathlib import Path -from typing import Mapping, Optional, Sequence +from typing import Optional, Sequence import pandas as pd from databricks.sql.client import Connection -from kumo_core.object_store_sink import FilesystemSink +from kumo_core.object_store_sink import FileStoreRouting, FilesystemSink from kumo_databricks.registry import DatabricksModelRegistryMixin from kumo_databricks.session import _target_config from kumo_databricks.volume import is_volume_path, volume_filesystem from kumo_databricks.write import RunVolume from kumoapi.data_source import DataSourceType from kumoapi.jobs import WriteMode -from kumoapi.plugin import ( - ModelRegistryLocation, - PublishedModel, - SourceConnector, -) +from kumoapi.plugin import SourceConnector from pyarrow.fs import FileSystem -class DatabricksSinkBackend(DatabricksModelRegistryMixin): +class DatabricksSinkBackend(FileStoreRouting, DatabricksModelRegistryMixin): r"""Writes a produced artifact to a Databricks Delta table. Registered by default beside @@ -201,49 +197,6 @@ def write_files(self, connector, table_name, files, prepare_file=prepare) return sql.qualified_name(config.catalog, schema, table) - def publish_model( - self, - connector: SourceConnector, - bundle_dir: str, - *, - path: Optional[str] = None, - location: Optional[ModelRegistryLocation] = None, - metadata: Mapping[str, str] = {}, - metrics: Mapping[str, float] = {}, - project_id: Optional[str] = None, - ) -> PublishedModel: - """Register in Unity Catalog, or copy to a volume ``path``.""" - if location is None and path and self.holds_files(connector, path): - return FilesystemSink(self).publish_model( - connector, bundle_dir, path=path, metadata=metadata) - return super().publish_model(connector, bundle_dir, path=path, - location=location, metadata=metadata, - metrics=metrics, project_id=project_id) - - def fetch_model( - self, - connector: SourceConnector, - local_dir: str, - *, - path: Optional[str] = None, - location: Optional[ModelRegistryLocation] = None, - project_id: Optional[str] = None, - ) -> PublishedModel: - """Download from Unity Catalog, or from a volume ``path``.""" - if location is None and path and self.holds_files(connector, path): - return FilesystemSink(self).fetch_model(connector, local_dir, - path=path) - return super().fetch_model(connector, local_dir, path=path, - location=location, project_id=project_id) - - def preflight_write_files(self, connector: SourceConnector, - table_name: str, - write_mode: Optional[WriteMode] = None) -> None: - """Refuse an occupied volume directory; tables are not checked.""" - if self.holds_files(connector, table_name): - FilesystemSink(self).preflight_write_files( - connector, table_name, write_mode) - def holds_files(self, connector: SourceConnector, path: str) -> bool: """Volume paths hold files; anything else is a table.""" return is_volume_path(path) From d60707731981348724d42ab043781e9c8f0ece74 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sat, 31 Oct 2026 21:47:28 +0000 Subject: [PATCH 12/25] Share one bounded put loop between a volume upload's writes and abort Co-Authored-By: Claude Opus 5.5 --- .../src/kumo_databricks/volume.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/packages/kumo-databricks/src/kumo_databricks/volume.py b/packages/kumo-databricks/src/kumo_databricks/volume.py index ae7def3fe..2e14c31d8 100644 --- a/packages/kumo-databricks/src/kumo_databricks/volume.py +++ b/packages/kumo-databricks/src/kumo_databricks/volume.py @@ -206,13 +206,18 @@ def write(self, data) -> int: self._put(piece) return len(piece) - def _put(self, piece: Optional[bytes]) -> None: + def _offer(self, piece: Any) -> bool: while self._thread.is_alive(): try: self._feed.pieces.put(piece, timeout=1) - return + return True except queue.Full: continue + return False + + def _put(self, piece: Optional[bytes]) -> None: + if self._offer(piece): + return self._thread.join() if self._error is not None: raise self._error @@ -222,12 +227,7 @@ def _put(self, piece: Optional[bytes]) -> None: def abort(self) -> None: if self.closed: return - while self._thread.is_alive(): - try: - self._feed.pieces.put(_CANCEL, timeout=1) - break - except queue.Full: - continue + self._offer(_CANCEL) self._thread.join() super().close() From cdfb75d860f4c50d45572a82c861b70511758313 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sat, 31 Oct 2026 21:47:28 +0000 Subject: [PATCH 13/25] Drop the fake volume's unused directory set, bad-request branch and stream log Co-Authored-By: Claude Opus 5.5 --- packages/kumo-databricks/tests/fakes.py | 37 ++++++------------------- 1 file changed, 9 insertions(+), 28 deletions(-) diff --git a/packages/kumo-databricks/tests/fakes.py b/packages/kumo-databricks/tests/fakes.py index b629e404d..d42545f71 100644 --- a/packages/kumo-databricks/tests/fakes.py +++ b/packages/kumo-databricks/tests/fakes.py @@ -11,7 +11,7 @@ from urllib.parse import unquote import pyarrow as pa -from databricks.sdk.errors import BadRequest, NotFound +from databricks.sdk.errors import NotFound from databricks.sql.exc import DataError from sources.recording import RecordingConnection, RecordingCursor @@ -199,29 +199,23 @@ def delete_directory(self, path: str) -> None: class FakeVolume(FakeFiles): - """A volume with real directories, uploads and ranged downloads.""" + """A volume with uploads and ranged downloads.""" def __init__(self, objects: Mapping[str, bytes] | None = None) -> None: super().__init__(objects) - self.directories: set[str] = set() self.ignore_range = False - def _is_directory(self, path: str) -> bool: - return path in self.directories or any( - name.startswith(path + "/") for name in self.objects) - def list_directory_contents(self, path: str) -> list[FakeEntry]: path = path.rstrip("/") - if not self._is_directory(path): - raise NotFound(f"No such directory: {path}") children: dict[str, FakeEntry] = {} - for name in sorted(self.objects) + sorted(self.directories): + for name in sorted(self.objects): if not name.startswith(path + "/"): continue - child = name[len(path) + 1:].split("/")[0] - full = f"{path}/{child}" - children[full] = (FakeEntry(full, len(self.objects[full])) - if full in self.objects else FakeEntry( - full + "/", 0, is_directory=True)) + child = f"{path}/{name[len(path) + 1:].split('/')[0]}" + children[child] = (FakeEntry(child, len(self.objects[child])) + if child in self.objects else FakeEntry( + child + "/", 0, is_directory=True)) + if not children: + raise NotFound(f"No such directory: {path}") return list(children.values()) def upload(self, path: str, contents: Any, *, overwrite: bool) -> None: @@ -229,22 +223,10 @@ def upload(self, path: str, contents: Any, *, overwrite: bool) -> None: self.objects[path] = contents.read() def get_metadata(self, path: str) -> Any: - if path.endswith("/"): - raise BadRequest("Bad Request") if path not in self.objects: raise NotFound(f"No such file: {path}") return FakeMetadata(len(self.objects[path])) - def get_directory_metadata(self, path: str) -> None: - if not self._is_directory(path): - raise NotFound(f"No such directory: {path}") - - def create_directory(self, path: str) -> None: - self.directories.add(path) - - def delete_directory(self, path: str) -> None: - self.directories.discard(path.rstrip("/")) - def do(self, method: str, path: str, *, raw: bool, response_headers: list[str], headers: Mapping[str, str]) -> dict: assert method == "GET" and raw @@ -258,7 +240,6 @@ def do(self, method: str, path: str, *, raw: bool, last = last or str(size - 1) stream = FakeStreamingResponse( self.objects[remote][int(first):int(last) + 1]) - self.streams.append(stream) return { "contents": stream, "content-range": f"bytes {first}-{last}/{size}" From 868becaae77599c90e55df51ec80f4a80e37aef7 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sat, 31 Oct 2026 22:32:20 +0000 Subject: [PATCH 14/25] Let FileStoreRouting send Databricks volume paths to file writers; keep only table writes Co-Authored-By: Claude Opus 5.5 --- .../src/kumo_databricks/sink.py | 113 +++++++++--------- 1 file changed, 54 insertions(+), 59 deletions(-) diff --git a/packages/kumo-databricks/src/kumo_databricks/sink.py b/packages/kumo-databricks/src/kumo_databricks/sink.py index fd0a7a353..80f32a6a6 100644 --- a/packages/kumo-databricks/src/kumo_databricks/sink.py +++ b/packages/kumo-databricks/src/kumo_databricks/sink.py @@ -26,7 +26,7 @@ import pandas as pd from databricks.sql.client import Connection -from kumo_core.object_store_sink import FileStoreRouting, FilesystemSink +from kumo_core.object_store_sink import FileStoreRouting from kumo_databricks.registry import DatabricksModelRegistryMixin from kumo_databricks.session import _target_config from kumo_databricks.volume import is_volume_path, volume_filesystem @@ -37,56 +37,7 @@ from pyarrow.fs import FileSystem -class DatabricksSinkBackend(FileStoreRouting, DatabricksModelRegistryMixin): - r"""Writes a produced artifact to a Databricks Delta table. - - Registered by default beside - :class:`~kumo_databricks.backend.DatabricksSourceBackend`, so a - run that read its graph from a catalog exports its predictions back to one - with no further setup. Installed as - ``pip install 'kumoai-studio[databricks]'``. - - A destination names its schema and its table in one string, joined by - ``__kumo__``, because that is the only field a Databricks output has -- - :func:`kumoai_studio.databricks.to_db_table_name` builds it from the - ``(schema, table)`` tuple an :class:`~kumoai_studio.artifact_export.OutputConfig` - takes. The schema must be the connector's own, which - :func:`kumo_databricks.output.within_schema` checks and the - hosted connector checks through as well. - - Every write stages through the connector's managed volume, - ``.._kumo_staging``, created on first use. That is not a - local choice: ``COPY INTO`` is how a Delta table is loaded, and hosted - states the alternatives -- S3, ADLS, an arbitrary volume, inline cloud - credentials -- are unsupported on purpose. The volume is emptied of what - this export put there once the table has landed. - - Args: - identity: A fallback identity, used only for connectors that name none - themselves. As with the reader, leaving it unset is the common - case, since a bound - :class:`~kumoai_studio.connector.DatabricksConnector` records the - identity it was created with. That identity needs ``CREATE TABLE`` - on the connector's schema and ``READ``/``WRITE VOLUME`` on its - staging volume, which a read-only role used for ingest will not - have. - - Example: - >>> import kumoai_studio as kumo # doctest: +SKIP - >>> from kumo_databricks import ( # doctest: +SKIP - ... DatabricksSinkBackend) - >>> kumo.set_backend(kumo.LocalBackend( # doctest: +SKIP - ... storage=paths, - ... sinks={kumo.DataSourceType.DATABRICKS: - ... DatabricksSinkBackend(identity='prod-dbx')}, - ... )) - """ - - source_type = DataSourceType.DATABRICKS - - _default_hint = ('LocalBackend(sinks={DataSourceType.DATABRICKS: ' - 'DatabricksSinkBackend(identity=identity)})') - +class _DatabricksTables(DatabricksModelRegistryMixin): def write_table( self, connector: SourceConnector, @@ -128,9 +79,6 @@ def write_table( names a schema other than the connector's, or if the staging volume is neither reachable nor creatable. """ - if self.holds_files(connector, table_name): - return FilesystemSink(self).write_table( - connector, table_name, frame, write_mode, export_id=export_id) del project_id, export_id from kumo_databricks import output, sql, write @@ -159,11 +107,7 @@ def write_table( def write_files(self, connector, table_name, files, write_mode=WriteMode.OVERWRITE, project_id=None, *, export_id=None, extra_columns=None): - """Copy Parquet to a volume path, or COPY it into a table.""" - if self.holds_files(connector, table_name): - return FilesystemSink(self).write_files( - connector, table_name, files, write_mode, - export_id=export_id, extra_columns=extra_columns) + """COPY selected Parquet into a table.""" del project_id, export_id from functools import partial @@ -197,6 +141,57 @@ def write_files(self, connector, table_name, files, prepare_file=prepare) return sql.qualified_name(config.catalog, schema, table) + +class DatabricksSinkBackend(FileStoreRouting, _DatabricksTables): + r"""Writes a produced artifact to a Databricks Delta table. + + Registered by default beside + :class:`~kumo_databricks.backend.DatabricksSourceBackend`, so a + run that read its graph from a catalog exports its predictions back to one + with no further setup. Installed as + ``pip install 'kumoai-studio[databricks]'``. + + A destination names its schema and its table in one string, joined by + ``__kumo__``, because that is the only field a Databricks output has -- + :func:`kumoai_studio.databricks.to_db_table_name` builds it from the + ``(schema, table)`` tuple an :class:`~kumoai_studio.artifact_export.OutputConfig` + takes. The schema must be the connector's own, which + :func:`kumo_databricks.output.within_schema` checks and the + hosted connector checks through as well. + + Every write stages through the connector's managed volume, + ``.._kumo_staging``, created on first use. That is not a + local choice: ``COPY INTO`` is how a Delta table is loaded, and hosted + states the alternatives -- S3, ADLS, an arbitrary volume, inline cloud + credentials -- are unsupported on purpose. The volume is emptied of what + this export put there once the table has landed. + + Args: + identity: A fallback identity, used only for connectors that name none + themselves. As with the reader, leaving it unset is the common + case, since a bound + :class:`~kumoai_studio.connector.DatabricksConnector` records the + identity it was created with. That identity needs ``CREATE TABLE`` + on the connector's schema and ``READ``/``WRITE VOLUME`` on its + staging volume, which a read-only role used for ingest will not + have. + + Example: + >>> import kumoai_studio as kumo # doctest: +SKIP + >>> from kumo_databricks import ( # doctest: +SKIP + ... DatabricksSinkBackend) + >>> kumo.set_backend(kumo.LocalBackend( # doctest: +SKIP + ... storage=paths, + ... sinks={kumo.DataSourceType.DATABRICKS: + ... DatabricksSinkBackend(identity='prod-dbx')}, + ... )) + """ + + source_type = DataSourceType.DATABRICKS + + _default_hint = ('LocalBackend(sinks={DataSourceType.DATABRICKS: ' + 'DatabricksSinkBackend(identity=identity)})') + def holds_files(self, connector: SourceConnector, path: str) -> bool: """Volume paths hold files; anything else is a table.""" return is_volume_path(path) From b970b76e64479993d44feca0cb62941346909d01 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sun, 1 Nov 2026 01:03:14 +0000 Subject: [PATCH 15/25] Cut the Databricks sink's class docstring to a summary and its argument Co-Authored-By: Claude Opus 5.5 --- .../src/kumo_databricks/sink.py | 42 +------------------ 1 file changed, 2 insertions(+), 40 deletions(-) diff --git a/packages/kumo-databricks/src/kumo_databricks/sink.py b/packages/kumo-databricks/src/kumo_databricks/sink.py index 80f32a6a6..e29027771 100644 --- a/packages/kumo-databricks/src/kumo_databricks/sink.py +++ b/packages/kumo-databricks/src/kumo_databricks/sink.py @@ -143,48 +143,10 @@ def write_files(self, connector, table_name, files, class DatabricksSinkBackend(FileStoreRouting, _DatabricksTables): - r"""Writes a produced artifact to a Databricks Delta table. - - Registered by default beside - :class:`~kumo_databricks.backend.DatabricksSourceBackend`, so a - run that read its graph from a catalog exports its predictions back to one - with no further setup. Installed as - ``pip install 'kumoai-studio[databricks]'``. - - A destination names its schema and its table in one string, joined by - ``__kumo__``, because that is the only field a Databricks output has -- - :func:`kumoai_studio.databricks.to_db_table_name` builds it from the - ``(schema, table)`` tuple an :class:`~kumoai_studio.artifact_export.OutputConfig` - takes. The schema must be the connector's own, which - :func:`kumo_databricks.output.within_schema` checks and the - hosted connector checks through as well. - - Every write stages through the connector's managed volume, - ``.._kumo_staging``, created on first use. That is not a - local choice: ``COPY INTO`` is how a Delta table is loaded, and hosted - states the alternatives -- S3, ADLS, an arbitrary volume, inline cloud - credentials -- are unsupported on purpose. The volume is emptied of what - this export put there once the table has landed. + """Writes artifacts to Delta tables, staged through the schema's volume. Args: - identity: A fallback identity, used only for connectors that name none - themselves. As with the reader, leaving it unset is the common - case, since a bound - :class:`~kumoai_studio.connector.DatabricksConnector` records the - identity it was created with. That identity needs ``CREATE TABLE`` - on the connector's schema and ``READ``/``WRITE VOLUME`` on its - staging volume, which a read-only role used for ingest will not - have. - - Example: - >>> import kumoai_studio as kumo # doctest: +SKIP - >>> from kumo_databricks import ( # doctest: +SKIP - ... DatabricksSinkBackend) - >>> kumo.set_backend(kumo.LocalBackend( # doctest: +SKIP - ... storage=paths, - ... sinks={kumo.DataSourceType.DATABRICKS: - ... DatabricksSinkBackend(identity='prod-dbx')}, - ... )) + identity: Fallback identity for connectors that name none. """ source_type = DataSourceType.DATABRICKS From 4c99876b4f6b451f5c357f8f9cebbbd0b970a2f4 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sun, 1 Nov 2026 01:49:00 +0000 Subject: [PATCH 16/25] Drop duplicate dockerignore entries the base already has Co-Authored-By: Claude Opus 5.5 --- tests/public/Dockerfile.dockerignore | 4 ---- 1 file changed, 4 deletions(-) diff --git a/tests/public/Dockerfile.dockerignore b/tests/public/Dockerfile.dockerignore index 00eca5c99..c6d21c61a 100644 --- a/tests/public/Dockerfile.dockerignore +++ b/tests/public/Dockerfile.dockerignore @@ -14,10 +14,6 @@ !packages/kumo-core/src/kumo_core/artifact_staging.py # Model registries validate bundles through the portable artifact contract. !packages/kumo-core/src/kumo_core/artifacts.py -# Volume and stage sinks hold files through the shared chunked filesystem. -!packages/kumo-core/src/kumo_core/chunked_fs.py -!packages/kumo-core/src/kumo_core/object_store_sink.py -!packages/kumo-core/src/kumo_core/artifact_export.py !packages/kumo-core/src/kumo_core/config.py !packages/kumo-core/src/kumo_core/manifest.py !packages/kumo-core/src/kumo_core/io/__init__.py From d63748e4f28548aeb780f383c363d56761aa5479 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sun, 1 Nov 2026 01:49:10 +0000 Subject: [PATCH 17/25] Share one upload-finish path between write and close Co-Authored-By: Claude Opus 5.5 --- .../src/kumo_databricks/volume.py | 16 ++++++---------- 1 file changed, 6 insertions(+), 10 deletions(-) diff --git a/packages/kumo-databricks/src/kumo_databricks/volume.py b/packages/kumo-databricks/src/kumo_databricks/volume.py index 2e14c31d8..1f0657de1 100644 --- a/packages/kumo-databricks/src/kumo_databricks/volume.py +++ b/packages/kumo-databricks/src/kumo_databricks/volume.py @@ -203,7 +203,9 @@ def writable(self) -> bool: def write(self, data) -> int: piece = bytes(data) - self._put(piece) + if not self._offer(piece): + self._finish() + raise OSError(f'{self._path}: the upload stopped reading') return len(piece) def _offer(self, piece: Any) -> bool: @@ -215,14 +217,10 @@ def _offer(self, piece: Any) -> bool: continue return False - def _put(self, piece: Optional[bytes]) -> None: - if self._offer(piece): - return + def _finish(self) -> None: self._thread.join() if self._error is not None: raise self._error - if piece is not None: - raise OSError(f'{self._path}: the upload stopped reading') def abort(self) -> None: if self.closed: @@ -235,9 +233,7 @@ def close(self) -> None: if self.closed: return try: - self._put(None) - self._thread.join() - if self._error is not None: - raise self._error + self._offer(None) + self._finish() finally: super().close() From 54ac96c700838a2716f881cb97fe192870230419 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sun, 1 Nov 2026 01:49:13 +0000 Subject: [PATCH 18/25] Share the fake volume sink setup across volume tests Co-Authored-By: Claude Opus 5.5 --- packages/kumo-databricks/tests/test_volume.py | 23 +++++++++---------- 1 file changed, 11 insertions(+), 12 deletions(-) diff --git a/packages/kumo-databricks/tests/test_volume.py b/packages/kumo-databricks/tests/test_volume.py index 15cef7d32..6bcdf8bcc 100644 --- a/packages/kumo-databricks/tests/test_volume.py +++ b/packages/kumo-databricks/tests/test_volume.py @@ -25,6 +25,13 @@ def filesystem(objects=None, part_bytes=None): return PyFileSystem(handler), volume +def volume_sink(): + sink = DatabricksSinkBackend() + volume = FakeVolume() + sink._client = lambda connector: FakeWorkspaceClient(volume) + return sink, volume + + def write(fs, path, body, piece=700): with fs.open_output_stream(path) as writer: for start in range(0, len(body), piece): @@ -108,9 +115,7 @@ def refuse(path, contents, *, overwrite): def test_sink_filesystem_only_takes_volume_paths(): - sink = DatabricksSinkBackend() - volume = FakeVolume() - sink._client = lambda connector: FakeWorkspaceClient(volume) + sink, volume = volume_sink() connector = SimpleNamespace() fs, path = sink.filesystem(connector, f'{ROOT}/exports/') assert path == f'{ROOT}/exports' @@ -125,9 +130,7 @@ def test_sink_filesystem_only_takes_volume_paths(): def test_bundle_publish_and_fetch_through_a_volume_path(tmp_path): - sink = DatabricksSinkBackend() - volume = FakeVolume() - sink._client = lambda connector: FakeWorkspaceClient(volume) + sink, volume = volume_sink() bundle = tmp_path / 'bundle' (bundle / 'model').mkdir(parents=True) (bundle / 'training-spec.json').write_text('{}') @@ -145,9 +148,7 @@ def test_bundle_publish_and_fetch_through_a_volume_path(tmp_path): def test_predictions_and_frames_written_to_a_volume_path(tmp_path): - sink = DatabricksSinkBackend() - volume = FakeVolume() - sink._client = lambda connector: FakeWorkspaceClient(volume) + sink, volume = volume_sink() part = tmp_path / 'part-0.parquet' pq.write_table(pa.table({'users': [1, 2], 'score': [0.1, 0.9]}), part) written = sink.write_files(SimpleNamespace(), f'{ROOT}/predictions/', @@ -164,9 +165,7 @@ def test_an_occupied_volume_directory_is_refused_before_scoring(tmp_path): from kumo_core.artifact_export import ArtifactDestinationConflict from kumoapi.jobs import WriteMode - sink = DatabricksSinkBackend() - volume = FakeVolume() - sink._client = lambda connector: FakeWorkspaceClient(volume) + sink, volume = volume_sink() target = f'{ROOT}/predictions/' sink.preflight_write_files(SimpleNamespace(), target) part = tmp_path / 'part-0.parquet' From 849848fa04aa4cb0139f791de16346bfb38c6f14 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sun, 1 Nov 2026 01:49:17 +0000 Subject: [PATCH 19/25] Reflow the volume README paragraph Co-Authored-By: Claude Opus 5.5 --- packages/kumo-databricks/README.md | 20 +++++++++----------- 1 file changed, 9 insertions(+), 11 deletions(-) diff --git a/packages/kumo-databricks/README.md b/packages/kumo-databricks/README.md index 2c0a1a53f..2874ff510 100644 --- a/packages/kumo-databricks/README.md +++ b/packages/kumo-databricks/README.md @@ -86,17 +86,15 @@ trainer.predict(prediction_table, snapshot=snapshot, output_config=PredictionOut For an MLOps artifact root on a volume, see `packages/kumo-engine/docs/mlops.md`. -A file whose size is known when it is written (bundle files, exported -Parquet parts) is stored as one object up to 4 GiB, so Spark -and other tools read it directly. A larger one streams into -`.chunks//00000`, `00001`, ... (4 GiB each, under the Files -API's 5 GiB upload cap) with an empty `_SUCCESS` written last, and Kumo reads -the newest committed generation back as one file. A file written without a -known size stays one object up to -256 MiB, staged on local disk under the engine scratch directory, and streams -into parts above that. The identity needs `USE CATALOG`, `USE SCHEMA`, and -`READ VOLUME` (loads) and `WRITE VOLUME` (exports) on the volume; no -warehouse is used. +A file whose size is known when it is written (bundle files, exported Parquet +parts) is stored as one object up to 4 GiB, so Spark and other tools read it +directly. A larger one streams into `.chunks//00000`, +`00001`, ... (4 GiB each, under the Files API's 5 GiB upload cap) with an empty +`_SUCCESS` written last, and Kumo reads the newest committed generation back as +one file. A file written without a known size stays one object up to 256 MiB, +staged on local disk under the engine scratch directory, and streams into parts +above that. The identity needs `USE CATALOG`, `USE SCHEMA`, and `READ VOLUME` +(loads) and `WRITE VOLUME` (exports) on the volume; no warehouse is used. ## Model registry From a75c11b4c2acb018f6c2314e537895de3086f3c9 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Mon, 2 Nov 2026 20:30:34 +0000 Subject: [PATCH 20/25] Reopen a volume download after a failed read Co-Authored-By: Claude Opus 5.5 --- .../src/kumo_databricks/volume.py | 20 +++++++++------- packages/kumo-databricks/tests/test_volume.py | 24 +++++++++++++++++++ 2 files changed, 36 insertions(+), 8 deletions(-) diff --git a/packages/kumo-databricks/src/kumo_databricks/volume.py b/packages/kumo-databricks/src/kumo_databricks/volume.py index 1f0657de1..e5face782 100644 --- a/packages/kumo-databricks/src/kumo_databricks/volume.py +++ b/packages/kumo-databricks/src/kumo_databricks/volume.py @@ -118,14 +118,18 @@ def read_at(self, offset: int, length: int) -> bytes: return b'' if self._stream is None or offset != self._at: self._open(offset) - with _translated(self._path): - pieces, got = [], 0 - while got < length: - piece = self._stream.read(min(_BUFFER, length - got)) - if not piece: - break - pieces.append(piece) - got += len(piece) + try: + with _translated(self._path): + pieces, got = [], 0 + while got < length: + piece = self._stream.read(min(_BUFFER, length - got)) + if not piece: + break + pieces.append(piece) + got += len(piece) + except BaseException: + self.close() + raise self._at += got return b''.join(pieces) diff --git a/packages/kumo-databricks/tests/test_volume.py b/packages/kumo-databricks/tests/test_volume.py index 6bcdf8bcc..4ae15d24d 100644 --- a/packages/kumo-databricks/tests/test_volume.py +++ b/packages/kumo-databricks/tests/test_volume.py @@ -89,6 +89,30 @@ def test_a_download_that_ignores_the_range_is_refused(): reader.read_at(10, 100) +def test_a_read_after_a_dropped_download_starts_where_asked(): + fs, volume = filesystem({f'{ROOT}/a': BODY}) + do = volume.do + + def dropping(*args, **kwargs): + volume.do = do + response = do(*args, **kwargs) + stream = response['contents'] + + def read(size, read=stream.read): + read(size // 2) + stream.read = read + raise ConnectionError('connection reset') + + stream.read = read + return response + + volume.do = dropping + with fs.open_input_file(f'{ROOT}/a') as reader: + with pytest.raises(OSError, match='ConnectionError'): + reader.read_at(10, 100) + assert reader.read_at(10, 100) == BODY[100:110] + + def test_directories_and_missing_paths(): fs, volume = filesystem({ f'{ROOT}/s/kumo_snapshot.json': b'{}', From 15dcf523d4fa1d181c351a0788748620da8f7684 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Mon, 2 Nov 2026 21:57:45 +0000 Subject: [PATCH 21/25] Keep unsized volume exports as plain objects up to a part CSV exports open the stream without a size, so anything over 256 MiB rolled over to .chunks parts that Spark cannot read as one file. Co-Authored-By: Claude Opus 5.5 --- packages/kumo-databricks/src/kumo_databricks/volume.py | 2 +- packages/kumo-databricks/tests/test_volume.py | 9 +++++++++ 2 files changed, 10 insertions(+), 1 deletion(-) diff --git a/packages/kumo-databricks/src/kumo_databricks/volume.py b/packages/kumo-databricks/src/kumo_databricks/volume.py index e5face782..6d85765d9 100644 --- a/packages/kumo-databricks/src/kumo_databricks/volume.py +++ b/packages/kumo-databricks/src/kumo_databricks/volume.py @@ -51,7 +51,7 @@ def _translated(path: str) -> Iterator[None]: class VolumeHandler(ChunkedFileHandler): part_bytes = 4 << 30 # under the Files API's 5 GiB upload cap - plain_bytes = 256 << 20 + plain_bytes = part_bytes def __init__(self, client: WorkspaceClient) -> None: self._client = client diff --git a/packages/kumo-databricks/tests/test_volume.py b/packages/kumo-databricks/tests/test_volume.py index 4ae15d24d..71de4a7eb 100644 --- a/packages/kumo-databricks/tests/test_volume.py +++ b/packages/kumo-databricks/tests/test_volume.py @@ -62,6 +62,15 @@ def test_sized_writes_stay_plain_up_to_a_part(): and path.endswith('/_SUCCESS') for path in volume.objects) +def test_unsized_writes_stay_plain_through_a_part(): + assert VolumeHandler.plain_bytes == VolumeHandler.part_bytes == 4 << 30 + fs, volume = filesystem(part_bytes=1000) + write(fs, f'{ROOT}/export.csv', BODY[:1000]) + assert volume.objects[f'{ROOT}/export.csv'] == BODY[:1000] + write(fs, f'{ROOT}/over.csv', BODY[:1001]) + assert f'{ROOT}/over.csv' not in volume.objects + + def test_a_failed_upload_raises_without_the_signed_url(): from kumo_core.object_store_sink import open_sized From 0e1bf22cbc8f9b1c98241c8c0670f03f1c263995 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Mon, 2 Nov 2026 21:57:45 +0000 Subject: [PATCH 22/25] Document that volume exports need READ VOLUME too The export preflight lists the destination before writing. Co-Authored-By: Claude Opus 5.5 --- packages/kumo-databricks/README.md | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/packages/kumo-databricks/README.md b/packages/kumo-databricks/README.md index 2874ff510..8de659f9d 100644 --- a/packages/kumo-databricks/README.md +++ b/packages/kumo-databricks/README.md @@ -91,10 +91,11 @@ parts) is stored as one object up to 4 GiB, so Spark and other tools read it directly. A larger one streams into `.chunks//00000`, `00001`, ... (4 GiB each, under the Files API's 5 GiB upload cap) with an empty `_SUCCESS` written last, and Kumo reads the newest committed generation back as -one file. A file written without a known size stays one object up to 256 MiB, +one file. A file written without a known size stays one object up to 4 GiB, staged on local disk under the engine scratch directory, and streams into parts above that. The identity needs `USE CATALOG`, `USE SCHEMA`, and `READ VOLUME` -(loads) and `WRITE VOLUME` (exports) on the volume; no warehouse is used. +(loads and exports) and `WRITE VOLUME` (exports) on the volume; no warehouse +is used. ## Model registry From 2bafa47246adea737021542c82e374663d61f5dc Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Tue, 3 Nov 2026 00:02:25 +0000 Subject: [PATCH 23/25] Trim volume and sink docstrings and redundant path stripping Co-Authored-By: Claude Opus 5.5 --- packages/kumo-databricks/src/kumo_databricks/sink.py | 2 -- packages/kumo-databricks/src/kumo_databricks/volume.py | 8 +------- 2 files changed, 1 insertion(+), 9 deletions(-) diff --git a/packages/kumo-databricks/src/kumo_databricks/sink.py b/packages/kumo-databricks/src/kumo_databricks/sink.py index e29027771..0b82aca01 100644 --- a/packages/kumo-databricks/src/kumo_databricks/sink.py +++ b/packages/kumo-databricks/src/kumo_databricks/sink.py @@ -155,12 +155,10 @@ class DatabricksSinkBackend(FileStoreRouting, _DatabricksTables): 'DatabricksSinkBackend(identity=identity)})') def holds_files(self, connector: SourceConnector, path: str) -> bool: - """Volume paths hold files; anything else is a table.""" return is_volume_path(path) def filesystem(self, connector: SourceConnector, path: str) -> tuple[FileSystem, str]: - """A volume path's filesystem, as the connector's identity.""" if not self.holds_files(connector, path): raise ValueError( f"{path!r} is not a Unity Catalog volume path; Databricks " diff --git a/packages/kumo-databricks/src/kumo_databricks/volume.py b/packages/kumo-databricks/src/kumo_databricks/volume.py index 6d85765d9..55e486107 100644 --- a/packages/kumo-databricks/src/kumo_databricks/volume.py +++ b/packages/kumo-databricks/src/kumo_databricks/volume.py @@ -1,7 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. # All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Unity Catalog volumes as a pyarrow filesystem, over the Files API.""" +"""Unity Catalog volumes as a pyarrow filesystem.""" from __future__ import annotations @@ -36,7 +36,6 @@ def volume_filesystem(client: WorkspaceClient) -> FileSystem: @contextmanager def _translated(path: str) -> Iterator[None]: - """OS errors, unchained; non-SDK text may hold a signed URL.""" try: yield except NotFound as error: @@ -60,7 +59,6 @@ def get_type_name(self) -> str: return 'databricks-volume' def _size(self, path: str) -> Optional[int]: - path = path.rstrip('/') with _translated(path): try: return int(self._client.files.get_metadata(path).content_length) @@ -68,7 +66,6 @@ def _size(self, path: str) -> Optional[int]: return None def _list(self, directory: str) -> Dict[str, int]: - directory = directory.rstrip('/') with _translated(directory): try: entries = list( @@ -95,7 +92,6 @@ def _part(self, path: str, size: int) -> _Part: return _Part(self._client, path) def _remove(self, path: str, *, tree: bool) -> None: - path = path.rstrip('/') with _translated(path): if tree: clear_volume(self._client, path) @@ -107,7 +103,6 @@ def _remove(self, path: str, *, tree: bool) -> None: class _Part: - """Serves consecutive reads from one open-ended ranged download.""" def __init__(self, client: WorkspaceClient, path: str) -> None: self._client, self._path = client, path self._stream = None @@ -185,7 +180,6 @@ def readinto(self, buffer) -> int: class _Upload(io.RawIOBase): - """One ``files.upload`` fed from writes; ``close`` commits it.""" def __init__(self, client: WorkspaceClient, path: str) -> None: self._path = path self._feed = _Feed() From dfdf8017e88568efc0e59ecd28cdcc146c629443 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Tue, 3 Nov 2026 00:15:50 +0000 Subject: [PATCH 24/25] Tighten volume README and hoist test imports Co-Authored-By: Claude Opus 5.5 --- packages/kumo-databricks/README.md | 19 +++++++++---------- packages/kumo-databricks/tests/test_volume.py | 10 +++------- 2 files changed, 12 insertions(+), 17 deletions(-) diff --git a/packages/kumo-databricks/README.md b/packages/kumo-databricks/README.md index 8de659f9d..bfa1b9c4a 100644 --- a/packages/kumo-databricks/README.md +++ b/packages/kumo-databricks/README.md @@ -86,16 +86,15 @@ trainer.predict(prediction_table, snapshot=snapshot, output_config=PredictionOut For an MLOps artifact root on a volume, see `packages/kumo-engine/docs/mlops.md`. -A file whose size is known when it is written (bundle files, exported Parquet -parts) is stored as one object up to 4 GiB, so Spark and other tools read it -directly. A larger one streams into `.chunks//00000`, -`00001`, ... (4 GiB each, under the Files API's 5 GiB upload cap) with an empty -`_SUCCESS` written last, and Kumo reads the newest committed generation back as -one file. A file written without a known size stays one object up to 4 GiB, -staged on local disk under the engine scratch directory, and streams into parts -above that. The identity needs `USE CATALOG`, `USE SCHEMA`, and `READ VOLUME` -(loads and exports) and `WRITE VOLUME` (exports) on the volume; no warehouse -is used. +A file written with a known size (bundle files, exported Parquet parts) is one +object up to 4 GiB, readable directly by Spark and other tools. A larger one +streams into `.chunks//00000`, `00001`, ... (4 GiB each, +under the Files API's 5 GiB upload cap) with an empty `_SUCCESS` written last; +Kumo reads the newest committed generation back as one file. A file written +without a known size is staged on local disk under the engine scratch +directory and stored the same way. The identity needs `USE CATALOG`, +`USE SCHEMA`, and `READ VOLUME` (loads and exports) and `WRITE VOLUME` +(exports) on the volume; no warehouse is used. ## Model registry diff --git a/packages/kumo-databricks/tests/test_volume.py b/packages/kumo-databricks/tests/test_volume.py index 71de4a7eb..266403c43 100644 --- a/packages/kumo-databricks/tests/test_volume.py +++ b/packages/kumo-databricks/tests/test_volume.py @@ -9,8 +9,11 @@ import pytest from databricks.sdk.errors import PermissionDenied from fakes import FakeVolume, FakeWorkspaceClient +from kumo_core.artifact_export import ArtifactDestinationConflict +from kumo_core.object_store_sink import open_sized from kumo_databricks.sink import DatabricksSinkBackend from kumo_databricks.volume import VolumeHandler +from kumoapi.jobs import WriteMode from pyarrow.fs import FileType, PyFileSystem ROOT = '/Volumes/main/analytics/artifacts' @@ -47,8 +50,6 @@ def test_reads_cross_part_boundaries(): def test_sized_writes_stay_plain_up_to_a_part(): - from kumo_core.object_store_sink import open_sized - fs, volume = filesystem(part_bytes=1000) fs.handler.plain_bytes = 100 for name, body in (('plain', BODY[:1000]), ('parts', BODY[:1500])): @@ -72,8 +73,6 @@ def test_unsized_writes_stay_plain_through_a_part(): def test_a_failed_upload_raises_without_the_signed_url(): - from kumo_core.object_store_sink import open_sized - fs, volume = filesystem() signed = 'https://bucket.example/part?X-Amz-Signature=secret' @@ -195,9 +194,6 @@ def test_predictions_and_frames_written_to_a_volume_path(tmp_path): def test_an_occupied_volume_directory_is_refused_before_scoring(tmp_path): - from kumo_core.artifact_export import ArtifactDestinationConflict - from kumoapi.jobs import WriteMode - sink, volume = volume_sink() target = f'{ROOT}/predictions/' sink.preflight_write_files(SimpleNamespace(), target) From 01ca25a465b0c0ead38a84d716d5e7fa0157f1e0 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Tue, 3 Nov 2026 18:26:59 +0000 Subject: [PATCH 25/25] Bound Databricks volume upload buffering Pass an explicit 16 MiB part size and 4 upload threads to files.upload when the SDK accepts them (databricks-sdk 0.70+), so a large plain object no longer buffers up to parallelism x part_size with the SDK's defaults. Older SDKs upload single-threaded in 10 MiB chunks and get no extra args. Co-Authored-By: Claude Opus 5.5 --- .../src/kumo_databricks/volume.py | 13 +++++++++++-- packages/kumo-databricks/tests/test_volume.py | 17 +++++++++++++++++ 2 files changed, 28 insertions(+), 2 deletions(-) diff --git a/packages/kumo-databricks/src/kumo_databricks/volume.py b/packages/kumo-databricks/src/kumo_databricks/volume.py index 55e486107..02b635197 100644 --- a/packages/kumo-databricks/src/kumo_databricks/volume.py +++ b/packages/kumo-databricks/src/kumo_databricks/volume.py @@ -5,6 +5,7 @@ from __future__ import annotations +import inspect import io import posixpath import queue @@ -22,6 +23,7 @@ _FILES = '/api/2.0/fs/files' _BUFFER = 8 << 20 +_UPLOAD_BOUNDS = {'part_size': 16 << 20, 'parallelism': 4} def is_volume_path(path: str) -> bool: @@ -30,6 +32,13 @@ def is_volume_path(path: str) -> bool: and all(parts[1:]) and posixpath.normpath(path) == path.rstrip('/')) +def _upload(client: WorkspaceClient, path: str, source: BinaryIO) -> None: + upload = client.files.upload + bounds = (_UPLOAD_BOUNDS if 'parallelism' in inspect.signature( + upload).parameters else {}) + upload(path, source, overwrite=True, **bounds) + + def volume_filesystem(client: WorkspaceClient) -> FileSystem: return PyFileSystem(VolumeHandler(client)) @@ -83,7 +92,7 @@ def _list(self, directory: str) -> Dict[str, int]: def _put(self, local: str, path: str) -> None: with _translated(path), open(local, 'rb') as contents: - self._client.files.upload(path, contents, overwrite=True) + _upload(self._client, path, contents) def _stream(self, path: str) -> BinaryIO: return _Upload(self._client, path) @@ -189,7 +198,7 @@ def __init__(self, client: WorkspaceClient, path: str) -> None: def upload() -> None: try: with _translated(path): - client.files.upload(path, source, overwrite=True) + _upload(client, path, source) except BaseException as error: self._error = error diff --git a/packages/kumo-databricks/tests/test_volume.py b/packages/kumo-databricks/tests/test_volume.py index 266403c43..fde116ede 100644 --- a/packages/kumo-databricks/tests/test_volume.py +++ b/packages/kumo-databricks/tests/test_volume.py @@ -72,6 +72,23 @@ def test_unsized_writes_stay_plain_through_a_part(): assert f'{ROOT}/over.csv' not in volume.objects +def test_uploads_bound_the_sdk_part_buffers(): + fs, volume = filesystem(part_bytes=1000) + calls = [] + + def upload(path, contents, *, overwrite, part_size=None, + use_parallel=True, parallelism=None): + calls.append((part_size, parallelism)) + volume.objects[path] = contents.read() + + volume.upload = upload + write(fs, f'{ROOT}/export.csv', BODY[:1000]) + with open_sized(fs, f'{ROOT}/sized.csv', 1000) as writer: + writer.write(BODY[:1000]) + assert calls == [(16 << 20, 4)] * 2 + assert volume.objects[f'{ROOT}/export.csv'] == BODY[:1000] + + def test_a_failed_upload_raises_without_the_signed_url(): fs, volume = filesystem() signed = 'https://bucket.example/part?X-Amz-Signature=secret'