Hold MLOps artifacts on Unity Catalog volumes through the Files API

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/<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 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
+`<name>.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

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

@@ -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/<catalog>/<schema>/<volume>/...``; 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/<catalog>/<schema>/<volume>/.")
+        return volume_filesystem(self._client(connector)), path.rstrip('/')
+
     def __repr__(self) -> str:
         return f'DatabricksSinkBackend(identity={self._identity!r})'