Hold MLOps artifacts on Unity Catalog volumes through the Files API

packages/kumo-databricks/README.md

@@ -67,6 +67,35 @@ events = Table.from_source_table(
 snapshot = Graph(tables=[events]).snapshot()
 ```
 
+## Artifacts on a volume
+
+A `/Volumes/<catalog>/<schema>/<volume>/...` 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 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 `<name>.chunks/<generation>/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
 
 The sink publishes Kumo bundles directly to Unity Catalog using the connector's

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/<catalog>/<schema>/<volume>/... '
+                         'path.')
     if location is None:
         raise ValueError('Pass a ModelRegistryLocation naming the model.')
     config = _target_config(connector)

packages/kumo-databricks/src/kumo_databricks/sink.py

@@ -26,64 +26,18 @@
 
 import pandas as pd
 from databricks.sql.client import Connection
+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
 from kumo_databricks.write import RunVolume
 from kumoapi.data_source import DataSourceType
 from kumoapi.jobs import WriteMode
 from kumoapi.plugin import SourceConnector
+from pyarrow.fs import FileSystem
 
 
-class DatabricksSinkBackend(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,
-    ``<catalog>.<schema>._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,
@@ -153,7 +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):
-        """Pass selected Parquet directly to the existing hosted COPY owner."""
+        """COPY selected Parquet into a table."""
         del project_id, export_id
         from functools import partial
 
@@ -187,6 +141,30 @@ def write_files(self, connector, table_name, files,
                     prepare_file=prepare)
         return sql.qualified_name(config.catalog, schema, table)
 
+
+class DatabricksSinkBackend(FileStoreRouting, _DatabricksTables):
+    """Writes artifacts to Delta tables, staged through the schema's volume.
+
+    Args:
+        identity: Fallback identity for connectors that name none.
+    """
+
+    source_type = DataSourceType.DATABRICKS
+
+    _default_hint = ('LocalBackend(sinks={DataSourceType.DATABRICKS: '
+                     'DatabricksSinkBackend(identity=identity)})')
+
+    def holds_files(self, connector: SourceConnector, path: str) -> bool:
+        return is_volume_path(path)
+
+    def filesystem(self, connector: SourceConnector,
+                   path: str) -> tuple[FileSystem, str]:
+        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/<catalog>/<schema>/<volume>/.")
+        return volume_filesystem(self._client(connector)), path.rstrip('/')
+
     def __repr__(self) -> str:
         return f'DatabricksSinkBackend(identity={self._identity!r})'
 

packages/kumo-databricks/src/kumo_databricks/volume.py

@@ -0,0 +1,237 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES.
+# All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Unity Catalog volumes as a pyarrow filesystem."""
+
+from __future__ import annotations
+
+import io
+import posixpath
+import queue
+import threading
+from contextlib import contextmanager
+from typing import Any, 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 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'
+_BUFFER = 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]:
+    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):
+    part_bytes = 4 << 30  # under the Files API's 5 GiB upload cap
+    plain_bytes = part_bytes
+
+    def __init__(self, client: WorkspaceClient) -> None:
+        self._client = client
+
+    def get_type_name(self) -> str:
+        return 'databricks-volume'
+
+    def _size(self, path: str) -> Optional[int]:
+        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]:
+        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:
+        with _translated(path):
+            if tree:
+                clear_volume(self._client, path)
+                return
+            try:
+                self._client.files.delete(path)
+            except NotFound:
+                pass
+
+
+class _Part:
+    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)
+        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)
+
+    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(_BUFFER)
+        self._at = offset
+
+    def close(self) -> None:
+        if self._stream is not None:
+            self._stream.close()
+            self._stream = None
+
+
+_CANCEL = object()
+
+
+class _Feed(io.RawIOBase):
+    def __init__(self) -> None:
+        self.pieces: queue.Queue[Any] = 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 _CANCEL:
+                raise OSError('upload cancelled')
+            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):
+    def __init__(self, client: WorkspaceClient, path: str) -> None:
+        self._path = path
+        self._feed = _Feed()
+        self._error: Optional[BaseException] = None
+        source = io.BufferedReader(self._feed, _BUFFER)
+
+        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)
+        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:
+        while self._thread.is_alive():
+            try:
+                self._feed.pieces.put(piece, timeout=1)
+                return True
+            except queue.Full:
+                continue
+        return False
+
+    def _finish(self) -> None:
+        self._thread.join()
+        if self._error is not None:
+            raise self._error
+
+    def abort(self) -> None:
+        if self.closed:
+            return
+        self._offer(_CANCEL)
+        self._thread.join()
+        super().close()
+
+    def close(self) -> None:
+        if self.closed:
+            return
+        try:
+            self._offer(None)
+            self._finish()
+        finally:
+            super().close()