From 014431e155ae7f39fd705a2e540c3165c240d13a Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Tue, 27 Oct 2026 09:50:58 +0000 Subject: [PATCH 1/7] Hold MLOps artifacts in Snowflake internal stages SnowflakeSinkBackend now serves named internal stages (@DB.SCHEMA.STAGE/ prefix) as a FileStoreSink: a chunked pyarrow filesystem over PUT, LIST, REMOVE, GET and presigned-URL range reads, so snapshots, training tables, models, prediction exports and the MLOps artifact root can live on a stage. Live run on a Snowflake account fixed two things the in-memory stage missed. REMOVE and GET match PATTERN against the stage's internal storage path, two segments longer than LIST shows, so the exact-file pattern anchors only the stage-relative suffix (.*/); the fake stage now matches against such a path. A SNOWFLAKE_FULL stage lists padded ciphertext sizes, so its files cannot be read back exactly: the one-time probe detects that from the probe's listed size and reads raise, naming SNOWFLAKE_SSE, instead of failing a size check. Verified live: writes, exact deletes with LIST checks, ranged reads on an SSE stage, the GET fallback on SSE, the FULL refusal, and no warehouse needed for PUT/LIST/GET/REMOVE/GET_PRESIGNED_URL; an MLOps fit, predict and model load with the artifact root on a stage. Unit tests: the kumo-snowflake non-live suite passes. Co-Authored-By: Claude Opus 5.5 --- packages/kumo-snowflake/README.md | 35 +++ .../src/kumo_snowflake/artifacts.py | 2 +- .../kumo-snowflake/src/kumo_snowflake/sink.py | 69 ++++- .../src/kumo_snowflake/stage_fs.py | 256 +++++++++++++++++ packages/kumo-snowflake/tests/fakes.py | 150 ++++++++++ .../kumo-snowflake/tests/test_stage_fs.py | 270 ++++++++++++++++++ tests/public/Dockerfile.dockerignore | 4 + 7 files changed, 783 insertions(+), 3 deletions(-) create mode 100644 packages/kumo-snowflake/src/kumo_snowflake/stage_fs.py create mode 100644 packages/kumo-snowflake/tests/test_stage_fs.py diff --git a/packages/kumo-snowflake/README.md b/packages/kumo-snowflake/README.md index 8c47d90d5..58b3d7e8c 100644 --- a/packages/kumo-snowflake/README.md +++ b/packages/kumo-snowflake/README.md @@ -7,6 +7,41 @@ The reader implements `kumo_core.sources.protocol.SourceReader` in resulting files over presigned URLs. A temporary stage needs no grant beyond the `USAGE` and `SELECT` a read-only analyst already holds. +`SnowflakeSinkBackend` also holds files in a named internal stage +(`@DB.SCHEMA.STAGE/prefix`). Output configs take stage paths, such as +`ModelOutputConfig(snowflake_connector, path="@ML.ARTIFACTS.KUMO/models/m1")` +or `PredictionOutputConfig(snowflake_connector, +output_paths={"predictions": "@ML.ARTIFACTS.KUMO/scores/"})`, and +`Trainer.load_from_connector` reads a model back from one. For an MLOps +artifact root on a stage, see `packages/kumo-engine/docs/mlops.md`. File names may +use only `A-Z a-z 0-9 _ $ = + . , / -`. + +A file of up to 4 GiB is one staged object. A larger one is stored as numbered +parts under `.chunks/`, which only kumo reads back. To load exported +prediction Parquet into a table, skip the marker and manifest files: + +```sql +COPY INTO SCORES FROM @ML.ARTIFACTS.KUMO/scores/ + PATTERN = '.*[.]parquet' FILE_FORMAT = (TYPE = PARQUET) + MATCH_BY_COLUMN_NAME = CASE_INSENSITIVE; +``` + +The role needs `USAGE` on the database and schema and `READ` and `WRITE` on the +stage. No warehouse runs: `PUT`, `LIST`, `GET`, `REMOVE` and +`GET_PRESIGNED_URL` complete while the session's warehouse stays suspended. +Create the stage with server-side encryption so that reads fetch byte ranges +over presigned URLs: + +```sql +CREATE STAGE ML.ARTIFACTS.KUMO ENCRYPTION = (TYPE = 'SNOWFLAKE_SSE'); +``` + +A stage with the default `SNOWFLAKE_FULL` (client-side) encryption is written +to, but reading from it raises: its `LIST` reports encrypted sizes (padded to +16 bytes), so files cannot be read back exactly. With a role that cannot write +the probe file, each file is downloaded once with `GET` into the spool +directory instead, and a warning says so. + From the repository root, run the public tests without cloud credentials: ```sh diff --git a/packages/kumo-snowflake/src/kumo_snowflake/artifacts.py b/packages/kumo-snowflake/src/kumo_snowflake/artifacts.py index 0ae19f1a7..cd374c1dd 100644 --- a/packages/kumo-snowflake/src/kumo_snowflake/artifacts.py +++ b/packages/kumo-snowflake/src/kumo_snowflake/artifacts.py @@ -29,7 +29,6 @@ from pathlib import Path from typing import Callable, Dict, List, Optional, Sequence, Tuple, TypeVar -from kumo_core.diskgraph.snapshot import rebased # Aliased: three methods below take a `SnapshotStore` parameter named # `store`, and one name meaning two stores inside one class is how a line # reaching for the identity store quietly gets the snapshot store instead. @@ -550,6 +549,7 @@ def restore(self, store: SnapshotStore[_Handle], copy_tree=files.copy_tree) if restored is None: return None + from kumo_core.diskgraph.snapshot import rebased return rebased(store, handle) def checkpoint(self, store: SnapshotStore[_Handle], diff --git a/packages/kumo-snowflake/src/kumo_snowflake/sink.py b/packages/kumo-snowflake/src/kumo_snowflake/sink.py index 330f52a72..2e89209ff 100644 --- a/packages/kumo-snowflake/src/kumo_snowflake/sink.py +++ b/packages/kumo-snowflake/src/kumo_snowflake/sink.py @@ -13,6 +13,8 @@ import numpy as np import pandas as pd +from kumo_core.object_store_sink import FilesystemSink +from kumo_snowflake.artifacts import _split_stage_location from kumo_snowflake.errors import is_lapsed_credential from kumo_snowflake.registry import SnowflakeModelRegistryMixin from kumo_snowflake.session import ( @@ -133,6 +135,10 @@ def write_table( TypeError: If the registered factory returned something that is not a Snowpark session. """ + if self.holds_files(connector, table_name): + return FilesystemSink(self).write_table(connector, table_name, + frame, write_mode, + export_id=export_id) # `export_id` identifies a partial result for a destination that can # hold one. A warehouse write is a single statement, so there is # nothing here for a rerun to recognise. @@ -161,8 +167,14 @@ def write_table( def write_files(self, connector, table_name, files, write_mode=WriteMode.OVERWRITE, project_id=None, *, - export_id=None, extra_columns=None): - """Load selected Parquet through the same upload/COPY owner as hosted.""" + export_id=None, extra_columns=None, **options): + """Copy Parquet to a stage 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 kumo_core.artifact_staging import selected_local_parquet_files from kumo_snowflake.write import export_table_files @@ -192,6 +204,59 @@ def export(*, retry: bool = False) -> None: export(retry=True) return self._qualified(connector, table_name) + def preflight_write_files(self, connector, table_name, write_mode=None): + """Refuse an occupied stage 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, path): + """Whether ``path`` names a named internal stage location.""" + try: + _split_stage_location(path) + except (TypeError, ValueError): + return False + return True + + def filesystem(self, connector, path): + """A pyarrow filesystem for an ``@DB.SCHEMA.STAGE/prefix`` path.""" + from kumo_snowflake.stage_fs import stage_filesystem + with self._session(connector) as session: + _require_snowpark(session) + return stage_filesystem(session, path) + + def publish_model(self, connector, bundle_dir, *, path=None, location=None, + metadata={}, metrics={}, project_id=None): + """Registers a model version, or uploads to a stage ``path``.""" + if location is None and path is not None 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, local_dir, *, path=None, location=None, + project_id=None): + """Downloads a registry version, or the bundle at a stage ``path``.""" + if location is None and path is not None 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 __repr__(self) -> str: return f'SnowflakeSinkBackend(identity={self._identity!r})' diff --git a/packages/kumo-snowflake/src/kumo_snowflake/stage_fs.py b/packages/kumo-snowflake/src/kumo_snowflake/stage_fs.py new file mode 100644 index 000000000..f859a1241 --- /dev/null +++ b/packages/kumo-snowflake/src/kumo_snowflake/stage_fs.py @@ -0,0 +1,256 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +r"""A pyarrow filesystem over a Snowflake named internal stage. + +Writes are ``PUT``, listings ``LIST``, deletes ``REMOVE``. Reads fetch byte +ranges from ``GET_PRESIGNED_URL`` links when the stage serves plaintext +(``ENCRYPTION = (TYPE = 'SNOWFLAKE_SSE')``); otherwise each file is +downloaded once with ``GET`` and read locally. +""" + +import logging +import posixpath +import shutil +import tempfile +import threading +import uuid +import weakref +from pathlib import Path +from typing import Dict, Optional, Tuple + +from kumo_core.chunked_fs import ChunkedFileHandler, spool_root +from kumo_core.sources.diagnostics import safe_summary +from kumo_snowflake.artifacts import ( + _split_stage_location, + _stage_join, + _StageFiles, +) +from kumo_snowflake.session import SnowflakeSession, _rows +from kumo_snowflake.staging import ( + DOWNLOAD_TIMEOUT, + _download_session, + presigned_urls, +) +from pyarrow.fs import PyFileSystem + +logger = logging.getLogger(__name__) + +_PROBE = b'kumo-stage-range-probe\n' +_ranged_stages: Dict[Tuple[Optional[str], str], str] = {} +_RANGED, _GET, _ENCRYPTED = 'ranged', 'get', 'encrypted' +_ranged_lock = threading.Lock() + + +def _pattern(relative: str) -> str: + # REMOVE and GET match PATTERN against the stage's internal storage path, + # whose prefix LIST does not show, so the full stage-relative path is the + # anchored suffix. + escaped = ''.join(f'[{c}]' if c in '.$+' else c for c in relative) + return f'.*/{escaped}' + + +class StageFileSystemHandler(ChunkedFileHandler): + # PUT has no small size cap, so a file stays one object (readable by + # COPY INTO) up to the same bound as a Databricks volume. + part_bytes = 4 << 30 + + def __init__(self, session: SnowflakeSession, root: str) -> None: + self._stage = _split_stage_location(root)[0] + self._root = root + self._session = session + self._files = _StageFiles(session, parallel=1) + self._urls: Dict[str, str] = {} + self._http = threading.local() + self._downloads: Dict[str, Tuple[threading.Lock, list]] = {} + self._downloads_lock = threading.Lock() + self._download_dir: Optional[str] = None + + def get_type_name(self): + return 'snowflake-stage' + + def _relative(self, path: str) -> str: + return path[len(self._stage) + 1:] + + def _size(self, path: str) -> Optional[int]: + found = self._files._list(path, tree=False).get('') + return None if found is None else found.size + + def _list(self, directory: str) -> Dict[str, int]: + return { + f'{directory}/{relative}': obj.size + for relative, obj in self._files._list(directory).items() + } + + def _put(self, local: str, path: str) -> None: + _split_stage_location(path) + self._files._put(Path(local), posixpath.dirname(path)) + + def _remove(self, path: str, *, tree: bool) -> None: + if tree: + self._files.remove(path) + return + parent = posixpath.dirname(path) + _rows( + self._session, f'REMOVE {_stage_join(parent, directory=True)} ' + f"PATTERN = '{_pattern(self._relative(path))}'") + if self._size(path) is not None: + raise OSError(f'REMOVE left {path} in place.') + + def _forget(self, path: str) -> None: + super()._forget(path) + path = path.rstrip('/') + with self._downloads_lock: + for name in list(self._downloads): + if name == path or name.startswith(path): + del self._downloads[name] + + def _part(self, path: str, size: int): + mode = self._read_mode() + if mode == _ENCRYPTED: + raise OSError( + f'Snowflake stage {self._stage} encrypts files client-side ' + "(SNOWFLAKE_FULL), so LIST reports encrypted sizes and kumo " + 'cannot read its files back. Use a stage created with ' + "ENCRYPTION = (TYPE = 'SNOWFLAKE_SSE').") + if mode == _RANGED: + return _RangedPart(self, path) + return _DownloadedPart(self, path, size) + + def _read_mode(self) -> str: + key = (getattr(self._session.connection, 'account', None), self._stage) + with _ranged_lock: + if key not in _ranged_stages: + _ranged_stages[key] = self._probe() + return _ranged_stages[key] + + def _probe(self) -> str: + root = _stage_join(self._stage, '.kumo-preflight', uuid.uuid4().hex) + try: + with tempfile.TemporaryDirectory() as tmp: + probe = Path(tmp) / 'probe' + probe.write_bytes(_PROBE) + self._files._put(probe, root) + except Exception as error: + logger.warning( + 'Could not verify ranged reads on Snowflake stage %s (%s; ' + 'the role may lack WRITE), so each file is downloaded whole.', + self._stage, safe_summary(error)) + return _GET + try: + if self._size(f'{root}/probe') != len(_PROBE): + return _ENCRYPTED + ranged = self._get(f'{root}/probe', 1, 4) == _PROBE[1:5] + except OSError: + ranged = False + finally: + try: + self._files.remove(root) + except Exception as error: + logger.warning('Could not remove the range probe: %s', + safe_summary(error)) + if not ranged: + logger.warning( + 'Snowflake stage %s does not serve byte ranges of plaintext ' + 'over presigned URLs, so each file is downloaded whole. ' + "Create artifact stages with ENCRYPTION = (TYPE = " + "'SNOWFLAKE_SSE') for ranged reads.", self._stage) + return _RANGED if ranged else _GET + + def _url(self, path: str) -> str: + url = self._urls.get(path) + if url is None: + relative = self._relative(path) + cursor = self._session.connection.cursor() + try: + url = presigned_urls(cursor, self._stage[1:], + [relative])[relative] + finally: + cursor.close() + self._urls[path] = url + return url + + def _get(self, path: str, offset: int, length: int) -> bytes: + import requests + if length == 0: + return b'' + http = getattr(self._http, 'session', None) + if http is None: + http = self._http.session = _download_session() + try: + response = http.get( + self._url(path), timeout=DOWNLOAD_TIMEOUT, + headers={'Range': f'bytes={offset}-{offset + length - 1}'}) + response.raise_for_status() + except requests.RequestException as error: + self._urls.pop(path, None) + status = getattr(getattr(error, 'response', None), 'status_code', + None) + raise OSError(f'Reading {path} failed: ' + f'{safe_summary(error, status=status)}') from None + if (response.status_code != 206 or not response.headers.get( + 'Content-Range', '').startswith(f'bytes {offset}-')): + raise OSError(f'Reading {path} returned no byte range ' + f'(status={response.status_code}).') + return response.content + + def _downloaded(self, path: str, size: int) -> Path: + with self._downloads_lock: + if self._download_dir is None: + self._download_dir = tempfile.mkdtemp(prefix='kumo-stage-get-', + dir=spool_root(self)) + weakref.finalize(self, shutil.rmtree, self._download_dir, True) + lock, local = self._downloads.setdefault(path, + (threading.Lock(), [])) + with lock: + if not local: + target = Path(self._download_dir) / uuid.uuid4().hex / ( + posixpath.basename(path)) + target.parent.mkdir() + parent = posixpath.dirname(path) + self._session.file.get(_stage_join(parent, directory=True), + str(target.parent), parallel=1, + pattern=_pattern(self._relative(path))) + if not target.is_file() or target.stat().st_size != size: + shutil.rmtree(target.parent, ignore_errors=True) + raise OSError(f'Downloading {path} did not produce its ' + f'{size} bytes.') + local.append(target) + return local[0] + + +class _RangedPart: + def __init__(self, handler: StageFileSystemHandler, path: str) -> None: + self._handler = handler + self._path = path + + def read_at(self, offset: int, length: int) -> bytes: + return self._handler._get(self._path, offset, length) + + def close(self) -> None: + pass + + +class _DownloadedPart: + def __init__(self, handler: StageFileSystemHandler, path: str, + size: int) -> None: + self._handler = handler + self._path = path + self._size = size + + def read_at(self, offset: int, length: int) -> bytes: + with open(self._handler._downloaded(self._path, self._size), + 'rb') as reader: + reader.seek(offset) + return reader.read(length) + + def close(self) -> None: + pass + + +def stage_filesystem(session: SnowflakeSession, + location: str) -> Tuple[PyFileSystem, str]: + r"""A filesystem for ``@DB.SCHEMA.STAGE/prefix`` and its normalized path.""" + stage, prefix, _ = _split_stage_location(location) + root = f'{stage}/{prefix}' if prefix else stage + return PyFileSystem(StageFileSystemHandler(session, root)), root diff --git a/packages/kumo-snowflake/tests/fakes.py b/packages/kumo-snowflake/tests/fakes.py index ddcd53bb9..9aae929a4 100644 --- a/packages/kumo-snowflake/tests/fakes.py +++ b/packages/kumo-snowflake/tests/fakes.py @@ -102,3 +102,153 @@ def __call__(self) -> FakeSession: @property def released(self) -> bool: return all(session.closed for session in self.opened) + + +class FakeStage: + """An in-memory named internal stage behind a Snowpark-shaped session. + + LIST and REMOVE match by string prefix, as Snowflake does, and report + names under the lower-cased stage name. ``ciphertext`` makes LIST report + padded sizes and presigned fetches return encrypted bytes, as a + ``SNOWFLAKE_FULL`` stage does. + """ + def __init__(self, stage: str = '@DB.SCHEMA.KUMO_ARTIFACTS') -> None: + self.stage = stage + self.objects: dict[str, bytes] = {} + self.statements: list[str] = [] + self.puts: list[dict[str, Any]] = [] + self.ranges: list[str] = [] + self.ciphertext = False + self.ignores_range = False + self.remove_pattern_ignored = False + self.gets: list[str] = [] + + def _relative(self, location: str) -> str: + location = location.strip("'") + assert location.startswith(self.stage), location + return location[len(self.stage):].lstrip('/') + + def _listed(self, relative: str) -> str: + return f"{self.stage.rsplit('.', 1)[1].lower()}/{relative}" + + def _matched(self, relative: str) -> str: + # What Snowflake matches REMOVE/GET PATTERN against: the listed name + # behind further internal path segments. + return f'internal/{self._listed(relative)}' + + def session(self) -> Any: + return _stage_session(self) + + def sql(self, sql: str) -> list[tuple]: + import re + self.statements.append(sql) + verb, _, rest = sql.partition(' ') + if verb == 'LIST': + prefix = self._relative(rest) + # A client-side encrypted stage lists padded ciphertext sizes. + return [(self._listed(name), (len(data) // 16 + 1) * + 16 if self.ciphertext else len(data), 'md5', 'date') + for name, data in sorted(self.objects.items()) + if name.startswith(prefix)] + if verb == 'REMOVE': + location, _, pattern = rest.partition(' PATTERN = ') + prefix = self._relative(location) + for name in [n for n in self.objects if n.startswith(prefix)]: + if not pattern or (not self.remove_pattern_ignored + and re.fullmatch(pattern.strip("'"), + self._matched(name))): + del self.objects[name] + return [] + raise AssertionError(f'unexpected statement {sql!r}') + + def put(self, local: str, target: str, **options: Any) -> list[str]: + import os + self.puts.append(options) + name = self._relative(target).rstrip('/') + key = f'{name}/{os.path.basename(local)}' if name else \ + os.path.basename(local) + with open(local, 'rb') as reader: + self.objects[key] = reader.read() + return ['UPLOADED'] + + def get(self, source: str, directory: str, *, pattern: str, + **options: Any) -> list[str]: + import os + import re + prefix = self._relative(source) + got = [] + for name, data in self.objects.items(): + if name.startswith(prefix) and re.fullmatch(pattern, + self._matched(name)): + with open(os.path.join(directory, os.path.basename(name)), + 'wb') as writer: + writer.write(data) + got.append(name) + self.gets += got + return got + + def presign(self, sql: str) -> list[tuple]: + import re + assert sql.startswith('SELECT'), sql + paths = re.findall(r"\('((?:[^']|'')*)'\)", sql) or re.findall( + r", '((?:[^']|'')*)', \d+\)", sql) + return [(path, f'https://stage.invalid/{path}?signature=secret') + for path in paths] + + def http_get(self, url: str, headers: Mapping[str, str], + **kwargs: Any) -> Any: + import re + + import requests + name = url.split('https://stage.invalid/', 1)[1].split('?')[0] + response = requests.Response() + response.url = url + if name not in self.objects: + response.status_code = 404 + return response + data = self.objects[name] + if self.ciphertext: + data = bytes(b ^ 0xFF for b in data) + self.ranges.append(headers['Range']) + start, end = map( + int, + re.fullmatch(r'bytes=(\d+)-(\d+)', headers['Range']).groups()) + if self.ignores_range: + response.status_code, response._content = 200, data + else: + response.status_code = 206 + response._content = data[start:end + 1] + response.headers['Content-Range'] = ( + f'bytes {start}-{start + len(response._content) - 1}/' + f'{len(data)}') + return response + + +def _stage_session(stage: FakeStage) -> Any: + from types import SimpleNamespace + + from snowflake.snowpark import Session + + class _Cursor: + def execute(self, sql: str) -> None: + self.rows = stage.presign(sql) + + def fetchall(self) -> list[tuple]: + return self.rows + + def close(self) -> None: + pass + + class _Session(Session): + def __init__(self) -> None: + pass + + file = SimpleNamespace(put=lambda *a, **k: stage.put(*a, **k), + get=lambda *a, **k: stage.get(*a, **k)) + connection = SimpleNamespace(cursor=_Cursor) + + def sql(self, sql: str, params: Any = None) -> Any: + rows = stage.sql(sql) + return SimpleNamespace(collect=lambda **_: rows) + + return _Session() diff --git a/packages/kumo-snowflake/tests/test_stage_fs.py b/packages/kumo-snowflake/tests/test_stage_fs.py new file mode 100644 index 000000000..c217750fd --- /dev/null +++ b/packages/kumo-snowflake/tests/test_stage_fs.py @@ -0,0 +1,270 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""A Snowflake stage seen through pyarrow, against an in-memory stage.""" + +import gc +import os +from contextlib import contextmanager +from types import SimpleNamespace + +import pytest +from fakes import FakeStage +from kumo_snowflake import stage_fs +from kumo_snowflake.sink import SnowflakeSinkBackend +from kumoapi.plugin import FileStoreSink +from pyarrow.fs import FileSelector, FileType + +ROOT = '@DB.SCHEMA.KUMO_ARTIFACTS/runs' +PART = 32 + + +@pytest.fixture +def stage(monkeypatch, tmp_path): + monkeypatch.setattr('kumo_snowflake.staging.SPCS_TOKEN_PATH', + str(tmp_path / 'absent')) + monkeypatch.setattr(stage_fs, '_ranged_stages', {}) + monkeypatch.setattr(stage_fs.StageFileSystemHandler, 'part_bytes', PART) + fake = FakeStage() + monkeypatch.setattr(stage_fs, '_download_session', + lambda: SimpleNamespace(get=fake.http_get)) + return fake + + +@pytest.fixture +def sink(stage, monkeypatch): + session = stage.session() + sink = SnowflakeSinkBackend() + + @contextmanager + def connected(connector): + yield session + + monkeypatch.setattr(sink, '_session', connected) + return sink + + +@pytest.fixture +def filesystem(sink): + return sink.filesystem(object(), ROOT + '/') + + +def _write(fs, path, data): + with fs.open_output_stream(path) as writer: + writer.write(data) + + +def _read(fs, path): + with fs.open_input_stream(path) as reader: + return reader.readall() + + +def test_the_sink_resolves_a_stage_path(filesystem, stage): + fs, root = filesystem + assert root == ROOT + _write(fs, f'{root}/model/weights.bin', b'weights') + assert stage.objects == {'runs/model/weights.bin': b'weights'} + assert stage.puts == [{ + 'parallel': 1, + 'auto_compress': False, + 'source_compression': 'NONE', + 'overwrite': True + }] + assert _read(fs, f'{root}/model/weights.bin') == b'weights' + + +@pytest.mark.parametrize('path,held', [(ROOT, True), (ROOT + '/m/1', True), + ('USER_SCORES', False), + ('@~/runs', False), + ('@%USERS/runs', False)]) +def test_only_named_stage_paths_are_file_locations(path, held): + sink = SnowflakeSinkBackend() + assert isinstance(sink, FileStoreSink) + assert sink.holds_files(object(), path) is held + + +def test_a_bundle_round_trips_through_a_stage_path(sink, stage, tmp_path): + bundle = tmp_path / 'bundle' + (bundle / 'model').mkdir(parents=True) + (bundle / 'model' / 'weights.pt').write_bytes(os.urandom(2 * PART + 3)) + (bundle / 'bundle.json').write_text('{}') + published = sink.publish_model(object(), str(bundle), + path=f'{ROOT}/models/b1') + assert published.uri == f'{ROOT}/models/b1' + fetched = sink.fetch_model(object(), str(tmp_path / 'fetched'), + path=f'{ROOT}/models/b1') + assert fetched.version == f'{ROOT}/models/b1' + for name in ('model/weights.pt', 'bundle.json'): + assert (tmp_path / 'fetched' / + name).read_bytes() == (bundle / name).read_bytes() + with pytest.raises(ValueError, match='Nothing is published'): + sink.fetch_model(object(), str(tmp_path / 'x'), + path=f'{ROOT}/models/none') + + +def test_parquet_files_are_written_to_a_stage_path( + sink, stage, tmp_path): + import pandas as pd + parts = [] + for index in range(2): + parts.append(tmp_path / f'part-{index}.parquet') + pd.DataFrame({'ID': [index]}).to_parquet(parts[-1]) + written = sink.write_files(object(), f'{ROOT}/predictions/r1/', + [str(path) for path in parts], export_id='r1') + assert written == f'{ROOT}/predictions/r1/' + fs, root = sink.filesystem(object(), f'{ROOT}/predictions/r1') + names = [i.base_name for i in fs.get_file_info(FileSelector(root))] + assert names == [ + '_KUMO_EXPORT', '_SUCCESS', 'part-0.parquet', 'part-1.parquet' + ] + assert sink.write_table(object(), f'{ROOT}/tables/t1.parquet', + pd.DataFrame({'ID': [1] + })) == (f'{ROOT}/tables/t1.parquet') + + +def test_a_user_stage_is_refused(filesystem): + with pytest.raises(ValueError, match='fully-qualified'): + stage_fs.stage_filesystem(None, '@~/runs') + + +def test_listing_reports_sizes_and_logical_files(filesystem, stage): + fs, root = filesystem + _write(fs, f'{root}/a/small.json', b'{}') + _write(fs, f'{root}/a/big.bin', os.urandom(2 * PART + 1)) + infos = fs.get_file_info(FileSelector(f'{root}/a', recursive=True)) + assert [(i.base_name, i.size) for i in infos] == [('big.bin', 2 * PART + 1), + ('small.json', 2)] + assert fs.get_file_info(f'{root}/a').type is FileType.Directory + assert fs.get_file_info(f'{root}/a/small').type is FileType.NotFound + + +def test_removing_a_file_spares_names_it_prefixes(filesystem, stage): + fs, root = filesystem + for name in ('part-1', 'part-10', 'part-1.x', 'd/part-1'): + _write(fs, f'{root}/{name}', name.encode()) + fs.delete_file(f'{root}/part-1') + assert sorted(stage.objects) == [ + 'runs/d/part-1', 'runs/part-1.x', 'runs/part-10' + ] + fs.delete_dir_contents(root) + assert stage.objects == {} + + +def test_a_chunked_file_reads_exact_ranges_over_presigned_urls( + filesystem, stage): + fs, root = filesystem + data = os.urandom(3 * PART) + _write(fs, f'{root}/index.bin', data) + assert sorted(stage.objects) == [ + f'runs/index.bin.chunks/{name}' + for name in ('00000', '00001', '00002', '_SUCCESS') + ] + with fs.open_input_file(f'{root}/index.bin') as reader: + assert reader.read_at(PART, PART // 2) == data[PART // 2:PART * 3 // 2] + assert _read(fs, f'{root}/index.bin') == data + assert len(stage.ranges) > 1 and not stage.gets + assert any( + s.startswith('REMOVE @DB.SCHEMA.KUMO_ARTIFACTS/.kumo-preflight/') + for s in stage.statements) + assert not any('.kumo-preflight' in name for name in stage.objects) + + +def test_a_stage_without_ranged_plaintext_reads_each_file_once_by_get( + sink, stage, caplog, tmp_path, monkeypatch): + spool = tmp_path / 'spool' + monkeypatch.setattr(stage_fs.StageFileSystemHandler, 'spool_dir', + str(spool)) + fs, root = sink.filesystem(object(), ROOT) + stage.ignores_range = True + data = os.urandom(2 * PART) + _write(fs, f'{root}/a.bin', data) + _write(fs, f'{root}/a.bin2', b'sibling') + for _ in range(2): + with fs.open_input_file(f'{root}/a.bin') as reader: + assert reader.read_at(PART, + PART // 2) == data[PART // 2:PART * 3 // 2] + assert _read(fs, f'{root}/a.bin') == data + assert sorted( + stage.gets) == ['runs/a.bin.chunks/00000', 'runs/a.bin.chunks/00001'] + assert [ + r.levelname for r in caplog.records + if 'SNOWFLAKE_SSE' in r.getMessage() + ] == ['WARNING'] + del fs, reader + gc.collect() + assert not list(spool.iterdir()) + + +def test_a_client_side_encrypted_stage_is_refused_on_read(filesystem, stage): + fs, root = filesystem + stage.ciphertext = True + _write(fs, f'{root}/a.bin', b'0123456789') + with pytest.raises(OSError, match="ENCRYPTION = \\(TYPE = 'SNOWFLAKE_SSE'\\)"): + _read(fs, f'{root}/a.bin') + assert stage.gets == [] + + +def test_a_role_that_cannot_write_reads_by_get_without_blaming_encryption( + filesystem, stage, caplog, monkeypatch): + fs, root = filesystem + _write(fs, f'{root}/a.bin', b'0123456789') + + def refused(*args, **kwargs): + raise PermissionError('insufficient privileges') + + monkeypatch.setattr(stage, 'put', refused) + with fs.open_input_file(f'{root}/a.bin') as reader: + assert reader.read_at(4, 3) == b'3456' + messages = [r.getMessage() for r in caplog.records] + assert any('lack WRITE' in m for m in messages) + assert not any('SNOWFLAKE_SSE' in m for m in messages) + + +def test_a_remove_that_matches_nothing_raises(filesystem, stage): + fs, root = filesystem + _write(fs, f'{root}/a.bin', b'old') + stage.remove_pattern_ignored = True + with pytest.raises(OSError, match='left'): + _write(fs, f'{root}/a.bin', os.urandom(2 * PART)) + + +def test_a_name_outside_the_stage_character_set_fails_at_write(filesystem): + fs, root = filesystem + with pytest.raises(ValueError): + _write(fs, f'{root}/my table.json', b'{}') + + +def test_a_failed_read_names_no_signed_url(filesystem, stage): + fs, root = filesystem + _write(fs, f'{root}/a.bin', b'data') + reader = fs.open_input_file(f'{root}/a.bin') + del stage.objects['runs/a.bin'] + with pytest.raises(OSError) as error: + reader.read_at(2, 0) + assert 'signature' not in str(error.value) + assert 'status=404' in str(error.value) + + +def test_a_registry_path_is_not_sent_to_the_stage(sink): + with pytest.raises(ValueError, match='ModelRegistryLocation'): + sink.publish_model(object(), '/tmp/none', path='models/m1') + + +def test_an_occupied_stage_directory_is_refused_before_scoring( + sink, stage, tmp_path): + import pandas as pd + from kumo_core.artifact_export import ArtifactDestinationConflict + from kumoapi.jobs import WriteMode + + target = f'{ROOT}/predictions/r1/' + sink.preflight_write_files(object(), target) + part = tmp_path / 'part-0.parquet' + pd.DataFrame({'ID': [1]}).to_parquet(part) + sink.write_files(object(), target, [str(part)], export_id='r1') + + with pytest.raises(ArtifactDestinationConflict, match='/predictions/r1/'): + sink.preflight_write_files(object(), target) + sink.preflight_write_files(object(), target, + write_mode=WriteMode.OVERWRITE) + sink.preflight_write_files(object(), f'{ROOT}/predictions/r2/') + sink.preflight_write_files(object(), 'SCORES') diff --git a/tests/public/Dockerfile.dockerignore b/tests/public/Dockerfile.dockerignore index ddae86d2c..6a91b5f4d 100644 --- a/tests/public/Dockerfile.dockerignore +++ b/tests/public/Dockerfile.dockerignore @@ -16,6 +16,10 @@ !packages/kumo-core/src/kumo_core/artifact_export.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 60d1efd4c51441cbd093347eedd3db0c2c4f0533 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Tue, 27 Oct 2026 19:13:16 +0000 Subject: [PATCH 2/7] Score Snowflake Model Registry versions with run_batch A Kumo model registered in the Snowflake Model Registry was a carrier whose predict raised, so Snowflake ML batch inference jobs could not run it, and re-registering the same bundle logged a duplicate version. kumo-engine now registers, under the kumo_snowflake.registry_models entry point, a builder that kumo_snowflake.registry uses when present: the version is a CustomModel that loads the bundle and an exported stage snapshot once per process and scores an ENTITY column in Snowpark Container Services. ModelOutputConfig gains serving_snapshot and serving_options: cuda_version (default 12.6), python_version (the registering Python), function_type (TABLE_FUNCTION) and pip_repository. Without the builder, or if it fails, the carrier is registered as before. A bundle whose kumo_bundle_id is already a Kumo version of the model returns that version. The scoring helpers shared with the MLflow pyfunc move to integrations/_serving.py. kumo's wheels are not on PyPI, so the version embeds kumo's packages (and torchmetrics, whose pin the CUDA index the image build prefers does not carry) as code_paths, with their third-party requirements in a form the image build accepts (no ==X.* pins, torch at this installation's minor version, faiss-cpu, and kumo-glove-6b-50d for text models). pip_repository makes those requirements install from a Snowflake pip artifact repository instead of PyPI. snowflake-ml pickles the model class by value, so the container must run the registering Python: python_version defaults to it, and another value raises (the version is then registered as a carrier, with a warning). Verified live on a Snowflake account, as a role holding only the documented grants (CREATE MODEL and the Model Monitor grants added to the docs, CREATE IMAGE REPOSITORY removed as unneeded): binary and regression versions registered through kumo.MLOps scored 300 entities with run_batch on a GPU_NV_S pool, matching local scores; the regression version installed its requirements from Snowflake's shared PyPI repository; Model Monitors on both refreshed with holdout_df() of the registry-loaded model as the baseline. Unit tests: snowflake-ml packager round-trip and registry fakes. Co-Authored-By: Claude Opus 5.5 --- examples/mlops/snowflake/run_batch_smoke.py | 132 ++++++++++ packages/kumo-engine/docs/mlops.md | 64 ++++- .../kumo/engine/integrations/_mlflow_model.py | 121 +-------- .../kumo/engine/integrations/_serving.py | 214 ++++++++++++++++ .../engine/integrations/_snowflake_model.py | 238 ++++++++++++++++++ .../kumo/engine/trainer/output_config.py | 26 +- .../kumo/engine/trainer/trainer.py | 7 + packages/kumo-engine/pyproject.toml | 5 + .../kumo-engine/test/test_mlflow_model.py | 2 +- .../kumo-engine/test/test_snowflake_model.py | 225 +++++++++++++++++ packages/kumo-snowflake/README.md | 9 + .../src/kumo_snowflake/registry.py | 141 ++++++++--- .../kumo-snowflake/tests/test_registry.py | 104 +++++++- 13 files changed, 1122 insertions(+), 166 deletions(-) create mode 100644 examples/mlops/snowflake/run_batch_smoke.py create mode 100644 packages/kumo-engine/kumo/engine/integrations/_serving.py create mode 100644 packages/kumo-engine/kumo/engine/integrations/_snowflake_model.py create mode 100644 packages/kumo-engine/test/test_snowflake_model.py diff --git a/examples/mlops/snowflake/run_batch_smoke.py b/examples/mlops/snowflake/run_batch_smoke.py new file mode 100644 index 000000000..87d4d169f --- /dev/null +++ b/examples/mlops/snowflake/run_batch_smoke.py @@ -0,0 +1,132 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Live smoke: fit a small model on synthetic data, register it in the +Snowflake Model Registry with a serving snapshot on a stage, score it with +``ModelVersion.run_batch`` on a GPU compute pool, and compare the job's output +with local scores. The steps were verified live; this script was not run +itself. See "Batch inference in Snowflake" in +packages/kumo-engine/docs/mlops.md for the pool, grants and stage it needs. + + python run_batch_smoke.py --connection-name \\ + --warehouse WH --database ML --schema MODELS \\ + --stage @ML.MODELS.KUMO --compute-pool KUMO_GPU_POOL +""" + +import argparse +import tempfile +from pathlib import Path + +import numpy as np +import pandas as pd +from kumoapi.graph import Edge +from kumoapi.model_plan import ModelPlan +from kumoapi.typing import Stype + +import kumo.engine as kumo + + +def fit(data: Path) -> tuple: + rng = np.random.default_rng(0) + pd.DataFrame({'ID': np.arange(200)}).to_parquet(data / 'users.parquet') + pd.DataFrame({ + 'OID': np.arange(5000), + 'USER_ID': rng.integers(0, 200, 5000), + 'TS': pd.Timestamp('2024-01-01') + + pd.to_timedelta(rng.integers(0, 120, 5000), unit='D'), + }).to_parquet(data / 'orders.parquet') + files = kumo.FileConnector(root_dir=str(data)) + graph = kumo.Graph( + tables={ + 'users': + kumo.Table.from_source_table(files['users.parquet'], + primary_key='ID').infer_metadata(), + 'orders': + kumo.Table.from_source_table( + files['orders.parquet'], primary_key='OID', + time_column='TS').infer_metadata().with_column( + 'USER_ID', stype=Stype.ID), + }, edges=[Edge(src_table='orders', fkey='USER_ID', dst_table='users')]) + snapshot = graph.snapshot() + query = kumo.PredictiveQuery( + graph=graph, query='PREDICT COUNT(orders.*, 0, 7, days) > 0 ' + 'FOR EACH users.ID') + return snapshot, query + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__.split('\n\n')[0]) + parser.add_argument('--connection-name', required=True) + parser.add_argument('--warehouse', required=True) + parser.add_argument('--database', required=True) + parser.add_argument('--schema', required=True) + parser.add_argument('--stage', required=True, + help='named internal stage, @DB.SCHEMA.STAGE') + parser.add_argument('--compute-pool', required=True) + parser.add_argument('--model-name', default='KUMO_RUN_BATCH_SMOKE') + parser.add_argument('--pip-repository', default=None) + args = parser.parse_args() + + from snowflake.ml.model.batch_inference import ( + InferenceSpec, + InputSpec, + OutputSpec, + ResourcesSpec, + SaveMode, + ) + from snowflake.ml.registry import Registry + + sf = kumo.SnowflakeConnector.from_connection_toml( + connection_name=args.connection_name, warehouse=args.warehouse, + database=args.database, schema_name=args.schema) + stage = args.stage.rstrip('/') + with tempfile.TemporaryDirectory() as tmp: + snapshot, query = fit(Path(tmp)) + exported = snapshot.export( + kumo.SnapshotOutputConfig(sf, f'{stage}/smoke/snapshots')) + trainer = kumo.Trainer(model_plan=ModelPlan()) + result = trainer.fit( + query.generate_training_table(snapshot), snapshot=snapshot, + export_config=kumo.ModelOutputConfig( + sf, location=kumo.ModelRegistryLocation(args.model_name), + serving_snapshot=exported, + serving_options={'pip_repository': args.pip_repository} + if args.pip_repository else None)) + published = result.published_model + print(f'registered {published.name} version {published.version}') + + table = query.generate_prediction_table(snapshot) + time = pd.Timestamp(table.timestamps[0]) + expected = trainer.predict(table, snapshot=snapshot, + binary_classification_threshold=0.5).frame + + session = sf.session() + version = Registry(session, database_name=args.database, + schema_name=args.schema).get_model( + published.name).version(published.version) + rows = session.create_dataframe( + pd.DataFrame({'ENTITY': expected['ENTITY'].astype(str)})) + output = f'{stage}/smoke/scores/' + job = version.run_batch( + rows, compute_pool=args.compute_pool, + output_spec=OutputSpec(stage_location=output, mode=SaveMode.OVERWRITE), + input_spec=InputSpec(params={ + 'prediction_time': str(time), + 'binary_classification_threshold': 0.5, + }), resources_spec=ResourcesSpec(gpu_requests='1'), + inference_spec=InferenceSpec(num_workers=1)) + job.wait() + job_name = job.id.split('.')[-1].strip('"') + scored = session.read.parquet(f'{output}{job_name}/').to_pandas() + scored.columns = [column.strip('"') for column in scored.columns] + + local = expected.assign(ENTITY=expected['ENTITY'].astype(str)).sort_values( + 'ENTITY').reset_index(drop=True) + remote = scored[list(local.columns)].sort_values('ENTITY').reset_index( + drop=True) + pd.testing.assert_frame_equal(local, remote, check_dtype=False, atol=1e-5) + print(f'PASS: {len(remote)} rows match local scores') + + +if __name__ == '__main__': + main() diff --git a/packages/kumo-engine/docs/mlops.md b/packages/kumo-engine/docs/mlops.md index f179d9202..853bb7b05 100644 --- a/packages/kumo-engine/docs/mlops.md +++ b/packages/kumo-engine/docs/mlops.md @@ -71,14 +71,14 @@ graph = kumo.Graph(tables={"users": users, "orders": orders.with_column("USER_ID ### Snowflake -Tracking and the Snowflake Model Registry through one connector (`kumo-engine[snowflake,s3]`): +With a named internal stage as the artifact root (`kumo-engine[snowflake]`): ```python import kumo.engine as kumo from kumo.engine.integrations.snowflake import SnowflakeTracking sf = kumo.SnowflakeConnector(warehouse="WH", database="ML", schema_name="MODELS") -cfg = kumo.MLOps("s3:///kumo", tracking=SnowflakeTracking(sf, experiment_name="CHURN"), +cfg = kumo.MLOps((sf, "@ML.MODELS.KUMO"), tracking=SnowflakeTracking(sf, experiment_name="CHURN"), registry=kumo.ModelRegistryLocation("CHURN_GNN", alias="candidate")) with kumo.mlops(cfg): snapshot = graph.snapshot() @@ -87,7 +87,7 @@ with kumo.mlops(cfg): run = result.predict(snapshot, output_config=kumo.PredictionOutputConfig(sf, path="CHURN_SCORES")) ``` -To serve only a better model, move its alias yourself: [Promote a model](#promote-a-model). +The registered version also scores inside Snowflake: see [Batch inference in Snowflake](#batch-inference-in-snowflake). To serve only a better model, move its alias yourself: [Promote a model](#promote-a-model). ## How it fits together @@ -116,7 +116,7 @@ with kumo.mlops(cfg): trainer.fit(table, snapshot=snapshot, mlops=cfg) ``` -- **`artifact_root`** (required): a URI (`s3://`, `gs://`, `file://`, or a local path, including `/Volumes/...` on a Databricks cluster) or a `(connector, path)` pair for a store a URI cannot reach: a volume through a `DatabricksConnector` (`"/Volumes////kumo"`) or a named stage through a `SnowflakeConnector` (`"@DB.SCHEMA.STAGE/kumo"`) once the installed sink holds files there. +- **`artifact_root`** (required): a URI (`s3://`, `gs://`, `file://`, or a local path, including `/Volumes/...` on a Databricks cluster) or a `(connector, path)` pair for a store a URI cannot reach: a volume through a `DatabricksConnector` (`"/Volumes////kumo"`) or a named stage through a `SnowflakeConnector` (`"@DB.SCHEMA.STAGE/kumo"`). - **`tracking`**: `MlflowTracking` or `SnowflakeTracking` (below), or your own [`Tracking`](#custom-trackers). - **`registry`**: a `ModelRegistryLocation(name, alias=...)`; every fit registers its model there, in the tracker's registry. - **`store`**: the artifact kinds stored when produced: `snapshot`, `training_table`, `model` and `prediction` (default: all). @@ -145,7 +145,7 @@ table = kumo.TrainingTable.load_from_connector(cfg, path=f"training_tables/{tabl - `result.registered_model` is the version `MLOps(registry=)` registered; `result.published_model` the one `fit(export_config=...)` published. - `result.tracking` is the tracker's run: `TrackingRef(platform="databricks", run_id="...", url="https://.../ml/experiments/.../runs/...")`. The model's `_provenance.json` records the same. -- `fit(scoring_defaults=kumo.ScoringDefaults(...))` stores prediction defaults in the bundle before it is stored and registered: `binary_classification_threshold` (binary and multilabel classification) and `num_classes_to_return` (multiclass, multilabel, link prediction). `predict`, the MLflow pyfunc and the scoring jobs apply them when a call omits them; explicit arguments and a query's `TOP k` win. `result.with_scoring_defaults(...)` retunes a fitted model as a new bundle whose manifest names the original in `parents.derived_from`. +- `fit(scoring_defaults=kumo.ScoringDefaults(...))` stores prediction defaults in the bundle before it is stored and registered: `binary_classification_threshold` (binary and multilabel classification) and `num_classes_to_return` (multiclass, multilabel, link prediction). `predict`, the MLflow pyfunc, Snowflake `run_batch` and the scoring jobs apply them when a call omits them; explicit arguments and a query's `TOP k` win. `result.with_scoring_defaults(...)` retunes a fitted model as a new bundle whose manifest names the original in `parents.derived_from`. - `result.export(cfg)` (or `trainer.export(cfg)` for a loaded model) registers a model at `cfg.registry` through its tracker. On every platform, registering a bundle the registry already holds under that name returns its version instead of adding one. - If registration fails, `fit` warns with the full error and returns with `registered_model=None` (with `on_error="raise"`, it raises), and the fit run ends failed with `kumo.error`. Registering a bundle the registry already holds under the name returns that version (logged at info) and moves the alias to it. @@ -187,10 +187,10 @@ Jobs that load `alias="champion"` pick up the new version on their next run. - With no `target`, it records where MLflow's own configuration points (`MLFLOW_TRACKING_URI`, `mlflow.set_tracking_uri`; automatic on Databricks). `target` names another server with its credentials: a `DatabricksConnector`, a `SageMakerMLflow(arn)`, or a tracking URI. With a `target`, `experiment_name` is required to record runs. - `log_model=True` logs the model on the fit run without registering it; `MLOps(registry=)` implies it. -`SnowflakeTracking(source, *, experiment_name, database=None, schema_name=None, run_name=None, tags=None, native_model_lineage=False)`: +`SnowflakeTracking(source, *, experiment_name, database=None, schema_name=None, run_name=None, tags=None, native_model_lineage=False, serving_options=None)`: - `source` is a `SnowflakeConnector`, whose identity opens the session and whose Snowflake Model Registry `MLOps(registry=)` registers in, or a Snowpark session, which tracks only. - `database` and `schema_name` default to the session's (`PUBLIC` when only `database` is given). The fit is in Snowsight under AI & ML ยป Experiments. -- The registered version carries the bundle; load it with `Trainer.load_from_connector`. +- `serving_options` sets how a registered version runs in Snowflake (see [Batch inference in Snowflake](#batch-inference-in-snowflake)). ## Score @@ -224,6 +224,26 @@ scores = model.predict(pd.DataFrame({"ENTITY": [1, 2, 3]}), params={ - Each replica loads a snapshot on its first `predict` with that URI and keeps it. `log_model(trainer, snapshot=..., serving_snapshot=...)` logs a model whose replicas load that snapshot at startup and use it when no `snapshot` param is given; `serving_snapshot` takes a URI, or a `GraphSnapshot` or snapshot id stored under the block's artifact root. - Serving needs an image with kumo-engine installed: `mlflow models serve -m --env-manager local`. Its pip requirement is `kumo-engine[s3,gcs]` of the same version (plus `[text]` for text columns). +### Batch inference in Snowflake + +A version registered in the Snowflake Model Registry with kumo-engine installed is runnable: Snowflake ML [batch inference jobs](https://docs.snowflake.com/en/developer-guide/snowflake-ml/inference/batch-inference-jobs) (`ModelVersion.run_batch`) score entity ids with it on a GPU compute pool. With a stage as the artifact root, the version scores against the fit's snapshot stored there, and carries the model's scoring defaults: + +```python +from snowflake.ml.model.batch_inference import InferenceSpec, InputSpec, OutputSpec, ResourcesSpec +from snowflake.ml.registry import Registry + +mv = Registry(session, database_name="ML", schema_name="MODELS").get_model("CHURN_GNN").version("champion") +job = mv.run_batch(session.sql("SELECT ID::STRING AS ENTITY FROM ML.MODELS.USERS"), compute_pool="KUMO_GPU_POOL", + output_spec=OutputSpec(stage_location="@ML.MODELS.KUMO_OUT/churn/"), + input_spec=InputSpec(params={"prediction_time": "2026-10-01"}), + resources_spec=ResourcesSpec(gpu_requests="1"), inference_spec=InferenceSpec(num_workers=1)) +job.wait() +``` + +- **Input:** one `ENTITY` column of ids as strings. Params: `prediction_time` (temporal queries), `binary_classification_threshold` and `num_classes_to_return` (override the model's defaults), and `snapshot`, another stage snapshot to score against. +- **Output:** Parquet under `//` with the columns `predict` returns. Load it with `COPY INTO FROM @... FILE_FORMAT = (TYPE = PARQUET) MATCH_BY_COLUMN_NAME = CASE_INSENSITIVE`. +- Pool, grants, environment and serving options: [Snowflake serving](#snowflake-serving). + ## Monitor `score_drift` compares a run's scores with the model's holdout and a previous run; `evaluate` scores a run once its outcomes are in a newer snapshot. With tracking on, each call logs a monitoring run named `drift-` or `evaluate-`, tagged `kumo.kind=monitoring`, `kumo.monitors=` and `kumo.bundle_id`. @@ -339,7 +359,7 @@ Every artifact has one way out and one way back. A `kumo.MLOps` stands in for th | Prediction run | `run.export(PredictionOutputConfig(connector, path=...))`, or `predict(output_config=...)` | `PredictionResult.load_from_connector(connector, path=...)` | | Evaluation | `evaluation.export(PredictionOutputConfig(connector, path=...))` | `EvaluationResult.load_from_connector(connector, path=...)` | -Loading needs a file destination; a table export is for consumers outside kumo-engine. Exporting again to a destination that holds a complete copy of the same artifact reuses it (logged at info). `ModelOutputConfig(connector, location=...)` publishes to a connector's registry outside a block. +Loading needs a file destination; a table export is for consumers outside kumo-engine. Exporting again to a destination that holds a complete copy of the same artifact reuses it (logged at info). `ModelOutputConfig(sf, location=..., serving_snapshot=..., serving_options=...)` publishes to the Snowflake Model Registry outside a block; `serving_snapshot` takes a stage path, or a `GraphSnapshot` or snapshot id stored under the block's stage root. ### Lineage and manifests @@ -353,7 +373,7 @@ Every artifact carries `_provenance.json`, metadata only (the underscore keeps S ### Artifacts, bundles and dependencies - A model's **bundle** holds only its own state: weights, encoders, the training spec and the holdout. Its manifest's **`parents`** are lineage; its **`dependencies`** are what `predict` needs besides the bundle and a snapshot: the `predictor` entry point, the `kumo-engine` `package` (with `[text]` when a text column is encoded) and content-addressed `asset`s. -- **Assets** are stored once under `/assets///` and shared by every model that uses them. `Trainer.load`, the pyfunc and the jobs look for each at the location recorded, then under the enclosing block's root, then in the embedding-model store, check its sha256 and cache non-local copies under `~/.cache/kumo/assets/`. A stored asset that cannot be found raises `MissingDependencyError`. +- **Assets** are stored once under `/assets///` and shared by every model that uses them. `Trainer.load`, the pyfunc, `run_batch` and the jobs look for each at the location recorded, then under the enclosing block's root, then in the embedding-model store, check its sha256 and cache non-local copies under `~/.cache/kumo/assets/`. A stored asset that cannot be found raises `MissingDependencyError`. - **Text encoders** are the first assets: word-vector models (`GloVeVectors`) are stored; other `TokenEmbeddingModel`s are recorded by fingerprint and must be registered wherever the model loads: ```python @@ -402,6 +422,32 @@ The statistics are read off the recorded percentile curves (min, q01..q99, max), - Databricks keys the monitor on `EVALUATED_AT` (anchors are older than its lookback by the label window). For regression use `INFERENCE_PROBLEM_TYPE_REGRESSION`; for multiclass `prediction_column="CLASS"`. `True_PROB` is profiled as an ordinary column. - Snowflake: binary classification uses `PREDICTION_SCORE_COLUMNS` (in [0, 1]) with `ACTUAL_CLASS_COLUMNS`; multiclass `PREDICTION_CLASS_COLUMNS=(CLASS)`; regression `PREDICTION_SCORE_COLUMNS=(TARGET_PRED)` with `ACTUAL_SCORE_COLUMNS`. The timestamp is `TIMESTAMP_NTZ`, columns free of nulls, NaN and infinities, each used once; mixed-case names such as `"True_PROB"` are double-quoted. Without `BASELINE` there is no drift. Registration records the version's task type, which Model Monitor requires. +### Snowflake serving + +- **Snapshot:** a version `MLOps(registry=)` registers with a stage root scores against the fit's snapshot stored under `/snapshots/`; pass `InputSpec(params={"snapshot": ...})` to score a newer one. The container reads it with its own Snowflake identity, in the registering connector's warehouse, over presigned URLs: create the stage with `ENCRYPTION = (TYPE = 'SNOWFLAKE_SSE')` (a `SNOWFLAKE_FULL` stage cannot be read back in ranges). +- **Environment:** Snowflake builds the inference image on the first job (about 5 minutes). The version embeds kumo's own packages (kumo-engine, kumo-core, kumo-api, kumo-pql, kumo-snowflake, DiskGraph; about 100 MB) as code; its pip requirements are their third-party dependencies, with torch at this installation's minor version and `faiss-cpu`. snowflake-ml pickles the model by value, so the container runs the registering Python; another `python_version` raises at registration. Assets are read from `/assets/` on the stage. A version without kumo-engine (from the hosted service) carries the bundle but refuses to score; `Trainer.load_from_connector` loads either. +- **`SnowflakeTracking(serving_options={...})`** (or `ModelOutputConfig(serving_options=)` outside a block): + + | Option | Default | | + |---|---|---| + | `cuda_version` | `12.6` | must match a torch CUDA build the pool's driver runs | + | `python_version` | the registering Python | must equal it | + | `function_type` | `TABLE_FUNCTION` | `FUNCTION` suits only one output row per input row | + | `pip_repository` | PyPI | a Snowflake pip artifact repository, e.g. `SNOWFLAKE.SNOWPARK.PYPI_SHARED_REPOSITORY` | + +- **Compute pool:** a GPU pool (`INSTANCE_FAMILY = GPU_NV_S`, `GPU_NV_M`, `GPU_NV_SM` on Azure). Each worker loads the bundle and snapshot once: use one worker per GPU and large batches. The first job, including the image build, takes about 12 minutes. +- **Grants:** + + | For | Grants | + |---|---| + | Artifact root and output stages | `USAGE` on the database and schema; `READ` and `WRITE` on the stages | + | Registering | `CREATE MODEL` on the schema | + | `SnowflakeTracking` | `CREATE EXPERIMENT` on the schema | + | `run_batch` | `USAGE` on the compute pool and warehouse; `CREATE SERVICE` on the schema; ownership or `USAGE` on the model | + | Model Monitor | `CREATE MODEL MONITOR` and `CREATE TABLE` on the schema | + +- Snowflake COPY and external tables read `_provenance.json` too: use `PATTERN='.*[.]parquet'`. [`examples/mlops/snowflake/run_batch_smoke.py`](../../../examples/mlops/snowflake/run_batch_smoke.py) scripts registration and a `run_batch` job. + ### Custom trackers Other systems plug in through `kumo.engine.tracking.Tracking`, which receives ready-built `RunRecord`s so every tracker records the same names: `start(record)` with the fit run, `on_trial(record, event)` with a trial attempt after each of its events, and `finish(record, error)` with the final record (`None` when tracking stopped early). Optional members: `prepare()`, `tags`, `check(registry)` (when `kumo.MLOps` is built), `log_model(bundle_dir, manifest, snapshot, registry)` (logs and registers, returning the version) or `connector` (a connector whose model registry `registry` registers in), `load_model(location)` and `register(result, location)` (for loading and `result.export(cfg)`), `platform`/`run_id`/`run_url` (recorded as `result.tracking`), and the best-effort `log_prediction_run(record)` and `log_monitoring(record)`. A tracker that raises is not called again for that fit and its run ends failed. `RunRecord.inputs` maps each source table to its `connector`, `location`, `version` and `version_kind`; a tracker's `tags` are added except `kumo.*` and param names. diff --git a/packages/kumo-engine/kumo/engine/integrations/_mlflow_model.py b/packages/kumo-engine/kumo/engine/integrations/_mlflow_model.py index bc85aa556..e8b80c607 100644 --- a/packages/kumo-engine/kumo/engine/integrations/_mlflow_model.py +++ b/packages/kumo-engine/kumo/engine/integrations/_mlflow_model.py @@ -1,23 +1,18 @@ # SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. # All rights reserved. # SPDX-License-Identifier: Apache-2.0 -import logging -import tempfile -from pathlib import Path -from typing import Any, Dict, Mapping, Optional, Tuple +from typing import Any, Dict, Optional, Tuple import pandas as pd from kumo.engine.connector._extras import import_provider -from kumo.engine.connector.base import Connector from kumo.engine.graph import GraphSnapshot +from kumo.engine.integrations._serving import predict, snapshot_location _pyfunc = import_provider('mlflow.pyfunc', 'mlflow') _schema = import_provider('mlflow.types.schema', 'mlflow') _signature = import_provider('mlflow.models.signature', 'mlflow') -logger = logging.getLogger(__name__) -_ANCHOR = 'ANCHOR_TIMESTAMP' _DTYPES = { 'PREDICTED': 'boolean', 'SCORE': 'double', @@ -58,118 +53,6 @@ def _snapshot(self, uri: str) -> GraphSnapshot: return self._snapshots[uri] -def predict(trainer: Any, snapshot: GraphSnapshot, rows: pd.DataFrame, - params: Mapping[str, Any]) -> pd.DataFrame: - from kumo.engine._mlops import resolve - from kumo.engine._scratch import scoped_scratch_dir - from kumo.engine.connector.file import FileConnector - - time = params.get('prediction_time') - time = None if time is None else pd.Timestamp(time) - options = dict( - binary_classification_threshold=params.get( - 'binary_classification_threshold'), - num_classes_to_return=params.get('num_classes_to_return'), - ) - _require_inputs(trainer, snapshot, rows, params, time) - trainer.scratch_dir.mkdir(parents=True, exist_ok=True) - with tempfile.TemporaryDirectory(dir=trainer.scratch_dir) as tmp, \ - scoped_scratch_dir(Path(tmp) / 'scratch'): - if len(rows) == 0: - return trainer.predict( - snapshot, prediction_time=None - if time is None else time.to_pydatetime(), **options).frame - rows = rows.rename(columns={'ENTITY': 'ENTITY1'}) - if time is not None and 'TIMESTAMP' not in rows: - rows = rows.assign(TIMESTAMP=time) - (Path(tmp) / 'rows').mkdir() - rows.to_parquet(Path(tmp) / 'rows' / 'rows.parquet', index=False) - entities = FileConnector(root_dir=str(Path(tmp) / 'rows'))[ - 'rows.parquet'] - timed = 'TIMESTAMP' in rows and _temporal(trainer, snapshot) - scored = trainer._predict( - trainer._prediction_table(snapshot, entities, None), - snapshot=snapshot, extra_fields=[_ANCHOR] if timed else [], - settings=resolve(), **options).frame - return _in_input_order(rows, scored) - - -def _in_input_order(rows: pd.DataFrame, scored: pd.DataFrame) -> pd.DataFrame: - r"""``scored`` in the order of ``rows``, which scoring does not keep; - rows scoring skipped stay out, and repeated rows pair up in order.""" - keys = ['ENTITY'] - left = pd.DataFrame({'ENTITY': rows['ENTITY1'].to_numpy()}) - right = scored[['ENTITY']].copy() - if _ANCHOR in scored: - keys.append(_ANCHOR) - left[_ANCHOR] = _seconds(rows['TIMESTAMP']) - right[_ANCHOR] = _seconds(scored[_ANCHOR]) - try: - left = left.astype({'ENTITY': right['ENTITY'].dtype}) - left['_n'] = left.groupby(keys, sort=False, dropna=False).cumcount() - right['_n'] = right.groupby(keys, sort=False, dropna=False).cumcount() - right['_row'] = range(len(right)) - order = left.merge(right, on=[*keys, '_n'])['_row'] - except (TypeError, ValueError): - order = None - scored = scored.drop(columns=_ANCHOR, errors='ignore') - if order is None or len(order) != len(scored): - logger.warning('Could not match the scores to the input rows; they ' - 'are returned in scoring order: join them on ENTITY') - return scored - return scored.iloc[order.to_numpy()].reset_index(drop=True) - - -def _temporal(trainer: Any, snapshot: GraphSnapshot) -> bool: - from kumo_core.pquery import query_is_temporal - - from kumo.engine.pquery import PredictiveQuery - - return query_is_temporal( - PredictiveQuery._from_snapshot(snapshot, - trainer._training_query)._parsed, - snapshot._native_graph) - - -def _seconds(values: pd.Series) -> pd.Series: - stamps = pd.to_datetime(values) - if stamps.dt.tz is not None: - stamps = stamps.dt.tz_convert(None) - return stamps.astype('datetime64[s]') - - -def _require_inputs(trainer: Any, snapshot: GraphSnapshot, rows: pd.DataFrame, - params: Mapping[str, Any], - time: Optional[pd.Timestamp]) -> None: - from kumo_core.pquery import query_is_temporal - from kumo_core.types import TaskType - - from kumo.engine.pquery import PredictiveQuery - - untimed = len(rows) > 0 and time is None and 'TIMESTAMP' not in rows - unthresholded = ( - params.get('binary_classification_threshold') is None and - trainer.result.scoring_defaults.binary_classification_threshold is None) - if not (untimed or unthresholded): - return - parsed = PredictiveQuery._from_snapshot(snapshot, - trainer._training_query)._parsed - if untimed and query_is_temporal(parsed, snapshot._native_graph): - raise ValueError("temporal query: pass params={'prediction_time': " - "...} or a TIMESTAMP column") - if unthresholded and (TaskType(parsed.task_type) - is TaskType.BINARY_CLASSIFICATION): - raise ValueError("binary classification: pass " - "params={'binary_classification_threshold': ...}") - - -def snapshot_location(uri: str) -> Tuple[Connector, str]: - from kumo.engine.connector import connector_for_root - - parent, slash, name = uri.rstrip('/').rpartition('/') - return connector_for_root(parent or slash or '.'), name - - def signature(snapshot: GraphSnapshot, query: str) -> Tuple[Any, pd.DataFrame]: from kumo_core.train.formatting import prediction_output_columns from kumo_core.types import TaskType diff --git a/packages/kumo-engine/kumo/engine/integrations/_serving.py b/packages/kumo-engine/kumo/engine/integrations/_serving.py new file mode 100644 index 000000000..bb5d99bb6 --- /dev/null +++ b/packages/kumo-engine/kumo/engine/integrations/_serving.py @@ -0,0 +1,214 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +r"""Scoring rows against a registered bundle, shared by the MLflow pyfunc +and the Snowflake Model Registry model.""" + +import json +import logging +import tempfile +import threading +from pathlib import Path +from typing import Any, Callable, Dict, List, Mapping, Optional, Tuple + +import pandas as pd + +from kumo.engine.connector.base import Connector +from kumo.engine.graph import GraphSnapshot + +logger = logging.getLogger(__name__) +_ANCHOR = 'ANCHOR_TIMESTAMP' +_lock = threading.Lock() +_trainers: Dict[str, Any] = {} +_snapshots: Dict[str, GraphSnapshot] = {} + + +def predict(trainer: Any, snapshot: GraphSnapshot, rows: pd.DataFrame, + params: Mapping[str, Any]) -> pd.DataFrame: + from kumo.engine._mlops import resolve + from kumo.engine._scratch import scoped_scratch_dir + from kumo.engine.connector.file import FileConnector + + time = params.get('prediction_time') + time = None if time is None else pd.Timestamp(time) + options = dict( + binary_classification_threshold=params.get( + 'binary_classification_threshold'), + num_classes_to_return=params.get('num_classes_to_return'), + ) + _require_inputs(trainer, snapshot, rows, params, time) + trainer.scratch_dir.mkdir(parents=True, exist_ok=True) + with tempfile.TemporaryDirectory(dir=trainer.scratch_dir) as tmp, \ + scoped_scratch_dir(Path(tmp) / 'scratch'): + if len(rows) == 0: + return trainer.predict( + snapshot, prediction_time=None + if time is None else time.to_pydatetime(), **options).frame + rows = rows.rename(columns={'ENTITY': 'ENTITY1'}) + if time is not None and 'TIMESTAMP' not in rows: + rows = rows.assign(TIMESTAMP=time) + (Path(tmp) / 'rows').mkdir() + rows.to_parquet(Path(tmp) / 'rows' / 'rows.parquet', index=False) + entities = FileConnector(root_dir=str(Path(tmp) / 'rows'))[ + 'rows.parquet'] + timed = 'TIMESTAMP' in rows and _temporal(trainer, snapshot) + scored = trainer._predict( + trainer._prediction_table(snapshot, entities, None), + snapshot=snapshot, extra_fields=[_ANCHOR] if timed else [], + settings=resolve(), **options).frame + return _in_input_order(rows, scored) + + +def _in_input_order(rows: pd.DataFrame, scored: pd.DataFrame) -> pd.DataFrame: + r"""``scored`` in the order of ``rows``, which scoring does not keep; + rows scoring skipped stay out, and repeated rows pair up in order.""" + keys = ['ENTITY'] + left = pd.DataFrame({'ENTITY': rows['ENTITY1'].to_numpy()}) + right = scored[['ENTITY']].copy() + if _ANCHOR in scored: + keys.append(_ANCHOR) + left[_ANCHOR] = _seconds(rows['TIMESTAMP']) + right[_ANCHOR] = _seconds(scored[_ANCHOR]) + try: + left = left.astype({'ENTITY': right['ENTITY'].dtype}) + left['_n'] = left.groupby(keys, sort=False, dropna=False).cumcount() + right['_n'] = right.groupby(keys, sort=False, dropna=False).cumcount() + right['_row'] = range(len(right)) + order = left.merge(right, on=[*keys, '_n'])['_row'] + except (TypeError, ValueError): + order = None + scored = scored.drop(columns=_ANCHOR, errors='ignore') + if order is None or len(order) != len(scored): + logger.warning('Could not match the scores to the input rows; they ' + 'are returned in scoring order: join them on ENTITY') + return scored + return scored.iloc[order.to_numpy()].reset_index(drop=True) + + +def _temporal(trainer: Any, snapshot: GraphSnapshot) -> bool: + from kumo_core.pquery import query_is_temporal + + from kumo.engine.pquery import PredictiveQuery + + return query_is_temporal( + PredictiveQuery._from_snapshot(snapshot, + trainer._training_query)._parsed, + snapshot._native_graph) + + +def _seconds(values: pd.Series) -> pd.Series: + stamps = pd.to_datetime(values) + if stamps.dt.tz is not None: + stamps = stamps.dt.tz_convert(None) + return stamps.astype('datetime64[s]') + + +def _require_inputs(trainer: Any, snapshot: GraphSnapshot, rows: pd.DataFrame, + params: Mapping[str, Any], + time: Optional[pd.Timestamp]) -> None: + from kumo_core.pquery import query_is_temporal + from kumo_core.types import TaskType + + from kumo.engine.pquery import PredictiveQuery + + untimed = len(rows) > 0 and time is None and 'TIMESTAMP' not in rows + unthresholded = ( + params.get('binary_classification_threshold') is None and + trainer.result.scoring_defaults.binary_classification_threshold is None) + if not (untimed or unthresholded): + return + parsed = PredictiveQuery._from_snapshot(snapshot, + trainer._training_query)._parsed + if untimed and query_is_temporal(parsed, snapshot._native_graph): + raise ValueError("temporal query: pass params={'prediction_time': " + "...} or a TIMESTAMP column") + if unthresholded and (TaskType(parsed.task_type) + is TaskType.BINARY_CLASSIFICATION): + raise ValueError("binary classification: pass " + "params={'binary_classification_threshold': ...}") + + +def snapshot_location(uri: str) -> Tuple[Connector, str]: + from kumo.engine.connector import connector_for_root + + parent, slash, name = uri.rstrip('/').rpartition('/') + return connector_for_root(parent or slash or '.'), name + + +def task_type(bundle: Path) -> Any: + from kumo_core.types import TaskType + + schema = json.loads((bundle / 'prediction-schema.json').read_text()) + return TaskType(schema['task_type']) + + +def output_columns(bundle: Path) -> List[str]: + from kumo_core.train.formatting import prediction_output_columns + + return ['ENTITY', *prediction_output_columns(task_type(bundle))] + + +def score_registered(bundle: str, serving: str, rows: pd.DataFrame, + params: Mapping[str, Any]) -> pd.DataFrame: + r"""Scores ``rows`` with the bundle at ``bundle`` against the snapshot + in ``params`` or the ``serving`` record. The bundle and each snapshot + load once per process.""" + from kumo.engine.trainer.trainer import Trainer + + record = json.loads(Path(serving).read_text()) + columns = output_columns(Path(bundle)) + if rows.empty: + return pd.DataFrame(columns=columns) + uri = params.get('snapshot') or record.get('snapshot') + if not uri: + raise ValueError( + "No snapshot to score against: register with ModelOutputConfig(" + "..., serving_snapshot='@DB.SCHEMA.STAGE/...') or pass " + "InputSpec(params={'snapshot': ...})") + trainer = _cached(_trainers, bundle, lambda: Trainer.load(bundle)) + snapshot = _cached(_snapshots, uri, lambda: load_snapshot(uri, record)) + rows = rows.assign(ENTITY=_entities(rows['ENTITY'], trainer, snapshot)) + frame = predict(trainer, snapshot, rows, + {k: v for k, v in params.items() if k != 'snapshot'}) + frame = frame.reindex(columns=columns) + for column in ('ENTITY', 'CLASS'): + if column in frame: + frame[column] = frame[column].astype(str) + return frame + + +def load_snapshot(uri: str, target: Mapping[str, Any]) -> GraphSnapshot: + r"""Loads an exported snapshot; an ``@`` stage path is read with the + Snowpark Container Services identity, in ``target``'s warehouse.""" + if uri.startswith('@'): + from kumo.engine.connector.snowflake import SnowflakeConnector + + connector: Connector = SnowflakeConnector.from_spcs_env( + warehouse=target.get('warehouse'), + database=target.get('database'), + schema_name=target.get('schema')) + return GraphSnapshot.load_from_connector(connector, path=uri) + connector, path = snapshot_location(uri) + return GraphSnapshot.load_from_connector(connector, path=path) + + +def _cached(cache: Dict[str, Any], key: str, load: Callable[[], Any]) -> Any: + with _lock: + if key not in cache: + cache[key] = load() + return cache[key] + + +def _entities(values: pd.Series, trainer: Any, + snapshot: GraphSnapshot) -> pd.Series: + from kumo.engine.pquery import PredictiveQuery + + parsed = PredictiveQuery._from_snapshot(snapshot, + trainer._training_query)._parsed + entity = snapshot._native_graph.table(parsed.entity_table) + dtype = str(entity.column(entity.primary_key).dtype) + if dtype.startswith('int') or dtype.startswith('uint'): + return pd.to_numeric(values).astype('int64') + if dtype.startswith('float'): + return pd.to_numeric(values).astype('float64') + return values.astype(str) diff --git a/packages/kumo-engine/kumo/engine/integrations/_snowflake_model.py b/packages/kumo-engine/kumo/engine/integrations/_snowflake_model.py new file mode 100644 index 000000000..4219f31c8 --- /dev/null +++ b/packages/kumo-engine/kumo/engine/integrations/_snowflake_model.py @@ -0,0 +1,238 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +r"""A Kumo model bundle as a Snowflake Model Registry version that scores in +Snowpark Container Services, e.g. with ``ModelVersion.run_batch``. + +snowflake-ml pickles this module by value into the version, so the container +must run the Python the version was registered from. kumo itself is embedded +in the version as code, on the path before the model loads; its third-party +requirements install from PyPI or a pip artifact repository. +""" + +import json +import shutil +import sys +from pathlib import Path +from typing import Any, Dict, List, Mapping, Optional + +import pandas as pd + +from kumo.engine.connector._extras import import_provider + +custom_model = import_provider('snowflake.ml.model.custom_model', 'snowflake') + +DEFAULTS = { + 'cuda_version': '12.6', + 'python_version': f'{sys.version_info.major}.{sys.version_info.minor}', + 'function_type': 'TABLE_FUNCTION', +} +#: Packages embedded as code: kumo's own, which are not on PyPI, and +#: torchmetrics, whose pinned version the CUDA index the image installs from +#: first does not carry. +EMBEDDED = ('kumo.engine', 'kumo_core', 'kumoapi', 'kumopql', 'kumo_snowflake', + 'diskgraph', 'torchmetrics') +_EMBEDDED_DISTS = ('diskgraph', 'torchmetrics') +#: Modules whose distributions' requirements the embedded code needs, with +#: the distribution of a workspace install; a bundled kumo-engine wheel +#: provides both modules itself. +REQUIRING = (('kumo_core', 'kumo-core'), ('kumopql', 'kumo-pql')) +#: kumo-engine's ``text`` extra, for models that embed text columns. +TEXT_REQUIREMENT = 'kumo-glove-6b-50d>=1.0.0,<2' +#: Installed by the Snowflake inference image; their own pins there win. +PROVIDED = ('snowflake-snowpark-python', 'snowflake-connector-python') + + +class KumoModel(custom_model.CustomModel): + r"""Scores entity ids (an ``ENTITY`` column) against a graph snapshot.""" + @custom_model.inference_api + def predict(self, df: pd.DataFrame, *, snapshot: str = '', + prediction_time: str = '', + binary_classification_threshold: float = -1.0, + num_classes_to_return: int = 0) -> pd.DataFrame: + from kumo.engine.integrations._serving import score_registered + + params = { + 'snapshot': snapshot or None, + 'prediction_time': prediction_time or None, + 'binary_classification_threshold': + (binary_classification_threshold + if binary_classification_threshold >= 0 else None), + 'num_classes_to_return': num_classes_to_return or None, + } + return score_registered(self.context.path('bundle'), + self.context.path('serving'), df, params) + + +def log_model_args(bundle_dir: str, metadata: Mapping[str, str], + target: Mapping[str, Optional[str]], + workdir: str) -> Dict[str, Any]: + r"""``Registry.log_model`` arguments that register the bundle as a + :class:`KumoModel` for Snowpark Container Services on GPU. ``target`` is + the registering connector's warehouse, database and schema.""" + from kumo.engine.trainer._spec import read_bundle_artifact_paths + from kumo.engine.trainer.output_config import ( + SERVING_OPTION_PREFIX, + SERVING_SNAPSHOT_KEY, + ) + + bundle = Path(bundle_dir) + serving = Path(workdir) / 'serving.json' + serving.write_text( + json.dumps({ + 'snapshot': metadata.get(SERVING_SNAPSHOT_KEY), + **target + })) + options = { + **DEFAULTS, + **{ + key[len(SERVING_OPTION_PREFIX):]: value + for key, value in metadata.items() + if key.startswith(SERVING_OPTION_PREFIX) + } + } + repository = options.pop('pip_repository', None) + python = options.pop('python_version') + if python != DEFAULTS['python_version']: + raise ValueError( + f"serving_options python_version={python!r} differs from the " + f"registering Python {DEFAULTS['python_version']}: snowflake-ml " + f"pickles the model by value, so the container must run the same " + f"Python. Register from Python {python} instead.") + paths, _ = read_bundle_artifact_paths(bundle) + environment: Dict[str, Any] = dict( + pip_requirements=pip_requirements() + + text_requirements(paths.get('encoders')), + code_paths=embedded_code(Path(workdir) / 'code')) + if repository: + environment['artifact_repository_map'] = {'pip': repository} + return dict( + model=KumoModel( + custom_model.ModelContext(artifacts={ + 'bundle': str(bundle), + 'serving': str(serving), + })), + signatures={'predict': signature(bundle)}, + target_platforms=['SNOWPARK_CONTAINER_SERVICES'], + python_version=python, + options=options, + **environment, + ) + + +def embedded_code(directory: Path) -> List[str]: + r"""Copies the :data:`EMBEDDED` packages into ``directory``, without + caches and tests, and returns the package directories.""" + import importlib + + paths = [] + for name in EMBEDDED: + source = Path(importlib.import_module(name).__file__).parent + if name == 'kumo.engine': + source, name = source.parent, 'kumo' + target = directory / name + shutil.copytree( + source, target, + ignore=shutil.ignore_patterns('__pycache__', 'test', 'tests')) + paths.append(str(target)) + return paths + + +def pip_requirements() -> List[str]: + r"""The third-party requirements of the distributions providing + :data:`REQUIRING`, in a form the Snowflake image build accepts: torch at + this installation's minor version, faiss-cpu for kumo's own faiss build, + one line per package, and no ``==X.*`` pins.""" + from importlib.metadata import packages_distributions, requires, version + + from packaging.requirements import Requirement + from packaging.utils import canonicalize_name + + provided = packages_distributions() + dists = sorted({ + dist + for module, fallback in REQUIRING + for dist in provided.get(module) or [fallback] + }) + found: Dict[str, Any] = {} + faiss = None + for dist in dists: + for line in requires(dist) or []: + req = Requirement(line) + if req.marker is not None and not req.marker.evaluate( + {'extra': 'faiss-cpu'}): + continue + name = canonicalize_name(req.name) + if name.startswith('faiss'): + faiss = req.specifier + elif name in found: + found[name].specifier &= req.specifier + elif not name.startswith('kumo') and name not in _EMBEDDED_DISTS: + found[name] = req + major, minor = (int(part) for part in version('torch').split('.')[:2]) + lines = [_pinned(req) for name, req in sorted(found.items()) + if name != 'torch'] + if faiss is not None: + lines.append(f'faiss-cpu{faiss}') + return [*lines, f'torch>={major}.{minor},<{major}.{minor + 1}', *PROVIDED] + + +def text_requirements(encoders: Optional[str]) -> List[str]: + r"""The word-embedding package, from PyPI, when ``encoders`` embeds + text columns.""" + from kumo_core.encoders.config import EncodersConfig + from kumo_core.io.json import load_json + from kumoapi.encoder import TextEmbedding + + config = (load_json(encoders, EncodersConfig) + if encoders and Path(encoders).is_file() else None) + if config is not None and any( + isinstance(encoder, TextEmbedding) + for table in config.encoders.values() + for encoder in table.values()): + return [TEXT_REQUIREMENT] + return [] + + +def _pinned(req: Any) -> str: + specs = [] + for spec in req.specifier: + if spec.operator == '==' and spec.version.endswith('.*'): + parts = spec.version[:-2].split('.') + upper = [*parts[:-1], str(int(parts[-1]) + 1)] + specs += [f'>={".".join(parts)}', f'<{".".join(upper)}'] + else: + specs.append(str(spec)) + return req.name + ','.join(sorted(specs)) + + +def signature(bundle: Path) -> Any: + from kumo_core.types import TaskType + ms = import_provider('snowflake.ml.model.model_signature', 'snowflake') + + from kumo.engine.integrations._serving import output_columns, task_type + + binary = task_type(bundle) is TaskType.BINARY_CLASSIFICATION + dtypes = { + 'PREDICTED': ms.DataType.BOOL, + 'CLASS': ms.DataType.STRING, + 'TIMESTAMP': ms.DataType.TIMESTAMP_NTZ, + 'TARGET_PRED': + ms.DataType.BOOL if binary else ms.DataType.DOUBLE, + } + return ms.ModelSignature( + inputs=[ms.FeatureSpec('ENTITY', ms.DataType.STRING)], + outputs=[ + ms.FeatureSpec( + name, ms.DataType.STRING if name == 'ENTITY' else dtypes.get( + name, ms.DataType.DOUBLE)) + for name in output_columns(bundle) + ], + params=[ + ms.ParamSpec('snapshot', ms.DataType.STRING, ''), + ms.ParamSpec('prediction_time', ms.DataType.STRING, ''), + ms.ParamSpec('binary_classification_threshold', + ms.DataType.DOUBLE, -1.0), + ms.ParamSpec('num_classes_to_return', ms.DataType.INT64, 0), + ], + ) diff --git a/packages/kumo-engine/kumo/engine/trainer/output_config.py b/packages/kumo-engine/kumo/engine/trainer/output_config.py index 160e9c834..dac3933c9 100644 --- a/packages/kumo-engine/kumo/engine/trainer/output_config.py +++ b/packages/kumo-engine/kumo/engine/trainer/output_config.py @@ -158,17 +158,41 @@ class SnapshotOutputConfig: path: str +#: Publication metadata a registry may act on: the serving snapshot, and +#: each serving option under ``kumo.serving.
FROM @... FILE_FORMAT = (TYPE = PARQUET) MATCH_BY_COLUMN_NAME = CASE_INSENSITIVE`. - Pool, grants, environment and serving options: [Snowflake serving](#snowflake-serving). @@ -425,12 +425,12 @@ The statistics are read off the recorded percentile curves (min, q01..q99, max), ### Snowflake serving - **Snapshot:** a version `MLOps(registry=)` registers with a stage root scores against the fit's snapshot stored under `/snapshots/`; pass `InputSpec(params={"snapshot": ...})` to score a newer one. The container reads it with its own Snowflake identity, in the registering connector's warehouse, over presigned URLs: create the stage with `ENCRYPTION = (TYPE = 'SNOWFLAKE_SSE')` (a `SNOWFLAKE_FULL` stage cannot be read back in ranges). -- **Environment:** Snowflake builds the inference image on the first job (about 5 minutes). The version embeds kumo's own packages (kumo-engine, kumo-core, kumo-api, kumo-pql, kumo-snowflake, DiskGraph; about 100 MB) as code; its pip requirements are their third-party dependencies, with torch at this installation's minor version and `faiss-cpu`. snowflake-ml pickles the model by value, so the container runs the registering Python; another `python_version` raises at registration. Assets are read from `/assets/` on the stage. A version without kumo-engine (from the hosted service) carries the bundle but refuses to score; `Trainer.load_from_connector` loads either. +- **Environment:** Snowflake builds the inference image on the first job (about 5 minutes). The version embeds kumo's own packages (kumo-engine, kumo-core, kumo-api, kumo-pql, kumo-snowflake, DiskGraph; about 100 MB) as code; its pip requirements are their third-party dependencies, with torch at this installation's minor version for its CUDA, and `faiss-cpu`. Registering the same bundle again returns its version only if it was registered for the same environment (Python, torch, CUDA, requirements); otherwise it adds a version. snowflake-ml pickles the model by value, so the container runs the registering Python; another `python_version` raises at registration. Assets are read from `/assets/` on the stage. A version without kumo-engine (from the hosted service) carries the bundle but refuses to score; `Trainer.load_from_connector` loads either. - **`SnowflakeTracking(serving_options={...})`** (or `ModelOutputConfig(serving_options=)` outside a block): | Option | Default | | |---|---|---| - | `cuda_version` | `12.6` | must match a torch CUDA build the pool's driver runs | + | `cuda_version` | this installation's torch CUDA (`torch.version.cuda`; `12.6` for a CPU torch) | another value than the installed torch's raises at registration | | `python_version` | the registering Python | must equal it | | `function_type` | `TABLE_FUNCTION` | `FUNCTION` suits only one output row per input row | | `pip_repository` | PyPI | a Snowflake pip artifact repository, e.g. `SNOWFLAKE.SNOWPARK.PYPI_SHARED_REPOSITORY` | diff --git a/packages/kumo-engine/kumo/engine/integrations/_snowflake_model.py b/packages/kumo-engine/kumo/engine/integrations/_snowflake_model.py index 020fd7dda..908ec1e70 100644 --- a/packages/kumo-engine/kumo/engine/integrations/_snowflake_model.py +++ b/packages/kumo-engine/kumo/engine/integrations/_snowflake_model.py @@ -22,8 +22,9 @@ custom_model = import_provider('snowflake.ml.model.custom_model', 'snowflake') +#: The CUDA a version serves with when this installation's torch has none. +CPU_TORCH_CUDA = '12.6' DEFAULTS = { - 'cuda_version': '12.6', 'python_version': f'{sys.version_info.major}.{sys.version_info.minor}', 'function_type': 'TABLE_FUNCTION', } @@ -93,10 +94,14 @@ def log_model_args(bundle_dir: str, metadata: Mapping[str, str], if key.startswith(SERVING_OPTION_PREFIX) } } + ServingEnvironmentError = import_provider( + 'kumo_snowflake.registry', 'snowflake').ServingEnvironmentError + repository = options.pop('pip_repository', None) python = options.pop('python_version') + options['cuda_version'] = cuda_version(options.get('cuda_version')) if python != DEFAULTS['python_version']: - raise ValueError( + raise ServingEnvironmentError( f"serving_options python_version={python!r} differs from the " f"registering Python {DEFAULTS['python_version']}: snowflake-ml " f"pickles the model by value, so the container must run the same " @@ -122,6 +127,36 @@ def log_model_args(bundle_dir: str, metadata: Mapping[str, str], ) +def cuda_version(requested: Optional[str]) -> str: + r"""The CUDA the version's image installs torch for: this installation's + torch build's, which ``requested`` must equal, so the image resolves the + same torch minor version for it.""" + import torch + ServingEnvironmentError = import_provider( + 'kumo_snowflake.registry', 'snowflake').ServingEnvironmentError + + def minor(version: str) -> str: + return '.'.join(version.split('.')[:2]) + + built = minor(torch.version.cuda) if torch.version.cuda else None + if requested is None: + return built or CPU_TORCH_CUDA + if built is not None and minor(requested) != built: + release = minor(torch.__version__.split('+')[0]) + raise ServingEnvironmentError( + f"serving_options cuda_version={requested!r}, but this " + f"installation's torch {torch.__version__} is built for CUDA " + f"{built}; the version's image would need torch {release} for " + f"CUDA {minor(requested)}, which jobs fail to install when it does " + f"not exist. Register from a torch built for CUDA " + f"{minor(requested)}, e.g. pip install --index-url " + f"https://download.pytorch.org/whl/cu" + f"{minor(requested).replace('.', '')} \"torch=={release}.*\" " + f"(another release if that one has no such build), or drop " + f"cuda_version to serve with CUDA {built}.") + return requested + + def embedded_code(directory: Path) -> List[str]: r"""Copies the :data:`EMBEDDED` packages into ``directory``, without caches and tests, and returns the package directories.""" diff --git a/packages/kumo-engine/kumo/engine/trainer/output_config.py b/packages/kumo-engine/kumo/engine/trainer/output_config.py index ff241cb87..a975f8d7a 100644 --- a/packages/kumo-engine/kumo/engine/trainer/output_config.py +++ b/packages/kumo-engine/kumo/engine/trainer/output_config.py @@ -175,7 +175,8 @@ class ModelOutputConfig: ``@DB.SCHEMA.STAGE/...`` path, or a `GraphSnapshot` or snapshot id stored under the stage artifact root of the enclosing `kumo.mlops` block. ``serving_options`` overrides how it runs: - ``cuda_version`` (default ``12.6``), ``python_version`` (the registering + ``cuda_version`` (the installed torch's, which it must equal), + ``python_version`` (the registering Python, which it must equal), ``function_type`` (``TABLE_FUNCTION``), and ``pip_repository``, a Snowflake pip artifact repository the version's third-party requirements install from instead of PyPI. diff --git a/packages/kumo-engine/test/test_snowflake_model.py b/packages/kumo-engine/test/test_snowflake_model.py index d96d7bb5c..8182beded 100644 --- a/packages/kumo-engine/test/test_snowflake_model.py +++ b/packages/kumo-engine/test/test_snowflake_model.py @@ -50,7 +50,9 @@ def _package(trainer, metadata, tmp_path): return args, loaded.model -def test_a_packaged_version_scores_like_the_trainer(served, tmp_path): +def test_a_packaged_version_scores_like_the_trainer(served, tmp_path, + monkeypatch): + monkeypatch.setattr('torch.version.cuda', '12.6.3') trainer, table, snapshot, uri = served args, model = _package(trainer, {'kumo.serving_snapshot': uri}, tmp_path) @@ -97,9 +99,10 @@ def test_a_packaged_version_scores_like_the_trainer(served, tmp_path): def test_serving_options_override_the_defaults_and_a_snapshot_is_required( - served, tmp_path): + served, tmp_path, monkeypatch): from kumo.engine.integrations._snowflake_model import pip_requirements + monkeypatch.setattr('torch.version.cuda', None) trainer, *_ = served args, model = _package( trainer, { @@ -114,13 +117,36 @@ def test_serving_options_override_the_defaults_and_a_snapshot_is_required( 'cuda_version': '12.4', 'function_type': 'FUNCTION' } - with pytest.raises(ValueError, match='serving_snapshot='): + with pytest.raises(ValueError, match='No snapshot to score against'): model.predict(pd.DataFrame({'ENTITY': ['1']})) with pytest.raises(ValueError, match='same Python'): _package(trainer, {'kumo.serving.python_version': '2.7'}, tmp_path / 'other') +def test_the_version_serves_with_the_cuda_of_the_installed_torch( + monkeypatch): + from kumo_snowflake.registry import ServingEnvironmentError + + from kumo.engine.integrations._snowflake_model import cuda_version + + monkeypatch.setattr('torch.__version__', '2.15.0+cu132') + monkeypatch.setattr('torch.version.cuda', '13.2') + assert cuda_version(None) == cuda_version('13.2') == '13.2' + with pytest.raises(ServingEnvironmentError) as raised: + cuda_version('12.6') + assert str(raised.value) == ( + "serving_options cuda_version='12.6', but this installation's torch " + "2.15.0+cu132 is built for CUDA 13.2; the version's image would need " + "torch 2.15 for CUDA 12.6, which jobs fail to install when it does " + "not exist. Register from a torch built for CUDA 12.6, e.g. pip " + "install --index-url https://download.pytorch.org/whl/cu126 " + "\"torch==2.15.*\" (another release if that one has no such build), " + "or drop cuda_version to serve with CUDA 13.2.") + monkeypatch.setattr('torch.version.cuda', None) + assert cuda_version(None) == '12.6' and cuda_version('12.4') == '12.4' + + def test_a_bundled_wheel_supplies_the_requirements(monkeypatch): import importlib.metadata as metadata diff --git a/packages/kumo-snowflake/src/kumo_snowflake/registry.py b/packages/kumo-snowflake/src/kumo_snowflake/registry.py index 1a21e1233..75c267dac 100644 --- a/packages/kumo-snowflake/src/kumo_snowflake/registry.py +++ b/packages/kumo-snowflake/src/kumo_snowflake/registry.py @@ -30,6 +30,7 @@ snowflake-ml's private layout; the marker is ours. """ +import hashlib import json import logging import re @@ -87,6 +88,16 @@ _LOG = logging.getLogger(__name__) +#: Version metric: a digest of the serving environment a runnable version +#: was registered with. +SERVING_KEY = 'kumo_serving' + + +class ServingEnvironmentError(ValueError): + r"""A runnable version's serving environment could not work, raised at + registration rather than when a job builds its image.""" + + #: Entry-point group of the callable that turns a bundle into ``log_model`` #: arguments for a runnable version; kumo-engine registers one. MODEL_BUILDER_GROUP = 'kumo_snowflake.registry_models' @@ -204,6 +215,8 @@ def _model_args(bundle: Path, metadata: Mapping[str, str], if build is not None: return build(str(bundle), metadata, _serving_target(connector), workdir) + except ServingEnvironmentError: + raise except Exception as error: _LOG.warning( 'Registering the bundle as a version that does not score in ' @@ -235,9 +248,26 @@ def _serving_target(connector: SourceConnector) -> Dict[str, Optional[str]]: } +def _serving(args: Mapping[str, Any]) -> Optional[str]: + r"""A digest of the environment ``log_model`` arguments serve with; + `None` for a version that does not score.""" + if 'pip_requirements' not in args: + return None + environment = { + name: args.get(name) + for name in ('python_version', 'options', 'pip_requirements', + 'artifact_repository_map') + } + return hashlib.sha256( + json.dumps(environment, sort_keys=True, + default=str).encode()).hexdigest()[:16] + + def _version_holding(registry: "Registry", model_name: str, - bundle_id: Optional[str]) -> Optional["ModelVersion"]: - r"""The Kumo version of ``model_name`` that records ``bundle_id``.""" + bundle_id: Optional[str], + serving: Optional[str]) -> Optional["ModelVersion"]: + r"""The Kumo version of ``model_name`` that records ``bundle_id`` and + was registered for the ``serving`` environment.""" if not bundle_id: return None try: @@ -252,7 +282,13 @@ def _version_holding(registry: "Registry", model_name: str, continue if (recorded.get(ENGINE_MARKER) == '1' and recorded.get(BUNDLE_ID_KEY) == bundle_id): - return version + if recorded.get(SERVING_KEY) == serving: + return version + _LOG.warning( + '%s version %s holds bundle %s for another serving ' + 'environment (torch, CUDA, Python or requirements); ' + 'registering a new version', model_name, + version.version_name, bundle_id) return None @@ -444,8 +480,12 @@ def publish_model( record = _version_record(metadata) with tempfile.TemporaryDirectory() as workdir, \ self._registry(connector) as registry: + args = _model_args(bundle, metadata, connector, workdir) + serving = _serving(args) + if serving is not None: + record[SERVING_KEY] = serving version = _version_holding(registry, location.model_name, - metadata.get(BUNDLE_ID_KEY)) + metadata.get(BUNDLE_ID_KEY), serving) if version is not None: _LOG.info('%s version %s already holds bundle %s; reused it', location.model_name, version.version_name, @@ -456,7 +496,7 @@ def publish_model( version_name=location.version, comment=location.comment, metrics=record, - **_model_args(bundle, metadata, connector, workdir), + **args, **_task(bundle), ) return _published(connector, version, metadata=metadata, diff --git a/packages/kumo-snowflake/tests/test_registry.py b/packages/kumo-snowflake/tests/test_registry.py index 8d0e31c73..0d0ff6c89 100644 --- a/packages/kumo-snowflake/tests/test_registry.py +++ b/packages/kumo-snowflake/tests/test_registry.py @@ -386,6 +386,50 @@ def log_model_args(bundle, metadata, target, workdir): 'warehouse': 'WH', 'database': 'DB', 'schema': 'ML'})] +def test_a_serving_environment_that_cannot_work_fails_registration( + tmp_path, monkeypatch): + def refuse(*_): + raise registry_module.ServingEnvironmentError('torch has no cu126') + + monkeypatch.setattr(registry_module, '_model_builder', lambda: refuse) + calls = [] + sink = _publishing_sink(monkeypatch, [], calls) + connector = SimpleNamespace(config=SnowflakeConnectorResourceConfig( + warehouse='WH', database='DB', schema_name='ML')) + with pytest.raises(ValueError, match='torch has no cu126'): + sink.publish_model(connector, str(_bundle(tmp_path)), + location=ModelRegistryLocation('CHURN')) + assert calls == [] + + +def test_a_bundle_registered_for_another_environment_gets_a_new_version( + tmp_path, monkeypatch, caplog): + environments = iter([['torch>=2.15,<2.16'], ['torch>=2.12,<2.13'], + ['torch>=2.15,<2.16']]) + monkeypatch.setattr( + registry_module, '_model_args', lambda *_: { + 'model': 'runnable', 'pip_requirements': next(environments)}) + versions, calls = [], [] + sink = _publishing_sink(monkeypatch, versions, calls) + bundle = str(_bundle(tmp_path)) + + def publish(): + logged = len(calls) + sink.publish_model(SimpleNamespace(), bundle, + location=ModelRegistryLocation('CHURN'), + metadata={'kumo_bundle_id': 'b1'}) + if len(calls) > logged: + versions.append(SimpleNamespace( + version_name=f'V{len(versions) + 1}', model_name='CHURN', + fully_qualified_model_name='DB.S.CHURN', + show_metrics=lambda metrics=calls[-1]['metrics']: metrics)) + return len(calls) > logged + + with caplog.at_level('WARNING', logger=registry_module.__name__): + assert [publish(), publish(), publish()] == [True, True, False] + assert 'V1 holds bundle b1 for another serving environment' in caplog.text + + @pytest.mark.parametrize('engine', [None, 'raises']) def test_without_a_usable_engine_the_bundle_travels_in_the_carrier( tmp_path, monkeypatch, engine): From 707ba8f93817232b1ecaea7955fc578bd6ebb6e2 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Thu, 29 Oct 2026 03:55:11 +0000 Subject: [PATCH 6/7] Load a Snowflake version by downloading only its bundle, in parallel ModelVersion.export GETs every file of a version one at a time, including the ~740 files of kumo code embedded for run_batch, so an alias load took about 2 minutes. fetch_model now lists the version once and GETs the files outside model/code/ on a pool of 8 cursors; the registered version and run_batch are unchanged. A live alias load takes about 6 s. Co-Authored-By: Claude Opus 5.5 --- .../src/kumo_snowflake/registry.py | 38 ++++++- .../kumo-snowflake/tests/test_registry.py | 103 ++++++++++++------ 2 files changed, 103 insertions(+), 38 deletions(-) diff --git a/packages/kumo-snowflake/src/kumo_snowflake/registry.py b/packages/kumo-snowflake/src/kumo_snowflake/registry.py index 75c267dac..a0bcaac9d 100644 --- a/packages/kumo-snowflake/src/kumo_snowflake/registry.py +++ b/packages/kumo-snowflake/src/kumo_snowflake/registry.py @@ -36,8 +36,9 @@ import re import shutil import tempfile +from concurrent.futures import ThreadPoolExecutor from contextlib import contextmanager -from pathlib import Path +from pathlib import Path, PurePosixPath from typing import ( TYPE_CHECKING, Any, @@ -106,6 +107,7 @@ class ServingEnvironmentError(ValueError): #: rather than a temporary directory so a half-finished import leaves something #: a reader can look at. _EXPORT_DIR = '_registry-export' +_EXPORT_WORKERS = 8 def _require_registry() -> type["Registry"]: @@ -313,6 +315,35 @@ def _locate_bundle(root: Path) -> Path: f"published by kumoai_studio.export_model().") +def _export(session: Any, version: "ModelVersion", directory: Path) -> None: + r"""Downloads the version's files but the code embedded in it, which + loading the bundle back does not import: one LIST, then GETs on a + bounded pool (``ModelVersion.export`` GETs every file in turn).""" + from kumo_core.sources.layout import confined + + root = (f'snow://model/{version.fully_qualified_model_name}/versions/' + f'{resolved(str(version.version_name))}/') + names = [ + PurePosixPath(*PurePosixPath(row['name']).parts[2:]) + for row in session.sql(f"LIST {_quoted(root + 'model/')}").collect() + ] + + def get(name: PurePosixPath) -> None: + local = directory / confined(str(name.parent)) + local.mkdir(parents=True, exist_ok=True) + with session.connection.cursor() as cursor: + cursor.execute(f"GET {_quoted(root + str(name))} " + f"{_quoted(f'file://{local}')}") + + with ThreadPoolExecutor(_EXPORT_WORKERS) as pool: + list(pool.map(get, [name for name in names + if name.parts[:2] != ('model', 'code')])) + + +def _quoted(url: str) -> str: + return "'" + url.replace('\\', '\\\\').replace("'", "\\'") + "'" + + def _model_uri(connector: SourceConnector, version: "ModelVersion") -> str: r"""Returns the fully qualified name of one model version.""" name = getattr(version, 'fully_qualified_model_name', None) @@ -568,7 +599,8 @@ def fetch_model( 'ModelRegistryLocation.') destination = Path(local_dir) - with self._registry(connector) as registry: + with self._registry(connector) as registry, \ + self._session(connector) as session: # Resolved before anything on disk is touched: a name the registry # does not know must not cost the caller the bundle already there. model = registry.get_model(location.model_name) @@ -591,7 +623,7 @@ def fetch_model( if staged.exists(): shutil.rmtree(staged) staged.mkdir(parents=True) - version.export(str(staged)) + _export(session, version, staged) bundle = _locate_bundle(staged) metrics_path = bundle / 'metrics.json' diff --git a/packages/kumo-snowflake/tests/test_registry.py b/packages/kumo-snowflake/tests/test_registry.py index 0d0ff6c89..163d6b1be 100644 --- a/packages/kumo-snowflake/tests/test_registry.py +++ b/packages/kumo-snowflake/tests/test_registry.py @@ -4,6 +4,7 @@ """Metrics returned by model downloads come from the exported bundle.""" import json +import re from contextlib import contextmanager from pathlib import Path from types import SimpleNamespace @@ -16,46 +17,85 @@ from kumoapi.plugin import ModelRegistryLocation -@pytest.fixture -def fetch(tmp_path, monkeypatch): - metrics_payload = None +class _VersionFiles: + r"""A session over model version files, as LIST and GET see them.""" + def __init__(self, files): + self.files, self.got = files, [] + + def sql(self, query): + version = re.search(r'/versions/([^/]+)/model/', query).group(1) + return SimpleNamespace(collect=lambda: [{ + 'name': f'versions/{version}/{name}' + } for name in self.files[version]]) + + @property + def connection(self): + return SimpleNamespace(cursor=self._cursor) + + @contextmanager + def _cursor(self): + yield SimpleNamespace(execute=self._get) + + def _get(self, query): + version, name, local = re.fullmatch( + r"GET 'snow://model/DB\.SCHEMA\.CHURN/versions/([^/]+)/(\S+)' " + r"'file://(\S+)'", query).groups() + self.got.append(name) + (Path(local) / Path(name).name).write_bytes(self.files[version][name]) + + +def _version_files(**files): + return { + f'model/models/CHURN/artifacts/bundle/{name}': data + for name, data in { + 'training-spec.json': b'{}', + 'publication-manifest.json': b'{}', + **files + }.items() + } | {'model/code/kumo/__init__.py': b''} + - def export(directory): - bundle = Path(directory) / 'artifacts' / 'bundle' - bundle.mkdir(parents=True) - (bundle / 'training-spec.json').write_text('{}') - (bundle / 'publication-manifest.json').write_text('{}') - (bundle / 'weights.bin').write_bytes(b'model weights') - if metrics_payload is not None: - (bundle / 'metrics.json').write_text(metrics_payload) +def _sink(monkeypatch, model, files): + @contextmanager + def registry(_): + yield SimpleNamespace(get_model=lambda _: model) + + @contextmanager + def session(_): + yield files + + sink = SnowflakeModelRegistryMixin() + monkeypatch.setattr(sink, '_registry', registry) + monkeypatch.setattr(sink, '_session', session, raising=False) + return sink + +@pytest.fixture +def fetch(tmp_path, monkeypatch): + files = _VersionFiles({}) version = SimpleNamespace( model_name='CHURN', version_name='V1', - fully_qualified_model_name='DB.SCHEMA.CHURN', export=export, + fully_qualified_model_name='DB.SCHEMA.CHURN', show_metrics=lambda: { 'kumo_engine': '1', 'kumo_bundle_id': 'b1', 'auc': 0.5 }) model = SimpleNamespace(default=version, version=lambda _: version) - - @contextmanager - def registry(_): - yield SimpleNamespace(get_model=lambda _: model) - - sink = SnowflakeModelRegistryMixin() - monkeypatch.setattr(sink, '_registry', registry) + sink = _sink(monkeypatch, model, files) destination = tmp_path / 'restored' destination.mkdir() (destination / 'previous').write_bytes(b'previous bundle') def run(payload, requested_version): - nonlocal metrics_payload - metrics_payload = payload + extra = {} if payload is None else {'metrics.json': payload.encode()} + files.files['V1'] = _version_files(**{'weights.bin': b'model weights'}, + **extra) return sink.fetch_model( SimpleNamespace(), str(destination), location=ModelRegistryLocation('CHURN', version=requested_version)) + run.files = files return run, destination @@ -70,6 +110,7 @@ def test_fetch_returns_bundle_metrics(fetch, requested_version, metrics): assert published.version == 'V1' assert (destination / 'weights.bin').read_bytes() == b'model weights' assert not (destination / 'previous').exists() + assert not any(name.startswith('model/code/') for name in run.files.got) if payload is not None: assert (destination / 'metrics.json').read_text() == payload @@ -166,15 +207,9 @@ def version(self, version_or_alias): if version_or_alias.upper() in held), version_or_alias) held = self.aliases[name] - def export(directory): - bundle = Path(directory) / 'bundle' - bundle.mkdir() - (bundle / 'training-spec.json').write_text(name) - (bundle / 'publication-manifest.json').write_text('{}') - return SimpleNamespace( version_name=name, model_name='CHURN', - fully_qualified_model_name='DB.SCHEMA.CHURN', export=export, + fully_qualified_model_name='DB.SCHEMA.CHURN', set_alias=lambda alias: held.append(alias.upper()), unset_alias=lambda alias: held.remove(alias.upper()), show_metrics=lambda: { @@ -184,13 +219,11 @@ def export(directory): def _aliasing_sink(monkeypatch, model): - @contextmanager - def registry(_): - yield SimpleNamespace(get_model=lambda _: model) - - sink = SnowflakeModelRegistryMixin() - monkeypatch.setattr(sink, '_registry', registry) - return sink + files = _VersionFiles({ + name: _version_files(**{'training-spec.json': name.encode()}) + for name in model.aliases + }) + return _sink(monkeypatch, model, files) def test_set_alias_moves_it_and_refuses_what_snowflake_cannot(monkeypatch): From 79edaa9f25b7b3720ad4587e2786384622447130 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Thu, 29 Oct 2026 03:55:11 +0000 Subject: [PATCH 7/7] Read stored assets with the identity the loading call was given A model whose text encoder was stored on a /Volumes root loaded inside a kumo.mlops block but raised MissingDependencyError outside it. Assets on a stage, volume or bucket are now read through the connector the call has: load_model's target (a DatabricksConnector, or SageMakerMLflow's identity as an S3Connector), Trainer.load_from_connector's connector, the block root's connector for any location it reads, or the pyfunc's model_config={'assets_from': connector}. When none can, the error names the argument to pass. Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/docs/mlops.md | 2 +- packages/kumo-engine/kumo/engine/_deps.py | 57 ++++++++++++++++--- .../kumo/engine/integrations/_mlflow_model.py | 14 ++++- .../engine/integrations/_mlflow_target.py | 17 ++++-- .../kumo/engine/integrations/mlflow.py | 19 ++++++- .../kumo/engine/trainer/trainer.py | 5 +- .../kumo-engine/test/test_mlflow_target.py | 2 + .../test/test_model_agnostic_mlops.py | 45 +++++++++++++++ 8 files changed, 143 insertions(+), 18 deletions(-) diff --git a/packages/kumo-engine/docs/mlops.md b/packages/kumo-engine/docs/mlops.md index 04fa15613..09ef3c863 100644 --- a/packages/kumo-engine/docs/mlops.md +++ b/packages/kumo-engine/docs/mlops.md @@ -373,7 +373,7 @@ Every artifact carries `_provenance.json`, metadata only (the underscore keeps S ### Artifacts, bundles and dependencies - A model's **bundle** holds only its own state: weights, encoders, the training spec and the holdout. Its manifest's **`parents`** are lineage; its **`dependencies`** are what `predict` needs besides the bundle and a snapshot: the `predictor` entry point, the `kumo-engine` `package` (with `[text]` when a text column is encoded) and content-addressed `asset`s. -- **Assets** are stored once under `/assets///` and shared by every model that uses them. `Trainer.load`, the pyfunc, `run_batch` and the jobs look for each at the location recorded, then under the enclosing block's root, then in the embedding-model store, check its sha256 and cache non-local copies under `~/.cache/kumo/assets/`. A stored asset that cannot be found raises `MissingDependencyError`. +- **Assets** are stored once under `/assets///` and shared by every model that uses them. `Trainer.load`, the pyfunc, `run_batch` and the jobs look for each at the location recorded, then under the enclosing block's root, then in the embedding-model store, check its sha256 and cache non-local copies under `~/.cache/kumo/assets/`. An asset on a stage, a volume or a bucket is read with the identity the loading call was given: `Trainer.load_from_connector(connector, ...)`, `load_model(uri, target=DatabricksConnector(...) or SageMakerMLflow(arn, identity=...))`, the block root's connector, or `mlflow.pyfunc.load_model(uri, model_config={"assets_from": connector})`. A stored asset that cannot be found raises `MissingDependencyError` naming the argument that passes such a connector. - **Text encoders** are the first assets: word-vector models (`GloVeVectors`) are stored; other `TokenEmbeddingModel`s are recorded by fingerprint and must be registered wherever the model loads: ```python diff --git a/packages/kumo-engine/kumo/engine/_deps.py b/packages/kumo-engine/kumo/engine/_deps.py index eab991b4a..60d8fdbbe 100644 --- a/packages/kumo-engine/kumo/engine/_deps.py +++ b/packages/kumo-engine/kumo/engine/_deps.py @@ -15,6 +15,7 @@ from pathlib import Path from typing import ( Any, + Callable, Iterator, List, Mapping, @@ -40,18 +41,30 @@ Location = Union[str, Tuple[Any, str]] +#: How to name the connector that reads an asset where it is stored, by +#: location prefix; and, per loading call, which argument passes it. +Hint = Callable[[str, str], str] +READERS = (('@', 'SnowflakeConnector'), ('/Volumes/', 'DatabricksConnector'), + ('s3://', 'S3Connector'), ('gs://', 'GCSConnector')) + _LOCATIONS: ContextVar[Tuple[Location, ...]] = ContextVar( 'kumo_asset_locations', default=()) +_HINT: ContextVar[Optional[Hint]] = ContextVar('kumo_asset_hint', + default=None) @contextmanager -def asset_locations(locations: Sequence[Location]) -> Iterator[None]: +def asset_locations(locations: Sequence[Location], + hint: Optional[Hint] = None) -> Iterator[None]: r"""Asset roots that dependencies resolved in the block look in - first.""" + first; ``hint(reader, location)`` says which argument passes a + connector that reads an asset none of them holds.""" token = _LOCATIONS.set((*locations, *_LOCATIONS.get())) + hinted = _HINT.set(hint or _HINT.get()) try: yield finally: + _HINT.reset(hinted) _LOCATIONS.reset(token) @@ -146,6 +159,7 @@ def __init__(self, records: Sequence[Dependency] = (), self.errors: List[BaseException] = [] self._settings = settings self._locations = [*locations, *_LOCATIONS.get()] + self._hint = _HINT.get() @classmethod def of(cls, manifest: Optional[ModelManifest], settings: Any = None, @@ -196,6 +210,14 @@ def resolve(self, asset_kind: str, name: str) -> Path: continue if directory is not None and _verified(directory, record): return directory + location = record.location or '' + reader = next((name for prefix, name in READERS + if location.startswith(prefix)), None) + if reader is not None and self._hint is not None: + raise MissingDependencyError( + f"The model needs {asset_kind} {name!r} (sha256 " + f"{record.sha256}), stored at {location}, which this " + f"process cannot read: {self._hint(reader, location)}.") raise MissingDependencyError( f"The model needs {asset_kind} {name!r} (sha256 " f"{record.sha256}), which is in none of " @@ -225,20 +247,39 @@ def candidates(self, record: Dependency) -> List[Tuple[Any, str]]: if settings is not None: found.append((settings.root.connector, settings.root.location(f'assets/{relative}'))) + if record.location and reads(settings.root.connector, + record.location): + found.append((settings.root.connector, record.location)) return found +def reads(connector: Any, location: str) -> bool: + r"""Whether ``connector`` reads the stored ``location`` with its own + identity: a stage, a volume, or its object store's URI.""" + from kumo.engine.connector import ( + FileConnector, + GCSConnector, + S3Connector, + is_file_destination, + ) + + for kind, prefix in ((S3Connector, 's3://'), (GCSConnector, 'gs://'), + (FileConnector, None)): + if isinstance(connector, kind): + return prefix is not None and location.startswith(prefix) + return (location.startswith(('@', '/Volumes/')) + and is_file_destination(connector, location)) + + def connector_roots(manifest: Optional[ModelManifest], connector: Any) -> List[Tuple[Any, str]]: - r"""The stage and volume asset roots of ``manifest`` that ``connector`` - reads, which no URI reaches.""" - from kumo.engine.connector import is_file_destination - + r"""The asset roots of ``manifest`` that ``connector`` reads.""" + if connector is None: + return [] return [(connector, root) for root in sorted({ record.location.rsplit('/', 2)[0] for record in Dependencies.of(manifest).records - if (record.location or '').startswith(('@', '/Volumes/')) - and is_file_destination(connector, record.location) + if record.location and reads(connector, record.location) })] diff --git a/packages/kumo-engine/kumo/engine/integrations/_mlflow_model.py b/packages/kumo-engine/kumo/engine/integrations/_mlflow_model.py index b5589cfda..311340196 100644 --- a/packages/kumo-engine/kumo/engine/integrations/_mlflow_model.py +++ b/packages/kumo-engine/kumo/engine/integrations/_mlflow_model.py @@ -24,9 +24,21 @@ class KumoModel(_pyfunc.PythonModel): def load_context(self, context: Any) -> None: + from pathlib import Path + + from kumo.engine._deps import asset_locations, connector_roots from kumo.engine._predictor import load_predictor + from kumo.engine.trainer._provenance import read_provenance - self._predictor = load_predictor(context.artifacts['bundle']) + bundle = context.artifacts['bundle'] + with asset_locations( + connector_roots(read_provenance(Path(bundle)), + context.model_config.get('assets_from')), + lambda reader, _: ( + f"pass model_config={{'assets_from': {reader}(...)}} to " + f"mlflow.pyfunc.load_model(model_uri, ...), with an " + f"identity that reads it")): + self._predictor = load_predictor(bundle) self._columns = list(context.model_config['columns']) self._snapshots: Dict[str, GraphSnapshot] = {} self._default = context.model_config.get('snapshot') diff --git a/packages/kumo-engine/kumo/engine/integrations/_mlflow_target.py b/packages/kumo-engine/kumo/engine/integrations/_mlflow_target.py index 62ecac2aa..8fc564fe3 100644 --- a/packages/kumo-engine/kumo/engine/integrations/_mlflow_target.py +++ b/packages/kumo-engine/kumo/engine/integrations/_mlflow_target.py @@ -57,8 +57,10 @@ class Target: def __init__(self, tracking_uri: str, registry_uri: Optional[str], swap: Optional[Callable[[], ContextManager[Any]]] = None, - aws_session: Any = None, host: Optional[str] = None) -> None: + aws_session: Any = None, host: Optional[str] = None, + assets: Any = None) -> None: self.tracking_uri = tracking_uri + self.assets = assets self.registry_uri = registry_uri self._swap = swap self._aws_session = aws_session @@ -169,12 +171,13 @@ def _databricks(connector: Any) -> Target: # Unity Catalog registration downloads the logged model through the # process-wide tracking URI, so it must point at the workspace. return Target('databricks', 'databricks-uc', - swap=lambda: global_tracking_uri('databricks')) + swap=lambda: global_tracking_uri('databricks'), + assets=connector) client = import_provider('kumo_databricks.session', 'databricks')._client_for(principal) return Target('databricks', 'databricks-uc', swap=lambda: _databricks_token(client), - host=client.config.host) + host=client.config.host, assets=connector) @contextmanager @@ -204,15 +207,19 @@ def _databricks_token(client: Any) -> Iterator[None]: def _sagemaker(target: SageMakerMLflow) -> Target: from kumoapi.data_source import S3AmbientAuthentication + from kumo.engine.connector import S3Connector + arn = target.tracking_server_arn + assets = S3Connector(identity=target.identity) principal = _principal(target.identity) if principal is None or isinstance(principal.config.authentication, S3AmbientAuthentication): - return Target(arn, arn, swap=lambda: global_tracking_uri(arn)) + return Target(arn, arn, swap=lambda: global_tracking_uri(arn), + assets=assets) session = import_provider('kumo_s3.backend', 's3')._session_for(principal) return Target(arn, arn, swap=lambda: _sagemaker_session(arn, session), - aws_session=session) + aws_session=session, assets=assets) @contextmanager diff --git a/packages/kumo-engine/kumo/engine/integrations/mlflow.py b/packages/kumo-engine/kumo/engine/integrations/mlflow.py index 58d7a769f..8ed7f722b 100644 --- a/packages/kumo-engine/kumo/engine/integrations/mlflow.py +++ b/packages/kumo-engine/kumo/engine/integrations/mlflow.py @@ -674,7 +674,9 @@ def load_model( r"""Restore the `Trainer` of a model `log_model` logged, from the server ``target`` reaches, as in `MlflowTracking` (`None`: MLflow's own tracking URI).""" + from kumo.engine._deps import asset_locations, connector_roots from kumo.engine._scratch import new_run_id, trainer_scratch_dir + from kumo.engine.trainer._provenance import read_provenance from kumo.engine.trainer.trainer import Trainer, warn_if_other_bundle mlflow = import_provider('mlflow', 'mlflow') @@ -685,12 +687,24 @@ def load_model( local = Path(_download(mlflow, target, model_uri, destination)) model = mlflow.models.Model.load(str(local)) bundle = model.flavors['python_function']['artifacts']['bundle']['path'] - trainer = Trainer.load(local / bundle, embedding_models, - accelerator=accelerator) + with asset_locations( + connector_roots(read_provenance(local / bundle), target.assets), + _target_hint): + trainer = Trainer.load(local / bundle, embedding_models, + accelerator=accelerator) warn_if_other_bundle(model.metadata, trainer) return trainer +def _target_hint(reader: str, location: str) -> str: + target = { + 'DatabricksConnector': 'DatabricksConnector(identity=...)', + 'S3Connector': 'SageMakerMLflow(arn, identity=S3Identity(...))', + }.get(reader, f'a {reader}') + return (f'pass target={target} to load_model(model_uri, ...), with an ' + f'identity that reads {location.rsplit("/", 2)[0]}') + + def _download(mlflow: Any, target: Any, model_uri: str, destination: Path) -> str: if target.explicit and model_uri.startswith('models:/'): @@ -762,6 +776,7 @@ def _log_model(target: Any, client: Any, bundle_dir: Path, 'columns': [] if signature is None else signature.outputs.input_names(), 'snapshot': serving_snapshot, + 'assets_from': None, }, artifacts={ 'bundle': str(bundle_dir), diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index bb052523f..7dabee110 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -1735,7 +1735,10 @@ def load_from_connector( from kumo.engine.trainer._provenance import read_provenance with asset_locations( - connector_roots(read_provenance(bundle_dir), model_registry)): + connector_roots(read_provenance(bundle_dir), model_registry), + lambda reader, _: ( + f'pass a {reader}(...) with an identity that reads it ' + f'as Trainer.load_from_connector(model_registry, ...)')): trainer = cls._load_bundle(bundle_dir, embedding_models, accelerator=accelerator) warn_if_other_bundle(getattr(published, 'metadata', None), trainer) diff --git a/packages/kumo-engine/test/test_mlflow_target.py b/packages/kumo-engine/test/test_mlflow_target.py index 4b9b360a8..79453009a 100644 --- a/packages/kumo-engine/test/test_mlflow_target.py +++ b/packages/kumo-engine/test/test_mlflow_target.py @@ -70,8 +70,10 @@ def test_ambient_identities_set_no_credentials(monkeypatch, tracking_uri): assert mlflow.get_tracking_uri() == tracking_uri assert dict(os.environ) == before + assert resolve(_connector(DatabricksIdentity())).assets.catalog == 'main' for target in (resolve(SageMakerMLflow(ARN, S3Identity())), resolve(SageMakerMLflow(ARN))): + assert type(target.assets).__name__ == 'S3Connector' with target.scope(): assert mlflow.get_tracking_uri() == ARN assert target.artifact_repository('s3://b/run') is None diff --git a/packages/kumo-engine/test/test_model_agnostic_mlops.py b/packages/kumo-engine/test/test_model_agnostic_mlops.py index 2d7b851e0..bc76b0dba 100644 --- a/packages/kumo-engine/test/test_model_agnostic_mlops.py +++ b/packages/kumo-engine/test/test_model_agnostic_mlops.py @@ -423,6 +423,51 @@ def test_staged_assets_load_through_the_loading_connector(tmp_path): (stage, f'@DB.S.ART/root/assets/toy/{sha}')] +def test_volume_assets_load_through_the_connector_the_loader_was_given( + tmp_path, monkeypatch): + from pyarrow.fs import LocalFileSystem, SubTreeFileSystem + + from kumo.engine.integrations import _mlflow_model + from kumo.engine.integrations.mlflow import _target_hint + + monkeypatch.setenv('XDG_CACHE_HOME', str(tmp_path / 'cache')) + (tmp_path / 'volume').mkdir() + files = SubTreeFileSystem(str(tmp_path / 'volume'), LocalFileSystem()) + volume = SimpleNamespace(sink_backend=SimpleNamespace( + holds_files=lambda connector, path: path.startswith('/Volumes/'), + filesystem=lambda connector, path: (files, path.lstrip('/')))) + (tmp_path / 'vocab').mkdir() + (tmp_path / 'vocab' / 'words.txt').write_text('a b c') + recorded = Dependencies(settings=kumo.MLOps((volume, '/Volumes/c/s/v/a'))) + record = recorded.require(FileAsset('toy', 'vocab', tmp_path / 'vocab')) + assert record.location.startswith('/Volumes/c/s/v/a/assets/toy/') + manifest = ModelManifest(id='m', dependencies=[record]) + + with _deps.asset_locations([], _target_hint): + with pytest.raises(MissingDependencyError, + match=r'pass target=DatabricksConnector\(' + r'identity=\.\.\.\) to load_model\(model_uri, ' + r'\.\.\.\), with an identity that reads ' + r'/Volumes/c/s/v/a/assets\.$'): + Dependencies.of(manifest).resolve('toy', 'vocab') + with _deps.asset_locations(_deps.connector_roots(manifest, volume)): + assert (Dependencies.of(manifest).resolve('toy', 'vocab') / + 'words.txt').read_text() == 'a b c' + other_root = kumo.MLOps((volume, '/Volumes/c/s/v/b')) + assert Dependencies.of(manifest, other_root).resolve('toy', 'vocab') + + seen = [] + monkeypatch.setattr( + 'kumo.engine._predictor.load_predictor', lambda bundle: seen.append( + Dependencies.of(manifest).candidates(record)[0])) + monkeypatch.setattr('kumo.engine.trainer._provenance.read_provenance', + lambda bundle: manifest) + _mlflow_model.KumoModel().load_context(SimpleNamespace( + artifacts={'bundle': str(tmp_path)}, + model_config={'columns': [], 'assets_from': volume})) + assert seen == [(volume, record.location)] + + def test_a_stage_root_registers_versions_that_serve_its_snapshot(tmp_path): from kumo.engine.trainer.trainer import _connector_registry, _publish