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 355ca19ed..6e9ca4454 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). It serves with the torch and CUDA build installed where it registers, so install one the GPU pool runs first, e.g. `pip install --index-url https://download.pytorch.org/whl/cu126 "torch==2.12.*"`. 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. A binary model always returns `True_PROB`; with no threshold passed or stored, its `TARGET_PRED` is null. `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. A binary model always returns `True_PROB`; with no threshold passed or stored, its `TARGET_PRED` is null. `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. @@ -188,10 +188,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 @@ -218,13 +218,34 @@ A model registered through `MlflowTracking` is an MLflow pyfunc: it serves with model = mlflow.pyfunc.load_model("models:/main.ml.users_churn@champion") scores = model.predict(pd.DataFrame({"ENTITY": [1, 2, 3]}), params={ "snapshot": cfg.snapshot_uri(snapshot), # "/Volumes/main/ml/kumo/snapshots//" - "prediction_time": "2024-06-01"}) # temporal queries + "prediction_time": "2024-06-01"}) # temporal queries need it, or a TIMESTAMP column ``` -- Input rows have an `ENTITY` column and optionally a `TIMESTAMP`; scores come back in input order, one per scored row. Binary classification returns `True_PROB`; `TARGET_PRED` needs `params={"binary_classification_threshold": ...}` or a stored threshold, and is null without one. +- Input rows have an `ENTITY` column and optionally a `TIMESTAMP`; a temporal query needs one of `TIMESTAMP` or `prediction_time` (the anchor of a batch run is `run.manifest.prediction_table.anchors[0]`). Scores come back in input order, one per scored row. Binary classification returns `True_PROB`; `TARGET_PRED` needs `params={"binary_classification_threshold": ...}` or a stored threshold, and is null without one. - 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. +- A model with stored assets (word vectors for text columns) on a volume, stage or bucket reads them with the identity it is given: `mlflow.pyfunc.load_model(uri, model_config={"assets_from": connector})`, a connector for that store with its credentials; an asset it cannot find raises, naming this option. - 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; `predict(snapshot)` defaults to the snapshot's latest time, so pass `run.manifest.prediction_table.anchors[0]` to reproduce a local run), `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`. @@ -327,7 +348,7 @@ snowflake = SnowflakeTracking(SnowflakeConnector(warehouse="", database="ML" ### Install -- A bare `kumo-engine[mlflow]` with `--find-links ` can resolve to an unrelated `kumo-engine` on PyPI; install the engine wheel by path: `pip install --find-links "/[mlflow]"`. +- Install the engine wheel by path, `pip install "/[mlflow]"`: a bare `kumo-engine[mlflow]` can resolve to an unrelated `kumo-engine` on PyPI. The wheel `scripts/build_wheel.sh` builds bundles kumo-core, kumo-api, kumo-pql, DiskGraph and the connector packages; with separately built wheels, add `--find-links `. - **GPU:** on a plain GPU VM or SageMaker image, install a torch build that matches the NVIDIA driver first, e.g. `pip install --index-url https://download.pytorch.org/whl/cu126 "torch==2.12.*" torchvision`. PyPI's default torch is built for a recent CUDA (`torch.version.cuda`; CUDA 13 needs driver ≥ 580); cu126 runs on ≥ 525.60. With `accelerator="auto"`, a torch that cannot use the driver runs on CPU: check `torch.cuda.is_available()`. - **Databricks:** Runtime 17.3 LTS ML or later keeps its torch, torchvision and pyarrow (kumo-engine accepts pyarrow 16 through 24; a custom CA bundle for S3 needs 21 or later). On 16.4 LTS ML pip replaces torch and breaks torchvision. The parent folder of an experiment path must exist. Registering to Unity Catalog from outside Databricks also needs `boto3`. - **Disk:** `kumo.configure(scratch_dir=...)` holds snapshots, tables and bundles; Python's own temporary files (multiprocessing) follow `TMPDIR`, so point both at a large local disk. @@ -346,7 +367,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 @@ -360,7 +381,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/`. 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 @@ -409,6 +430,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; columns are listed as string literals in their stored case, `('True_PROB')` (`'"True_PROB"'` is not found). 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 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` | 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` | + +- **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/_deps.py b/packages/kumo-engine/kumo/engine/_deps.py index f8d6b5713..60d8fdbbe 100644 --- a/packages/kumo-engine/kumo/engine/_deps.py +++ b/packages/kumo-engine/kumo/engine/_deps.py @@ -10,9 +10,13 @@ import os import shutil import tempfile +from contextlib import contextmanager +from contextvars import ContextVar from pathlib import Path from typing import ( Any, + Callable, + Iterator, List, Mapping, Optional, @@ -37,6 +41,33 @@ 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], + hint: Optional[Hint] = None) -> Iterator[None]: + r"""Asset roots that dependencies resolved in the block look in + 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) + + class Asset(Protocol): r"""A file or directory a model needs at predict time, addressed by its content. ``files`` maps names to local files; the class's @@ -127,7 +158,8 @@ def __init__(self, records: Sequence[Dependency] = (), self.records: List[Dependency] = list(records) self.errors: List[BaseException] = [] self._settings = settings - self._locations = list(locations) + self._locations = [*locations, *_LOCATIONS.get()] + self._hint = _HINT.get() @classmethod def of(cls, manifest: Optional[ModelManifest], settings: Any = None, @@ -178,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 " @@ -197,7 +237,7 @@ def candidates(self, record: Dependency) -> List[Tuple[Any, str]]: else: connector, path = location found.append((connector, f"{path.rstrip('/')}/{relative}")) - if record.location: + if record.location and not record.location.startswith('@'): parent, _, leaf = record.location.rstrip('/').rpartition('/') try: found.append((connector_for_root(parent or '/'), leaf)) @@ -207,9 +247,42 @@ 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 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 and reads(connector, record.location) + })] + + def asset_cache() -> Path: base = os.environ.get('XDG_CACHE_HOME') or Path.home() / '.cache' return Path(base) / 'kumo' / 'assets' diff --git a/packages/kumo-engine/kumo/engine/_predictor.py b/packages/kumo-engine/kumo/engine/_predictor.py index b65ebb906..402e5f34a 100644 --- a/packages/kumo-engine/kumo/engine/_predictor.py +++ b/packages/kumo-engine/kumo/engine/_predictor.py @@ -56,7 +56,7 @@ def load(cls, bundle_dir: Union[str, Path]) -> 'TrainerPredictor': def predict(self, rows: Any, snapshot: Any, params: Mapping[str, Any]) -> Union[pd.DataFrame, Tuple[str, ...]]: - from kumo.engine.integrations._mlflow_model import predict + from kumo.engine.integrations._serving import predict from kumo.engine.pquery.prediction_table import PredictionTable if isinstance(rows, PredictionTable): diff --git a/packages/kumo-engine/kumo/engine/integrations/_mlflow_model.py b/packages/kumo-engine/kumo/engine/integrations/_mlflow_model.py index 638d0a313..311340196 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 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', @@ -29,9 +24,21 @@ class KumoModel(_pyfunc.PythonModel): def load_context(self, context: Any) -> None: - from kumo.engine._predictor import load_predictor + from pathlib import Path - self._predictor = load_predictor(context.artifacts['bundle']) + from kumo.engine._deps import asset_locations, connector_roots + from kumo.engine._predictor import load_predictor + from kumo.engine.trainer._provenance import read_provenance + + 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') @@ -58,108 +65,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, 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, - time: Optional[pd.Timestamp]) -> None: - from kumo_core.pquery import query_is_temporal - - from kumo.engine.pquery import PredictiveQuery - - if len(rows) == 0 or time is not None or 'TIMESTAMP' in rows: - return - parsed = PredictiveQuery._from_snapshot(snapshot, - trainer._training_query)._parsed - if query_is_temporal(parsed, snapshot._native_graph): - raise ValueError("temporal query: pass params={'prediction_time': " - "...} or a TIMESTAMP column") - - -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/_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/_serving.py b/packages/kumo-engine/kumo/engine/integrations/_serving.py new file mode 100644 index 000000000..42ddbdd3f --- /dev/null +++ b/packages/kumo-engine/kumo/engine/integrations/_serving.py @@ -0,0 +1,233 @@ +# 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() +_predictors: 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, 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, + time: Optional[pd.Timestamp]) -> None: + from kumo_core.pquery import query_is_temporal + + from kumo.engine.pquery import PredictiveQuery + + if len(rows) == 0 or time is not None or 'TIMESTAMP' in rows: + return + parsed = PredictiveQuery._from_snapshot(snapshot, + trainer._training_query)._parsed + if query_is_temporal(parsed, snapshot._native_graph): + raise ValueError("temporal query: pass params={'prediction_time': " + "...} or a TIMESTAMP column") + + +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 predictor the bundle at ``bundle`` names, + against the snapshot in ``params`` or the ``serving`` record, with the + bundle's assets read from their stage. The bundle and each snapshot + load once per process.""" + from kumo.engine._predictor import load_predictor + + record = json.loads(Path(serving).read_text()) + columns = record.get('columns') or 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: fit with a stage artifact root, " + "which stores the training snapshot there, or pass " + "InputSpec(params={'snapshot': '@DB.SCHEMA.STAGE/...'})") + predictor = _cached(_predictors, bundle, + lambda: _staged(bundle, record, load_predictor)) + snapshot = _cached(_snapshots, uri, lambda: load_snapshot(uri, record)) + trainer = getattr(predictor, 'trainer', None) + if trainer is not None: + rows = rows.assign( + ENTITY=_entities(rows['ENTITY'], trainer, snapshot)) + frame = predictor.predict( + rows, snapshot, {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 _staged(bundle: str, target: Mapping[str, Any], + load: Callable[[str], Any]) -> Any: + r"""``load(bundle)``, with the stages holding the bundle's assets + readable with the Snowpark Container Services identity.""" + from kumo.engine._deps import Dependencies, asset_locations + from kumo.engine.trainer._provenance import read_provenance + + roots = sorted({ + record.location.rsplit('/', 2)[0] + for record in Dependencies.of(read_provenance(Path(bundle))).records + if (record.location or '').startswith('@') + }) + if not roots: + return load(bundle) + from kumo.engine.connector.snowflake import SnowflakeConnector + + connector = SnowflakeConnector.from_spcs_env( + warehouse=target.get('warehouse'), database=target.get('database'), + schema_name=target.get('schema')) + with asset_locations([(connector, root) for root in roots]): + return load(bundle) + + +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..908ec1e70 --- /dev/null +++ b/packages/kumo-engine/kumo/engine/integrations/_snowflake_model.py @@ -0,0 +1,275 @@ +# 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') + +#: The CUDA a version serves with when this installation's torch has none. +CPU_TORCH_CUDA = '12.6' +DEFAULTS = { + '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.integrations._serving import output_columns + 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), + 'columns': output_columns(bundle), + **target + })) + options = { + **DEFAULTS, + **{ + key[len(SERVING_OPTION_PREFIX):]: value + for key, value in metadata.items() + 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 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 " + 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 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.""" + 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/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/integrations/snowflake.py b/packages/kumo-engine/kumo/engine/integrations/snowflake.py index 080df1d3c..148501af6 100644 --- a/packages/kumo-engine/kumo/engine/integrations/snowflake.py +++ b/packages/kumo-engine/kumo/engine/integrations/snowflake.py @@ -58,6 +58,10 @@ class SnowflakeTracking: tags: extra tags for the fit run. native_model_lineage: keep the run open as a context for the whole fit, so models registered meanwhile are linked to it. + serving_options: how a version ``kumo.MLOps(registry=...)`` + registers runs in Snowflake (``cuda_version``, + ``python_version``, ``function_type``, ``pip_repository``; see + ``ModelOutputConfig``). ``run_id`` is the most recent tracked run's name (Snowflake runs are identified by name), once it has started. @@ -76,6 +80,7 @@ def __init__( run_name: Optional[str] = None, tags: Optional[Mapping[str, str]] = None, native_model_lineage: bool = False, + serving_options: Optional[Mapping[str, str]] = None, ) -> None: from kumo.engine.connector.base import Connector @@ -92,6 +97,7 @@ def __init__( self.run_name = run_name self.tags = dict(tags or {}) self.native_model_lineage = native_model_lineage + self.serving_options = dict(serving_options or {}) or None self.run_id: Optional[str] = None self._reset() diff --git a/packages/kumo-engine/kumo/engine/trainer/output_config.py b/packages/kumo-engine/kumo/engine/trainer/output_config.py index 160e9c834..a975f8d7a 100644 --- a/packages/kumo-engine/kumo/engine/trainer/output_config.py +++ b/packages/kumo-engine/kumo/engine/trainer/output_config.py @@ -4,7 +4,7 @@ """Prediction, training table and model export destinations.""" from dataclasses import dataclass, field -from typing import List, Mapping, Optional, cast +from typing import Any, List, Mapping, Optional, cast from kumo_core.sources import layout from kumo_core.sources.object_store_sessions import ObjectStoreSessions @@ -158,17 +158,44 @@ class SnapshotOutputConfig: path: str +#: Publication metadata a registry may act on: the serving snapshot, and +#: each serving option under ``kumo.serving.