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 4bedfe32e..b5e0341fb 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