diff --git a/docs/development/testing.md b/docs/development/testing.md index 844411696..e6f3d5531 100644 --- a/docs/development/testing.md +++ b/docs/development/testing.md @@ -23,6 +23,14 @@ Run from the repository root unless a component guide says otherwise. Python package test suites have different roots and fixtures; inspect the package's configuration before invoking pytest over multiple packages. +kumo-engine and kumo-core tests load the DiskGraph native module +(`diskgraph/_c.abi3.so`) and its `bin/ingest` binary from the installed +`diskgraph` package. Both must be built from the same tree as the Python code +under test: a build from an older checkout fails with missing attributes or +arguments (for example `RuntimeIndex.iter(projections=...)` or +`compute_col_stats_for_parquet(percentiles=...)`). After pulling DiskGraph +changes, rebuild with `make install-diskgraph`, which builds and installs both. + For connector changes, start with the [public test command](public-tests.md). It builds its prerequisites and runs without cloud credentials. The [test matrix](../connector-test-matrix.md) lists coverage and CI limits. diff --git a/examples/mlops/README.md b/examples/mlops/README.md new file mode 100644 index 000000000..a63be837c --- /dev/null +++ b/examples/mlops/README.md @@ -0,0 +1,69 @@ +# Scheduled kumo-engine jobs + +Two plain Python job scripts that run the normal `kumo.engine` API inside +`with kumo.mlops(kumo.MLOps(...)):`, and the wiring that runs them on a +schedule. Both take their settings as command-line flags. + +- `fit_job.py` builds a snapshot, generates the users-churn training table and + trains the model. Inside the block each artifact is stored under the + artifact root (`snapshots/…`, `training_tables//`, `models//`). + The model carries its `--threshold` (0.5 by default) as scoring defaults. + With a registry, the model is registered and the `--alias` alias + (`champion` by default) moves to it. It prints `model: models/`. +- `score_job.py` loads the version `--alias` points at (or `--model-path`) + through `Trainer.load_from_connector(cfg, ...)`, scores a fresh snapshot + (or the stored one `--snapshot` names) and writes `--predictions-table` when + given. Inside the block the run is stored under `predictions//` with its + `_provenance.json`. +- `jobs.py` holds what both share: the flags, the graph (users, items, orders), + the connectors and the `kumo.MLOps` settings. Replace it with your own module; + list text columns in `TEXT_COLUMNS` (e.g. `('users.BIO',)`), since inference + reads short text as categorical. With `--on-error raise`, a score job whose + graph differs from the one the model was trained on fails before scoring. + +| Directory | Platform | Registry | Schedule | +|---|---|---|---| +| [databricks](databricks/README.md) | Databricks Jobs (Asset Bundle) | Unity Catalog through MLflow | train on demand, score daily | +| [sagemaker](sagemaker/README.md) | SageMaker Pipelines | SageMaker managed MLflow | fit step, then score step | +| [snowflake](snowflake/README.md) (experimental, preview APIs) | Snowflake ML Jobs + Tasks | Snowflake Model Registry, Snowflake Experiment Tracking | train weekly, score daily | + +## Flags + +| Flag | Meaning | +|---|---| +| `--artifact-root` | `s3://`, `gs://`, `file://`, a local or `/Volumes` path, or with `--source snowflake` an `@DB.SCHEMA.STAGE/...` path; required | +| `--source` | `databricks`, `snowflake`, or a directory holding `users`, `items` and `orders` Parquet files | +| `--catalog`, `--schema`, `--warehouse` | Databricks catalog, schema and SQL warehouse id, or Snowflake database, schema and warehouse | +| `--table-prefix` | prefix of the `users`, `items` and `orders` table or file names, e.g. `demo_` for `demo_users`; default none | +| `--registry` | `mlflow` (`MLOps(tracking=MlflowTracking(), registry=ModelRegistryLocation(...))`) or `snowflake` (`MLOps(tracking=SnowflakeTracking(connector, ...), registry=ModelRegistryLocation(...))`); omit for no registry | +| `--model-name` | registered model name, with `--registry`: `users-churn` on SageMaker (letters, digits and hyphens only), `..users_churn` in Unity Catalog, `USERS_CHURN` in Snowflake | +| `--alias` | with `--registry`: the alias the fit moves to its new version and the score job loads; default `champion`. With MLflow, register under `candidate` and promote yourself to gate scoring | +| `--model-path` | score only, without a registry: `models/` under the artifact root | +| `--snapshot` | score only: the id of a snapshot stored under the artifact root | +| `--predictions-table` | score only: a table of the source connector to write | +| `--query`, `--num-experiments`, `--threshold` | fit only: the predictive query, the number of experiments and the classification threshold the model scores with | +| `--scratch-dir` | local disk for snapshots, tables and bundles | +| `--on-error` | `raise` (default): a failed tracking call or registration fails the job after its result is stored (`kumo.MLOps(on_error="raise")`); `warn`: the job warns and exits 0 | + +## Run locally + +With `users.parquet`, `items.parquet` and `orders.parquet` in a directory: + +```bash +cd examples/mlops +python fit_job.py --source --artifact-root +python score_job.py --source --artifact-root \ + --model-path models/ +``` + +The query needs several 30-day windows of order history; for a small sample +pass a shorter `--query` and `--num-experiments 1`. + +## What every scheduler sees + +The scripts are ordinary Python: they exit 0 on success and non-zero with a +traceback on failure, and log to stderr. A failed registration or tracking +call is a failure too (`--on-error raise`), so a fit that stored its model but +did not register it does not leave the score job on the previous version +unnoticed. Databricks Jobs, SageMaker Pipelines +and Snowflake Tasks keep their own run ids, status and logs. diff --git a/examples/mlops/databricks/README.md b/examples/mlops/databricks/README.md new file mode 100644 index 000000000..9f9f41ef8 --- /dev/null +++ b/examples/mlops/databricks/README.md @@ -0,0 +1,103 @@ +# Databricks Jobs + +`databricks.yml` is an Asset Bundle with two jobs on a single-node job +cluster, a GPU node with the GPU ML runtime by default. Each task is a `spark_python_task` running a job script with its +flags as `parameters`; the cluster's `spark_env_vars` hold only MLflow's own +settings. + +- **`kumo_train`** runs `fit_job.py` on the `users`, `items` and `orders` + tables (each prefixed with `table_prefix`) in `.` inside + `kumo.mlops(...)` with the artifact root + `/Volumes////kumo`: the snapshot, training table + and model are stored under it. `--registry mlflow` logs the fit to + `MLFLOW_EXPERIMENT_NAME`, registers the model as a new version of + `..` (default `users_churn`) and moves the + `alias` variable's alias (default `candidate`) to it; move `champion` to it + once it is better (see + [Promote a model](../../../packages/kumo-engine/docs/mlops.md#promote-a-model)). + Set `alias` to `champion` to serve every fit without a gate. +- **`kumo_score`** runs `score_job.py` daily at 03:00 UTC (created paused). + It loads the version `score_alias` (default `champion`) points at in + `..`, builds a fresh snapshot and appends to + `` (default `users_churn_predictions`) in the + connector's schema. Add `--snapshot ` to score a stored + snapshot instead. + +## Before deploying + +1. Put the kumo-engine wheel and a `requirements.txt` for its dependencies + (kumo-core, kumo-api, kumo-pql, diskgraph, kumo-databricks, mlflow-skinny) + on a volume and set `kumo_engine_wheel` and `requirements`. Each task + installs the requirements, then the wheel. Start the requirements file with + `--prefer-binary`, so pip installs wheels rather than building sdists. + - **Runtime:** Databricks Runtime 17.3 LTS ML or later: the GPU flavour + (`17.3.x-gpu-ml-scala2.13`) on a GPU node such as `g5.xlarge`, or the + CPU flavour (`17.3.x-cpu-ml-scala2.13`) on a node type with local NVMe + disks such as `m5d.xlarge`. These runtimes already carry a torch that + kumo-core accepts, so pip keeps the runtime's torch, torchvision, + MLflow and cryptography; leave torch out of the requirements file. + kumo-engine accepts pyarrow 16 through 24, so the runtime's pyarrow + (19 on 17.3) is kept too. On 16.4 LTS ML, whose torch 2.6 is + older than kumo-core needs, pip replaces torch and breaks the + preinstalled torchvision. +2. Replace the `<...>` placeholders in `databricks.yml` (or pass + `--var catalog=...`). `model_name`, `predictions_table`, `table_prefix`, + `experiment` and `on_error` have defaults; set them to run beside existing + objects in the same schema. The parent folder of the MLflow experiment + path must already exist in the workspace. +3. Adjust the `dev` target (workspace `root_path`, under your user folder by + default) or add your own, then + `databricks bundle validate` and `databricks bundle deploy -t `. The bundle syncs this + whole `examples/mlops` directory (`sync.paths: [..]`). + +## Privileges for the run-as principal + +- `USE CATALOG`, `USE SCHEMA` and `SELECT` on the three tables; `CREATE TABLE` + and `MODIFY` on the schema for predictions. +- `READ VOLUME`/`WRITE VOLUME` on ``, and `CREATE VOLUME` on the + schema (or a pre-created `.._kumo_staging` volume) for + staged table writes. +- `CAN USE` on the SQL warehouse; edit permission on the MLflow experiment. +- `CREATE MODEL` on the schema and `EXECUTE` on the model; moving an alias + needs rights on the model itself. + +## Notes + +- **Strict mode.** Both jobs pass `--on-error raise` (variable `on_error`): + if MLflow tracking or Unity Catalog registration fails, the task fails after + the model or predictions are stored, and `champion` stays where it was. + +- **Script, not `python_wheel_task`.** Databricks overrides a custom + `PYTHONPATH` at driver startup, so a wheel task cannot import `jobs.py`; the + scripts put their own directory on `sys.path` (from `sys.argv[0]`: a + `spark_python_task` runs the file without `__file__`). +- **Dedicated access mode** (`SINGLE_USER`): standard access mode runs as a + low-privilege user and has no GPUs. `--scratch-dir` is + `/local_disk0/kumo`. +- **MLflow.** `MLFLOW_TRACKING_URI=databricks` and + `MLFLOW_REGISTRY_URI=databricks-uc` are set explicitly; + `MLFLOW_EXPERIMENT_NAME` is required in jobs. +- **Job clusters.** Both jobs repeat the cluster block; YAML anchors would + remove the copy but are untested with `databricks bundle validate`. + +## Verified + +This bundle was deployed and both jobs were run on a single-node cluster +against synthetic `users`, `items` and `orders` tables, on a GPU ML runtime +on a `g5.xlarge` (A10G; the fit trained on the GPU): the fit registered a +new `users_churn` version and moved `champion`, and the score job loaded it and +wrote the predictions table. `test_job_templates.py` checks the scripts +locally. + +## References + +- Task types and parameters: + https://docs.databricks.com/aws/en/dev-tools/bundles/job-task-types +- Custom `PYTHONPATH` overridden: + https://kb.databricks.com/en_US/clusters/cannot-set-a-custom-pythonpath +- Dynamic value references: + https://docs.databricks.com/aws/en/jobs/dynamic-value-references +- Standard compute limitations: + https://docs.databricks.com/aws/en/compute/standard-limitations +- Unity Catalog models: + https://docs.databricks.com/aws/en/machine-learning/manage-model-lifecycle/ diff --git a/examples/mlops/databricks/databricks.yml b/examples/mlops/databricks/databricks.yml new file mode 100644 index 000000000..bd0722327 --- /dev/null +++ b/examples/mlops/databricks/databricks.yml @@ -0,0 +1,173 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +# Train and daily score jobs for kumo-engine; see README.md. + +bundle: + name: kumo-engine-users-churn + +sync: + paths: + - .. + +variables: + catalog: + description: Unity Catalog catalog holding the users, items and orders tables + default: + schema: + description: schema holding the tables, the registered model and the predictions + default: + volume: + description: Unity Catalog volume for exported models and snapshots + default: + warehouse_id: + description: SQL warehouse that reads the source tables and writes predictions + default: + kumo_engine_wheel: + description: kumo-engine wheel on a volume or workspace path + default: /Volumes////wheels/.whl + requirements: + description: requirements.txt for kumo-engine's dependencies; see README.md + default: /Volumes////wheels/requirements.txt + spark_version: + description: Databricks Runtime 17.3 LTS ML or later, for example 17.3.x-gpu-ml-scala2.13 (or 17.3.x-cpu-ml-scala2.13 on a CPU node) + default: + node_type_id: + description: a GPU node type, for example g5.xlarge (or a CPU node type with local NVMe disks, such as m5d.xlarge) + default: + experiment: + description: MLflow experiment (workspace path) for train and score runs + default: /Shared/kumo-engine/users-churn + model_name: + description: registered model, created as .. + default: users_churn + predictions_table: + description: table in . the score job appends to + default: users_churn_predictions + table_prefix: + description: prefix of the users, items and orders table names, e.g. demo_ for demo_users + default: "" + on_error: + description: raise fails a job whose tracking or registration fails; warn lets it succeed + default: raise + alias: + description: alias kumo_train moves to its new version; promote it to score_alias yourself + default: candidate + score_alias: + description: alias kumo_score loads + default: champion + +resources: + jobs: + kumo_train: + name: kumo-engine-train + max_concurrent_runs: 1 + job_clusters: + - job_cluster_key: kumo + new_cluster: + spark_version: ${var.spark_version} + node_type_id: ${var.node_type_id} + num_workers: 0 + data_security_mode: SINGLE_USER + spark_conf: + spark.databricks.cluster.profile: singleNode + spark.master: local[*] + custom_tags: + ResourceClass: SingleNode + spark_env_vars: + MLFLOW_TRACKING_URI: databricks + MLFLOW_REGISTRY_URI: databricks-uc + MLFLOW_EXPERIMENT_NAME: ${var.experiment} + tasks: + - task_key: fit + job_cluster_key: kumo + libraries: + - requirements: ${var.requirements} + - whl: ${var.kumo_engine_wheel} + spark_python_task: + python_file: ${workspace.file_path}/fit_job.py + parameters: + - --artifact-root + - /Volumes/${var.catalog}/${var.schema}/${var.volume}/kumo + - --source + - databricks + - --catalog + - ${var.catalog} + - --schema + - ${var.schema} + - --warehouse + - ${var.warehouse_id} + - --registry + - mlflow + - --model-name + - ${var.catalog}.${var.schema}.${var.model_name} + - --table-prefix=${var.table_prefix} + - --on-error + - ${var.on_error} + - --alias + - ${var.alias} + - --scratch-dir + - /local_disk0/kumo + + kumo_score: + name: kumo-engine-score + max_concurrent_runs: 1 + schedule: + quartz_cron_expression: 0 0 3 * * ? + timezone_id: UTC + pause_status: PAUSED + job_clusters: + - job_cluster_key: kumo + new_cluster: + spark_version: ${var.spark_version} + node_type_id: ${var.node_type_id} + num_workers: 0 + data_security_mode: SINGLE_USER + spark_conf: + spark.databricks.cluster.profile: singleNode + spark.master: local[*] + custom_tags: + ResourceClass: SingleNode + spark_env_vars: + MLFLOW_TRACKING_URI: databricks + MLFLOW_REGISTRY_URI: databricks-uc + MLFLOW_EXPERIMENT_NAME: ${var.experiment} + tasks: + - task_key: score + job_cluster_key: kumo + libraries: + - requirements: ${var.requirements} + - whl: ${var.kumo_engine_wheel} + spark_python_task: + python_file: ${workspace.file_path}/score_job.py + parameters: + - --artifact-root + - /Volumes/${var.catalog}/${var.schema}/${var.volume}/kumo + - --source + - databricks + - --catalog + - ${var.catalog} + - --schema + - ${var.schema} + - --warehouse + - ${var.warehouse_id} + - --registry + - mlflow + - --model-name + - ${var.catalog}.${var.schema}.${var.model_name} + - --table-prefix=${var.table_prefix} + - --on-error + - ${var.on_error} + - --alias + - ${var.score_alias} + - --scratch-dir + - /local_disk0/kumo + - --predictions-table + - ${var.predictions_table} + +targets: + dev: + default: true + workspace: + root_path: /Workspace/Users/${workspace.current_user.userName}/.bundle/${bundle.name}/${bundle.target} diff --git a/examples/mlops/fit_job.py b/examples/mlops/fit_job.py new file mode 100644 index 000000000..8160bb59c --- /dev/null +++ b/examples/mlops/fit_job.py @@ -0,0 +1,38 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Train the users-churn model: ``python fit_job.py --artifact-root ROOT +--source SOURCE [...]``; see README.md.""" + +import os +import sys +from typing import Optional, Sequence + +sys.path.insert(0, os.path.dirname(os.path.abspath( + globals().get('__file__') or sys.argv[0]))) + +import jobs # noqa: E402 + +import kumo.engine as kumo # noqa: E402 + + +def main(argv: Optional[Sequence[str]] = None) -> kumo.TrainingResult: + args = jobs.parse_args(argv, score=False) + connector = jobs.source(args) + graph = jobs.build_graph(connector, args.table_prefix) + with kumo.mlops(jobs.mlops(args, connector)): + snapshot = graph.snapshot() + query = kumo.PredictiveQuery(graph=graph, query=args.query) + table = query.generate_training_table(snapshot) + plan = query.suggest_model_plan(snapshot) + if args.num_experiments: + plan.training_job.num_experiments = args.num_experiments + result = kumo.Trainer(plan).fit( + table, snapshot=snapshot, scoring_defaults=kumo.ScoringDefaults( + binary_classification_threshold=args.threshold)) + print(f'model: models/{result.id}') + return result + + +if __name__ == '__main__': + main() diff --git a/examples/mlops/jobs.py b/examples/mlops/jobs.py new file mode 100644 index 000000000..1dcc865c8 --- /dev/null +++ b/examples/mlops/jobs.py @@ -0,0 +1,170 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The flags, graph, connectors and MLOps settings both job scripts share.""" + +import argparse +from typing import Optional, Sequence + +from kumoapi.jobs import MetadataField +from kumoapi.typing import Stype + +import kumo.engine as kumo +from kumo.engine import ( + Edge, + Graph, + ModelRegistryLocation, + PredictionOutputConfig, + Table, + Trainer, +) +from kumo.engine.connector import ( + Connector, + DatabricksConnector, + FileConnector, + S3Connector, + SnowflakeConnector, + connector_for_root, +) + +# Needs several 30-day windows of order history; shorten it for small data. +QUERY = 'PREDICT COUNT(orders.*, 0, 30, days) > 0 FOR EACH users.USER_ID' +# 'table.COLUMN' names to encode as text: inference reads short text as +# categorical, and the score job must declare the graph the fit trained on. +TEXT_COLUMNS: Sequence[str] = () + + +def parse_args(argv: Optional[Sequence[str]], *, + score: bool) -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument( + '--artifact-root', required=True, + help='s3://, gs://, file://, a local or /Volumes path, or with ' + '--source snowflake an @DB.SCHEMA.STAGE path') + parser.add_argument( + '--source', required=True, + help='databricks, snowflake, or a directory holding users, items ' + 'and orders parquet files') + parser.add_argument('--catalog', + help='Databricks catalog or Snowflake database') + parser.add_argument('--schema') + parser.add_argument( + '--table-prefix', default='', + help='prefix of the users, items and orders table or file names') + parser.add_argument( + '--warehouse', help='Databricks SQL warehouse id or Snowflake ' + 'warehouse') + parser.add_argument('--registry', choices=('mlflow', 'snowflake')) + parser.add_argument('--model-name') + parser.add_argument( + '--alias', default='champion', + help='alias the fit moves to its registered version and the score ' + 'job loads') + parser.add_argument('--scratch-dir') + parser.add_argument( + '--on-error', choices=('raise', 'warn'), default='raise', + help='raise (fail the job) or warn when tracking or registration ' + 'fails') + if score: + parser.add_argument( + '--model-path', + help='models/ under the artifact root, without a ' + 'registry') + parser.add_argument( + '--snapshot', help='the id of a snapshot stored under the ' + 'artifact root, instead of a new one') + parser.add_argument('--predictions-table') + else: + parser.add_argument('--query', default=QUERY) + parser.add_argument('--num-experiments', type=int) + parser.add_argument( + '--threshold', type=float, default=0.5, + help='binary classification threshold the model scores with') + args = parser.parse_args(argv) + if args.registry and not args.model_name: + parser.error('--registry needs --model-name') + if score and not args.registry and not args.model_path: + parser.error('without --registry, pass --model-path') + if args.scratch_dir: + kumo.configure(scratch_dir=args.scratch_dir) + return args + + +def source(args: argparse.Namespace) -> Connector: + if args.source == 'databricks': + return DatabricksConnector(catalog=args.catalog, + schema_name=args.schema, + warehouse_id=args.warehouse) + if args.source == 'snowflake': + return SnowflakeConnector.from_spcs_env(warehouse=args.warehouse, + database=args.catalog, + schema_name=args.schema) + return connector_for_root(args.source) + + +def mlops(args: argparse.Namespace, connector: Connector) -> kumo.MLOps: + root = ((connector, args.artifact_root) + if args.artifact_root.startswith('@') else args.artifact_root) + if args.registry == 'mlflow': + from kumo.engine.integrations.mlflow import MlflowTracking + + tracking = MlflowTracking() + elif args.registry == 'snowflake': + from kumo.engine.integrations.snowflake import SnowflakeTracking + + tracking = SnowflakeTracking(connector, + experiment_name=args.model_name) + else: + return kumo.MLOps(root, on_error=args.on_error) + return kumo.MLOps(root, tracking=tracking, registry=registry(args), + on_error=args.on_error) + + +def registry(args: argparse.Namespace) -> ModelRegistryLocation: + return ModelRegistryLocation(args.model_name, alias=args.alias) + + +def build_graph(connector: Connector, prefix: str = '') -> Graph: + + def table(alias: str, **kwargs: str) -> Table: + alias = f'{prefix}{alias}' + name = { + S3Connector: f'{alias}.parquet', + FileConnector: f'{alias}.parquet', + SnowflakeConnector: alias.upper(), + }.get(type(connector), alias) + return Table.from_source_table(connector[name], **kwargs) + + orders = table('orders', primary_key='ORDER_ID', + time_column='ORDER_TIME').with_column( + 'USER_ID', + stype=Stype.ID).with_column('ITEM_ID', stype=Stype.ID) + tables = { + 'users': table('users', primary_key='USER_ID'), + 'items': table('items', primary_key='ITEM_ID'), + 'orders': orders, + } + for name in TEXT_COLUMNS: + alias, column = name.split('.') + tables[alias] = tables[alias].with_column(column, stype=Stype.text) + return Graph( + tables=tables, edges=[ + Edge(src_table='orders', fkey='USER_ID', dst_table='users'), + Edge(src_table='orders', fkey='ITEM_ID', dst_table='items'), + ]) + + +def load_model(args: argparse.Namespace, cfg: kumo.MLOps) -> Trainer: + if args.registry: + return Trainer.load_from_connector(cfg, location=registry(args)) + return Trainer.load_from_connector(cfg, path=args.model_path) + + +def prediction_output( + args: argparse.Namespace, + connector: Connector) -> Optional[PredictionOutputConfig]: + if not args.predictions_table: + return None + return PredictionOutputConfig( + connector, path=args.predictions_table, + extra_fields=[MetadataField.ANCHOR_TIMESTAMP]) diff --git a/examples/mlops/sagemaker/Dockerfile b/examples/mlops/sagemaker/Dockerfile new file mode 100644 index 000000000..c0b165c02 --- /dev/null +++ b/examples/mlops/sagemaker/Dockerfile @@ -0,0 +1,22 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +# Image for the pipeline's Fit and Predict steps; see README.md. +# Build context: a directory holding the kumo-engine wheel and its kumo-core, +# kumo-api, kumo-pql, kumo-s3 and diskgraph wheels. +# GPU: docker build -f examples/mlops/sagemaker/Dockerfile -t kumo-engine-sagemaker +# CPU: add --build-arg TORCH=cpu +FROM python:3.12-slim +ARG TORCH=cu126 +COPY *.whl /wheels/ +# The engine wheel is installed by path: PyPI's unrelated "kumo-engine" would +# win over a dev build. Once kumo-engine is on the index, install +# "kumo-engine[mlflow,sagemaker,text]" instead and drop the wheels. +RUN pip install --no-cache-dir --find-links /wheels \ + --extra-index-url "https://download.pytorch.org/whl/${TORCH}" \ + "$(ls /wheels/kumo_engine-*.whl)[mlflow,sagemaker,text]" \ + "torch==2.12.0+${TORCH}" \ + && python -c "import kumo.engine, torch; print(kumo.engine.__version__, torch.__version__)" \ + && rm -rf /wheels +ENV PYTHONUNBUFFERED=1 GIT_PYTHON_REFRESH=quiet MLFLOW_DISABLE_AGENT_HINT=1 diff --git a/examples/mlops/sagemaker/README.md b/examples/mlops/sagemaker/README.md new file mode 100644 index 000000000..c78cb62e0 --- /dev/null +++ b/examples/mlops/sagemaker/README.md @@ -0,0 +1,122 @@ +# SageMaker Pipelines + +`pipeline.py` defines and upserts a two-step pipeline: + +1. **`Fit`** (training job) runs `python fit_job.py` on `users.parquet`, + `items.parquet` and `orders.parquet` under `--data-root`, inside + `kumo.mlops(...)` with `--artifact-root`: the snapshot, training table and + model are stored under it. `--registry mlflow` logs the fit to the managed + MLflow tracking server, registers the model as a new version of + `users-churn` and moves the `--alias` alias (`champion` by default) to it. +2. **`Predict`** (processing job, after `Fit`) runs `python score_job.py`, + loading `models:/users-churn@` and scoring a fresh snapshot. + +`--fit-instance-type` and `--score-instance-type` size the two steps +separately. `--model-name` (default `users-churn`) and `--table-prefix` +(prefix of the three Parquet file names) are passed to both steps, as is +`--on-error` (default `raise`: a failed tracking call or registration fails +the step, so `Predict` never runs on the previous version unnoticed). With `--alias candidate`, the pipeline registers and scores a +candidate and leaves `champion` to be moved by hand (see +[Promote a model](../../../packages/kumo-engine/docs/mlops.md#promote-a-model)). +On a shared tracking server, pass your own `--name` and `--model-name`: the +defaults would update another team's pipeline or registered model. + +To score on a schedule without retraining, put `Predict` in its own pipeline +with a `PipelineSchedule` trigger (`pipeline.put_triggers`). + +## Before running + +1. Build and push the image. [`Dockerfile`](Dockerfile) installs + `kumo-engine[mlflow,sagemaker,text]` (MLflow, S3, `sagemaker-mlflow` and the default word vectors for text columns) and torch 2.12 for CUDA 12.6 + on `python:3.12-slim`, from a directory holding the kumo-engine wheel and + its kumo-core, kumo-api, kumo-pql, kumo-s3 and diskgraph wheels + (`--find-links`; the engine wheel goes in by path, since PyPI's unrelated + `kumo-engine` would otherwise win). Add `--build-arg TORCH=cpu` for a CPU + image. The images are about 8 GB (GPU) and 3.5 GB (CPU). + ```bash + docker build -f examples/mlops/sagemaker/Dockerfile -t kumo-engine-sagemaker + aws ecr create-repository --repository-name kumo-engine-sagemaker + aws ecr get-login-password | docker login --username AWS \ + --password-stdin .dkr.ecr..amazonaws.com + docker tag kumo-engine-sagemaker .dkr.ecr..amazonaws.com/kumo-engine-sagemaker: + docker push .dkr.ecr..amazonaws.com/kumo-engine-sagemaker: + ``` +2. Upload the code channel: `aws s3 sync examples/mlops s3:///kumo/code/`. +3. With the SageMaker Python SDK v2 (`pip install 'sagemaker<3'`): + ```bash + python examples/mlops/sagemaker/pipeline.py --role \ + --image .dkr.ecr..amazonaws.com/: \ + --code s3:///kumo/code/ --data-root s3:///kumo/data/ \ + --artifact-root s3:///kumo/artifacts/ \ + --tracking-arn \ + --fit-instance-type ml.g5.xlarge --score-instance-type ml.g5.xlarge \ + --start + ``` + SDK v2 reads the region only from `AWS_DEFAULT_REGION` (or the AWS + config file), not `AWS_REGION`. + The name (`--name`, default `kumo-engine-users-churn`) may not contain + `-aws-`: SageMaker reserves that token, and the execution then fails with + no steps (`The token -aws- is reserved`). + SDK v3 replaced `Estimator` with `ModelTrainer` and moved the workflow + classes (`sagemaker.mlops.workflow.*`); this script is not ported. + +## Privileges + +- **Caller:** `sagemaker:CreatePipeline`, `UpdatePipeline`, + `StartPipelineExecution`, and `iam:PassRole` on ``. +- **Execution role** (`--role` also runs the steps): `sagemaker:CreateTrainingJob`, + `DescribeTrainingJob`, `StopTrainingJob`, `CreateProcessingJob`, + `DescribeProcessingJob`, `StopProcessingJob` and `AddTags`; `iam:PassRole` + on itself; S3 read on the code and data prefixes and read/write on + the artifact prefix; ECR pull; CloudWatch Logs; the `sagemaker-mlflow:*` + actions for tracking and the registry (`CreateRun`, `LogBatch`, + `CreateRegisteredModel`, `CreateModelVersion`, `SetRegisteredModelAlias`, + `GetModelVersionByAlias`, ...). + +## Notes + +- **GPU quota for `Predict`.** The `Predict` step is a SageMaker Processing + job, which draws on the ` for processing job usage` quota + (for example `ml.g5.xlarge for processing job usage`), separate from the + training job quota and often 0 by default. Request an increase in Service + Quotas before scoring on a GPU type, or pass a CPU type as + `--score-instance-type` with a CPU image. A GPU `Predict` step was verified + on `ml.g5.xlarge` with that quota at 2. +- **Flags.** Both steps pass the job scripts' flags as plain strings + (`container_arguments` for the training step, `arguments` for the + processing step); only MLflow's own settings are environment variables. +- **MLflow.** `MLFLOW_TRACKING_URI` is the tracking server or MLflow App ARN. + Managed MLflow servers run MLflow 3.0; the `mlflow` extra installs + mlflow-skinny 3.0.1 or later, which works against them, so do not pin the + client to the server's version. Aliases are letters, digits and underscores, so + SageMaker's reserved `SageMakerLifecycle-` aliases cannot be set. Model names + are letters, digits and hyphens (no `_` or `.`), checked when `kumo.MLOps` + is built. +- **Scratch.** The fit step uses `/tmp/kumo` (on NVMe instance types + SageMaker puts `/tmp` on the instance store; on EBS-only types raise + `volume_size`); the score step uses `/opt/ml/processing/scratch`. +- **Retries.** Step `retry_policies` key on SageMaker failure reasons, never + on the exit code; a non-zero exit is not retried. + +## Verified + +This pipeline was upserted and run on `ml.m5.xlarge` with a CPU image, the +tracking server ARN as `MLFLOW_TRACKING_URI` and 180 days of synthetic orders +(the 30-day query needs about 150 days of history): both steps exit 0, the +fit registers a new `users-churn` version and moves `champion`, and the score +step loads it. It also ran with both steps on `ml.g5.xlarge` with a GPU image +built like `Dockerfile` (python 3.12-slim, torch 2.12 for CUDA 12.6, the +wheels and `sagemaker-mlflow`), using the GPU in both. The separate +instance-type flags, `--alias` and this exact `Dockerfile` have not been run +on SageMaker. An MLflow App ARN as tracking URI is unverified. +`test_job_templates.py` checks the scripts locally. + +## References + +- Pipeline steps: https://docs.aws.amazon.com/sagemaker/latest/dg/build-and-manage-steps-types.html +- Retry policies: https://docs.aws.amazon.com/sagemaker/latest/dg/pipelines-retry-policy.html +- Instance store and `/tmp`: + https://docs.aws.amazon.com/sagemaker/latest/dg/model-train-storage-tips-considerations.html +- Managed MLflow and IAM: + https://docs.aws.amazon.com/sagemaker/latest/dg/mlflow-track-experiments.html, + https://docs.aws.amazon.com/sagemaker/latest/dg/mlflow-create-tracking-server-iam.html diff --git a/examples/mlops/sagemaker/pipeline.py b/examples/mlops/sagemaker/pipeline.py new file mode 100644 index 000000000..12d97b50f --- /dev/null +++ b/examples/mlops/sagemaker/pipeline.py @@ -0,0 +1,112 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Upsert a SageMaker Pipeline (SDK v2) that fits and scores with kumo-engine; +see README.md.""" + +import argparse +from typing import Any, Optional, Sequence + +TRAIN_CODE = '/opt/ml/input/data/code' +PROCESS_CODE = '/opt/ml/processing/input/code' + + +def build_pipeline( + *, + name: str, + role: str, + image: str, + code: str, + data_root: str, + artifact_root: str, + tracking_arn: str, + experiment: str, + fit_instance_type: str, + score_instance_type: str, + alias: str, + model_name: str = 'users-churn', + table_prefix: str = '', + on_error: str = 'raise', +) -> Any: + from sagemaker.estimator import Estimator + from sagemaker.inputs import TrainingInput + from sagemaker.processing import ProcessingInput, Processor + from sagemaker.workflow.pipeline import Pipeline + from sagemaker.workflow.pipeline_context import PipelineSession + from sagemaker.workflow.steps import ProcessingStep, TrainingStep + + session = PipelineSession() + env = { + 'MLFLOW_TRACKING_URI': tracking_arn, + 'MLFLOW_EXPERIMENT_NAME': experiment, + } + flags = [ + '--artifact-root', artifact_root, '--source', data_root, + '--registry', 'mlflow', '--model-name', model_name, '--alias', + alias, f'--table-prefix={table_prefix}', '--on-error', on_error + ] + + estimator = Estimator( + image_uri=image, + role=role, + instance_count=1, + instance_type=fit_instance_type, + container_entry_point=['python', f'{TRAIN_CODE}/fit_job.py'], + container_arguments=[*flags, '--scratch-dir', '/tmp/kumo'], + environment=env, + output_path=f'{artifact_root.rstrip("/")}/sagemaker-output', + sagemaker_session=session, + ) + fit = TrainingStep( + name='Fit', + step_args=estimator.fit(inputs={'code': TrainingInput(code)}), + ) + + processor = Processor( + role=role, + image_uri=image, + instance_count=1, + instance_type=score_instance_type, + entrypoint=['python', f'{PROCESS_CODE}/score_job.py'], + env=env, + sagemaker_session=session, + ) + predict = ProcessingStep( + name='Predict', + step_args=processor.run( + inputs=[ProcessingInput(source=code, destination=PROCESS_CODE)], + arguments=[ + *flags, '--scratch-dir', '/opt/ml/processing/scratch' + ], + ), + depends_on=[fit], + ) + return Pipeline(name=name, steps=[fit, predict], + sagemaker_session=session) + + +def main(argv: Optional[Sequence[str]] = None) -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument('--name', default='kumo-engine-users-churn') + for flag in ('--role', '--image', '--code', '--data-root', + '--artifact-root', '--tracking-arn'): + parser.add_argument(flag, required=True) + parser.add_argument('--experiment', default='kumo-engine-users-churn') + parser.add_argument('--fit-instance-type', required=True) + parser.add_argument('--score-instance-type', required=True) + parser.add_argument('--alias', default='champion') + parser.add_argument('--model-name', default='users-churn') + parser.add_argument('--table-prefix', default='') + parser.add_argument('--on-error', choices=('raise', 'warn'), + default='raise') + parser.add_argument('--start', action='store_true') + args = vars(parser.parse_args(argv)) + start = args.pop('start') + pipeline = build_pipeline(**args) + pipeline.upsert(role_arn=args['role']) + if start: + print(pipeline.start().arn) + + +if __name__ == '__main__': + main() diff --git a/examples/mlops/score_job.py b/examples/mlops/score_job.py new file mode 100644 index 000000000..e836ac244 --- /dev/null +++ b/examples/mlops/score_job.py @@ -0,0 +1,37 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Score with the champion model, or the one ``--model-path`` names: +``python score_job.py --artifact-root ROOT --source SOURCE [...]``; see +README.md.""" + +import os +import sys +from typing import Optional, Sequence + +sys.path.insert(0, os.path.dirname(os.path.abspath( + globals().get('__file__') or sys.argv[0]))) + +import jobs # noqa: E402 + +import kumo.engine as kumo # noqa: E402 + + +def main(argv: Optional[Sequence[str]] = None) -> kumo.PredictionResult: + args = jobs.parse_args(argv, score=True) + connector = jobs.source(args) + graph = jobs.build_graph(connector, args.table_prefix) + cfg = jobs.mlops(args, connector) + with kumo.mlops(cfg): + trainer = jobs.load_model(args, cfg) + snapshot = (kumo.GraphSnapshot.load_from_connector( + cfg, id=args.snapshot, graph=graph) + if args.snapshot else graph.snapshot()) + result = trainer.predict( + snapshot, output_config=jobs.prediction_output(args, connector)) + print(f'predictions: {result.num_rows} rows') + return result + + +if __name__ == '__main__': + main() diff --git a/examples/mlops/snowflake/README.md b/examples/mlops/snowflake/README.md new file mode 100644 index 000000000..e206295fb --- /dev/null +++ b/examples/mlops/snowflake/README.md @@ -0,0 +1,95 @@ +# Snowflake ML Jobs and Tasks (experimental) + +**Experimental.** This template is built on preview Snowflake APIs +(`MLJobDefinition` is in preview and the `SYSTEM$EXECUTE_ML_JOB` statement it +generates is undocumented) and has not been run end to end; see +[Unverified](#unverified). They may change without notice. + +`register_jobs.py` registers two ML Job definitions whose payload is this +`examples/mlops` directory, and one Task per definition that runs it. + +- **`KUMO_TRAIN`** (`KUMO_TRAIN_TASK`, Sundays 02:00 UTC) runs `fit_job.py` + on `USERS`, `ITEMS` and `ORDERS` in `.`, publishes a new + version of model `USERS_CHURN` in that schema's Model Registry and moves the + `champion` alias to it (`--registry snowflake`, which builds + `kumo.MLOps(tracking=SnowflakeTracking(connector, ...), + registry=ModelRegistryLocation(...))`). The snapshot, training table and + model are stored under `--artifact-root`. The fit is recorded in + Snowflake Experiment Tracking through the connector's own Snowpark session + (`SnowflakeConnector.session()`, the container's OAuth token). +- **`KUMO_SCORE`** (`KUMO_SCORE_TASK`, daily 03:00 UTC) runs `score_job.py`, + loading the version `champion` points at and writing + `USERS_CHURN_PREDICTIONS`. + +To review a model before it is scored, drop the alias from the registry +location in `jobs.mlops` and move `champion` yourself with +`ALTER MODEL ..USERS_CHURN VERSION SET ALIAS = champion` +(see [Promote a model](../../../packages/kumo-engine/docs/mlops.md#promote-a-model) +for the statements that free both aliases first). +If moving the alias fails, the fit still publishes, warns, and scoring keeps +the previous `champion`. + +Each Task runs its job synchronously with the statement +`MLJobDefinition.to_sql(job_args=, use_async=False)` generates. `USER_TASK_TIMEOUT_MS` is 24 h for train and 6 h for score +(default 1 h); a job that outlives it, or the warehouse's +`STATEMENT_TIMEOUT_IN_SECONDS`, is cancelled. + +## Before running + +1. Make `kumo-engine[snowflake]` and its dependency wheels installable in the + job container (package index via an external access integration, or a + custom image). +2. Create the GPU compute pool `` and the stage ``. The + artifact root must outlive the job container: an `s3://` or `gs://` URI + the external access integration reaches. +3. Run: + ```bash + python examples/mlops/snowflake/register_jobs.py --connection \ + --compute-pool --stage --database \ + --schema --warehouse \ + --artifact-root s3:///kumo/ \ + --requirement 'kumo-engine[snowflake]' --external-access-integration + ``` + `--model-name` (default `USERS_CHURN`), `--predictions-table` (default + `USERS_CHURN_PREDICTIONS`) and `--table-prefix` (prefix of `USERS`, + `ITEMS` and `ORDERS`) are passed to the jobs, as is `--on-error` + (default `raise`: a failed registration or tracking call fails the job and + its Task run). +4. Tasks are created suspended: `ALTER TASK RESUME`, or + `EXECUTE TASK ` to run one now. + +## Privileges for the role that owns the tasks + +- `USAGE` on the database, schema, warehouse and compute pool; + `CREATE SERVICE`, `CREATE TASK`, `CREATE TABLE` and `CREATE MODEL` on the + schema; `EXECUTE TASK` on the account. +- Read and write on the stage; `USAGE` on the external access integration (or + `SNOWFLAKE.PYPI_REPOSITORY_USER`). +- `SELECT` on the source tables, `INSERT` for predictions, and ownership of + the model to move its alias. + +## Notes + +- The ML Jobs launcher runs the job script from the payload directory, so + `jobs.py` imports; a non-zero exit fails the job (its `mljob_result` is + still `None`, so judge a run by its status). +- The Task's synchronous `CALL` holds the task warehouse while the job runs; + a serverless task avoids that. +- `--scratch-dir` is `/tmp/kumo`; add a volume through `spec_overrides` if the + graph needs more disk. Tasks retry only via `TASK_AUTO_RETRY_ATTEMPTS`. + +## Unverified + +Nothing has been registered or run. `SYSTEM$EXECUTE_ML_JOB` (what `to_sql` +generates) is not a documented SQL API and `MLJobDefinition` is in preview. +Unverified: reaching the reader's presigned URLs without an external access +integration, a failed job failing the Task, and creating experiments from a job's OAuth session. + +## References + +- `MLJobDefinition`: + https://docs.snowflake.com/en/developer-guide/snowpark-ml/reference/latest/api/jobs/snowflake.ml.jobs.MLJobDefinition +- ML Jobs: https://docs.snowflake.com/en/developer-guide/snowflake-ml/ml-jobs/overview +- `CREATE TASK`: https://docs.snowflake.com/en/sql-reference/sql/create-task +- Model versions: + https://docs.snowflake.com/en/developer-guide/snowflake-ml/model-registry/model-management diff --git a/examples/mlops/snowflake/register_jobs.py b/examples/mlops/snowflake/register_jobs.py new file mode 100644 index 000000000..89dd0ea99 --- /dev/null +++ b/examples/mlops/snowflake/register_jobs.py @@ -0,0 +1,121 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Experimental: register fit_job.py/score_job.py as Snowflake ML Jobs with +one Task each (snowflake-ml-python 2.6, preview API); see README.md.""" + +import argparse +from pathlib import Path +from typing import Any, Dict, List, Optional, Sequence + +PAYLOAD = Path(__file__).resolve().parent.parent +JOBS = { + 'train': { + 'entrypoint': 'fit_job.py', + 'schedule': 'USING CRON 0 2 * * 0 UTC', + 'timeout_ms': 24 * 3600 * 1000, + }, + 'score': { + 'entrypoint': 'score_job.py', + 'schedule': 'USING CRON 0 3 * * * UTC', + 'timeout_ms': 6 * 3600 * 1000, + }, +} + + +def task_sql(*, task_name: str, warehouse: str, schedule: str, + timeout_ms: int, execute_ml_job_sql: str) -> str: + return (f"CREATE OR REPLACE TASK {task_name}\n" + f" WAREHOUSE = {warehouse}\n" + f" SCHEDULE = '{schedule}'\n" + f" USER_TASK_TIMEOUT_MS = {timeout_ms}\n" + f" SUSPEND_TASK_AFTER_NUM_FAILURES = 3\n" + f"AS\n{execute_ml_job_sql}") + + +def register(session: Any, *, compute_pool: str, stage: str, database: str, + schema: str, warehouse: str, artifact_root: str, + requirements: List[str], + external_access_integrations: List[str], + model_name: str = 'USERS_CHURN', + predictions_table: str = 'USERS_CHURN_PREDICTIONS', + table_prefix: str = '', + on_error: str = 'raise') -> Dict[str, str]: + from snowflake.ml.jobs import MLJobDefinition + + flags = [ + '--artifact-root', artifact_root, '--source', 'snowflake', + '--catalog', database, '--schema', schema, '--warehouse', warehouse, + '--registry', 'snowflake', '--model-name', model_name, + f'--table-prefix={table_prefix}', '--on-error', on_error, + '--scratch-dir', '/tmp/kumo' + ] + job_args = { + 'train': flags, + 'score': [*flags, '--predictions-table', predictions_table], + } + tasks = {} + for kind, job in JOBS.items(): + definition = MLJobDefinition.register( + str(PAYLOAD), + compute_pool, + stage, + session=session, + name=f'KUMO_{kind.upper()}', + entrypoint=str(job['entrypoint']), + pip_requirements=requirements, + external_access_integrations=external_access_integrations, + query_warehouse=warehouse, + database=database, + schema=schema, + overwrite=True, + ) + name = f'{database}.{schema}.KUMO_{kind.upper()}_TASK' + session.sql( + task_sql( + task_name=name, warehouse=warehouse, + schedule=str(job['schedule']), + timeout_ms=int(job['timeout_ms']), + execute_ml_job_sql=definition.to_sql( + job_args=job_args[kind], use_async=False), + )).collect() + tasks[kind] = name + return tasks + + +def main(argv: Optional[Sequence[str]] = None) -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument('--connection', help='connections.toml entry') + for flag in ('--compute-pool', '--stage', '--database', '--schema', + '--warehouse', '--artifact-root'): + parser.add_argument(flag, required=True) + parser.add_argument('--requirement', action='append', required=True) + parser.add_argument('--external-access-integration', action='append', + default=[]) + parser.add_argument('--model-name', default='USERS_CHURN') + parser.add_argument('--predictions-table', + default='USERS_CHURN_PREDICTIONS') + parser.add_argument('--table-prefix', default='') + parser.add_argument('--on-error', choices=('raise', 'warn'), + default='raise') + args = parser.parse_args(argv) + + from snowflake.snowpark import Session + + builder = Session.builder + if args.connection is not None: + builder = builder.config('connection_name', args.connection) + tasks = register( + builder.create(), compute_pool=args.compute_pool, stage=args.stage, + database=args.database, schema=args.schema, warehouse=args.warehouse, + artifact_root=args.artifact_root, + requirements=args.requirement, + external_access_integrations=args.external_access_integration, + model_name=args.model_name, predictions_table=args.predictions_table, + table_prefix=args.table_prefix, on_error=args.on_error) + for kind, name in tasks.items(): + print(f'{kind}: {name} (suspended; ALTER TASK {name} RESUME)') + + +if __name__ == '__main__': + main() diff --git a/packages/diskgraph/diskgraph-ingest/src/col_stats.rs b/packages/diskgraph/diskgraph-ingest/src/col_stats.rs index b8bbd94b9..c8e9d3343 100644 --- a/packages/diskgraph/diskgraph-ingest/src/col_stats.rs +++ b/packages/diskgraph/diskgraph-ingest/src/col_stats.rs @@ -791,6 +791,17 @@ impl FileColStats { } } + /// Values of numerical column `col_idx` at quantiles `qs`, from its + /// t-digest; `None` for other columns or without values. + pub fn numerical_quantiles(&mut self, col_idx: usize, qs: &[f64]) -> Option> { + match self.cols.get_mut(col_idx)? { + ColumnStatsCollector::Numerical { + tdigest, welford, .. + } if welford.count > 0 => Some(qs.iter().map(|&q| tdigest.quantile_at(q)).collect()), + _ => None, + } + } + /// Serialize accumulated stats to a JSON object keyed by column name. /// Category counts are read (not consumed) from the shared state. /// Reference for the `col_stats.json` that [`write_col_stats`] streams. diff --git a/packages/diskgraph/diskgraph-pybind/python/diskgraph/__init__.py b/packages/diskgraph/diskgraph-pybind/python/diskgraph/__init__.py index 5b9f5d0b6..a07eca938 100644 --- a/packages/diskgraph/diskgraph-pybind/python/diskgraph/__init__.py +++ b/packages/diskgraph/diskgraph-pybind/python/diskgraph/__init__.py @@ -44,6 +44,7 @@ has_resumable_ingest, validated_index_resume_progress, compute_col_stats_for_parquet, + evaluate_prediction_parquet, resolve_temporal, detect_multicat_delimiter, multicat_delimiter_is_proven, @@ -146,6 +147,7 @@ def __getattr__(name): "has_resumable_ingest", "validated_index_resume_progress", "compute_col_stats_for_parquet", + "evaluate_prediction_parquet", "DiskGraphPlanCompiler", "SplitSpec", "DiskgraphConfigConverter", diff --git a/packages/diskgraph/diskgraph-pybind/python/diskgraph/tt_artifacts.py b/packages/diskgraph/diskgraph-pybind/python/diskgraph/tt_artifacts.py index 9fc993ee8..2b3d491df 100644 --- a/packages/diskgraph/diskgraph-pybind/python/diskgraph/tt_artifacts.py +++ b/packages/diskgraph/diskgraph-pybind/python/diskgraph/tt_artifacts.py @@ -348,7 +348,8 @@ def build_prediction_table_artifacts_for_existing_table( task_type=task_type, ingest_on_error=ingest_on_error, ) - _warn_skipped_rows(result) + # The caller reports skipped prediction rows: it knows they are entities + # to score, not training rows. return result diff --git a/packages/diskgraph/diskgraph-pybind/python/tests/test_col_stats_options.py b/packages/diskgraph/diskgraph-pybind/python/tests/test_col_stats_options.py new file mode 100644 index 000000000..c56c5c801 --- /dev/null +++ b/packages/diskgraph/diskgraph-pybind/python/tests/test_col_stats_options.py @@ -0,0 +1,67 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import hashlib +import json + +import diskgraph +import pyarrow as pa +import pyarrow.parquet as pq + +SPECS = json.dumps([ + {"name": "SCORE", "stype": "numerical", "dtype": "float64"}, + {"name": "CLASS", "stype": "categorical", "dtype": "string"}, +]) + + +def test_prediction_options_through_python(tmp_path): + path = tmp_path / "part-00000.parquet" + pq.write_table( + pa.table({ + "ENTITY": [1, 1, 2, 2, 3, 3], + "SCORE": [0.2, 0.9, 0.7, 0.6, 0.1, 0.4], + "CLASS": ["a", "b", "a", "c", "b", "a"], + }), path, row_group_size=3) + + full, _ = diskgraph.compute_col_stats_for_parquet( + [str(path)], SPECS, "v1", percentiles=True, + summary_stats_by_class=False, histograms=False, + category_counts_limit=1) + stats = json.loads(full) + assert len(stats["SCORE__kumo.numerical__PERCENTILES__True"]["value"]) == 99 + assert stats["CLASS__kumo.categorical__CATEGORY_COUNTS__False"][ + "value"] == [["a", 3]] + assert stats["CLASS__kumo.categorical__CATEGORY_COUNTS_CAPPED__False"][ + "value"] is False + assert not any("HISTOGRAM" in key or "BY_CLASS" in key for key in stats) + + full, _ = diskgraph.compute_col_stats_for_parquet( + [str(path)], json.dumps(json.loads(SPECS)[:1]), "v1", + summary_stats_by_class=False, histograms=False, max_per=["ENTITY"]) + stats = json.loads(full) + assert stats["SCORE__kumo.numerical__COUNT__False"]["value"] == 3 + assert stats["SCORE__kumo.numerical__MIN__False"]["value"] == 0.4 + + +def test_default_output_is_unchanged(tmp_path): + path = tmp_path / "part-00000.parquet" + pq.write_table( + pa.table({ + "SCORE": [0.9, 0.2, 0.7, 0.6, 0.4, 0.1, 0.35, 0.8], + "CLASS": ["a", "b", "a", "c", "b", "a", "d", "a"], + "AT": pa.array([1_600_000_000 + 86_400 * i for i in range(8)], + pa.timestamp("s")), + }), path) + specs = json.dumps(json.loads(SPECS) + [ + {"name": "AT", "stype": "timestamp", "dtype": "date"}]) + + full, min_ = diskgraph.compute_col_stats_for_parquet([str(path)], specs, + "v1") + + expected = ("7e7c026d875b5eb33be64227ee539fa4" + "fa558d0618f794a7ae3ddb25c6764d9a") + assert hashlib.sha256(full.encode()).hexdigest() == expected + assert hashlib.sha256(min_.encode()).hexdigest() == expected diff --git a/packages/diskgraph/diskgraph-pybind/python/tests/test_tt_artifacts_row_warnings.py b/packages/diskgraph/diskgraph-pybind/python/tests/test_tt_artifacts_row_warnings.py index b5d4761a1..41588f5a6 100644 --- a/packages/diskgraph/diskgraph-pybind/python/tests/test_tt_artifacts_row_warnings.py +++ b/packages/diskgraph/diskgraph-pybind/python/tests/test_tt_artifacts_row_warnings.py @@ -8,7 +8,7 @@ @pytest.mark.parametrize('prediction', [False, True]) -def test_builder_warns_with_distinct_rows_and_all_reasons( +def test_training_builder_warns_with_distinct_rows_and_all_reasons( prediction, monkeypatch, caplog): result = { 'row_count': 3, @@ -41,6 +41,10 @@ def build(**kwargs): kwargs['split_config'] = None assert builder(**kwargs) is result assert calls[0]['ingest_on_error'] == 'skip_row_on_error' + if prediction: + # The caller reports skipped entities to score. + assert not caplog.records + return assert '2 of 5 rows were skipped; training on 3 rows' in caplog.text assert 'null weight[WEIGHT]=2' in caplog.text assert 'empty target list[TARGET]=1' in caplog.text diff --git a/packages/diskgraph/diskgraph-pybind/src/bp_col_stats.rs b/packages/diskgraph/diskgraph-pybind/src/bp_col_stats.rs index 923e1798c..15b81c0a4 100644 --- a/packages/diskgraph/diskgraph-pybind/src/bp_col_stats.rs +++ b/packages/diskgraph/diskgraph-pybind/src/bp_col_stats.rs @@ -13,15 +13,20 @@ use std::path::Path; use std::sync::Arc; use anyhow::{bail, Context, Result}; -use arrow_array::RecordBatch; -use arrow_schema::{DataType, Schema}; +use arrow::array::{new_null_array, Array, Float64Array, UInt32Array}; +use arrow::compute::{cast, concat, take}; +use arrow::row::{OwnedRow, RowConverter, SortField}; +use arrow_array::{ArrayRef, RecordBatch}; +use arrow_schema::{DataType, Field, Schema}; use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; use parquet::arrow::ProjectionMask; use rayon::prelude::*; use serde_json::{json, Map as JsonMap, Value as JsonValue}; +#[cfg(test)] +use diskgraph_core::accumulators::MAX_CATEGORIES; use diskgraph_core::accumulators::{ - kumo_entry, kumo_key, CategoryCountStrategy, WelfordAccumulator, MAX_CATEGORIES, + kumo_entry, kumo_key, CategoryCountStrategy, WelfordAccumulator, }; use diskgraph_ingest::col_stats::{FileColStats, TableCategoryCounts}; use diskgraph_ingest::graph_config::{ColumnSpec, IngestedType, SemanticType, TimestampUnit}; @@ -78,6 +83,238 @@ impl IntegerRange { } } +/// Options beyond the ingest-compatible statistics; the default changes nothing. +#[derive(Clone, Debug)] +pub struct ColStatsOptions { + /// Emit `PERCENTILES` (q01..q99, approximate) for every numerical column. + pub percentiles: bool, + /// Emit `SUMMARY_STATS_BY_CLASS` for a `SCORE`/`CLASS` column pair. + pub summary_stats_by_class: bool, + /// Emit `HISTOGRAM` / `HISTOGRAM_HALF_MONTH` (a second read). + pub histograms: bool, + /// Keep at most this many `CATEGORY_COUNTS` entries in the full output, + /// and record `CATEGORY_COUNTS_CAPPED`: whether counting hit the + /// category cap, so `NUM_UNIQUE` is an estimate. + pub category_counts_limit: Option, + /// Compute every statistic over the maximum of each column within each + /// run of equal values of these columns, in path and row order. Needs + /// numerical columns and no histograms or by-class summaries. + pub max_per: Vec, +} + +impl Default for ColStatsOptions { + fn default() -> Self { + Self { + percentiles: false, + summary_stats_by_class: true, + histograms: true, + category_counts_limit: None, + max_per: Vec::new(), + } + } +} + +fn key_rows(batch: &RecordBatch, names: &[String]) -> Result { + let arrays: Vec = names + .iter() + .map(|name| Ok(batch.column(batch.schema().index_of(name)?).clone())) + .collect::>()?; + let converter = RowConverter::new( + arrays + .iter() + .map(|array| SortField::new(array.data_type().clone())) + .collect(), + )?; + Ok(converter.convert_columns(&arrays)?) +} + +/// Per-column maximum of one run of equal `max_per` keys, as one-row arrays. +struct RunMax { + key: OwnedRow, + best: Vec<(Option, Option)>, +} + +impl RunMax { + fn merge(&mut self, other: RunMax) { + for (mine, theirs) in self.best.iter_mut().zip(other.best) { + if larger(theirs.0, mine.0) { + *mine = theirs; + } + } + } +} + +fn larger(candidate: Option, current: Option) -> bool { + matches!(candidate, Some(v) if !v.is_nan() && current.is_none_or(|c| v > c)) +} + +/// The first and last runs of a row group, which may continue into the +/// neighbouring row groups. +#[derive(Default)] +struct RowGroupRuns { + head: Option, + tail: Option, +} + +/// The open run while scanning: its best values so far, and which of them +/// sit in the current batch. +struct OpenRun { + key: Option, + best: Vec<(Option, Option)>, + in_batch: Vec>, + started_in_batch: bool, +} + +impl OpenRun { + fn new(columns: usize, key: Option, started_in_batch: bool) -> Self { + Self { + key, + best: vec![(None, None); columns], + in_batch: vec![None; columns], + started_in_batch, + } + } + + /// Copy (not slice, which would pin the batch) the best values out. + fn settle(mut self, arrays: &[ArrayRef], key: impl FnOnce() -> OwnedRow) -> Result { + for (col, row) in self.in_batch.iter().enumerate() { + if let Some(row) = row { + let row = UInt32Array::from(vec![*row as u32]); + self.best[col].1 = Some(take(&arrays[col], &row, None)?); + } + } + Ok(RunMax { + key: self.key.unwrap_or_else(key), + best: self.best, + }) + } +} + +fn runs_batch(columns: &[ColumnSpec], arrays: Vec) -> Result { + let fields: Vec = columns + .iter() + .zip(&arrays) + .map(|(col, array)| Field::new(&col.name, array.data_type().clone(), true)) + .collect(); + Ok(RecordBatch::try_new(Arc::new(Schema::new(fields)), arrays)?) +} + +fn runs_to_batch(columns: &[ColumnSpec], types: &[DataType], runs: &[RunMax]) -> Result { + let arrays = (0..columns.len()) + .map(|col| { + let parts: Vec = runs + .iter() + .map(|run| { + run.best[col] + .1 + .clone() + .unwrap_or_else(|| new_null_array(&types[col], 1)) + }) + .collect(); + Ok(concat(&parts.iter().map(|a| a.as_ref()).collect::>())?) + }) + .collect::>>()?; + runs_batch(columns, arrays) +} + +/// Reduce a row group's batches to the per-column maximum of each run; +/// `emit` gets the maxima of every run that starts and ends in the group +/// except the first. +fn reduce_runs( + batches: impl Iterator>, + keys: &[String], + columns: &[ColumnSpec], + emit: &mut dyn FnMut(&RecordBatch) -> Result<()>, +) -> Result { + let mut runs = RowGroupRuns::default(); + let mut spilled: Vec = Vec::new(); + let mut types: Vec = Vec::new(); + let mut open: Option = None; + for batch in batches { + let batch = batch?; + if batch.num_rows() == 0 { + continue; + } + let key = key_rows(&batch, keys)?; + let arrays: Vec = columns + .iter() + .map(|col| Ok(batch.column(batch.schema().index_of(&col.name)?).clone())) + .collect::>()?; + types = arrays.iter().map(|array| array.data_type().clone()).collect(); + let values: Vec = arrays + .iter() + .map(|array| { + Ok(cast(array, &DataType::Float64)? + .as_any() + .downcast_ref::() + .context("numerical column")? + .clone()) + }) + .collect::>()?; + let mut picks: Vec>> = vec![Vec::new(); columns.len()]; + for row in 0..batch.num_rows() { + let new_run = match (row, open.as_ref()) { + (_, None) => true, + (0, Some(o)) => o.key.as_ref().is_none_or(|k| k.row() != key.row(0)), + _ => key.row(row) != key.row(row - 1), + }; + if new_run { + if let Some(done) = open.take() { + let done_key = || key.row(row - 1).owned(); + if runs.head.is_none() { + runs.head = Some(done.settle(&arrays, done_key)?); + } else if done.started_in_batch { + for (col, row) in done.in_batch.iter().enumerate() { + picks[col].push(row.map(|row| row as u32)); + } + } else { + spilled.push(done.settle(&arrays, done_key)?); + } + } + open = Some(OpenRun::new(columns.len(), None, true)); + } + let o = open.as_mut().expect("open run"); + for (col, values) in values.iter().enumerate() { + let value = (!values.is_null(row)).then(|| values.value(row)); + if larger(value, o.best[col].0) { + o.best[col].0 = value; + o.in_batch[col] = Some(row); + } + } + } + if !picks[0].is_empty() { + let arrays = arrays + .iter() + .zip(picks) + .map(|(array, rows)| Ok(take(array, &UInt32Array::from(rows), None)?)) + .collect::>>()?; + emit(&runs_batch(columns, arrays)?)?; + } + if let Some(o) = open.take() { + let last = batch.num_rows() - 1; + let run = o.settle(&arrays, || key.row(last).owned())?; + let mut next = OpenRun::new(columns.len(), Some(run.key), false); + next.best = run.best; + open = Some(next); + } + } + if let Some(o) = open { + let run = RunMax { + key: o.key.expect("settled key"), + best: o.best, + }; + if runs.head.is_none() { + runs.head = Some(run); + } else { + runs.tail = Some(run); + } + } + if !spilled.is_empty() { + emit(&runs_to_batch(columns, &types, &spilled)?)?; + } + Ok(runs) +} + fn validate_stats_dtype(col: &ColumnSpec) -> Result<()> { if col.semantic_type == Some(SemanticType::Timestamp) { return Ok(()); @@ -582,9 +819,7 @@ fn accumulate_distributions_for_row_group( timestamp_specs: &[TimestampHistogramSpec], out: &mut DistributionPartial, ) -> Result<()> { - let file = File::open(path)?; - let builder = ParquetRecordBatchReaderBuilder::try_new(file)?; - let mut projected_indices: Vec = floating_specs + let projected_indices: Vec = floating_specs .iter() .map(|spec| parquet_indices[spec.spec_idx]) .chain( @@ -598,14 +833,7 @@ fn accumulate_distributions_for_row_group( .map(|spec| parquet_indices[spec.spec_idx]), ) .collect(); - projected_indices.sort_unstable(); - projected_indices.dedup(); - let projection = ProjectionMask::roots(builder.parquet_schema(), projected_indices); - let reader = builder - .with_row_groups(vec![row_group]) - .with_projection(projection) - .build()?; - for batch_res in reader { + for batch_res in row_group_batches(path, row_group, &projected_indices, &[])? { let batch = batch_res?; let schema = batch.schema(); for (hist_idx, spec) in floating_specs.iter().enumerate() { @@ -787,30 +1015,50 @@ fn accumulate_by_class( Ok(()) } -fn stats_for_row_group( +type RowGroupStats = (FileColStats, ByClassStats, Vec); + +/// The row group's batches of the projected columns and `keys`. +pub(crate) fn row_group_batches( path: &Path, row_group: usize, - columns: &[ColumnSpec], parquet_indices: &[usize], - integer_spec_indices: &[usize], - table_cat: &Arc, - score_class: Option<(usize, usize)>, -) -> Result<(FileColStats, ByClassStats, Vec)> { - let mut fcs = FileColStats::new_with_categories(columns, Arc::clone(table_cat)); - let mut by_class = ByClassStats::new(); - let mut integer_ranges = vec![IntegerRange::default(); integer_spec_indices.len()]; + keys: &[String], +) -> Result>> { let file = File::open(path)?; let builder = ParquetRecordBatchReaderBuilder::try_new(file)?; - let projection = ProjectionMask::roots(builder.parquet_schema(), parquet_indices.to_vec()); + let mut indices = parquet_indices.to_vec(); + for name in keys { + indices.push( + builder + .schema() + .index_of(name) + .with_context(|| format!("schema error: missing parquet column '{name}'"))?, + ); + } + indices.sort_unstable(); + indices.dedup(); + let projection = ProjectionMask::roots(builder.parquet_schema(), indices); let reader = builder .with_row_groups(vec![row_group]) .with_projection(projection) .build()?; - for batch_res in reader { - let batch = batch_res?; + Ok(reader.map(|batch| Ok(batch?))) +} + +struct StatsSink<'a> { + columns: &'a [ColumnSpec], + integer_spec_indices: &'a [usize], + score_class: Option<(usize, usize)>, + stats: RowGroupStats, +} + +impl StatsSink<'_> { + fn drive(&mut self, batch: &RecordBatch) -> Result<()> { + let (fcs, by_class, integer_ranges) = &mut self.stats; let schema = batch.schema(); - for (i, col) in columns.iter().enumerate() { - let range = integer_spec_indices + for (i, col) in self.columns.iter().enumerate() { + let range = self + .integer_spec_indices .binary_search(&i) .ok() .map(|idx| &mut integer_ranges[idx]); @@ -818,7 +1066,7 @@ fn stats_for_row_group( i, batch.column(schema.index_of(&col.name)?), col, - &mut fcs, + fcs, range, ); if let Err(error) = result { @@ -831,42 +1079,75 @@ fn stats_for_row_group( ))); } } - if let Some((score, class)) = score_class { + if let Some((score, class)) = self.score_class { accumulate_by_class( - &batch, - schema.index_of(&columns[score].name)?, - schema.index_of(&columns[class].name)?, - &columns[score], - &columns[class], - &mut by_class, + batch, + schema.index_of(&self.columns[score].name)?, + schema.index_of(&self.columns[class].name)?, + &self.columns[score], + &self.columns[class], + by_class, )?; } + Ok(()) } - fcs.flush_categories(); - Ok((fcs, by_class, integer_ranges)) } -/// Compute full and min Kumo-format statistics for prediction parquet files. -pub fn compute_parquet_col_stats( - parquet_paths: &[&Path], +type RunEdges = Vec<(Option, Option)>; + +fn stats_for_row_group( + path: &Path, + row_group: usize, + parquet_indices: &[usize], + sink: &mut StatsSink<'_>, + max_per: &[String], +) -> Result<(Option, Option)> { + let batches = row_group_batches(path, row_group, parquet_indices, max_per)?; + let edges = if max_per.is_empty() { + for batch in batches { + sink.drive(&batch?)?; + } + (None, None) + } else { + let columns = sink.columns; + let runs = reduce_runs(batches, max_per, columns, &mut |batch| sink.drive(batch))?; + (runs.head, runs.tail) + }; + sink.stats.0.flush_categories(); + Ok(edges) +} + +/// Join runs that cross row groups, in order, into one batch of maxima. +fn stitch_runs( columns: &[ColumnSpec], - data_version: &str, -) -> Result<(String, String)> { - compute_parquet_col_stats_with_max_categories( - parquet_paths, - columns, - data_version, - MAX_CATEGORIES, - ) + types: &[DataType], + edges: RunEdges, +) -> Result> { + let mut done = Vec::new(); + let mut pending: Option = None; + for (head, tail) in edges { + let Some(head) = head else { continue }; + match pending.as_mut() { + Some(open) if open.key == head.key => open.merge(head), + _ => done.extend(pending.replace(head)), + } + if let Some(tail) = tail { + done.extend(pending.replace(tail)); + } + } + done.extend(pending); + if done.is_empty() { + return Ok(None); + } + Ok(Some(runs_to_batch(columns, types, &done)?)) } -/// Test-focused variant of [`compute_parquet_col_stats`] with an injectable -/// retained-vocabulary limit. Production callers always use -/// [`MAX_CATEGORIES`] through the public wrapper above. +/// Accumulate per-column statistics for [`compute_parquet_col_stats_with_options`]. fn accumulate_parquet_col_stats<'a>( parquet_paths: &'a [&'a Path], columns: &'a [ColumnSpec], max_categories: u64, + options: &ColStatsOptions, ) -> Result<(FileColStats, Arc, Accumulated<'a>)> { if max_categories == 0 { bail!("category count limit must be greater than zero"); @@ -906,50 +1187,84 @@ fn accumulate_parquet_col_stats<'a>( }) .collect(); - let score_class = find_score_class(columns); - - let (acc, by_class, integer_ranges) = if work.is_empty() { - let mut fcs = FileColStats::new_with_categories(columns, Arc::clone(&table_cat)); - fcs.flush_categories(); - ( - fcs, + let score_class = find_score_class(columns).filter(|_| options.summary_stats_by_class); + let max_per = options.max_per.as_slice(); + if !max_per.is_empty() + && (options.histograms + || score_class.is_some() + || columns + .iter() + .any(|col| col.semantic_type != Some(SemanticType::Numerical))) + { + bail!("max_per needs numerical columns and no histograms or by-class summaries"); + } + let sink = || StatsSink { + columns, + integer_spec_indices: &integer_spec_indices, + score_class, + stats: ( + FileColStats::new_with_categories(columns, Arc::clone(&table_cat)), ByClassStats::new(), vec![IntegerRange::default(); integer_spec_indices.len()], - ) - } else { - work.par_iter() - .map(|&(pi, rg)| { - stats_for_row_group( - parquet_paths[pi], - rg, - columns, - &parquet_indices[pi], - &integer_spec_indices, - &table_cat, - score_class, - ) - }) - .try_reduce( - || { - ( - FileColStats::new_with_categories(columns, Arc::clone(&table_cat)), - ByClassStats::new(), - vec![IntegerRange::default(); integer_spec_indices.len()], - ) - }, - |(mut acc, mut by_class, mut integer_ranges), - (mut partial, partial_by_class, partial_ranges)| { - acc.merge(&mut partial); - for (class, stats) in partial_by_class { - by_class.entry(class).or_default().merge(&stats); - } - for (total, range) in integer_ranges.iter_mut().zip(partial_ranges) { - total.merge(range); - } - Ok((acc, by_class, integer_ranges)) - }, - )? + ), + }; + let merge = |total: &mut RowGroupStats, (mut partial, partial_by_class, partial_ranges): RowGroupStats| { + let (acc, by_class, integer_ranges) = total; + acc.merge(&mut partial); + for (class, stats) in partial_by_class { + by_class.entry(class).or_default().merge(&stats); + } + for (total, range) in integer_ranges.iter_mut().zip(partial_ranges) { + total.merge(range); + } }; + + let mut total = sink(); + if !work.is_empty() { + // Fixed chunks merged in order: the result never depends on scheduling. + let chunk = work.len().div_ceil(64); + let chunks = work + .par_chunks(chunk) + .map(|items| { + let mut out = sink(); + let mut edges = RunEdges::new(); + for &(pi, rg) in items { + edges.push(stats_for_row_group( + parquet_paths[pi], + rg, + &parquet_indices[pi], + &mut out, + max_per, + )?); + } + Ok((out.stats, edges)) + }) + .collect::>>()?; + let mut edges = RunEdges::new(); + for (stats, chunk_edges) in chunks { + merge(&mut total.stats, stats); + edges.extend(chunk_edges); + } + let types = if max_per.is_empty() { + Vec::new() + } else { + let builder = ParquetRecordBatchReaderBuilder::try_new(File::open(parquet_paths[0])?)?; + columns + .iter() + .map(|col| Ok(builder.schema().field_with_name(&col.name)?.data_type().clone())) + .collect::>>()? + }; + if let Some(batch) = stitch_runs(columns, &types, edges)? { + let mut boundary = sink(); + boundary.drive(&batch)?; + boundary.stats.0.flush_categories(); + merge(&mut total.stats, boundary.stats); + } + } + let (mut acc, by_class, integer_ranges) = total.stats; + if work.is_empty() { + acc.flush_categories(); + } Ok(( acc, table_cat, @@ -978,17 +1293,53 @@ struct Accumulated<'a> { score_class: Option<(usize, usize)>, } -/// Test-focused variant of [`compute_parquet_col_stats`] with an injectable -/// retained-vocabulary limit. Production callers always use -/// [`MAX_CATEGORIES`] through the public wrapper above. -pub(crate) fn compute_parquet_col_stats_with_max_categories( +/// [`compute_parquet_col_stats_with_options`] with default options and a given +/// retained-vocabulary limit. +#[cfg(test)] +fn compute_parquet_col_stats_with_max_categories( + parquet_paths: &[&Path], + columns: &[ColumnSpec], + data_version: &str, + max_categories: u64, +) -> Result<(String, String)> { + compute_parquet_col_stats_with_options( + parquet_paths, + columns, + data_version, + max_categories, + &ColStatsOptions::default(), + ) +} + +/// Compute full and min Kumo-format statistics for prediction parquet files, +/// with the optional statistics in `options`. +pub fn compute_parquet_col_stats_with_options( parquet_paths: &[&Path], columns: &[ColumnSpec], data_version: &str, max_categories: u64, + options: &ColStatsOptions, ) -> Result<(String, String)> { - let (acc, table_cat, run) = - accumulate_parquet_col_stats(parquet_paths, columns, max_categories)?; + let (mut acc, table_cat, run) = + accumulate_parquet_col_stats(parquet_paths, columns, max_categories, options)?; + let percentiles: Vec<(usize, Vec)> = if options.percentiles { + let levels: Vec = (1..100).map(|q| q as f64 / 100.0).collect(); + (0..columns.len()) + .filter_map(|idx| Some((idx, acc.numerical_quantiles(idx, &levels)?))) + .collect() + } else { + Vec::new() + }; + let capped: Vec<(usize, bool)> = if options.category_counts_limit.is_some() { + table_cat + .columns + .iter() + .enumerate() + .filter_map(|(idx, counts)| Some((idx, counts.as_ref()?.is_capped()))) + .collect() + } else { + Vec::new() + }; let names: Vec = columns.iter().map(|c| c.name.clone()).collect(); let types: Vec = columns.iter().map(|c| c.ingested_dtype.clone()).collect(); let mut kumo = acc.into_kumo_stats( @@ -1000,12 +1351,42 @@ pub(crate) fn compute_parquet_col_stats_with_max_categories( data_version, )?; let mut full = std::mem::take(&mut kumo.stats); - add_distribution_stats(&run, &mut full, data_version)?; + add_distribution_stats(&run, &mut full, data_version, options.histograms)?; // Category arrays are filled in, and trimmed for min, when serialized. let mut min = full.clone(); + for (idx, capped) in capped { + let name = &columns[idx].name; + let prefix = format!("{name}__"); + let stype = full.keys().find_map(|key| { + key.strip_prefix(&prefix)? + .strip_suffix("__CATEGORY_COUNTS__False") + .map(str::to_string) + }); + if let Some(stype) = stype { + let count = full + .get(&kumo_key(name, &stype, "COUNT", false)) + .and_then(|entry| entry.get("value")?.as_u64()) + .unwrap_or_default(); + full.insert( + kumo_key(name, &stype, "CATEGORY_COUNTS_CAPPED", false), + kumo_entry(json!(capped), false, data_version, count), + ); + } + } + for (idx, values) in percentiles { + let name = &columns[idx].name; + let count = full + .get(&kumo_key(name, "kumo.numerical", "COUNT", false)) + .and_then(|entry| entry.get("value")?.as_u64()) + .unwrap_or_default(); + full.insert( + kumo_key(name, "kumo.numerical", "PERCENTILES", true), + kumo_entry(json!(values), true, data_version, count), + ); + } add_summary_stats_by_class(&run, &mut full, &mut min, data_version); Ok(( - kumo.to_json_string(&full, usize::MAX)?, + kumo.to_json_string(&full, options.category_counts_limit.unwrap_or(usize::MAX))?, kumo.to_json_string(&min, MIN_COL_STATS_CATEGORY_LIMIT)?, )) } @@ -1030,6 +1411,7 @@ fn add_distribution_stats( run: &Accumulated<'_>, full: &mut JsonMap, data_version: &str, + histograms: bool, ) -> Result<()> { let Accumulated { parquet_paths, @@ -1059,6 +1441,9 @@ fn add_distribution_stats( } } + if !histograms { + return Ok(()); + } let floating_specs: Vec = columns .iter() .enumerate() @@ -1273,7 +1658,7 @@ mod tests { use arrow_array::builder::{Float64Builder, ListBuilder}; use arrow_array::{ - ArrayRef, Date32Array, Decimal128Array, Decimal256Array, Float64Array, Int32Array, + ArrayRef, Date32Array, Decimal128Array, Decimal256Array, Float32Array, Float64Array, Int32Array, Int64Array, RecordBatch, StringArray, TimestampMicrosecondArray, TimestampSecondArray, }; use arrow::datatypes::i256; @@ -1367,12 +1752,17 @@ mod tests { max_categories: u64, ) -> (String, String) { use diskgraph_ingest::col_stats::minify_kumo_stats; - let (mut acc, table_cat, run) = - accumulate_parquet_col_stats(paths, columns, max_categories).unwrap(); + let (mut acc, table_cat, run) = accumulate_parquet_col_stats( + paths, + columns, + max_categories, + &ColStatsOptions::default(), + ) + .unwrap(); let names: Vec = columns.iter().map(|c| c.name.clone()).collect(); let types: Vec = columns.iter().map(|c| c.ingested_dtype.clone()).collect(); let mut full = acc.to_kumo_json(&names, &types, Some(&table_cat), None, "v1"); - add_distribution_stats(&run, &mut full, "v1").unwrap(); + add_distribution_stats(&run, &mut full, "v1", true).unwrap(); let mut min = minify_kumo_stats(&full, MIN_COL_STATS_CATEGORY_LIMIT); add_summary_stats_by_class(&run, &mut full, &mut min, "v1"); ( @@ -1490,7 +1880,13 @@ mod tests { col("SCORE", IngestedType::Float64, SemanticType::Numerical), col("CLASS", IngestedType::String, SemanticType::Categorical), ]; - let (full, min) = compute_parquet_col_stats(&[path.as_path()], &columns, "v1").unwrap(); + let (full, min) = compute_parquet_col_stats_with_options( + &[path.as_path()], + &columns, + "v1", + MAX_CATEGORIES, + &ColStatsOptions::default(), + ).unwrap(); // COUNT == n across columns (merge across row groups correct). assert_eq!(val(&full, "entity__kumo.ID__COUNT__False"), n as u64); @@ -1601,7 +1997,13 @@ mod tests { ts_col("month", None), ts_col("day", None), ]; - let (full, _) = compute_parquet_col_stats(&[path.as_path()], &columns, "v1").unwrap(); + let (full, _) = compute_parquet_col_stats_with_options( + &[path.as_path()], + &columns, + "v1", + MAX_CATEGORIES, + &ColStatsOptions::default(), + ).unwrap(); for (name, stype) in [ ("rank", "kumo.numerical"), ("revenue", "kumo.numerical"), @@ -1649,7 +2051,13 @@ mod tests { IngestedType::Float32Array, SemanticType::Latlong, )]; - let (full, _) = compute_parquet_col_stats(&[path.as_path()], &columns, "v1").unwrap(); + let (full, _) = compute_parquet_col_stats_with_options( + &[path.as_path()], + &columns, + "v1", + MAX_CATEGORIES, + &ColStatsOptions::default(), + ).unwrap(); assert_eq!( val(&full, "location__kumo.latlong__SEQUENCE_MAX_LENGTH__False"), 2 @@ -1683,7 +2091,13 @@ mod tests { col("SCORE", IngestedType::Float32, SemanticType::Numerical), col("CLASS", IngestedType::Bool, SemanticType::Categorical), ]; - let (full, _min) = compute_parquet_col_stats(&[path.as_path()], &columns, "v1").unwrap(); + let (full, _min) = compute_parquet_col_stats_with_options( + &[path.as_path()], + &columns, + "v1", + MAX_CATEGORIES, + &ColStatsOptions::default(), + ).unwrap(); // Numerical scalar stats populated (not empty-Welford defaults). let mx = val(&full, "SCORE__kumo.numerical__MAX__False") .as_f64() @@ -1732,7 +2146,13 @@ mod tests { col("SCORE", IngestedType::Float64, SemanticType::Numerical), col("CLASS", IngestedType::String, SemanticType::Categorical), ]; - let (full, _min) = compute_parquet_col_stats(&[path.as_path()], &columns, "v1").unwrap(); + let (full, _min) = compute_parquet_col_stats_with_options( + &[path.as_path()], + &columns, + "v1", + MAX_CATEGORIES, + &ColStatsOptions::default(), + ).unwrap(); let bc = val( &full, "SCORE__kumo.numerical__SUMMARY_STATS_BY_CLASS__False", @@ -1783,7 +2203,13 @@ mod tests { col("score", IngestedType::Float64, SemanticType::Numerical), ]; let (full, _) = - compute_parquet_col_stats(&[p1.as_path(), p2.as_path()], &columns, "v1").unwrap(); + compute_parquet_col_stats_with_options( + &[p1.as_path(), p2.as_path()], + &columns, + "v1", + MAX_CATEGORIES, + &ColStatsOptions::default(), + ).unwrap(); assert_eq!(val(&full, "entity__kumo.ID__COUNT__False"), 4); assert_eq!(val(&full, "score__kumo.numerical__MIN__False"), 1.0); assert_eq!(val(&full, "score__kumo.numerical__MAX__False"), 4.0); @@ -1814,10 +2240,12 @@ mod tests { vec![StdArc::new(Float64Array::from(vec![2.0])) as ArrayRef], 1, ); - let err = compute_parquet_col_stats( + let err = compute_parquet_col_stats_with_options( &[p1.as_path(), p2.as_path()], &[col("score", IngestedType::Float64, SemanticType::Numerical)], "v1", + MAX_CATEGORIES, + &ColStatsOptions::default(), ) .unwrap_err(); assert!(format!("{err:#}").contains("column 'score'")); @@ -1882,7 +2310,13 @@ mod tests { ); let columns = vec![ts_col("created_at", None), ts_col("updated_at", None)]; - let (full, _min) = compute_parquet_col_stats(&[path.as_path()], &columns, "v1").unwrap(); + let (full, _min) = compute_parquet_col_stats_with_options( + &[path.as_path()], + &columns, + "v1", + MAX_CATEGORIES, + &ColStatsOptions::default(), + ).unwrap(); assert_eq!( val( &full, @@ -1937,7 +2371,13 @@ mod tests { ); let (full, _min) = - compute_parquet_col_stats(&[path.as_path()], &[ts_col("ts", None)], "v1").unwrap(); + compute_parquet_col_stats_with_options( + &[path.as_path()], + &[ts_col("ts", None)], + "v1", + MAX_CATEGORIES, + &ColStatsOptions::default(), + ).unwrap(); let histogram = val(&full, "ts__kumo.timestamp__HISTOGRAM_HALF_MONTH__False"); let histogram = histogram.as_array().unwrap(); assert_eq!(histogram.len(), TIMESTAMP_HISTOGRAM_BINS); @@ -1981,7 +2421,13 @@ mod tests { ); let (full, _) = - compute_parquet_col_stats(&[path.as_path()], &[ts_col("ts", None)], "v1").unwrap(); + compute_parquet_col_stats_with_options( + &[path.as_path()], + &[ts_col("ts", None)], + "v1", + MAX_CATEGORIES, + &ColStatsOptions::default(), + ).unwrap(); let histogram = val(&full, "ts__kumo.timestamp__HISTOGRAM_HALF_MONTH__False"); let histogram = histogram.as_array().unwrap(); assert_eq!(histogram.len(), TIMESTAMP_HISTOGRAM_BINS); @@ -2033,7 +2479,13 @@ mod tests { ts_col("null_ts", None), col("score", IngestedType::Float64, SemanticType::Numerical), ]; - let (full, _min) = compute_parquet_col_stats(&[path.as_path()], &columns, "v1").unwrap(); + let (full, _min) = compute_parquet_col_stats_with_options( + &[path.as_path()], + &columns, + "v1", + MAX_CATEGORIES, + &ColStatsOptions::default(), + ).unwrap(); let stats: JsonMap = serde_json::from_str(&full).unwrap(); assert!(!stats.contains_key("ts__kumo.timestamp__HISTOGRAM_HALF_MONTH__False")); assert!(!stats.contains_key("null_ts__kumo.timestamp__HISTOGRAM_HALF_MONTH__False")); @@ -2077,7 +2529,13 @@ mod tests { let mut timestamp = ts_col("ts", Some(TimestampUnit::Millis)); timestamp.ingested_dtype = IngestedType::Int64; let columns = vec![timestamp]; - let (full, _min) = compute_parquet_col_stats(&[path.as_path()], &columns, "v1").unwrap(); + let (full, _min) = compute_parquet_col_stats_with_options( + &[path.as_path()], + &columns, + "v1", + MAX_CATEGORIES, + &ColStatsOptions::default(), + ).unwrap(); assert_eq!(val(&full, "ts__kumo.timestamp__MIN__False"), expected); assert_eq!(val(&full, "ts__kumo.timestamp__MAX__False"), expected); assert_eq!( @@ -2088,7 +2546,13 @@ mod tests { let mut timestamp = ts_col("ts", None); timestamp.ingested_dtype = IngestedType::Int64; let (full, _min) = - compute_parquet_col_stats(&[path.as_path()], &[timestamp], "v1").unwrap(); + compute_parquet_col_stats_with_options( + &[path.as_path()], + &[timestamp], + "v1", + MAX_CATEGORIES, + &ColStatsOptions::default(), + ).unwrap(); assert_eq!(val(&full, "ts__kumo.timestamp__MIN__False"), expected); } @@ -2106,7 +2570,13 @@ mod tests { 64, ); let columns = vec![ts_col("ts", Some(TimestampUnit::Seconds))]; - let (full, _min) = compute_parquet_col_stats(&[path.as_path()], &columns, "v1").unwrap(); + let (full, _min) = compute_parquet_col_stats_with_options( + &[path.as_path()], + &columns, + "v1", + MAX_CATEGORIES, + &ColStatsOptions::default(), + ).unwrap(); assert_eq!( val(&full, "ts__kumo.timestamp__MIN__False"), "kumo_timestamp_2022-07-31 22:00:01" @@ -2129,7 +2599,13 @@ mod tests { IngestedType::Float64, SemanticType::Numerical, )]; - let err = compute_parquet_col_stats(&[path.as_path()], &columns, "v1").unwrap_err(); + let err = compute_parquet_col_stats_with_options( + &[path.as_path()], + &columns, + "v1", + MAX_CATEGORIES, + &ColStatsOptions::default(), + ).unwrap_err(); let msg = format!("{err:#}"); assert!(msg.contains("NOT_THERE") && msg.contains("schema error")); @@ -2142,7 +2618,13 @@ mod tests { 1024, ); let columns = vec![col("SCORE", IngestedType::Float64, SemanticType::Numerical)]; - let (full, _) = compute_parquet_col_stats(&[path.as_path()], &columns, "v1").unwrap(); + let (full, _) = compute_parquet_col_stats_with_options( + &[path.as_path()], + &columns, + "v1", + MAX_CATEGORIES, + &ColStatsOptions::default(), + ).unwrap(); assert_eq!( val(&full, "SCORE__kumo.numerical__MIN__False"), JsonValue::Null @@ -2175,7 +2657,13 @@ mod tests { col("cnt", IngestedType::Int32, SemanticType::Numerical), col("b", IngestedType::UInt8, SemanticType::Numerical), ]; - let (full, _min) = compute_parquet_col_stats(&[path.as_path()], &columns, "v1").unwrap(); + let (full, _min) = compute_parquet_col_stats_with_options( + &[path.as_path()], + &columns, + "v1", + MAX_CATEGORIES, + &ColStatsOptions::default(), + ).unwrap(); assert_eq!(val(&full, "cnt__kumo.numerical__COUNT__False"), 3u64); assert_eq!(val(&full, "b__kumo.numerical__COUNT__False"), 3u64); let mfull: JsonMap = serde_json::from_str(&full).unwrap(); @@ -2199,10 +2687,12 @@ mod tests { 1, ); - let (full, _) = compute_parquet_col_stats( + let (full, _) = compute_parquet_col_stats_with_options( &[path.as_path()], &[col("v", IngestedType::Int64, SemanticType::Numerical)], "v1", + MAX_CATEGORIES, + &ColStatsOptions::default(), ) .unwrap(); assert_eq!(val(&full, "v__kumo.numerical__MIN__False"), i64::MAX - 2); @@ -2227,10 +2717,12 @@ mod tests { 1, ); - let (full, _) = compute_parquet_col_stats( + let (full, _) = compute_parquet_col_stats_with_options( &[path.as_path()], &[col("v", IngestedType::Int64, SemanticType::Numerical)], "v1", + MAX_CATEGORIES, + &ColStatsOptions::default(), ) .unwrap(); let histogram = val(&full, "v__kumo.numerical__HISTOGRAM__False"); @@ -2259,14 +2751,12 @@ mod tests { 1, ); - let (full, _) = compute_parquet_col_stats( + let (full, _) = compute_parquet_col_stats_with_options( &[path.as_path()], - &[col( - "wide_int", - IngestedType::Int64, - SemanticType::Numerical, - )], + &[col( "wide_int", IngestedType::Int64, SemanticType::Numerical, )], "v1", + MAX_CATEGORIES, + &ColStatsOptions::default(), ) .unwrap(); @@ -2299,14 +2789,12 @@ mod tests { 1, ); - let (full, _) = compute_parquet_col_stats( + let (full, _) = compute_parquet_col_stats_with_options( &[path.as_path()], - &[col( - "integer", - IngestedType::Int64, - SemanticType::Numerical, - )], + &[col( "integer", IngestedType::Int64, SemanticType::Numerical, )], "v1", + MAX_CATEGORIES, + &ColStatsOptions::default(), ) .unwrap(); @@ -2369,14 +2857,12 @@ mod tests { vec![values], 1, ); - let error = compute_parquet_col_stats( + let error = compute_parquet_col_stats_with_options( &[path.as_path()], - &[col( - "wide_int", - IngestedType::Int64, - SemanticType::Numerical, - )], + &[col( "wide_int", IngestedType::Int64, SemanticType::Numerical, )], "v1", + MAX_CATEGORIES, + &ColStatsOptions::default(), ) .unwrap_err(); @@ -2411,13 +2897,12 @@ mod tests { 1, ); - let (full, _) = compute_parquet_col_stats( + let (full, _) = compute_parquet_col_stats_with_options( &[path.as_path()], - &[ - col("integer", IngestedType::Int64, SemanticType::Numerical), - col("float", IngestedType::Float64, SemanticType::Numerical), - ], + &[ col("integer", IngestedType::Int64, SemanticType::Numerical), col("float", IngestedType::Float64, SemanticType::Numerical), ], "v1", + MAX_CATEGORIES, + &ColStatsOptions::default(), ) .unwrap(); let stats: JsonMap = serde_json::from_str(&full).unwrap(); @@ -2452,7 +2937,13 @@ mod tests { col("subnormal", IngestedType::Float64, SemanticType::Numerical), ]; - let (full, _min) = compute_parquet_col_stats(&[path.as_path()], &columns, "v1").unwrap(); + let (full, _min) = compute_parquet_col_stats_with_options( + &[path.as_path()], + &columns, + "v1", + MAX_CATEGORIES, + &ColStatsOptions::default(), + ).unwrap(); let histogram = val(&full, "extreme__kumo.numerical__HISTOGRAM__False"); let histogram = histogram.as_array().unwrap(); assert_eq!(histogram.len(), NUMERICAL_HISTOGRAM_BINS); @@ -2509,7 +3000,13 @@ mod tests { ); let columns = vec![ts_col("ts", None)]; let (full, _min) = - compute_parquet_col_stats(&[path.as_path()], &columns, "v1").unwrap(); + compute_parquet_col_stats_with_options( + &[path.as_path()], + &columns, + "v1", + MAX_CATEGORIES, + &ColStatsOptions::default(), + ).unwrap(); assert_eq!( val(&full, "ts__kumo.timestamp__MIN__False"), "kumo_timestamp_1969-12-31 23:59:59", @@ -2530,4 +3027,260 @@ mod tests { ); } } + + fn with_options(paths: &[&Path], columns: &[ColumnSpec], options: ColStatsOptions) -> String { + compute_parquet_col_stats_with_options(paths, columns, "v1", MAX_CATEGORIES, &options) + .unwrap() + .0 + } + + fn write_scores(path: &Path, scores: Vec, rows_per_group: usize) { + write_parquet( + path, + Schema::new(vec![Field::new("SCORE", DataType::Float64, true)]), + vec![StdArc::new(Float64Array::from(scores)) as ArrayRef], + rows_per_group, + ); + } + + #[test] + fn unset_options_change_nothing_and_options_are_additive() { + let dir = tempdir().unwrap(); + let path = dir.path().join("scores.parquet"); + let n = 3000usize; + write_parquet( + &path, + Schema::new(vec![ + Field::new("SCORE", DataType::Float64, false), + Field::new("CLASS", DataType::Utf8, false), + ]), + vec![ + StdArc::new(Float64Array::from( + (0..n).map(|i| (i % 97) as f64 / 97.0).collect::>(), + )) as ArrayRef, + StdArc::new(StringArray::from( + (0..n) + .map(|i| format!("c{}", (i * i) % 50)) + .collect::>(), + )) as ArrayRef, + ], + 256, + ); + let columns = vec![ + col("SCORE", IngestedType::Float64, SemanticType::Numerical), + col("CLASS", IngestedType::String, SemanticType::Categorical), + ]; + let paths = [path.as_path()]; + let parse = |json: &str| serde_json::from_str::>(json).unwrap(); + + let (default, default_min) = compute_parquet_col_stats_with_options( + &paths, + &columns, + "v1", + MAX_CATEGORIES, + &ColStatsOptions::default(), + ).unwrap(); + let unset = compute_parquet_col_stats_with_options( + &paths, + &columns, + "v1", + MAX_CATEGORIES, + &ColStatsOptions::default(), + ) + .unwrap(); + assert_eq!((default.clone(), default_min.clone()), unset); + + let (full, min) = compute_parquet_col_stats_with_options( + &paths, + &columns, + "v1", + MAX_CATEGORIES, + &ColStatsOptions { + percentiles: true, + summary_stats_by_class: false, + histograms: false, + category_counts_limit: Some(5), + ..Default::default() + }, + ) + .unwrap(); + let (mut full, mut default) = (parse(&full), parse(&default)); + let percentiles = full + .remove("SCORE__kumo.numerical__PERCENTILES__True") + .unwrap(); + assert_eq!(percentiles["value"].as_array().unwrap().len(), 99); + assert_eq!(percentiles["num_entries"], json!(n)); + for key in [ + "SCORE__kumo.numerical__SUMMARY_STATS_BY_CLASS__False", + "SCORE__kumo.numerical__HISTOGRAM__False", + ] { + assert!(default.remove(key).is_some()); + } + assert_eq!( + full.remove("CLASS__kumo.categorical__CATEGORY_COUNTS_CAPPED__False").unwrap()["value"], + json!(false) + ); + let counts_key = "CLASS__kumo.categorical__CATEGORY_COUNTS__False"; + let all = default.remove(counts_key).unwrap()["value"].clone(); + let top = full.remove(counts_key).unwrap()["value"].clone(); + assert_eq!(top.as_array().unwrap()[..], all.as_array().unwrap()[..5]); + assert_eq!(full, default); + let mut default_min = parse(&default_min); + default_min.remove("SCORE__kumo.numerical__SUMMARY_STATS_BY_CLASS__False"); + default_min.remove("SCORE__kumo.numerical__HISTOGRAM__False"); + assert_eq!(parse(&min), default_min); + } + + #[test] + fn capped_category_counts_are_flagged() { + let dir = tempdir().unwrap(); + let path = dir.path().join("classes.parquet"); + write_parquet( + &path, + Schema::new(vec![Field::new("CLASS", DataType::Utf8, false)]), + vec![StdArc::new(StringArray::from(vec!["a", "b", "c", "a"])) as ArrayRef], + 100, + ); + let columns = vec![col("CLASS", IngestedType::String, SemanticType::Categorical)]; + let options = ColStatsOptions { + category_counts_limit: Some(1), + ..Default::default() + }; + for (cap, capped) in [(2, true), (3, false)] { + let (full, _) = + compute_parquet_col_stats_with_options(&[path.as_path()], &columns, "v1", cap, &options) + .unwrap(); + assert_eq!( + val(&full, "CLASS__kumo.categorical__CATEGORY_COUNTS_CAPPED__False"), + json!(capped) + ); + assert_eq!( + val(&full, "CLASS__kumo.categorical__CATEGORY_COUNTS__False"), + json!([["a", 2]]) + ); + } + } + + #[test] + fn percentiles_track_exact_ranks() { + let dir = tempdir().unwrap(); + let path = dir.path().join("scores.parquet"); + let mut scores: Vec = (1..=20_000).map(|i| (i as f64 / 2000.0).exp()).collect(); + scores.extend((1..=5_000).map(|i| -(i as f64) / 10.0)); + scores.push(f64::NAN); + write_scores(&path, scores.clone(), 1000); + let columns = vec![col("SCORE", IngestedType::Float64, SemanticType::Numerical)]; + + let full = with_options( + &[path.as_path()], + &columns, + ColStatsOptions { + percentiles: true, + ..Default::default() + }, + ); + let values = val(&full, "SCORE__kumo.numerical__PERCENTILES__True"); + let mut finite: Vec = scores.into_iter().filter(|v| v.is_finite()).collect(); + finite.sort_by(f64::total_cmp); + for (idx, value) in values.as_array().unwrap().iter().enumerate() { + let rank = finite.partition_point(|v| *v < value.as_f64().unwrap()) as f64 + / finite.len() as f64; + assert!( + (rank - (idx + 1) as f64 / 100.0).abs() < 0.005, + "q{:02}", + idx + 1 + ); + } + } + + #[test] + fn max_per_takes_each_runs_maximum_across_row_groups_and_files() { + let dir = tempdir().unwrap(); + let schema = Schema::new(vec![ + Field::new("ENTITY", DataType::Int64, false), + Field::new("ANCHOR", DataType::Int64, false), + Field::new("SCORE", DataType::Float32, true), + ]); + // Runs of 3 rows (entity, anchor) with the maximum in any position; + // the same entity repeats with a new anchor, and runs straddle row + // groups (7 rows each) and the file boundary. + let runs = 40i64; + let entity: Vec = (0..runs * 3).map(|i| i / 6).collect(); + let anchor: Vec = (0..runs * 3).map(|i| (i / 3) % 2).collect(); + let score: Vec> = (0..runs * 3) + .map(|i| match (i % 3, (i / 3) % 3) { + (pos, top) if pos == top => Some((i / 3) as f32 + 1.0), + (0, _) => None, + _ => Some(f32::NAN), + }) + .collect(); + let split = 61usize; + let mut paths = Vec::new(); + for (idx, range) in [0..split, split..(runs * 3) as usize].into_iter().enumerate() { + let path = dir.path().join(format!("part-{idx}.parquet")); + write_parquet( + &path, + schema.clone(), + vec![ + StdArc::new(Int64Array::from(entity[range.clone()].to_vec())) as ArrayRef, + StdArc::new(Int64Array::from(anchor[range.clone()].to_vec())) as ArrayRef, + StdArc::new(Float32Array::from(score[range].to_vec())) as ArrayRef, + ], + 7, + ); + paths.push(path); + } + let paths: Vec<&Path> = paths.iter().map(|p| p.as_path()).collect(); + let columns = vec![col("SCORE", IngestedType::Float32, SemanticType::Numerical)]; + let options = |keys: &[&str]| ColStatsOptions { + percentiles: true, + histograms: false, + max_per: keys.iter().map(|key| key.to_string()).collect(), + ..Default::default() + }; + + let full = with_options(&paths, &columns, options(&["ENTITY", "ANCHOR"])); + assert_eq!(val(&full, "SCORE__kumo.numerical__COUNT__False"), json!(runs)); + assert_eq!(val(&full, "SCORE__kumo.numerical__NUM_NA__False"), json!(0)); + assert_eq!(val(&full, "SCORE__kumo.numerical__MIN__False"), json!(1.0)); + assert_eq!( + val(&full, "SCORE__kumo.numerical__MEAN__False"), + json!((runs + 1) as f64 / 2.0) + ); + assert_eq!( + val(&full, "SCORE__kumo.numerical__PERCENTILES__True").as_array().unwrap()[49], + json!(20.5) + ); + + let entity_only = with_options(&paths, &columns, options(&["ENTITY"])); + assert_eq!(val(&entity_only, "SCORE__kumo.numerical__COUNT__False"), json!(runs / 2)); + assert_eq!(val(&entity_only, "SCORE__kumo.numerical__MIN__False"), json!(2.0)); + } + + #[test] + fn output_does_not_depend_on_thread_count() { + let dir = tempdir().unwrap(); + let path = dir.path().join("scores.parquet"); + let scores: Vec = (0..50_000u64) + .map(|i| ((i * 2_654_435_761) % 1_000_003) as f64 / 1_000_003.0) + .collect(); + write_scores(&path, scores, 300); + let columns = vec![col("SCORE", IngestedType::Float64, SemanticType::Numerical)]; + let options = ColStatsOptions { + percentiles: true, + ..Default::default() + }; + let run = |threads| { + rayon::ThreadPoolBuilder::new() + .num_threads(threads) + .build() + .unwrap() + .install(|| with_options(&[path.as_path()], &columns, options.clone())) + }; + let single = run(1); + for threads in [2, 5, 16] { + assert_eq!(run(threads), single); + } + } + } diff --git a/packages/diskgraph/diskgraph-pybind/src/bp_evaluate.rs b/packages/diskgraph/diskgraph-pybind/src/bp_evaluate.rs new file mode 100644 index 000000000..8d0f933d3 --- /dev/null +++ b/packages/diskgraph/diskgraph-pybind/src/bp_evaluate.rs @@ -0,0 +1,937 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +// All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Join batch-prediction outputs with generated labels and compute realized +//! metrics. Both sides are hash-partitioned on their key into Arrow IPC +//! spill files of about `partition_rows` rows each, and at most +//! `MAX_CONCURRENT_JOINS` partitions are joined at once. Metrics follow the +//! fit-time definitions: exact AUROC and AP as in torchmetrics, MAE/RMSE, +//! top-1 accuracy, and MAP@k / precision@k as in PyG. + +use std::collections::hash_map::DefaultHasher; +use std::collections::{BTreeMap, HashMap}; +use std::fs::File; +use std::hash::{Hash, Hasher}; +use std::io::BufWriter; +use std::path::{Path, PathBuf}; +use std::sync::{Arc, Mutex}; + +use anyhow::{bail, Context, Result}; +use arrow::array::TimestampNanosecondArray; +use arrow::array::{Array, AsArray, BooleanArray, FixedSizeBinaryArray, UInt32Array}; +use arrow::compute::{cast, filter, interleave, take_record_batch}; +use arrow::datatypes::Float64Type; +use arrow::ipc::reader::FileReader; +use arrow::ipc::writer::FileWriter; +use arrow::row::{RowConverter, SortField}; +use arrow_array::{ArrayRef, RecordBatch}; +use arrow_schema::{DataType, Field, Schema, SchemaRef, TimeUnit}; +use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; +use parquet::arrow::ArrowWriter; +use rayon::prelude::*; +use serde::Deserialize; +use serde_json::{json, Value as JsonValue}; + +use crate::bp_col_stats::row_group_batches; + +#[derive(Clone, Copy, Debug, Deserialize, PartialEq)] +#[serde(rename_all = "snake_case")] +pub enum EvaluateTask { + Binary, + Regression, + Multiclass, + Link, +} + +#[derive(Clone, Debug, Deserialize)] +pub struct EvaluateSpec { + pub task: EvaluateTask, + pub prediction_keys: Vec, + pub label_keys: Vec, + pub label_target: String, + #[serde(default)] + pub actuals_dir: Option, + /// (label column, output column) pairs written to `actuals_dir`. + #[serde(default)] + pub actual_columns: Vec<(String, String)>, + /// (prediction column, output column) pairs written next to them, from + /// the matched prediction (the top-scored row of a grouped task). + #[serde(default)] + pub prediction_columns: Vec<(String, String)>, + /// (output column, UTC nanoseconds) pairs written as constant timestamps. + #[serde(default)] + pub constant_timestamps: Vec<(String, i64)>, + pub spill_dir: PathBuf, + #[serde(default = "default_partition_rows")] + pub partition_rows: usize, +} + +fn default_partition_rows() -> usize { + 1_000_000 +} + +const MAX_CONCURRENT_JOINS: usize = 8; +const KEY: &str = "__key"; + +type Batches = Box> + Send>; + +/// The type a pair of columns is compared as. +fn common_type(a: &DataType, b: &DataType) -> DataType { + let int = |t: &DataType| t.is_integer() || *t == DataType::Boolean; + if a.is_temporal() || b.is_temporal() { + DataType::Timestamp(TimeUnit::Nanosecond, None) + } else if int(a) && int(b) { + DataType::Int64 + } else if [a, b].iter().all(|t| int(t) || t.is_floating()) { + DataType::Float64 + } else { + DataType::Utf8 + } +} + +/// 128-bit hash of each row of `arrays` cast to `types`; `None` on a null. +fn hashes(arrays: &[&ArrayRef], types: &[DataType]) -> Result>> { + let arrays = arrays + .iter() + .zip(types) + .map(|(array, to)| match (to, array.data_type()) { + (DataType::Timestamp(unit, _), DataType::Timestamp(_, tz)) => { + cast(array, &DataType::Timestamp(*unit, tz.clone())) + } + _ => cast(array, to), + }) + .collect::, _>>()?; + let fields = arrays + .iter() + .map(|a| SortField::new(a.data_type().clone())) + .collect(); + let rows = RowConverter::new(fields)?.convert_columns(&arrays)?; + Ok((0..rows.num_rows()) + .map(|row| { + arrays.iter().all(|array| array.is_valid(row)).then(|| { + let half = |seed: u64| { + let mut hasher = DefaultHasher::new(); + seed.hash(&mut hasher); + rows.row(row).hash(&mut hasher); + hasher.finish() + }; + u128::from(half(1)) << 64 | u128::from(half(2)) + }) + }) + .collect()) +} + +fn classes(array: &ArrayRef, class: &DataType) -> Result>> { + let hashes = hashes(&[array], std::slice::from_ref(class))?; + Ok(hashes.into_iter().map(|h| h.map(|h| h as u64)).collect()) +} + +fn column<'a>(batch: &'a RecordBatch, name: &str) -> Result<&'a ArrayRef> { + batch + .column_by_name(name) + .with_context(|| format!("missing column '{name}'")) +} + +fn floats(array: &ArrayRef) -> Result>> { + Ok(cast(array, &DataType::Float64)? + .as_primitive::() + .iter() + .collect()) +} + +fn keys(batch: &RecordBatch) -> Result + '_> { + let keys = column(batch, KEY)?.as_fixed_size_binary(); + Ok((0..keys.len()).map(|row| u128::from_le_bytes(keys.value(row).try_into().unwrap()))) +} + +/// Sorted, distinct class hashes of each label. +fn label_classes(target: &ArrayRef, link: bool, class: &DataType) -> Result>>> { + if !link { + let classes = classes(target, class)?; + return Ok(classes.into_iter().map(|c| c.map(|c| vec![c])).collect()); + } + let item = Field::new_list_field(class.clone(), true); + let list = cast(target, &DataType::List(Arc::new(item)))?; + let list = list.as_list::(); + let items = classes(list.values(), class)?; + let offsets = list.value_offsets(); + Ok((0..list.len()) + .map(|row| { + list.is_valid(row).then(|| { + let range = offsets[row] as usize..offsets[row + 1] as usize; + let mut row: Vec = items[range].iter().flatten().copied().collect(); + row.sort_unstable(); + row.dedup(); + row + }) + }) + .collect()) +} + +type Sink = (PathBuf, Mutex>>>); + +fn spill((path, writer): &Sink, batch: &RecordBatch) -> Result<()> { + let mut writer = writer.lock().unwrap(); + if writer.is_none() { + let file = BufWriter::new(File::create(path)?); + *writer = Some(FileWriter::try_new(file, &batch.schema())?); + } + Ok(writer.as_mut().unwrap().write(batch)?) +} + +fn unspill((path, writer): Sink) -> Result { + let Some(mut writer) = writer.into_inner().unwrap() else { + return Ok(Box::new(std::iter::empty())); + }; + writer.finish()?; + let reader = FileReader::try_new(File::open(path)?, None)?; + Ok(Box::new(reader.map(|batch| Ok(batch?)))) +} + +type RowGroups = Vec<(usize, usize)>; + +fn row_groups(paths: &[&Path]) -> Result<(RowGroups, usize, SchemaRef)> { + let (mut groups, mut rows, mut schema) = (Vec::new(), 0, Arc::new(Schema::empty())); + for (file, path) in paths.iter().enumerate() { + let builder = ParquetRecordBatchReaderBuilder::try_new(File::open(path)?)?; + schema = builder.schema().clone(); + rows += builder.metadata().file_metadata().num_rows() as usize; + groups.extend((0..builder.metadata().num_row_groups()).map(|group| (file, group))); + } + Ok((groups, rows, schema)) +} + +/// Read `keep` columns, hash `keys`, and route rows with a key to `sinks`. +fn scatter( + paths: &[&Path], + groups: &RowGroups, + (keys, types): (&[String], &[DataType]), + keep: &[String], + sinks: &[Sink], +) -> Result<()> { + let columns: Vec = keys.iter().chain(keep).cloned().collect(); + groups.par_iter().try_for_each(|&(file, group)| { + for batch in row_group_batches(paths[file], group, &[], &columns)? { + let batch = batch?; + let key_columns = keys + .iter() + .map(|name| column(&batch, name)) + .collect::>>()?; + let hashes = hashes(&key_columns, types)?; + let mut parts = vec![Vec::new(); sinks.len()]; + for (row, key) in hashes.iter().enumerate() { + if let Some(key) = key { + parts[((key >> 64) as u64 % sinks.len() as u64) as usize].push(row as u32); + } + } + let key = hashes.into_iter().map(|key| key.map(u128::to_le_bytes)); + let key = FixedSizeBinaryArray::try_from_sparse_iter_with_size(key, 16)?; + let mut out = Vec::new(); + for name in keep { + out.push((name.as_str(), column(&batch, name)?.clone(), true)); + } + out.push((KEY, Arc::new(key), true)); + let batch = RecordBatch::try_from_iter_with_nullable(out)?; + for (sink, rows) in sinks.iter().zip(parts) { + if !rows.is_empty() { + spill(sink, &take_record_batch(&batch, &UInt32Array::from(rows))?)?; + } + } + } + Ok(()) + }) +} + +#[derive(Default)] +struct Accumulator { + rows: u64, + labelled: u64, + /// Binary: score order key shifted left, label in bit 0. + pairs: Vec, + count: u64, + sum: f64, + squares: f64, + /// Link: MAP numerators by label count. + map_by_labels: BTreeMap, + k: usize, +} + +impl Accumulator { + fn merge(mut self, mut other: Self) -> Self { + self.rows += other.rows; + self.labelled += other.labelled; + self.pairs.append(&mut other.pairs); + self.count += other.count; + self.sum += other.sum; + self.squares += other.squares; + for (labels, value) in other.map_by_labels { + *self.map_by_labels.entry(labels).or_default() += value; + } + self.k = self.k.max(other.k); + self + } +} + +fn order_key(score: f64) -> u64 { + let bits = (score as f32).to_bits(); + u64::from(if bits & 0x8000_0000 != 0 { + !bits + } else { + bits | 0x8000_0000 + }) +} + +/// A grouped prediction row: score, class and (batch, row). +type Ranked = (f64, Option, (usize, usize)); + +/// Score each label row against the predictions with its key; the first +/// label of a key counts. +fn join( + predictions: Batches, + labels: Batches, + spec: &EvaluateSpec, + (score, class): (&str, &DataType), + actuals: Option, +) -> Result { + let grouped = matches!(spec.task, EvaluateTask::Multiclass | EvaluateTask::Link); + let mut first = HashMap::new(); + let mut groups: HashMap> = HashMap::new(); + let mut kept: Vec> = Vec::new(); + for (index, batch) in predictions.enumerate() { + let batch = batch?; + kept.push( + spec.prediction_columns + .iter() + .map(|(name, _)| column(&batch, name).cloned()) + .collect::>()?, + ); + let scores = floats(column(&batch, score)?)?; + if !grouped { + for (row, (key, score)) in keys(&batch)?.zip(scores).enumerate() { + first.entry(key).or_insert((score, (index, row))); + } + continue; + } + let classes = classes(column(&batch, "CLASS")?, class)?; + for (row, ((key, score), class)) in keys(&batch)?.zip(scores).zip(classes).enumerate() { + let score = score.unwrap_or(f64::NEG_INFINITY); + groups + .entry(key) + .or_default() + .push((score, class, (index, row))); + } + } + let mut acc = Accumulator { + rows: (first.len() + groups.len()) as u64, + k: groups.values().map(Vec::len).max().unwrap_or(0), + ..Default::default() + }; + + let mut writer: Option> = None; + for batch in labels { + let batch = batch?; + let target = column(&batch, &spec.label_target)?; + let mut matched = vec![false; batch.num_rows()]; + let mut picked = Vec::new(); + if !grouped { + for ((row, key), label) in keys(&batch)?.enumerate().zip(floats(target)?) { + let Some(label) = label else { continue }; + let Some((score, at)) = first.remove(&key) else { + continue; + }; + matched[row] = true; + picked.push(at); + match score { + Some(score) if spec.task == EvaluateTask::Binary => { + if score.is_finite() { + acc.pairs + .push(order_key(score) << 1 | u64::from(label != 0.0)); + } + } + Some(prediction) => { + let error = prediction - label; + acc.count += 1; + acc.sum += error.abs(); + acc.squares += error * error; + } + None => {} + } + } + } else { + let link = spec.task == EvaluateTask::Link; + let labels = label_classes(target, link, class)?; + for ((row, key), items) in keys(&batch)?.enumerate().zip(labels) { + let Some(items) = items else { continue }; + let Some(mut ranked) = groups.remove(&key) else { + continue; + }; + matched[row] = true; + // Highest score first; ties keep their order. + ranked.sort_by(|a, b| b.0.total_cmp(&a.0)); + picked.push(ranked[0].2); + let hit = + |class: Option| class.is_some_and(|c| items.binary_search(&c).is_ok()); + if !link { + acc.count += 1; + acc.sum += f64::from(u8::from(hit(ranked[0].1))); + } else if !items.is_empty() { + let (mut hits, mut numerator) = (0usize, 0.0); + for (rank, (_, class, _)) in ranked.iter().enumerate() { + if hit(*class) { + hits += 1; + numerator += hits as f64 / (rank + 1) as f64; + } + } + acc.count += 1; + acc.sum += hits as f64; + *acc.map_by_labels.entry(items.len()).or_default() += numerator; + } + } + } + let matched = BooleanArray::from(matched); + acc.labelled += matched.true_count() as u64; + let Some(path) = actuals.as_ref().filter(|_| matched.true_count() > 0) else { + continue; + }; + let mut out = spec + .actual_columns + .iter() + .map(|(name, out)| Ok((out.as_str(), filter(column(&batch, name)?, &matched)?, true))) + .collect::>>()?; + for (name, ns) in &spec.constant_timestamps { + let values = vec![*ns; picked.len()]; + let array = TimestampNanosecondArray::from(values).with_timezone("UTC"); + out.push((name.as_str(), Arc::new(array), true)); + } + for (column, (_, name)) in spec.prediction_columns.iter().enumerate() { + let arrays: Vec<&dyn Array> = kept.iter().map(|batch| batch[column].as_ref()).collect(); + out.push((name.as_str(), interleave(&arrays, &picked)?, true)); + } + let out = RecordBatch::try_from_iter_with_nullable(out)?; + if writer.is_none() { + writer = Some(ArrowWriter::try_new( + File::create(path)?, + out.schema(), + None, + )?); + } + writer.as_mut().unwrap().write(&out)?; + } + if let Some(writer) = writer { + writer.close()?; + } + Ok(acc) +} + +/// Exact AUROC and average precision from (score, label) order keys. +fn ranking_metrics(mut pairs: Vec) -> (Option, Option) { + pairs.par_sort_unstable(); + let positives = pairs.iter().filter(|pair| *pair & 1 == 1).count() as f64; + let negatives = pairs.len() as f64 - positives; + if positives == 0.0 || negatives == 0.0 { + return (None, None); + } + let groups: Vec<(f64, f64)> = pairs + .chunk_by(|a, b| a >> 1 == b >> 1) + .map(|group| { + let positive = group.iter().filter(|pair| *pair & 1 == 1).count() as f64; + (positive, group.len() as f64 - positive) + }) + .collect(); + let (mut below, mut area) = (0.0, 0.0); + for &(positive, negative) in &groups { + area += positive * (below + negative / 2.0); + below += negative; + } + let (mut true_positive, mut false_positive, mut previous_recall, mut precision_sum) = + (0.0, 0.0, 0.0, 0.0); + for &(positive, negative) in groups.iter().rev() { + true_positive += positive; + false_positive += negative; + let recall = true_positive / positives; + precision_sum += + (recall - previous_recall) * true_positive / (true_positive + false_positive); + previous_recall = recall; + } + (Some(area / (positives * negatives)), Some(precision_sum)) +} + +fn metrics(acc: Accumulator, task: EvaluateTask) -> JsonValue { + let ratio = |num: f64| (acc.count > 0).then(|| num / acc.count as f64); + match task { + EvaluateTask::Binary => { + let (auroc, ap) = ranking_metrics(acc.pairs); + json!({"auroc": auroc, "ap": ap}) + } + EvaluateTask::Regression => { + json!({"mae": ratio(acc.sum), "rmse": ratio(acc.squares).map(f64::sqrt)}) + } + EvaluateTask::Multiclass => json!({"acc": ratio(acc.sum)}), + EvaluateTask::Link => { + let k = acc.k.max(1); + let map: f64 = acc + .map_by_labels + .iter() + .map(|(labels, numerator)| numerator / (*labels).min(k) as f64) + .sum(); + let precision = ratio(acc.sum / k as f64); + json!({format!("map@{k}"): ratio(map), format!("precision@{k}"): precision}) + } + } +} + +/// Join `prediction_paths` with `label_paths` and compute the realized +/// metrics; returns `{"rows", "labelled", "metrics"}`. +pub fn evaluate_prediction_parquet( + prediction_paths: &[&Path], + label_paths: &[&Path], + spec: &EvaluateSpec, +) -> Result { + if spec.prediction_keys.is_empty() || spec.prediction_keys.len() != spec.label_keys.len() { + bail!("prediction_keys and label_keys must pair up"); + } + let (prediction_groups, prediction_rows, prediction_schema) = row_groups(prediction_paths)?; + let (label_groups, label_rows, label_schema) = row_groups(label_paths)?; + let types = spec + .prediction_keys + .iter() + .zip(&spec.label_keys) + .map(|(p, l)| { + Ok(common_type( + prediction_schema.field_with_name(p)?.data_type(), + label_schema.field_with_name(l)?.data_type(), + )) + }) + .collect::>>()?; + let score = match spec.task { + EvaluateTask::Binary => "True_PROB", + EvaluateTask::Regression => "TARGET_PRED", + EvaluateTask::Multiclass | EvaluateTask::Link => "SCORE", + }; + let mut prediction_columns = vec![score.to_string()]; + prediction_columns.extend(spec.prediction_columns.iter().map(|(name, _)| name.clone())); + let mut class = DataType::Null; + if matches!(spec.task, EvaluateTask::Multiclass | EvaluateTask::Link) { + prediction_columns.push("CLASS".to_string()); + } + prediction_columns.sort(); + prediction_columns.dedup(); + if matches!(spec.task, EvaluateTask::Multiclass | EvaluateTask::Link) { + let mut target = label_schema + .field_with_name(&spec.label_target)? + .data_type(); + if spec.task == EvaluateTask::Link { + let (DataType::List(item) + | DataType::LargeList(item) + | DataType::FixedSizeList(item, _)) = target + else { + bail!("link-prediction labels must be lists, got {target}"); + }; + target = item.data_type(); + } + class = common_type( + prediction_schema.field_with_name("CLASS")?.data_type(), + target, + ); + } + let mut label_columns = vec![spec.label_target.clone()]; + label_columns.extend(spec.actual_columns.iter().map(|(name, _)| name.clone())); + label_columns.sort(); + label_columns.dedup(); + + let partitions = prediction_rows + .max(label_rows) + .div_ceil(spec.partition_rows.max(1)) + .max(1); + std::fs::create_dir_all(&spec.spill_dir)?; + if let Some(dir) = &spec.actuals_dir { + std::fs::create_dir_all(dir)?; + } + let sinks = |side: &str| -> Vec { + (0..partitions) + .map(|p| { + ( + spec.spill_dir.join(format!("{side}-{p:05}.arrow")), + Mutex::default(), + ) + }) + .collect() + }; + let (predictions, labels) = (sinks("predictions"), sinks("labels")); + let prediction_keys = (spec.prediction_keys.as_slice(), types.as_slice()); + scatter( + prediction_paths, + &prediction_groups, + prediction_keys, + &prediction_columns, + &predictions, + )?; + let label_keys = (spec.label_keys.as_slice(), types.as_slice()); + scatter( + label_paths, + &label_groups, + label_keys, + &label_columns, + &labels, + )?; + + let pool = rayon::ThreadPoolBuilder::new() + .num_threads(MAX_CONCURRENT_JOINS) + .build()?; + let acc = pool.install(|| { + predictions + .into_par_iter() + .zip(labels) + .enumerate() + .map(|(p, (predictions, labels))| { + let actuals = + (spec.actuals_dir.as_ref()).map(|dir| dir.join(format!("part-{p:05}.parquet"))); + join( + unspill(predictions)?, + unspill(labels)?, + spec, + (score, &class), + actuals, + ) + }) + .try_reduce(Accumulator::default, |a, b| Ok(a.merge(b))) + })?; + let _ = std::fs::remove_dir_all(&spec.spill_dir); + let (rows, labelled) = (acc.rows, acc.labelled); + Ok(json!({"rows": rows, "labelled": labelled, "metrics": metrics(acc, spec.task)})) +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::{ + BooleanArray, Float64Array, Int64Array, Int64Builder, ListBuilder, StringArray, + TimestampMicrosecondArray, TimestampNanosecondArray, + }; + use parquet::file::properties::WriterProperties; + use tempfile::tempdir; + + fn write(path: &Path, columns: Vec<(&str, ArrayRef)>, rows_per_group: usize) { + let batch = RecordBatch::try_from_iter(columns).unwrap(); + let props = WriterProperties::builder() + .set_max_row_group_row_count(Some(rows_per_group)) + .build(); + let mut writer = + ArrowWriter::try_new(File::create(path).unwrap(), batch.schema(), Some(props)).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + } + + fn spec(task: EvaluateTask, dir: &Path, partition_rows: usize) -> EvaluateSpec { + EvaluateSpec { + task, + prediction_keys: vec!["ENTITY".into(), "ANCHOR_TIMESTAMP".into()], + label_keys: vec!["ENTITY1".into(), "TIMESTAMP".into()], + label_target: "TARGET1".into(), + actuals_dir: Some(dir.join("actuals")), + actual_columns: vec![ + ("ENTITY1".into(), "ENTITY".into()), + ("TIMESTAMP".into(), "ANCHOR_TIMESTAMP".into()), + ("TARGET1".into(), "TARGET".into()), + ], + prediction_columns: vec![], + constant_timestamps: Vec::new(), + spill_dir: dir.join("spill"), + partition_rows, + } + } + + fn evaluate(dir: &Path, spec: &EvaluateSpec) -> JsonValue { + let predictions = dir.join("predictions.parquet"); + let labels = dir.join("labels.parquet"); + evaluate_prediction_parquet(&[predictions.as_path()], &[labels.as_path()], spec).unwrap() + } + + fn actuals(dir: &Path) -> Vec { + std::fs::read_dir(dir.join("actuals")) + .unwrap() + .flat_map(|entry| { + let file = File::open(entry.unwrap().path()).unwrap(); + ParquetRecordBatchReaderBuilder::try_new(file) + .unwrap() + .build() + .unwrap() + .map(Result::unwrap) + }) + .collect() + } + + fn actual_column( + dir: &Path, + name: &str, + ) -> Vec { + actuals(dir) + .iter() + .flat_map(|batch| { + let array = cast(column(batch, name).unwrap(), &T::DATA_TYPE).unwrap(); + array.as_primitive::().values().to_vec() + }) + .collect() + } + + /// Mann-Whitney AUROC with ties counted half, by brute force. + fn brute_auroc(scores: &[f32], labels: &[bool]) -> f64 { + let (mut wins, mut pairs) = (0.0, 0.0); + for (i, &si) in scores.iter().enumerate() { + for (j, &sj) in scores.iter().enumerate() { + if labels[i] && !labels[j] { + pairs += 1.0; + wins += if si > sj { + 1.0 + } else if si == sj { + 0.5 + } else { + 0.0 + }; + } + } + } + wins / pairs + } + + #[test] + fn binary_join_matches_brute_force_in_one_and_many_partitions() { + let n = 3000i64; + let entity: Vec = (0..n).collect(); + // Several anchors; tz-aware microseconds against naive nanoseconds. + let anchor_us: Vec = (0..n).map(|i| (i % 3) * 86_400_000_000).collect(); + let scores: Vec = (0..n).map(|i| ((i * 7919) % 101) as f32 / 100.0).collect(); + let labels: Vec = (0..n).map(|i| (i * 31 + i / 7) % 5 < 2).collect(); + // Predictions for every entity; labels for all but the last 500. + let m = (n - 500) as usize; + for partition_rows in [1_000_000, 700] { + let dir = tempdir().unwrap(); + write( + &dir.path().join("predictions.parquet"), + vec![ + ( + "ENTITY", + Arc::new(Int64Array::from(entity.clone())) as ArrayRef, + ), + ( + "ANCHOR_TIMESTAMP", + Arc::new( + TimestampMicrosecondArray::from(anchor_us.clone()).with_timezone("UTC"), + ) as ArrayRef, + ), + ( + "True_PROB", + Arc::new(arrow::array::Float32Array::from(scores.clone())) as ArrayRef, + ), + ], + 256, + ); + write( + &dir.path().join("labels.parquet"), + vec![ + ( + "ENTITY1", + Arc::new(Int64Array::from( + entity[..m].iter().rev().copied().collect::>(), + )) as ArrayRef, + ), + ( + "TIMESTAMP", + Arc::new(TimestampNanosecondArray::from( + anchor_us[..m] + .iter() + .rev() + .map(|t| t * 1000) + .collect::>(), + )) as ArrayRef, + ), + ( + "TARGET1", + Arc::new(BooleanArray::from( + labels[..m].iter().rev().copied().collect::>(), + )) as ArrayRef, + ), + ], + 300, + ); + let mut spec_ = spec(EvaluateTask::Binary, dir.path(), partition_rows); + spec_.prediction_columns = vec![("True_PROB".into(), "True_PROB".into())]; + let out = evaluate(dir.path(), &spec_); + assert_eq!(out["rows"], json!(n)); + assert_eq!(out["labelled"], json!(m)); + let auroc = out["metrics"]["auroc"].as_f64().unwrap(); + assert!((auroc - brute_auroc(&scores[..m], &labels[..m])).abs() < 1e-12); + assert!(out["metrics"]["ap"].as_f64().unwrap() > 0.0); + let entities = actual_column::(dir.path(), "ENTITY"); + let emitted = actual_column::(dir.path(), "True_PROB"); + assert_eq!(entities.len(), m); + for (entity, score) in entities.iter().zip(emitted) { + assert_eq!(score, scores[*entity as usize]); + } + assert!(!dir.path().join("spill").exists()); + } + } + + #[test] + fn average_precision_matches_step_definition() { + // Scores 0.9(+), 0.8(-), 0.8(+), 0.1(-): thresholds 0.9, 0.8, 0.1. + let pairs = vec![ + order_key(0.9) << 1 | 1, + order_key(0.8) << 1, + order_key(0.8) << 1 | 1, + order_key(0.1) << 1, + ]; + let (auroc, ap) = ranking_metrics(pairs); + assert!((auroc.unwrap() - 0.875).abs() < 1e-12); + assert!((ap.unwrap() - (0.5 * 1.0 + 0.5 * 2.0 / 3.0)).abs() < 1e-12); + assert_eq!(ranking_metrics(vec![order_key(0.5) << 1 | 1]), (None, None)); + } + + #[test] + fn regression_multiclass_and_link_metrics() { + let dir = tempdir().unwrap(); + let mut spec_ = spec(EvaluateTask::Regression, dir.path(), 1_000_000); + spec_.prediction_keys.truncate(1); + spec_.label_keys.truncate(1); + spec_.actual_columns = vec![ + ("ENTITY1".into(), "ENTITY".into()), + ("TARGET1".into(), "TARGET".into()), + ]; + spec_.constant_timestamps = + vec![("ANCHOR_TIMESTAMP".into(), 7), ("EVALUATED_AT".into(), 9)]; + write( + &dir.path().join("predictions.parquet"), + vec![ + ( + "ENTITY", + Arc::new(StringArray::from(vec!["a", "b", "c"])) as ArrayRef, + ), + ( + "TARGET_PRED", + Arc::new(Float64Array::from(vec![1.0, 3.0, 4.0])) as ArrayRef, + ), + ], + 2, + ); + write( + &dir.path().join("labels.parquet"), + vec![ + ( + "ENTITY1", + Arc::new(StringArray::from(vec!["c", "a", "b"])) as ArrayRef, + ), + ( + "TARGET1", + Arc::new(Float64Array::from(vec![5.0, 1.0, 2.0])) as ArrayRef, + ), + ], + 2, + ); + let out = evaluate(dir.path(), &spec_); + assert_eq!( + out["metrics"], + json!({"mae": 2.0 / 3.0, "rmse": (2.0f64 / 3.0).sqrt()}) + ); + assert_eq!( + actual_column::( + dir.path(), + "ANCHOR_TIMESTAMP" + ), + vec![7, 7, 7] + ); + assert_eq!( + actual_column::(dir.path(), "EVALUATED_AT"), + vec![9, 9, 9] + ); + + // Two anchors of entity 1 (a hit and a miss at rank 1) and entity 2. + spec_ = spec(EvaluateTask::Multiclass, dir.path(), 1_000_000); + spec_.prediction_columns = vec![ + ("CLASS".into(), "CLASS".into()), + ("SCORE".into(), "SCORE".into()), + ]; + let entity = Int64Array::from(vec![1, 1, 1, 1, 2, 2]); + let anchor = TimestampMicrosecondArray::from(vec![0, 0, 1, 1, 0, 0]); + write( + &dir.path().join("predictions.parquet"), + vec![ + ("ENTITY", Arc::new(entity.clone()) as ArrayRef), + ("ANCHOR_TIMESTAMP", Arc::new(anchor.clone()) as ArrayRef), + ( + "CLASS", + Arc::new(Int64Array::from(vec![10, 11, 12, 10, 11, 12])) as ArrayRef, + ), + ( + "SCORE", + Arc::new(Float64Array::from(vec![0.9, 0.5, 0.7, 0.2, 0.8, 0.1])) as ArrayRef, + ), + ], + 3, + ); + write( + &dir.path().join("labels.parquet"), + vec![ + ( + "ENTITY1", + Arc::new(Int64Array::from(vec![1, 1, 2])) as ArrayRef, + ), + ( + "TIMESTAMP", + Arc::new(TimestampNanosecondArray::from(vec![0, 1000, 0])) as ArrayRef, + ), + ( + "TARGET1", + Arc::new(Float64Array::from(vec![10.0, 10.0, 11.0])) as ArrayRef, + ), + ], + 3, + ); + let out = evaluate(dir.path(), &spec_); + assert_eq!(out["metrics"], json!({"acc": 2.0 / 3.0})); + let top: Vec<(i64, f64)> = + actual_column::(dir.path(), "CLASS") + .into_iter() + .zip(actual_column::( + dir.path(), + "SCORE", + )) + .collect(); + assert_eq!(top, vec![(10, 0.9), (12, 0.7), (11, 0.8)]); + assert_eq!( + (out["rows"].clone(), out["labelled"].clone()), + (json!(3), json!(3)) + ); + + spec_ = spec(EvaluateTask::Link, dir.path(), 1_000_000); + let mut targets = ListBuilder::new(Int64Builder::new()); + for items in [vec![11, 13], vec![], vec![11]] { + for item in items { + targets.values().append_value(item); + } + targets.append(true); + } + write( + &dir.path().join("labels.parquet"), + vec![ + ( + "ENTITY1", + Arc::new(Int64Array::from(vec![1, 1, 2])) as ArrayRef, + ), + ( + "TIMESTAMP", + Arc::new(TimestampNanosecondArray::from(vec![0, 1000, 0])) as ArrayRef, + ), + ("TARGET1", Arc::new(targets.finish()) as ArrayRef), + ], + 3, + ); + // (1, 0) ranks 10, 11 and hits 11 at rank 2 of two labels; (1, 1) + // has no labels and is skipped; entity 2 hits 11 at rank 1. + let out = evaluate(dir.path(), &spec_); + assert_eq!( + out["metrics"], + json!({"map@2": (0.5 / 2.0 + 1.0) / 2.0, "precision@2": 0.5}) + ); + } +} diff --git a/packages/diskgraph/diskgraph-pybind/src/lib.rs b/packages/diskgraph/diskgraph-pybind/src/lib.rs index a47daef4d..9b9f0704e 100644 --- a/packages/diskgraph/diskgraph-pybind/src/lib.rs +++ b/packages/diskgraph/diskgraph-pybind/src/lib.rs @@ -184,6 +184,7 @@ use diskgraph_engine::state::{SamplingState, StatePool}; mod artifact_helpers; mod bp_col_stats; +mod bp_evaluate; mod col_stats_export; mod custom_tt_entry; mod index_metadata; @@ -1387,7 +1388,7 @@ fn estimate_ingest_size(py: Python<'_>, config_json: &str) -> PyResult { /// Compute Kumo-format column statistics over one or more parquet files. /// -/// Wraps :func:`bp_col_stats::compute_parquet_col_stats` for kumo's +/// Wraps :func:`bp_col_stats::compute_parquet_col_stats_with_options` for kumo's /// batch-prediction stats path: after a BP run writes its output parquet, /// kumo invokes this to produce the same ``kumo_col_stats_v1`` JSON that /// ingest emits, so downstream readers (UI, encoder inference, @@ -1414,6 +1415,21 @@ fn estimate_ingest_size(py: Python<'_>, config_json: &str) -> PyResult { /// invalid. Semantic and logical types are mapped through the same /// ``map_kumo_stype`` / ``map_kumo_dtype`` helpers ingest uses. /// data_version: version string stamped into every emitted entry. +/// percentiles: also emit approximate ``PERCENTILES`` (q01 through q99, +/// from each numerical column's t-digest). +/// summary_stats_by_class: emit ``SUMMARY_STATS_BY_CLASS`` for a +/// ``SCORE``/``CLASS`` pair (the default); ``False`` skips its +/// per-class map. +/// histograms: emit ``HISTOGRAM`` statistics (the default), which take a +/// second read; ``False`` skips them. +/// category_counts_limit: keep at most this many ``CATEGORY_COUNTS`` +/// entries in the full output (all by default; the min output keeps +/// 20) and record ``CATEGORY_COUNTS_CAPPED``, whether counting hit +/// the category cap so ``NUM_UNIQUE`` is an estimate. +/// max_per: compute every statistic over the maximum of each column +/// within each run of equal values of these columns, in path and row +/// order (for example each prediction's top score). Needs numerical +/// columns, ``histograms=False`` and ``summary_stats_by_class=False``. /// _category_count_limit_for_testing: keyword-only test hook that lowers /// the retained category limit. ``None`` preserves the production /// 1,000,000-category limit; a supplied value must be greater than 0. @@ -1428,6 +1444,11 @@ fn estimate_ingest_size(py: Python<'_>, config_json: &str) -> PyResult { kumo_column_specs_json, data_version, *, + percentiles=false, + summary_stats_by_class=true, + histograms=true, + category_counts_limit=None, + max_per=None, _category_count_limit_for_testing=None, ))] #[tracing::instrument(level = "debug", skip(py, kumo_column_specs_json))] @@ -1436,6 +1457,11 @@ fn compute_col_stats_for_parquet( parquet_paths: Vec, kumo_column_specs_json: &str, data_version: &str, + percentiles: bool, + summary_stats_by_class: bool, + histograms: bool, + category_counts_limit: Option, + max_per: Option>, _category_count_limit_for_testing: Option, ) -> PyResult<(String, String)> { use diskgraph_ingest::config_converter::{ @@ -1443,9 +1469,9 @@ fn compute_col_stats_for_parquet( }; use diskgraph_ingest::graph_config::ColumnSpec; - use crate::bp_col_stats::{ - compute_parquet_col_stats, compute_parquet_col_stats_with_max_categories, - }; + use diskgraph_core::accumulators::MAX_CATEGORIES; + + use crate::bp_col_stats::{compute_parquet_col_stats_with_options, ColStatsOptions}; if _category_count_limit_for_testing == Some(0) { return Err(DiskgraphConfigError::new_err( @@ -1505,19 +1531,24 @@ fn compute_col_stats_for_parquet( .collect::>>()?; let data_version = data_version.to_string(); let paths_owned = parquet_paths; + let options = ColStatsOptions { + percentiles, + summary_stats_by_class, + histograms, + category_counts_limit, + max_per: max_per.unwrap_or_default(), + }; py.detach(|| -> anyhow::Result<(String, String)> { let paths: Vec<&std::path::Path> = paths_owned.iter().map(std::path::Path::new).collect(); - match _category_count_limit_for_testing { - Some(limit) => compute_parquet_col_stats_with_max_categories( - &paths, - &columns, - &data_version, - limit, - ), - None => compute_parquet_col_stats(&paths, &columns, &data_version), - } + compute_parquet_col_stats_with_options( + &paths, + &columns, + &data_version, + _category_count_limit_for_testing.unwrap_or(MAX_CATEGORIES), + &options, + ) }) // Route schema failures through the shared exception classifier rather // than reporting missing columns or invalid timestamp metadata as a @@ -1526,6 +1557,44 @@ fn compute_col_stats_for_parquet( .map_err(map_anyhow_err) } +/// Join batch-prediction Parquet outputs with generated label Parquet parts +/// and compute realized metrics, streaming and hash-partitioning both sides. +/// +/// Args: +/// prediction_paths: the run's prediction Parquet files. +/// label_paths: the label Parquet parts. +/// spec_json: JSON object with ``task`` (``"binary"``, ``"regression"``, +/// ``"multiclass"`` or ``"link"``), ``prediction_keys`` and +/// ``label_keys`` (paired by position), ``label_target``, +/// ``spill_dir``, and optional ``actuals_dir`` with +/// ``actual_columns`` (``[label name, output name]`` pairs) to write +/// the matched label rows, ``prediction_columns`` (pairs copied +/// from the matched prediction), ``constant_timestamps`` +/// (``[output name, UTC nanoseconds]`` pairs written on every row), +/// and ``partition_rows``. +/// +/// Returns a JSON string ``{"rows", "labelled", "metrics"}``: distinct +/// predictions, those with a label, and the metric values (``null`` when +/// undefined). +#[pyfunction] +#[pyo3(signature = (prediction_paths, label_paths, spec_json))] +fn evaluate_prediction_parquet( + py: Python<'_>, + prediction_paths: Vec, + label_paths: Vec, + spec_json: &str, +) -> PyResult { + let spec: bp_evaluate::EvaluateSpec = serde_json::from_str(spec_json) + .map_err(|e| DiskgraphConfigError::new_err(format!("invalid spec_json: {e}")))?; + py.detach(|| -> anyhow::Result { + let predictions: Vec<&std::path::Path> = + prediction_paths.iter().map(std::path::Path::new).collect(); + let labels: Vec<&std::path::Path> = label_paths.iter().map(std::path::Path::new).collect(); + Ok(bp_evaluate::evaluate_prediction_parquet(&predictions, &labels, &spec)?.to_string()) + }) + .map_err(map_anyhow_err) +} + /// Verify a local snapshot directory matches its ``_MANIFEST.json``. /// /// Called by the Python ``DiskgraphService`` after a dataplane→local @@ -1627,6 +1696,7 @@ fn diskgraph_pybind(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_function(wrap_pyfunction!(has_resumable_ingest, m)?)?; m.add_function(wrap_pyfunction!(validated_index_resume_progress, m)?)?; m.add_function(wrap_pyfunction!(compute_col_stats_for_parquet, m)?)?; + m.add_function(wrap_pyfunction!(evaluate_prediction_parquet, m)?)?; m.add_function(wrap_pyfunction!(col_stats_export::read_col_stats, m)?)?; m.add_function(wrap_pyfunction!(col_stats_export::read_edge_stats, m)?)?; m.add_class::()?; diff --git a/packages/kumo-core/src/kumo_core/object_store_sink.py b/packages/kumo-core/src/kumo_core/object_store_sink.py index b6750b69a..4114c36bb 100644 --- a/packages/kumo-core/src/kumo_core/object_store_sink.py +++ b/packages/kumo-core/src/kumo_core/object_store_sink.py @@ -164,7 +164,10 @@ def write_files(self, connector, table_name, files, if write_mode is not WriteMode.OVERWRITE: raise NotImplementedError( 'Object-store artifact output does not support APPEND') - from kumo_core.artifact_export import export_selected_files + from kumo_core.artifact_export import ( + ArtifactDestinationConflict, + export_selected_files, + ) from kumo_core.sources import layout filesystem, path = self._destination(connector, table_name) kind = layout.file_type(path.rstrip('/')) or layout.PARQUET @@ -173,13 +176,41 @@ def write_files(self, connector, table_name, files, 'it', path, export_id) return f'{self.scheme}://{path.rstrip("/")}/' storage = _ArrowExportFiles(filesystem, extra_columns) - _, output = export_selected_files(files, path, fs=storage, - identity=export_id or str(files), - file_type=kind, - convert_csv=storage.convert_csv, - sidecars=sidecars) + try: + _, output = export_selected_files( + files, path, fs=storage, identity=export_id or str(files), + file_type=kind, convert_csv=storage.convert_csv, + sidecars=sidecars) + except ArtifactDestinationConflict: + raise self._conflict(path, '') from None return f'{self.scheme}://{output}' + def preflight_write_files(self, connector, table_name, write_mode=None): + """Refuse an occupied directory unless ``write_mode`` is + ``OVERWRITE``, before any rows are produced.""" + from pyarrow.fs import FileSelector + + if write_mode is WriteMode.OVERWRITE or not table_name.endswith('/'): + return + filesystem, path = self._destination(connector, table_name) + if filesystem.get_file_info( + FileSelector(path.rstrip('/'), allow_not_found=True)): + raise self._conflict(path) + + def _uri(self, path: str) -> str: + return f'{self.scheme}://{path}' + + def _conflict( + self, path: str, + hint: str = ', or replace it with write_mode=WriteMode.OVERWRITE' + ) -> ValueError: + from kumo_core.artifact_export import ArtifactDestinationConflict + + return ArtifactDestinationConflict( + f"Output directory '{self._uri(path.rstrip('/'))}/' already " + f"contains data from another export. Choose an empty destination " + f"path{hint}.") + def publish_model( self, connector: SourceConnector, @@ -465,6 +496,9 @@ def _destination(self, connector, coordinate): def _bare(self, uri: str) -> str: return uri.removeprefix(f'{self.scheme}://') + def _uri(self, path: str) -> str: + return path + def write_table(self, *args, **kwargs): return self._bare(super().write_table(*args, **kwargs)) diff --git a/packages/kumo-core/src/kumo_core/pquery.py b/packages/kumo-core/src/kumo_core/pquery.py index 6954e2e08..704da4351 100644 --- a/packages/kumo-core/src/kumo_core/pquery.py +++ b/packages/kumo-core/src/kumo_core/pquery.py @@ -20,7 +20,7 @@ from datetime import datetime, timezone from enum import Enum from numbers import Integral -from typing import Any, Mapping, Sequence +from typing import Any, Callable, Iterable, Mapping, Sequence from diskgraph import read_col_stats from kumoapi.common import ValidationResponse @@ -984,8 +984,15 @@ def latest_temporal_timestamp(query: ValidatedNativeQuery, if query.query_type != "temporal": raise PredictiveQueryPlanError( "a source-derived timestamp requires a temporal query") - return max(int(maximum.value) - for _, maximum in _target_time_bounds(query, index_path)) + return _target_maximum(query, index_path, max) + + +def _target_maximum(query: ValidatedNativeQuery, index_path: str, + pick: Callable[[Iterable[int]], int]) -> int: + """``pick`` (``min`` or ``max``) of the target aggregations' source + maxima.""" + return pick(int(maximum.value) + for _, maximum in _target_time_bounds(query, index_path)) def _temporal_timestamps( diff --git a/packages/kumo-core/src/kumo_core/prediction.py b/packages/kumo-core/src/kumo_core/prediction.py index effb6b69e..7537bf6d6 100644 --- a/packages/kumo-core/src/kumo_core/prediction.py +++ b/packages/kumo-core/src/kumo_core/prediction.py @@ -70,6 +70,9 @@ class GeneratedPredictionTable: table_path: str row_count: int tt_id: str + #: DiskGraph's ``row_failures`` report for caller-supplied rows: rows + #: read, rows skipped, and per-condition counts with source locations. + row_failures: Mapping[str, Any] | None = None @dataclass(frozen=True) @@ -254,11 +257,13 @@ def generate_custom_prediction_table( "custom prediction-table builder result missing tt_parquet_path") if row_count <= 0: raise RuntimeError("custom prediction table contains no rows") + failures = result.get("row_failures") return GeneratedPredictionTable( staging_dir=staging_dir, table_path=table_path, row_count=row_count, tt_id=request.tt_id, + row_failures=failures if isinstance(failures, Mapping) else None, ) diff --git a/packages/kumo-core/src/kumo_core/training/table.py b/packages/kumo-core/src/kumo_core/training/table.py index 558e5b273..8c5fef997 100644 --- a/packages/kumo-core/src/kumo_core/training/table.py +++ b/packages/kumo-core/src/kumo_core/training/table.py @@ -453,6 +453,7 @@ def generate_train_table( request: TrainTableRequest, *, runtime: Any | None = None, + validate_targets: bool = True, ) -> GeneratedTrainTableResult: """Compile and execute the real DiskGraph TT-NG node/link producer.""" @@ -531,11 +532,12 @@ def generate_train_table( if not isinstance(table_path, str) or not table_path: raise RuntimeError("execute_tt_gen result missing tt_parquet_path") try: - _validate_classification_targets(request, staging_dir) + if validate_targets: + _validate_classification_targets(request, staging_dir) except ValueError: _discard_invalid_staging_dir(staging_dir, request) raise - top_k = getattr(request.query, "top_k", None) + top_k = getattr(request.query, "top_k", None) if validate_targets else None if (task_type in {"multilabel_classification", "multilabel_ranking"} and top_k is not None): class_count = _observed_target_class_count(staging_dir) diff --git a/packages/kumo-core/tests/test_pquery.py b/packages/kumo-core/tests/test_pquery.py index 1dd0e985f..df825b785 100644 --- a/packages/kumo-core/tests/test_pquery.py +++ b/packages/kumo-core/tests/test_pquery.py @@ -15,6 +15,7 @@ _column_stats, _default_split, _split_options, + _target_maximum, _temporal_timestamps, build_generated_training_task_dataset, build_native_pquery_task_plan, @@ -762,7 +763,7 @@ def test_latest_temporal_timestamp_uses_source_maximum(tmp_path, monkeypatch): ("first", "second"), ("second", "first"), ]) -def test_latest_temporal_timestamp_is_order_independent(monkeypatch, order): +def test_target_maxima_are_order_independent(monkeypatch, order): maxima = { "first": 1_700_000_000_000_000_000, "second": 1_800_000_000_000_000_000, @@ -790,6 +791,7 @@ def stats(_index_path, table_name, _scope): assert latest_temporal_timestamp(query, "/index") == maxima["second"] assert calls == list(order) + assert _target_maximum(query, "/index", min) == maxima["first"] def test_generation_plan_keeps_earliest_anchor_needed_for_filter_history( diff --git a/packages/kumo-core/tests/test_train_table.py b/packages/kumo-core/tests/test_train_table.py index 9d156dc1d..3dd497cc1 100644 --- a/packages/kumo-core/tests/test_train_table.py +++ b/packages/kumo-core/tests/test_train_table.py @@ -160,6 +160,10 @@ def test_generated_degenerate_labels_are_rejected_before_publication( }, }, })) + labels = generate_train_table( + request(tmp_path, task_type=task_type, target_stype=target_stype), + runtime=FakeRuntime(str(staging)), validate_targets=False) + assert labels.staging_dir == str(staging) and staging.is_dir() with pytest.raises(GenerationValidationError, match=message): generate_train_table( request(tmp_path, task_type=task_type, target_stype=target_stype), diff --git a/packages/kumo-databricks/tests/test_write.py b/packages/kumo-databricks/tests/test_write.py index 45ee32dc3..615236335 100644 --- a/packages/kumo-databricks/tests/test_write.py +++ b/packages/kumo-databricks/tests/test_write.py @@ -279,3 +279,30 @@ def test_a_warehouse_value_error_is_not_restated_as_a_refusal(artifact): with pytest.raises(ValueError) as raised: load(warehouse, str(artifact), files=[artifact]) assert not isinstance(raised.value, write.ArtifactValidationError) + + +def test_sink_appends_files_stamped_with_extra_columns(artifact, monkeypatch): + from contextlib import contextmanager + from types import SimpleNamespace + + from kumo_databricks.sink import DatabricksSinkBackend + from kumoapi.data_source import DatabricksConnectorResourceConfig + from kumoapi.jobs import WriteMode + + sink = DatabricksSinkBackend() + monkeypatch.setattr(sink, "_connected", + contextmanager(lambda *a, **k: (yield None))) + loads = [] + + def load(**kwargs): + with kwargs["prepare_file"](str(artifact)) as staged: + loads.append((kwargs["replace"], pq.read_table(staged).to_pydict())) + + monkeypatch.setattr(write, "load", load) + connector = SimpleNamespace(id="dbx", config=DatabricksConnectorResourceConfig( + catalog="main", schema_name="features", warehouse_id="w")) + sink.write_files(connector, "features__kumo__churn", [str(artifact)], + WriteMode.APPEND, extra_columns={"PREDICTION_ID": "p1"}) + + assert loads == [(False, {"ENTITY": ["a"], "SCORE": [1.0], + "PREDICTION_ID": ["p1"]})] diff --git a/packages/kumo-engine/docs/mlops.md b/packages/kumo-engine/docs/mlops.md index 7db3f3f77..f179d9202 100644 --- a/packages/kumo-engine/docs/mlops.md +++ b/packages/kumo-engine/docs/mlops.md @@ -1,12 +1,12 @@ # MLOps with Kumo Engine -Turn MLOps on with one `kumo.MLOps` and the normal `kumo.engine` calls store their snapshots, training tables and models under an artifact root, record fits in your experiment tracker, and register models in its registry. +Turn MLOps on with one `kumo.MLOps` and the normal `kumo.engine` calls store their snapshots, training tables, models and prediction runs under an artifact root, record fits and prediction runs in your experiment tracker, and register models in its registry. Install with `pip install "kumo-engine[mlflow]"` (add `s3`, `gcs`, `databricks`, `sagemaker` or `snowflake` as needed; Snowflake needs only `kumo-engine[snowflake]`). GPU, CUDA and Databricks Runtime notes are under [Install](#install). ## Quickstart -Each block assumes a `graph` and a `pquery` (see the README). It trains, stores and registers a model, then scores with it. +Each block assumes a `graph` and a `pquery` (see the README). It trains, stores and registers a model, scores with it, and checks score drift. ### Laptop or VM @@ -20,8 +20,8 @@ with kumo.mlops(cfg): snapshot = graph.snapshot() result = kumo.Trainer(pquery.suggest_model_plan(snapshot)).fit(pquery.generate_training_table(snapshot), snapshot=snapshot, scoring_defaults=kumo.ScoringDefaults(binary_classification_threshold=0.5)) - scores = result.predict(snapshot).frame - print(result.registered_model.version, result.tracking.run_id, len(scores)) + run = result.predict(snapshot, output_config=kumo.PredictionOutputConfig(kumo.FileConnector(root_dir="out"), path="scores/")) + print(result.registered_model.version, run.scored_rows, result.score_drift(run).psi) ``` ### Databricks @@ -34,11 +34,12 @@ from kumo.engine.integrations.mlflow import MlflowTracking cfg = kumo.MLOps("/Volumes/main/ml/kumo", tracking=MlflowTracking("/Users//churn"), registry=kumo.ModelRegistryLocation("main.ml.users_churn", alias="candidate")) +dbx = kumo.DatabricksConnector(catalog="main", schema_name="ml", warehouse_id="") with kumo.mlops(cfg): snapshot = graph.snapshot() result = kumo.Trainer(pquery.suggest_model_plan(snapshot)).fit(pquery.generate_training_table(snapshot), snapshot=snapshot, scoring_defaults=kumo.ScoringDefaults(binary_classification_threshold=0.5)) - print(result.registered_model.version, result.tracking.url) + run = result.predict(snapshot, output_config=kumo.PredictionOutputConfig(dbx, path="churn_scores")) ``` ### SageMaker @@ -55,7 +56,7 @@ with kumo.mlops(cfg): snapshot = graph.snapshot() result = kumo.Trainer(pquery.suggest_model_plan(snapshot)).fit(pquery.generate_training_table(snapshot), snapshot=snapshot, scoring_defaults=kumo.ScoringDefaults(binary_classification_threshold=0.5)) - print(result.registered_model.version, result.tracking.url) + run = result.predict(snapshot, output_config=kumo.PredictionOutputConfig(kumo.S3Connector(), path="s3:///scores/")) ``` Tables on S3 are read through an `S3Connector`: @@ -83,7 +84,7 @@ with kumo.mlops(cfg): snapshot = graph.snapshot() result = kumo.Trainer(pquery.suggest_model_plan(snapshot)).fit(pquery.generate_training_table(snapshot), snapshot=snapshot, scoring_defaults=kumo.ScoringDefaults(binary_classification_threshold=0.5)) - print(result.registered_model.version, result.tracking.url) + 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). @@ -91,21 +92,22 @@ To serve only a better model, move its alias yourself: [Promote a model](#promot ## How it fits together - The **artifact root** holds every artifact a block produces, each in its own directory with a `_provenance.json` manifest that embeds its parents' manifests. -- The **tracker** records each fit as a run of its own; the **registry** is the tracker's: MLflow (Unity Catalog, SageMaker MLflow, any MLflow server) or the Snowflake Model Registry. +- The **tracker** records each fit and prediction run as a run of its own; the **registry** is the tracker's: MLflow (Unity Catalog, SageMaker MLflow, any MLflow server) or the Snowflake Model Registry. | id | names | where to find it | |---|---|---| | `snapshot.id` | a graph snapshot | `/snapshots///`; `result.snapshot_id` | | `table.id` | a training table | `/training_tables//` | | `result.id` | a model (its bundle id) | `/models//`; the `kumo.bundle_id` run tag; `kumo_bundle_id` on the registered version | +| `run.id` | a prediction run | `/predictions//`; the `kumo.prediction_id` run tag; the `PREDICTION_ID` column of a table output | | `result.registered_model.version` | the registered version | the registry | | `result.tracking` | the tracker's own run of the fit: `platform`, `run_id`, `url` | the tracker's UI | -Find runs with the platform's UI, or by tag: `kumo.kind` (`fit`, `trial`) and `kumo.bundle_id`. +Find runs with the platform's UI, or by tag: `kumo.kind` (`fit`, `trial`, `prediction`, `monitoring`), `kumo.bundle_id`, `kumo.prediction_id`. ## Turning MLOps on -`kumo.MLOps` says where artifacts go and what records them. Use it as a block, or pass it as `mlops=` on a single call (`fit`, `graph.snapshot`, `generate_training_table`): +`kumo.MLOps` says where artifacts go and what records them. Use it as a block, or pass it as `mlops=` on a single call (`fit`, `predict`, `evaluate`, `score_drift`, `graph.snapshot`, `generate_training_table`): ```python cfg = kumo.MLOps("s3://bucket/kumo", tracking=..., registry=..., on_error="raise") @@ -117,10 +119,10 @@ 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. - **`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` and `model` (default: all). -- **`on_error`**: `"warn"` (default) warns about a failed store, tracking call or registration and returns; `"raise"`, for scheduled jobs, raises once the call has finished, with the failure as the `__cause__`. After `fit`, `trainer.result` still holds the model. A model whose bundle or assets could not be stored under the root is never registered. +- **`store`**: the artifact kinds stored when produced: `snapshot`, `training_table`, `model` and `prediction` (default: all). +- **`on_error`**: `"warn"` (default) warns about a failed store, tracking call or registration and returns; `"raise"`, for scheduled jobs, raises once the call has finished, with the failure as the `__cause__`. After `fit`, `trainer.result` still holds the model. A model whose bundle or assets could not be stored under the root is never registered. Under `"raise"`, `predict` also refuses a snapshot whose graph schema differs from the training graph's, and a run that could not be stored or tracked ends failed with `kumo.error`. - **Precedence:** a call's `mlops=` replaces the enclosing block, and an inner block replaces an outer one. A block covers the calling thread or asyncio task. -- **Off by default.** Without a block or `mlops=`, every call behaves as it always has; plain exports (`export(config)`, `export_config=`) never write to the artifact root. +- **Off by default.** Without a block or `mlops=`, every call behaves as it always has; plain exports (`export(config)`, `export_config=`, `output_config=`) never write to the artifact root. - `kumo.MLOps(...)` checks the tracker and registry when it is built (one experiment lookup, one registry listing, the model name's rules), so wrong credentials or a bad name fail before any work. ### Reading artifacts back @@ -131,6 +133,7 @@ Pass the `kumo.MLOps` where a connector goes, with paths relative to the artifac snapshot = kumo.GraphSnapshot.load_from_connector(cfg, id=result.snapshot_id) trainer = kumo.Trainer.load_from_connector(cfg, path=f"models/{result.id}") champion = kumo.Trainer.load_from_connector(cfg, location=kumo.ModelRegistryLocation("users-churn", alias="champion")) +run = kumo.PredictionResult.load_from_connector(cfg, path=f"predictions/{run_id}") table = kumo.TrainingTable.load_from_connector(cfg, path=f"training_tables/{table_id}", pq=pquery, snapshot=snapshot) ``` @@ -142,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` and the MLflow pyfunc 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 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. @@ -191,7 +194,20 @@ Jobs that load `alias="champion"` pick up the new version on their next run. ## Score -`predict(snapshot)` scores the trained query's population at its prediction time; `entities=` scores given rows; `predict(prediction_table, snapshot=...)` scores a table generated beforehand. A scoring job needs no `Graph` code: `GraphSnapshot.load_from_connector(cfg, id=...)` and `Trainer.load_from_connector(cfg, ...)`, then `predict`. It warns when the snapshot was taken of a different graph than the model's, naming the schema changes. +`predict(snapshot)` scores the trained query's population at its prediction time; `entities=` scores given rows; `predict(prediction_table, snapshot=...)` scores a table generated beforehand. + +```python +run = champion.predict(snapshot, output_config=kumo.PredictionOutputConfig(connector, path="...")) +run.id, run.bundle_id, run.snapshot_id, run.scored_rows, run.skipped_rows +``` + +- **`PredictionOutputConfig(connector, path=...)`**: a directory or file on a file destination (local, S3, GCS, a volume or stage), else a table. `output_paths={"predictions": ..., "embeddings": ...}` writes embeddings too. +- **Tables** append and get two columns, `PREDICTION_ID` (`run.id`) and `BUNDLE_ID`, so each row names its run and model. **Files** go to a directory per run, `//` (or `//scores.csv` for a file name), with the run's `_provenance.json` next to the Parquet parts. `write_mode=WriteMode.OVERWRITE` replaces the table or exactly `path`. +- With MLOps on, each run is stored under `/predictions//` and recorded as a tracker run named after its id: scoring options as params, row counts and score percentiles as metrics, the scoring snapshot's source tables as inputs. A failed export still records the run, ended failed with `kumo.error`. +- Entities the snapshot does not hold are skipped and counted (`run.skipped_rows`, the `predictions.skipped_fraction` metric), never refused. +- A scoring job needs no `Graph` code: `GraphSnapshot.load_from_connector(cfg, id=...)` and `Trainer.load_from_connector(cfg, ...)`, then `predict`. It warns when the snapshot was taken of a different graph than the model's, naming the schema changes (tables, columns, types, keys, edges); under `on_error="raise"` it raises instead, before scoring. + +**Scheduled jobs:** [`examples/mlops`](../../../examples/mlops/README.md) has a fit job and a score job with Databricks Jobs, SageMaker Pipelines and Snowflake Tasks wiring. ### MLflow pyfunc @@ -208,6 +224,73 @@ 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). +## 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`. + +```python +with kumo.mlops(cfg): + drift = champion.score_drift(run, previous=earlier_run) # psi, ks, wasserstein, features, input_changes + evaluation = champion.evaluate(run, labels_snapshot=newer) # metrics, matured, labelled_fraction +``` + +- Inside a block, the training and scoring snapshots come from the manifests and the artifact root (only their column statistics are fetched); outside one, pass `snapshot=` and `training_snapshot=`. Missing inputs leave a field `None` with a `reasons` entry. +- `evaluate` returns the metrics the fit reports (`auroc`/`ap`, `mae`/`rmse`, `acc`, `map@k`), `matured=False` while the label window is still open, and `actuals`, a local Parquet directory of the evaluated rows (`pd.read_parquet(evaluation.actuals)`). `evaluation.export(PredictionOutputConfig(...))` writes them to a table or directory for a warehouse monitor. + +A later job finds the champion's runs by tag and reads them from the root: + +```python +champion = kumo.Trainer.load_from_connector(cfg, location=kumo.ModelRegistryLocation("users-churn", alias="champion")) +with cfg.tracking.client(): # the tracker's MLflow server, as its identity + runs = mlflow.search_runs(experiment_names=["churn"], order_by=["attributes.start_time DESC"], max_results=2, + filter_string=f"tags.`kumo.kind` = 'prediction' AND tags.`kumo.bundle_id` = '{champion.result.id}' " + "AND attributes.status = 'FINISHED'") +latest, *earlier = [kumo.PredictionResult.load_from_connector(cfg, path=f"predictions/{prediction_id}") + for prediction_id in runs["tags.kumo.prediction_id"]] +with kumo.mlops(cfg): + drift = champion.score_drift(latest, previous=earlier[0] if earlier else None) +``` + +Snowflake runs have no tags, so read the ids from a table the runs wrote instead: + +```python +ids = [row[0] for row in sf.session().sql( + f"SELECT DISTINCT PREDICTION_ID FROM CHURN_SCORES WHERE BUNDLE_ID = '{champion.result.id}'").collect()] +latest, *earlier = sorted((kumo.PredictionResult.load_from_connector(cfg, path=f"predictions/{prediction_id}") + for prediction_id in ids), key=lambda run: run.manifest.created_at, reverse=True) +``` + +`load_from_connector` needs the run's `_SUCCESS` marker and fetches only its `_provenance.json`; the Parquet parts download when `evaluate` (or `files`) first reads them. + +### Warehouse model monitors + +Both platforms monitor the evaluated rows exported to a table; neither monitors link prediction or ranking outputs. + +**Databricks (Data Quality Monitoring, inference log; databricks-sdk ≥ 0.68):** + +```python +evaluation.export(kumo.PredictionOutputConfig(dbx, path="churn_evaluated")) +w = WorkspaceClient() +table_id = w.tables.get("main.ml.churn_evaluated").table_id +w.data_quality.create_monitor(monitor=Monitor(object_type="table", object_id=table_id, + data_profiling_config=DataProfilingConfig(output_schema_id=w.schemas.get("main.ml").schema_id, + assets_dir="/Workspace/Users//monitoring", inference_log=InferenceLogConfig( + problem_type=InferenceProblemType.INFERENCE_PROBLEM_TYPE_CLASSIFICATION, + prediction_column="TARGET_PRED", label_column="TARGET", model_id_column="BUNDLE_ID", + timestamp_column="EVALUATED_AT", granularities=[AggregationGranularity.AGGREGATION_GRANULARITY_1_DAY])))) +``` + +**Snowflake (Model Monitor)**, on a version `MLOps(registry=)` registered. `evaluation.export` to a Snowflake table writes `TARGET` as a number and `EVALUATED_AT` as `TIMESTAMP_NTZ`, as the monitor needs; the baseline is `result.holdout_df().astype({"TARGET": int}).assign(EVALUATED_AT=pd.Timestamp.now())` written with `session.write_pandas`: + +```sql +CREATE MODEL MONITOR churn_monitor WITH MODEL=churn_gnn VERSION=v3 FUNCTION=predict + SOURCE=churn_evaluated BASELINE=churn_holdout_evaluated WAREHOUSE=wh REFRESH_INTERVAL='1 day' + AGGREGATION_WINDOW='1 day' TIMESTAMP_COLUMN=EVALUATED_AT ID_COLUMNS=(ENTITY) + PREDICTION_SCORE_COLUMNS=("True_PROB") ACTUAL_CLASS_COLUMNS=(TARGET); +``` + +Details per task type: [Warehouse monitor columns](#warehouse-monitor-columns). + ## Credentials and running off-platform `MlflowTracking(target=...)` and `SnowflakeTracking(connector)` take the same identities as the connectors, so a fit on a laptop or any VM can record to Databricks, SageMaker or Snowflake. An identity with no credentials (or no identity) uses the platform's ambient chain (a notebook, job or instance role, `aws configure`, `databricks auth login`); one with credentials uses exactly those. @@ -253,6 +336,8 @@ Every artifact has one way out and one way back. A `kumo.MLOps` stands in for th | Snapshot | `snapshot.export(SnapshotOutputConfig(connector, path))`, returns the location | `GraphSnapshot.load_from_connector(connector, path=location)`, `(cfg, id=...)` | | Training table | `table.export(TrainingTableOutputConfig(connector, name))` | `TrainingTable.load_from_connector(connector, path=name, pq=pq, snapshot=snapshot)` | | Model | `result.export(ModelOutputConfig(connector, path=... or location=...))`, or `fit(export_config=...)` | `Trainer.load_from_connector(connector, path=... or location=...)` | +| 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. @@ -262,12 +347,13 @@ Every artifact carries `_provenance.json`, metadata only (the underscore keeps S - **Snapshot:** each source table's connector, location and version, and the engine build. `graph.snapshot(record_sources=False)` keeps only connector and table names, a SHA-256 of each location, and Custom SQL as a hash. - **Training table:** the query, generation plan, task type and split seed, its snapshot, and for `TrainingTable.from_source_table` the source table. `table.id` is deterministic over snapshot, query, plan, seed and engine build; a source without a version gets a random id. -- **Model** (`ModelManifest`, `result.manifest`): the bundle id (a content hash excluding the manifest), fit times and fit mode, the training table and warm-start model (an embedded model keeps only ids of its own model parents), `tracking` (the tracker's run) and `dependencies`. The bundle carries the holdout predictions (`holdout.parquet`), so `holdout_df()` works on a loaded model. +- **Model** (`ModelManifest`, `result.manifest`): the bundle id (a content hash excluding the manifest), fit times and fit mode, the training table and warm-start model (an embedded model keeps only ids of its own model parents), `tracking` (the tracker's run), `dependencies` and `holdout` (score statistics of the holdout rows: count, mean, min, max, q01..q99; the positive rate for binary tasks), the baseline `score_drift` compares against. The bundle carries the holdout predictions (`holdout.parquet`), so `holdout_df()` works on a loaded model. +- **Prediction run** (`PredictionManifest`, `run.manifest`): its id; the model and the scoring snapshot; the prediction table (query hash, anchors, `requested_rows`, `scored_rows`, `skipped_rows` and up to 5 skipped ids per reason); the scoring options and where each came from; `outputs`, the destinations `output_config` wrote (even after a partial failure); status, timings and score statistics per output. ### 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` and the pyfunc 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 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 @@ -281,6 +367,18 @@ trainer = kumo.Trainer.load_from_connector(cfg, path=f"models/{result.id}") # A text column `infer_metadata()` reads as categorical (short values from a small vocabulary) takes a text encoder once its stype is text: `users.with_column("bio", stype="text")`. +### Prediction output contract + +| Column | Task types | Meaning | +|---|---|---| +| `ENTITY` | all | the entity's primary key | +| `ANCHOR_TIMESTAMP` | temporal queries | the prediction time, when a run scores several | +| `TARGET_PRED` | binary classification, regression, forecasting | the predicted label or value | +| `False_PROB`, `True_PROB` | binary classification | class probabilities; `True_PROB` is summarized | +| `CLASS`, `SCORE` | multiclass, multilabel, ranking | one row per returned class; the top `SCORE` is summarized | + +`predict` checks each destination first: a directory holding another export raises `ArtifactDestinationConflict`. A table written without the id columns takes one `OVERWRITE` before it can append. Object stores replace a directory by deleting, then writing, so a reader in between finds no `_SUCCESS`; `APPEND` to files is not supported. + ### Moving snapshots between machines - **Layout:** `export` writes `///`, holding `kumo_snapshot.json` and one directory per export generation (`index/` and `_provenance.json`). @@ -289,9 +387,24 @@ A text column `infer_metadata()` reads as categorical (short values from a small - **Staleness:** `snapshot.source_changes(graph)` reports each table as `unchanged`, `changed` or `unknown`; `load_from_connector(..., graph=graph)` warns about another graph's schema or changed tables. - **Trust:** load only from paths whose writers you trust; checksums detect corruption, not tampering. A warehouse connector cannot hold a snapshot. +### Drift statistics + +| statistic | definition | error from the stored curves | +|---|---|---| +| `psi` | population stability over the baseline's deciles | — | +| `ks` | largest gap between the two CDFs | ≤ 0.02 + 2× the sketch's rank error | +| `wasserstein` | mean absolute gap between the quantile functions, in score units | ≤ 0.01 × (sum of both score ranges) + sketch error | + +The statistics are read off the recorded percentile curves (min, q01..q99, max), with no pass over the data. `DriftReport.features` compares, per column the model encodes (by its stype in the model plan), the statistics DiskGraph wrote when both snapshots were built: numeric deltas of mean, std, min, max, quartiles, skewness, kurtosis, null rate and distinct count plus `standardized_mean_shift`; categorical null rate, distinct count, PSI and Jensen-Shannon divergence (`category_counts_capped` flags truncated counts); timestamp min and max shift. Other columns, such as text, are not compared and are named in `reasons["features"]`. These are whole-snapshot statistics, so a slow change in a large event table is diluted. A tracked monitoring run logs the top ten columns by standardized shift and by PSI and the full comparison as `feature_drift.json`. + +### Warehouse monitor columns + +- 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. + ### 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`). 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. +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. ### Progress events diff --git a/packages/kumo-engine/kumo/engine/__init__.py b/packages/kumo-engine/kumo/engine/__init__.py index ca5a6b00f..78d7b8a3a 100644 --- a/packages/kumo-engine/kumo/engine/__init__.py +++ b/packages/kumo-engine/kumo/engine/__init__.py @@ -29,10 +29,12 @@ from kumo.engine._manifest import ( ModelManifest, + PredictionManifest, SnapshotManifest, TrainingTableManifest, ) from kumo.engine._mlops import MLOps, mlops +from kumo.engine._monitoring import DriftReport, EvaluationResult from kumo.engine._version import __version__ from kumo.engine.config import configure from kumo.engine.connector import ( @@ -69,6 +71,8 @@ __all__ = [ "__version__", "configure", + "EvaluationResult", + "DriftReport", "EmbeddingModelStore", "GloVeVectors", "TextEmbedding", @@ -100,6 +104,7 @@ "SnapshotManifest", "TrainingTableManifest", "ModelManifest", + "PredictionManifest", "PredictiveQuery", "TrainingTable", "PredictionTable", diff --git a/packages/kumo-engine/kumo/engine/_manifest.py b/packages/kumo-engine/kumo/engine/_manifest.py index 574e84836..f23a628ba 100644 --- a/packages/kumo-engine/kumo/engine/_manifest.py +++ b/packages/kumo-engine/kumo/engine/_manifest.py @@ -127,7 +127,7 @@ class SourceTableEntry(Record): #: The artifact kinds, as manifests and the artifact root name them. -KINDS = ('snapshot', 'training_table', 'model') +KINDS = ('snapshot', 'training_table', 'model', 'prediction') @dataclass(frozen=True, kw_only=True) @@ -271,6 +271,14 @@ class Dependency(Record): location: Optional[str] = omit_none() +@dataclass(frozen=True, kw_only=True) +class HoldoutStats(Record): + truncated: Optional[bool] + rows: int + columns: Dict[str, Dict[str, Any]] + label_positive_rate: Optional[float] = omit_none() + + @dataclass(frozen=True, kw_only=True) class ModelManifest(Manifest): """How a model bundle was fitted: from which training table (embedded, @@ -285,8 +293,79 @@ class ModelManifest(Manifest): fit_mode: Optional[str] = None tracking: Optional[TrackingRef] = omit_none() dependencies: Optional[List[Dependency]] = omit_none() + holdout: Optional[HoldoutStats] = omit_none() + + EMBED_OMITS = ('tracking', 'holdout') + + +@dataclass(frozen=True, kw_only=True) +class PredictionParents(Record): + model: ModelManifest + snapshot: SnapshotManifest + + +@dataclass(frozen=True, kw_only=True) +class PredictionTableRecord(Record): + query_sha256: str + source: Optional[str] + anchors: List[str] + rows: Optional[int] + requested_rows: Optional[int] = None + scored_rows: Optional[int] = None + skipped_rows: Optional[int] = None + skipped_sample: Optional[List[Any]] = None + skip_reason: Optional[str] = omit_none() + + +@dataclass(frozen=True, kw_only=True) +class RunScoring(Record): + binary_classification_threshold: Optional[float] + threshold_source: str + num_classes_to_return: Optional[int] + count_source: str + num_workers: Optional[int] + - EMBED_OMITS = ('tracking', ) +@dataclass(frozen=True, kw_only=True) +class OutputDestination(Record): + connector: str + location: Optional[str] + + +@dataclass(frozen=True, kw_only=True) +class PredictionOutput(Record): + artifact: str + destination: OutputDestination + + +@dataclass(frozen=True, kw_only=True) +class RunTimings(Record): + started_at: str + finished_at: Optional[str] + + +@dataclass(frozen=True, kw_only=True) +class ArtifactStats(Record): + rows: Optional[int] + columns: Optional[Dict[str, Dict[str, Any]]] = omit_none() + predicted_positive_rate: Optional[float] = omit_none() + + +@dataclass(frozen=True, kw_only=True) +class PredictionManifest(Manifest): + """A prediction run: its model and scoring snapshot (embedded), the + prediction table, scoring options, outputs and score statistics.""" + + kind: str = 'prediction' + parents: PredictionParents + prediction_table: PredictionTableRecord + scoring: RunScoring + outputs: Optional[List[PredictionOutput]] + status: Optional[str] + timings: Optional[RunTimings] + stats: Optional[Dict[str, ArtifactStats]] + + EMBED_OMITS = ('outputs', 'stats', 'status', 'timings') @dataclass(frozen=True, kw_only=True) diff --git a/packages/kumo-engine/kumo/engine/_mlops.py b/packages/kumo-engine/kumo/engine/_mlops.py index e225872bd..1fead5769 100644 --- a/packages/kumo-engine/kumo/engine/_mlops.py +++ b/packages/kumo-engine/kumo/engine/_mlops.py @@ -59,7 +59,7 @@ class MLOps: MLflow) for ``MlflowTracking``, the Snowflake Model Registry for ``SnowflakeTracking``. store: the artifact kinds stored under the root when produced - (``snapshot``, ``training_table``, ``model``); + (``snapshot``, ``training_table``, ``model``, ``prediction``); ``True`` for all, ``False`` for none. on_error: what a failure to track a run or register a model does: ``'warn'`` warns and the call returns; ``'raise'`` raises once @@ -160,7 +160,7 @@ def wants(self, kind: str) -> bool: @contextlib.contextmanager def mlops(settings: MLOps) -> Iterator[MLOps]: r"""Turns MLOps on for the calls in the block (in this thread or asyncio - task): each snapshot, training table and model is stored + task): each snapshot, training table, model and prediction run is stored under ``settings.artifact_root``. An inner block replaces an outer one; an ``mlops=`` argument replaces both for its call. Runs are tracked and models registered as ``settings`` says.""" diff --git a/packages/kumo-engine/kumo/engine/_monitoring.py b/packages/kumo-engine/kumo/engine/_monitoring.py new file mode 100644 index 000000000..0003ccad3 --- /dev/null +++ b/packages/kumo-engine/kumo/engine/_monitoring.py @@ -0,0 +1,898 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +r"""Realized performance and score drift of prediction runs.""" + +import dataclasses +import json +import math +import shutil +import time +import warnings +from dataclasses import dataclass, field, replace +from pathlib import Path +from typing import Any, Callable, Dict, Mapping, Optional, Tuple +from uuid import uuid4 + +import numpy as np +import pandas as pd + +from kumo.engine._manifest import PredictionManifest + +_EVALUATION_FILE = '_evaluation.json' + + +def _for_model_monitor(files: Any, directory: Path) -> list: + r"""``files`` with the types Snowflake Model Monitor accepts: a boolean + ``TARGET`` as 0/1 and timestamps without a time zone (UTC), which load as + NUMBER and TIMESTAMP_NTZ.""" + import pyarrow as pa + import pyarrow.parquet as parquet + + def monitored(field: Any) -> Any: + if field.name == 'TARGET' and pa.types.is_boolean(field.type): + return field.with_type(pa.int8()) + if pa.types.is_timestamp(field.type) and field.type.tz is not None: + return field.with_type(pa.timestamp(field.type.unit)) + return field + + typed = [] + for path in files: + source = parquet.ParquetFile(path) + schema = pa.schema([monitored(field) for field in source.schema_arrow]) + target = str(directory / Path(path).name) + with parquet.ParquetWriter(target, schema) as writer: + for batch in source.iter_batches(): + writer.write_table( + pa.Table.from_batches([batch]).cast(schema)) + typed.append(target) + return typed + + +@dataclass(frozen=True) +class EvaluationResult: + """Metrics of a prediction run against labels from a newer snapshot. + + ``actuals`` is a Parquet directory of the evaluated rows: ``ENTITY``, + ``ANCHOR_TIMESTAMP`` (the prediction time for static queries), the + prediction (``True_PROB`` and ``TARGET_PRED``; ``TARGET_PRED``; or the + top ``CLASS`` and its ``SCORE``), the label ``TARGET`` and + ``EVALUATED_AT``, when `evaluate` ran. + """ + + prediction_id: str + bundle_id: Optional[str] + snapshot_id: str + task_type: Optional[str] + anchors: Tuple[str, ...] + metrics: Mapping[str, Optional[float]] + rows: int + labelled: int + labelled_fraction: Optional[float] + matured: Optional[bool] + bundle_matches: Optional[bool] + reasons: Mapping[str, str] = field(default_factory=dict) + actuals: Optional[str] = None + + _exported: Any = field(default=None, compare=False, repr=False) + + def export(self, output_config: Any) -> None: + r"""Copy the evaluated rows to ``output_config``'s predictions + destination (a `PredictionOutputConfig`), e.g. a table to monitor. + A Parquet directory also gets ``_evaluation.json``, the other fields, + for `load_from_connector`. A Snowflake table gets the types Model + Monitor needs: ``TARGET`` as a number and TIMESTAMP_NTZ times.""" + import tempfile + + from kumo.engine._scratch import get_scratch_dir + from kumo.engine.connector import ( + SnowflakeConnector, + is_file_destination, + ) + from kumo.engine.trainer.trainer import _holds_parquet + + if self.actuals is None: + raise ValueError(f'No evaluated rows: {dict(self.reasons)}') + destination = output_config._targets( + self.prediction_id)['predictions'] + files = sorted(str(path) for path in Path(self.actuals).glob('*.parquet')) + written = max(Path(path).stat().st_mtime_ns for path in files) + extra: Dict[str, Any] = {} + columns = output_config._id_columns(destination, self.prediction_id, + self.bundle_id) + if columns: + extra['extra_columns'] = columns + if _holds_parquet(output_config.connector, destination): + fields = dataclasses.asdict(self) + for name in ('actuals', '_exported'): + fields.pop(name) + extra['sidecars'] = { + _EVALUATION_FILE: json.dumps(fields, sort_keys=True).encode() + } + get_scratch_dir().mkdir(parents=True, exist_ok=True) + with tempfile.TemporaryDirectory(dir=get_scratch_dir()) as typed: + if isinstance(output_config.connector, SnowflakeConnector + ) and not is_file_destination( + output_config.connector, destination): + files = _for_model_monitor(files, Path(typed)) + output_config._write(destination, files, + export_id=f'{self.prediction_id}-evaluation-' + f'{Path(self.actuals).name}-{written}', + **extra) + + @classmethod + def load_from_connector(cls, connector: Any, *, + path: str) -> 'EvaluationResult': + r"""Load an evaluation `export` wrote to a Parquet directory; its rows + are downloaded into ``actuals``, removed with the result.""" + import tempfile + import weakref + + from kumo_core.sources.transfer import StagedFile, fetch_all + from pyarrow.fs import FileSelector, FileType + + from kumo.engine._scratch import get_scratch_dir + from kumo.engine.connector import file_store + + store = file_store(connector, path) + if store is None: + raise TypeError(f'{path} is not a file destination of this ' + f'connector') + filesystem, root = store + root = root.rstrip('/') + try: + with filesystem.open_input_stream( + f'{root}/{_EVALUATION_FILE}') as reader: + fields = json.loads(reader.read()) + except (FileNotFoundError, OSError) as error: + raise FileNotFoundError( + f'{path} has no {_EVALUATION_FILE}: not an evaluation export' + ) from error + if filesystem.get_file_info( + f'{root}/_SUCCESS').type == FileType.NotFound: + raise FileNotFoundError( + f'{path} has no _SUCCESS marker: the export is incomplete or ' + f'still running') + scratch = get_scratch_dir() + scratch.mkdir(parents=True, exist_ok=True) + local = Path(tempfile.mkdtemp(prefix='fetched-evaluation-', + dir=scratch)) + + def fetch(staged: StagedFile, target: str) -> None: + with filesystem.open_input_stream(staged.path) as reader, \ + open(target, 'wb') as writer: + shutil.copyfileobj(reader, writer, 1 << 20) + + fetch_all([ + StagedFile(info.path, info.size) + for info in filesystem.get_file_info(FileSelector(root)) + if info.type == FileType.File and info.path.endswith('.parquet') + ], str(local / 'actuals'), fetch=fetch) + fields['anchors'] = tuple(fields['anchors']) + result = cls(**fields, actuals=str(local / 'actuals'), + _exported=local) + weakref.finalize(result, shutil.rmtree, str(local), True) + return result + + +@dataclass(frozen=True) +class DriftReport: + """Score and input drift of a prediction run; ``reasons`` explains Nones. + + ``psi``, ``ks`` and ``wasserstein`` compare the run's scores with the + holdout's; ``previous_*`` with ``previous_prediction_id``'s run. All are + computed from the recorded percentile curves (``wasserstein`` is in score + units). ``features`` compares, per table and column the model encodes, + the column statistics DiskGraph wrote when each snapshot was built + (training snapshot, then scoring snapshot). + """ + + prediction_id: Optional[str] + bundle_id: Optional[str] + score: Optional[str] + psi: Optional[float] + ks: Optional[float] + wasserstein: Optional[float] + holdout_truncated: Optional[bool] + bundle_matches: Optional[bool] + previous_prediction_id: Optional[str] + previous_psi: Optional[float] + previous_ks: Optional[float] + previous_wasserstein: Optional[float] + input_changes: Optional[Mapping[str, Tuple[Any, Any]]] + reasons: Mapping[str, str] + features: Mapping[str, Mapping[str, Mapping[str, Any]]] = field( + default_factory=dict) + + +def evaluate(trainer: Any, run: Any, snapshot: Any) -> EvaluationResult: + query = trainer._training_query + if query is None: + raise RuntimeError('Fit or load the Trainer before evaluate()') + reasons: Dict[str, str] = {} + out: Dict[str, Any] = dict(task_type=None, anchors=(), metrics={}, + rows=0, labelled=0, matured=None, + reasons=reasons, actuals=None, + step='importing') + try: + _evaluate(out, trainer, run, snapshot, query) + except Exception as error: + reasons['metrics'] = f"{out['step']} failed: {error}" + out.pop('step') + if out['matured'] is False: + warnings.warn( + f'Prediction run {run.id}: labels are not mature yet, since the ' + 'label window ends after the snapshot data; the metrics cover ' + 'only the outcomes observed so far.', UserWarning, stacklevel=3) + return EvaluationResult( + prediction_id=run.id, + bundle_id=run.bundle_id, + snapshot_id=snapshot.id, + labelled_fraction=(out['labelled'] / + out['rows'] if out['rows'] else None), + bundle_matches=_same(run.bundle_id, _bundle_id(trainer)), + **out, + ) + + +_TASKS = { + 'binary_classification': 'binary', + 'regression': 'regression', + 'multiclass_classification': 'multiclass', + 'link_prediction': 'link', + 'temporal_link_prediction': 'link', + 'static_link_prediction': 'link', +} + + +def _evaluate(out: Dict[str, Any], trainer: Any, run: Any, snapshot: Any, + query: str) -> None: + import diskgraph + from kumo_core.pquery import _target_maximum + from kumo_core.training import custom_table + from kumo_core.types import PredictionArtifactType + from pyarrow import parquet + + from kumo.engine.pquery.predictive_query import PredictiveQuery + + reasons = out['reasons'] + out['step'] = 'parsing the query' + pquery = PredictiveQuery._from_snapshot(snapshot, query) + parsed = pquery._parsed + static = parsed.query_type == 'static' + out['task_type'] = parsed.task_type + + out['step'] = 'reading the predictions' + files = [ + str(path) for path in replace( + run, artifact_type=PredictionArtifactType.PREDICTIONS).files + ] + columns = parquet.read_schema(files[0]).names + manifest = run.manifest + anchors = () if static or manifest is None else tuple( + sorted({_ns(value) for value in manifest.prediction_table.anchors})) + keys = ['ENTITY'] + if len(anchors) > 1: + keys.append('ANCHOR_TIMESTAMP') + out['anchors'] = tuple( + pd.Timestamp(anchor, tz='UTC').isoformat() for anchor in anchors) + if anchors: + window_end = parsed.query.get_combined_date_offset_range( + ).end_date_offset + try: + latest = _target_maximum(parsed, str(snapshot.index_path), min) + out['matured'] = bool( + pd.Timestamp(max(anchors)) + window_end <= pd.Timestamp(latest)) + except Exception as error: + reasons['matured'] = str(error) + task = _TASKS.get(parsed.task_type) + if task is None: + reasons['metrics'] = f'{parsed.task_type} metrics are not supported' + return + if not static and not anchors: + reasons['metrics'] = 'the run records no anchor timestamps' + return + if 'ANCHOR_TIMESTAMP' in keys and 'ANCHOR_TIMESTAMP' not in columns: + reasons['metrics'] = ('the run has several anchors but its ' + 'predictions have no ANCHOR_TIMESTAMP column') + return + + out['step'] = 'generating labels' + entity, timestamp, target = (custom_table.entity_column(), + custom_table.timestamp_column(), + custom_table.target_column()) + actuals = _actuals_dir(run, snapshot) + spill = actuals.parent / f'.spill-{uuid4().hex}' + try: + with pquery._label_parts( + snapshot, anchors, + trainer.training_table_generation_plan) as parts: + out['step'] = 'joining labels with predictions' + label_columns = parquet.read_schema(parts[0]).names + shutil.rmtree(actuals, ignore_errors=True) + emitted = [name for name in _EMITTED[task] if name in columns] + spec = { + 'task': task, + 'prediction_keys': keys, + 'label_keys': [entity, timestamp][:len(keys)], + 'label_target': target, + 'actuals_dir': str(actuals), + 'actual_columns': [[entity, 'ENTITY']] + + ([[timestamp, 'ANCHOR_TIMESTAMP']] + if timestamp in label_columns and not static else []) + + [[target, 'TARGET']], + 'prediction_columns': [[name, name] for name in emitted], + 'constant_timestamps': ([[ + 'ANCHOR_TIMESTAMP', _prediction_time(manifest) + ]] if static and manifest is not None else []) + + [['EVALUATED_AT', time.time_ns()]], + 'spill_dir': str(spill), + } + joined = json.loads( + diskgraph.evaluate_prediction_parquet(files, parts, + json.dumps(spec))) + except Exception as error: + if (out['step'] != 'generating labels' + or 'empty training table' not in str(error)): + raise + reasons['metrics'] = 'no entity has a label at these anchors' + return + finally: + shutil.rmtree(spill, ignore_errors=True) + out.update(rows=joined['rows'], labelled=joined['labelled'], + actuals=str(actuals) if joined['labelled'] else None) + if out['labelled'] == 0: + reasons['metrics'] = 'no prediction has a label' + return + metrics = out['metrics'] = { + name: _finite(value) + for name, value in joined['metrics'].items() + } + if None in metrics.values(): + reasons['metrics'] = 'a metric is undefined for these labels' + + +_EMITTED = { + 'binary': ('True_PROB', 'TARGET_PRED'), + 'regression': ('TARGET_PRED', ), + 'multiclass': ('CLASS', 'SCORE'), + 'link': (), +} + + +def _prediction_time(manifest: PredictionManifest) -> int: + anchors = manifest.prediction_table.anchors + return _ns(anchors[0] if anchors else manifest.timings.started_at) + + +def _bundle_id(trainer: Any) -> Optional[str]: + result = getattr(trainer, 'result', None) + return None if result is None else result.id + + +def _actuals_dir(run: Any, snapshot: Any) -> Path: + from kumo.engine._scratch import get_scratch_dir + + selection = getattr(run, 'selection', None) + root = (Path(selection.pointer_path).parent if selection is not None + else get_scratch_dir() / 'evaluations' / str(run.id)) + return root / 'actuals' / str(snapshot.id).replace('/', '_') + + +def score_drift(trainer: Any, run: Any, *, previous: Any, snapshot: Any, + training_snapshot: Any, root: Any = None) -> DriftReport: + reasons: Dict[str, str] = {} + score, summary = _score_summary(run) + result = getattr(trainer, 'result', None) + model = None if result is None else result.manifest + holdout = getattr(model, 'holdout', None) + baseline = None + if summary is None: + reasons['holdout'] = 'the run records no output score statistics' + elif holdout is None: + reasons['holdout'] = 'the model records no holdout score statistics' + else: + baseline = holdout.columns.get(score) + if _curves(summary, baseline) is None: + reasons['holdout'] = f'the holdout records no percentiles for {score}' + + earlier = None + if previous is not None: + previous_score, earlier = _score_summary(previous) + if summary is None or earlier is None: + reasons['previous'] = 'a run records no score statistics' + elif previous_score != score: + reasons['previous'] = 'the runs record different scores' + earlier = None + elif _curves(summary, earlier) is None: + reasons['previous'] = 'a run records no percentiles' + + sides = _snapshot_stats(trainer, run, snapshot, training_snapshot, root) + counts = [side['counts'] for side in sides.values()] + input_changes = None + if all(isinstance(count, dict) for count in counts): + input_changes = _count_changes(*counts) + else: + reasons['input_changes'] = _missing('ingest statistics', sides, + 'counts') + run_bundle = getattr(run, 'bundle_id', None) + stats = (_psi, _ks, _wasserstein) + against = [[None if b is None or summary is None else stat(summary, b) + for stat in stats] for b in (baseline, earlier)] + return DriftReport( + prediction_id=getattr(run, 'id', None), + bundle_id=run_bundle, + score=score, + psi=against[0][0], + ks=against[0][1], + wasserstein=against[0][2], + holdout_truncated=None if holdout is None else holdout.truncated, + bundle_matches=_same(run_bundle, _bundle_id(trainer)), + previous_prediction_id=getattr(previous, 'id', None), + previous_psi=against[1][0], + previous_ks=against[1][1], + previous_wasserstein=against[1][2], + input_changes=input_changes, + reasons=reasons, + features=_features(trainer, sides, reasons), + ) + + +def track(report: Any, run: Any, settings: Any) -> None: + r"""Log ``report`` as a new monitoring run of ``run``; a failure warns, + or raises under ``on_error='raise'``.""" + tracking = None if settings is None else settings.tracking + if tracking is None: + return + + from kumo.engine._mlops import full_error, raise_on_error, warn_failure + from kumo.engine._tracking import RunSpec, log_run, run_record + + try: + values = {**report.metrics} if isinstance( + report, EvaluationResult) else {} + for name, value in dataclasses.asdict(report).items(): + if name == 'input_changes': + for key, (before, after) in (value or {}).items(): + values[f'input.{key}.before'] = before + values[f'input.{key}.after'] = after + elif name == 'features': + values.update(feature_metrics(value or {})) + elif name not in ('metrics', 'reasons', 'actuals'): + values[name] = value + prediction_id = getattr(run, 'id', None) + record = run_record( + RunSpec( + kind='monitoring', + name=(f'evaluate-{prediction_id}' if isinstance( + report, EvaluationResult) else f'drift-{prediction_id}'), + params={ + key: value + for key, value in values.items() if isinstance(value, str) + }, + metrics={ + key: value + for key, value in values.items() + if isinstance(value, (int, float)) + and not isinstance(value, bool) + }, + parents={ + 'monitors': prediction_id, + 'bundle_id': getattr(run, 'bundle_id', None), + }, + artifacts=({ + 'feature_drift.json': + json.dumps(report.features, indent=2, sort_keys=True, + default=str) + } if getattr(report, 'features', None) else {}), + )) + except Exception as error: + if settings.on_error == 'warn': + warn_failure( + f'Could not track the monitoring run: {full_error(error)}', + error, 3) + failure: Optional[BaseException] = error + else: + failure = log_run(tracking, 'log_monitoring', record) + raise_on_error(settings, failure, f'Monitoring of prediction run ' + f'{getattr(run, "id", None)}', stacklevel=3) + + +def _count_changes(before: Mapping[str, Mapping[str, int]], + after: Mapping[str, Mapping[str, int]] + ) -> Dict[str, Tuple[Any, Any]]: + changes = {} + for kind, key in (('rows', 'tables.{}.rows'), ('edges', 'edges.{}')): + old, new = before.get(kind) or {}, after.get(kind) or {} + for name in sorted(set(old) | set(new)): + if not name.endswith('_reverse') and old.get(name) != new.get(name): + changes[key.format(name)] = (old.get(name), new.get(name)) + return changes + + +def _score_summary(run: Any) -> Tuple[Optional[str], Optional[Mapping]]: + try: + columns = run.manifest.stats['predictions'].columns + return next(iter(columns.items())) + except Exception: + return None, None + + +def _psi(actual: Mapping[str, Any], + expected: Optional[Mapping[str, Any]]) -> Optional[float]: + r"""PSI over ``expected``'s deciles; each side's bin shares are read off + its percentile curve (linear between min, q01..q99 and max).""" + curves = _curves(actual, expected) + if curves is None: + return None + edges = np.unique(curves[1][0][10:91:10]) + + def shares(values: np.ndarray, levels: np.ndarray) -> np.ndarray: + upper = np.searchsorted(values, edges, side='right') + lower = np.maximum(upper - 1, 0) + upper = np.minimum(upper, len(values) - 1) + width = values[upper] - values[lower] + step = np.divide(edges - values[lower], width, + out=np.ones_like(edges), where=width > 0) + cdf = levels[lower] + np.clip(step, 0, 1) * (levels[upper] - + levels[lower]) + return np.maximum(np.diff(np.concatenate([[0], cdf, [1]])), 1e-6) + + actual_shares, expected_shares = shares(*curves[0]), shares(*curves[1]) + return _finite( + np.sum((actual_shares - expected_shares) * + np.log(actual_shares / expected_shares))) + + +def _ks(a: Mapping[str, Any], b: Optional[Mapping[str, Any]]) -> Optional[float]: + r"""sup |F_a - F_b| of the percentile curves, at every knot from both + sides (point masses at repeated knots).""" + curves = _curves(a, b) + if curves is None: + return None + points = np.union1d(curves[0][0], curves[1][0]) + + def cdf(values: np.ndarray, levels: np.ndarray, side: str) -> np.ndarray: + index = np.searchsorted(values, points, side=side) + lower = np.clip(index - 1, 0, len(values) - 1) + upper = np.clip(index, 0, len(values) - 1) + width = values[upper] - values[lower] + step = np.divide(points - values[lower], width, + out=np.zeros_like(points), where=width > 0) + inside = levels[lower] + step * (levels[upper] - levels[lower]) + return np.where(index == 0, 0.0, + np.where(index == len(values), 1.0, inside)) + + return _finite( + max( + np.max(np.abs(cdf(*curves[0], side) - cdf(*curves[1], side))) + for side in ('left', 'right'))) + + +def _wasserstein(a: Mapping[str, Any], + b: Optional[Mapping[str, Any]]) -> Optional[float]: + r"""Integral of |Q_a - Q_b| over levels, exact for the linear curves.""" + curves = _curves(a, b) + if curves is None: + return None + d = curves[0][0] - curves[1][0] + left, right = d[:-1], d[1:] + size = np.abs(left) + np.abs(right) + crossing = np.divide(left**2 + right**2, 2 * size, out=np.zeros_like(size), + where=size > 0) + segments = np.where(left * right >= 0, size / 2, crossing) + return _finite(0.01 * np.sum(segments)) + + +def _curves(a: Mapping[str, Any], + b: Optional[Mapping[str, Any]]) -> Optional[tuple]: + try: + curves = (_percentile_curve(a), _percentile_curve(b)) + except Exception: + return None + return None if None in curves else curves + + +def _percentile_curve(summary: Optional[Mapping[str, Any]]) -> Optional[tuple]: + percentiles = (summary or {}).get('percentiles') + if not percentiles or len(percentiles) != 99: + return None + values = np.maximum.accumulate( + np.array([summary['min'], *percentiles, summary['max']], dtype=float)) + return values, np.linspace(0, 1, 101) + + +def _ns(value: Any) -> int: + timestamp = pd.Timestamp(value) + if timestamp.tzinfo is not None: + timestamp = timestamp.tz_convert(None) + return timestamp.value + + +def _same(a: Optional[str], b: Optional[str]) -> Optional[bool]: + return None if a is None or b is None else a == b + + +def _finite(value: Any) -> Optional[float]: + if value is None: + return None + value = float(value) + return value if math.isfinite(value) else None + + +#: Columns per ranked list, and tables, in a monitoring run's metrics; keeps +#: the run far below Snowflake's 200 metrics per run. +FEATURE_METRICS_TOP = 10 +_NUMERIC_STATS = ('mean', 'std', 'min', 'max', 'q25', 'q50', 'q75', + 'skewness', 'kurtosis', 'null_rate', 'num_unique') +_CATEGORY_STATS = ('null_rate', 'num_unique') +_TIME_STATS = ('min_secs', 'max_secs') +_SMOOTHING = 1e-6 + + +def _features(trainer: Any, sides: Mapping[str, Mapping[str, Any]], + reasons: Dict[str, str]) -> Dict[str, Dict[str, Any]]: + r"""Per-column statistic changes between the build-time ``col_stats.json`` + of the training and scoring snapshots, for the columns the model + encodes. Never raises; gaps go to ``reasons['features']``.""" + try: + before, after = (side['col_stats'] for side in sides.values()) + if not (isinstance(before, dict) and isinstance(after, dict)): + reasons['features'] = _missing('column statistics', sides, + 'col_stats') + return {} + used = _used_columns(getattr(trainer, 'result', None)) + if used is None: + reasons['features'] = 'the model records no encoded columns' + return {} + features: Dict[str, Dict[str, Any]] = {} + absent, uncompared = [], [] + for table, columns in sorted(used.items()): + for column in sorted(columns): + old = (before.get(table) or {}).get(column) + new = (after.get(table) or {}).get(column) + if old is None or new is None: + absent.append(f'{table}.{column}') + continue + change = _column_change(old, new) + if change is None: + uncompared.append( + f"{table}.{column} ({new.get('semantic_type')})") + else: + features.setdefault(table, {})[column] = change + notes = [] + if absent: + notes.append('column statistics missing on a side for ' + + ', '.join(absent)) + if uncompared: + notes.append('no drift statistics for the type of ' + + ', '.join(uncompared)) + if notes: + reasons['features'] = '; '.join(notes) + return features + except Exception as error: + reasons['features'] = f'feature statistics failed: {error}' + return {} + + +def _parent_snapshot_id(manifest: Any) -> Optional[str]: + parents = getattr(manifest, 'parents', None) + table = getattr(parents, 'training_table', None) + snapshot = getattr(getattr(table, 'parents', None), 'snapshot', None) + if snapshot is None: + snapshot = getattr(parents, 'snapshot', None) + return getattr(snapshot, 'id', None) + + +_ARGUMENTS = {'training': 'training_snapshot', 'scoring': 'snapshot'} + + +def _snapshot_stats(trainer: Any, run: Any, snapshot: Any, + training_snapshot: Any, + root: Any) -> Dict[str, Dict[str, Any]]: + r"""The build-time ``col_stats.json`` and ingest counts of the training + and scoring snapshots: those passed, else those the manifests name, from + the local snapshot store or the artifact ``root``. Never raises.""" + model = getattr(getattr(trainer, 'result', None), 'manifest', None) + sides = {} + for side, passed, manifest in ( + ('training', training_snapshot, model), + ('scoring', snapshot, getattr(run, 'manifest', None)), + ): + read = _reader(passed, _parent_snapshot_id(manifest), root) + ingest = read('ingest_stats.json') + stats = ingest.get('stats') if isinstance(ingest, dict) else None + sides[side] = { + 'col_stats': read('col_stats.json'), + 'counts': { + 'rows': dict(stats.get('source_rows_by_table', {})), + 'edges': dict(stats.get('edge_counts', {})), + } if isinstance(stats, dict) else None, + } + return sides + + +def _reader(snapshot: Any, snapshot_id: Optional[str], + root: Any) -> Callable[[str], Optional[Any]]: + index = getattr(snapshot, 'index_path', None) + if index is None and snapshot_id: + index = _local_index(snapshot_id) + if index is not None: + def local(name: str) -> Optional[Any]: + try: + return json.loads((Path(index) / name).read_text()) + except (OSError, ValueError): + return None + return local + if root is None or not snapshot_id: + return lambda name: None + from kumo.engine import _snapshot_io + + try: + path = root.location('snapshots/' + snapshot_id.replace('@', '/')) + generation, _ = _snapshot_io.read(root.connector, path) + filesystem, _ = _snapshot_io._location(root.connector, path) + except Exception: + return lambda name: None + + def stored(name: str) -> Optional[Any]: + try: + return _snapshot_io._read_json(filesystem, + f'{generation}/index/{name}') + except Exception: + return None + return stored + + +def _missing(what: str, sides: Mapping[str, Mapping[str, Any]], + name: str) -> str: + missing = [side for side, stats in sides.items() + if not isinstance(stats[name], dict)] + passes = ', '.join( + f'{_ARGUMENTS[side]}=GraphSnapshot.load_from_connector(store, ' + f'path=...)' for side in missing) + return (f"no {what} for the {' and '.join(missing)} snapshot: pass " + f"{passes}, or turn MLOps on with the artifact root that stores " + f"{'them' if len(missing) > 1 else 'it'}") + + +def _local_index(snapshot_id: str) -> Optional[Path]: + from kumo.engine._scratch import get_scratch_dir + + graph_key, _, version = snapshot_id.partition('@') + if not graph_key or not version: + return None + root = get_scratch_dir() / 'snapshots' / graph_key + found = sorted(path for path in root.glob(f'{version}*') + if (path / 'col_stats.json').is_file()) + return found[0] if len(found) == 1 else None + + +def _used_columns(result: Any) -> Optional[Dict[str, set]]: + from kumo_core.encoders.config import EncodersConfig + from kumo_core.io.json import load_json + + path = (getattr(result, 'artifact_paths', None) or {}).get('encoders') + if not path or not Path(path).is_file(): + return None + config = load_json(path, EncodersConfig) + if config is None: + return None + return {table: set(columns) for table, columns in config.encoders.items()} + + +def _column_change(old: Mapping[str, Any], + new: Mapping[str, Any]) -> Optional[Dict[str, Any]]: + kind = new.get('semantic_type') + if old.get('semantic_type') != kind: + return {'type': kind, 'type_changed': [old.get('semantic_type'), kind]} + if kind == 'numerical': + change = _deltas(old, new, _NUMERIC_STATS) + std = old.get('std') + shift = change['delta'].get('mean') + change['standardized_mean_shift'] = ( + shift / std if shift is not None and std else None) + elif kind in ('categorical', 'multicategorical'): + change = _deltas(old, new, _CATEGORY_STATS) + p, q = _distribution(old), _distribution(new) + change['psi'] = _category_psi(p, q) + change['js'] = _js(p, q) + change['category_counts_capped'] = bool( + old.get('category_counts_capped') + or new.get('category_counts_capped')) + elif kind == 'timestamp': + change = _deltas(old, new, _TIME_STATS) + else: + return None + return {'type': kind, **change} + + +def _stat(stats: Mapping[str, Any], name: str) -> Optional[float]: + if name == 'null_rate': + count = stats.get('count') + return stats.get('null_count', 0) / count if count else None + if name == 'num_unique': + value = stats.get('num_unique', stats.get('num_unique_multi')) + else: + value = stats.get(name) + return float(value) if isinstance(value, (int, float)) and not isinstance( + value, bool) and math.isfinite(value) else None + + +def _deltas(old: Mapping[str, Any], new: Mapping[str, Any], + names: Tuple[str, ...]) -> Dict[str, Any]: + out: Dict[str, Any] = {'training': {}, 'scoring': {}, 'delta': {}, + 'relative': {}} + for name in names: + before, after = _stat(old, name), _stat(new, name) + out['training'][name], out['scoring'][name] = before, after + if before is None or after is None: + continue + out['delta'][name] = after - before + if before: + out['relative'][name] = (after - before) / abs(before) + return out + + +def _distribution(stats: Mapping[str, Any]) -> Dict[str, float]: + counts: Dict[str, float] = {} + for value, count in stats.get('category_counts') or (): + counts[str(value)] = counts.get(str(value), 0.0) + float(count) + total = sum(counts.values()) + return {key: count / total for key, count in counts.items()} if total \ + else {} + + +def _smoothed(p: Mapping[str, float], q: Mapping[str, float] + ) -> Tuple[np.ndarray, np.ndarray]: + keys = sorted(set(p) | set(q)) + a = np.array([p.get(k, 0.0) for k in keys]) + _SMOOTHING + b = np.array([q.get(k, 0.0) for k in keys]) + _SMOOTHING + return a / a.sum(), b / b.sum() + + +def _category_psi(p: Mapping[str, float], + q: Mapping[str, float]) -> Optional[float]: + if not p or not q: + return None + a, b = _smoothed(p, q) + return float(np.sum((b - a) * np.log(b / a))) + + +def _js(p: Mapping[str, float], q: Mapping[str, float]) -> Optional[float]: + if not p or not q: + return None + a, b = _smoothed(p, q) + m = (a + b) / 2 + return float((np.sum(a * np.log2(a / m)) + np.sum(b * np.log2(b / m))) / 2) + + +def _metric_name(*parts: str) -> str: + import re + return '.'.join(re.sub(r'[^A-Za-z0-9_-]', '_', part) for part in parts) + + +def feature_metrics(features: Mapping[str, Mapping[str, Mapping[str, Any]]] + ) -> Dict[str, float]: + r"""A bounded summary of ``features`` for a monitoring run: the + `FEATURE_METRICS_TOP` columns by absolute standardized mean shift and by + PSI, and the largest of each per table (for the same number of + tables).""" + metrics: Dict[str, float] = {} + for key, name in (('standardized_mean_shift', 'shift'), ('psi', 'psi')): + ranked = sorted( + ((abs(change[key]), table, column) + for table, columns in features.items() + for column, change in columns.items() + if isinstance(change.get(key), float) + and math.isfinite(change[key])), + reverse=True) + for value, table, column in ranked[:FEATURE_METRICS_TOP]: + metrics[_metric_name('features', name, table, column)] = value + tables: Dict[str, float] = {} + for value, table, _ in ranked: + tables.setdefault(table, value) + for table, value in sorted(tables.items(), key=lambda item: -item[1] + )[:FEATURE_METRICS_TOP]: + metrics[_metric_name('features', f'max_{name}', table)] = value + return metrics diff --git a/packages/kumo-engine/kumo/engine/_score_stats.py b/packages/kumo-engine/kumo/engine/_score_stats.py new file mode 100644 index 000000000..5f83ca914 --- /dev/null +++ b/packages/kumo-engine/kumo/engine/_score_stats.py @@ -0,0 +1,110 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +r"""Score statistics of prediction Parquet files, computed by DiskGraph.""" + +import json +from pathlib import Path +from typing import Any, Dict, Optional, Sequence, Union + +from kumo_core.types import TaskType + +TOP_SCORE = 'top_score' +TOP_CLASSES = 100 +#: The column each task type's scores are summarised by. +SCORE_COLUMNS = { + TaskType.BINARY_CLASSIFICATION: 'True_PROB', + TaskType.REGRESSION: 'TARGET_PRED', + TaskType.FORECASTING: 'TARGET_PRED', + TaskType.MULTICLASS_CLASSIFICATION: TOP_SCORE, + TaskType.MULTILABEL_CLASSIFICATION: TOP_SCORE, + TaskType.MULTILABEL_RANKING: TOP_SCORE, + TaskType.STATIC_LINK_PREDICTION: TOP_SCORE, + TaskType.TEMPORAL_LINK_PREDICTION: TOP_SCORE, +} +_CLASS_COLUMNS = {'CLASS', 'SCORE', 'LABEL', 'PREDICTED'} + +Paths = Sequence[Union[str, Path]] + + +def parquet_stats(paths: Paths, columns: Sequence[str], *, + max_per: Optional[Sequence[str]] = None + ) -> Dict[str, Dict[str, Any]]: + r"""Summaries of ``columns``: percentiles of numbers, counts of labels.""" + import diskgraph + from pyarrow import parquet + + schema = parquet.read_schema(paths[0]) + specs = [_spec(name, schema.field(name).type) for name in columns] + full, _ = diskgraph.compute_col_stats_for_parquet( + [str(path) for path in paths], json.dumps(specs), '', + percentiles=True, summary_stats_by_class=False, histograms=False, + category_counts_limit=TOP_CLASSES, + max_per=list(max_per or [])) + stats = json.loads(full) + return { + spec['name']: _summary(stats, spec['name'], spec['stype']) + for spec in specs + } + + +def top_score_stats(paths: Paths) -> Dict[str, Any]: + r"""Summary of each prediction's highest ``SCORE``. + + A prediction is a run of consecutive rows with equal non-class columns + (entity, and anchor when present), in any score order. + """ + from pyarrow import parquet + + keys = [ + name for name in parquet.read_schema(paths[0]).names + if name not in _CLASS_COLUMNS and not name.startswith('__') + ] + return parquet_stats(paths, ['SCORE'], max_per=keys)['SCORE'] + + +def positive_rate(summary: Dict[str, Any]) -> Optional[float]: + r"""Share of true/1 values in a label column's summary.""" + counts = summary.get('top') or [] + total = sum(count for _, count in counts) + positive = sum(count for value, count in counts if value in (True, 1)) + return positive / total if total else None + + +def _spec(name: str, dtype: Any) -> Dict[str, str]: + import pyarrow as pa + + if pa.types.is_floating(dtype): + return {'name': name, 'stype': 'numerical', + 'dtype': 'float32' if pa.types.is_float32(dtype) else + 'float64'} + if pa.types.is_boolean(dtype): + return {'name': name, 'stype': 'categorical', 'dtype': 'bool'} + if pa.types.is_integer(dtype): + return {'name': name, 'stype': 'categorical', 'dtype': 'int64'} + return {'name': name, 'stype': 'categorical', 'dtype': 'string'} + + +def _summary(stats: Dict[str, Any], name: str, + stype: str) -> Dict[str, Any]: + def value(stat: str, approx: bool = False) -> Any: + entry = stats.get(f'{name}__kumo.{stype}__{stat}__{approx}') + return None if entry is None else entry['value'] + + if stype == 'numerical': + rows = value('COUNT') or 0 + non_finite = round((value('INVALID_FRACTION') or 0) * rows) + count = rows - (value('NUM_NA') or 0) - non_finite + return { + 'count': count, + 'mean': value('MEAN') if count else None, + 'min': value('MIN') if count else None, + 'max': value('MAX') if count else None, + 'percentiles': value('PERCENTILES', True) or [], + } + return { + 'count': (value('COUNT') or 0) - (value('NUM_NA') or 0), + 'distinct': value('NUM_UNIQUE'), + 'distinct_estimated': bool(value('CATEGORY_COUNTS_CAPPED')), + 'top': value('CATEGORY_COUNTS') or [], + } diff --git a/packages/kumo-engine/kumo/engine/_tracking.py b/packages/kumo-engine/kumo/engine/_tracking.py index 159c9a83a..dc17a416d 100644 --- a/packages/kumo-engine/kumo/engine/_tracking.py +++ b/packages/kumo-engine/kumo/engine/_tracking.py @@ -27,6 +27,7 @@ Tuple, ) +from kumo_core.artifacts import PROVENANCE_FILE from kumo_core.instrumentation import format_error from kumo.engine._events import Event, _finite_floats @@ -35,7 +36,11 @@ if TYPE_CHECKING: from kumoapi.plugin import ModelRegistryLocation, PublishedModel - from kumo.engine._manifest import ModelManifest, TrackingRef + from kumo.engine._manifest import ( + ModelManifest, + PredictionManifest, + TrackingRef, + ) logger = logging.getLogger(__name__) @@ -91,14 +96,17 @@ class RunSpec: def fit_inputs(manifest: Any) -> Dict[str, Dict[str, Optional[str]]]: - r"""Source tables of a training table or model manifest, read from its - embedded parents; stub parents add none.""" + r"""Source tables of a training table or model manifest, or a prediction + run's scoring snapshot, read from its embedded parents; stub parents add + none.""" try: - table = manifest - if getattr(manifest, 'kind', None) == 'model': + table, kind = manifest, getattr(manifest, 'kind', None) + if kind == 'model': table = getattr(_parents(manifest), 'training_table', None) - parents = _parents(table) - snapshot = _parents(getattr(parents, 'snapshot', None)) + parents = None if kind == 'prediction' else _parents(table) + snapshot = _parents( + getattr(_parents(manifest) if kind == 'prediction' else parents, + 'snapshot', None)) tables = dict(getattr(snapshot, 'source_tables', None) or {}) source = getattr(parents, 'source_table', None) if source is not None: @@ -188,6 +196,89 @@ def trial_spec(events: Sequence[Event]) -> RunSpec: ) +def prediction_record(manifest: 'PredictionManifest', + error: Optional[BaseException] = None) -> RunRecord: + r"""A prediction run's record, from its manifest: scoring options as + params, row counts and score percentiles as metrics, the scoring + snapshot's source tables as inputs, pointer tags (``kumo.error`` when + ``error`` ended the run), and the manifest as an artifact.""" + table = manifest.prediction_table + params = {**manifest.scoring.to_dict(), 'anchors': table.anchors} + metrics: Dict[str, Any] = {'prediction_table.rows': table.rows} + if table.skipped_rows is not None: + metrics['predictions.skipped_rows'] = table.skipped_rows + metrics['predictions.skipped_fraction'] = ( + table.skipped_rows / table.requested_rows + if table.requested_rows else 0.0) + for artifact, stats in (manifest.stats or {}).items(): + flat = stats.to_dict() + flat.update(flat.pop('columns', None) or {}) + for key, value in _flatten(flat, f'{artifact}.'): + if key.endswith('.percentiles'): + metrics.update(_logged_percentiles(key, value)) + else: + metrics[key] = value + return run_record( + RunSpec( + kind='prediction', + name=manifest.id, + params={ + key: value + for key, value in params.items() if value is not None + }, + metrics={ + key: value + for key, value in metrics.items() if isinstance( + value, (int, float)) and not isinstance(value, bool) + }, + parents={ + 'bundle_id': manifest.parents.model.id, + 'prediction_id': manifest.id, + }, + artifacts={ + PROVENANCE_FILE: + json.dumps(manifest.to_dict(), indent=2, sort_keys=True, + default=str) + }, + inputs=fit_inputs(manifest), + error=None if error is None else format_error(error), + )) + + +def record_failure(record: RunRecord) -> Optional[Exception]: + r"""The error a run record says ended its run (its ``kumo.error``).""" + error = record.tags.get('kumo.error') + return None if error is None else RuntimeError(error) + + +def track_prediction( + manifest: Optional['PredictionManifest'], tracking: Any, + error: Optional[BaseException] = None) -> Optional[Exception]: + r"""Record a prediction run in ``tracking``, failed if ``error`` ended + it; returns the tracking failure instead of raising it.""" + if getattr(tracking, 'log_prediction_run', None) is None: + return None + try: + if manifest is None: + raise ValueError('the run has no prediction manifest') + record = prediction_record(manifest, error) + except Exception as failure: + logger.warning('Experiment tracking failed: %s; the prediction run ' + 'is not tracked', full_error(failure), exc_info=_debug()) + return failure + return log_run(tracking, 'log_prediction_run', record) + + +def _logged_percentiles(key: str, values: Any) -> Dict[str, Any]: + if not isinstance(values, list) or len(values) != 99: + return {} + column = key[:-len('percentiles')] + return { + f'{column}q{q:02d}': values[q - 1] + for q in (1, 5, 25, 50, 75, 95, 99) + } + + class Tracking(Protocol): r"""Records one `Trainer.fit` in an experiment tracker. @@ -204,7 +295,9 @@ class Tracking(Protocol): ``ModelRegistryLocation``; ``connector``, a connector whose model registry ``registry`` registers in when there is no ``log_model``; and ``platform``, ``run_id`` and ``run_url``, the tracker's own run once - `start` returned, recorded as the model manifest's ``tracking``. + `start` returned, recorded as the model manifest's ``tracking``; and + the best-effort hooks ``log_prediction_run(record)`` and + ``log_monitoring(record)``. """ def start(self, record: RunRecord) -> None: ... @@ -440,6 +533,42 @@ def _fail(self, error: BaseException) -> None: self.error = error +def log_run(tracking: Any, hook: str, + record: RunRecord) -> Optional[Exception]: + r"""Log a standalone run with ``tracking.(record)`` on a tracking + thread, waiting up to `FINISH_SECONDS`; returns the tracking failure + instead of raising it.""" + method = getattr(tracking, hook, None) + if method is None or _rank() != 0: + return None + failed: List[Exception] = [] + + def log() -> None: + try: + method(record) + except Exception as error: + logger.warning('Experiment tracking failed: %s; the %s run is not ' + 'tracked', full_error(error), + record.tags.get('kumo.kind'), exc_info=_debug()) + failed.append(error) + + try: + getattr(tracking, 'prepare', _none)() + except Exception as error: + logger.warning('Experiment tracking failed: %s; the %s run is not ' + 'tracked', full_error(error), + record.tags.get('kumo.kind'), exc_info=_debug()) + return error + worker = threading.Thread(target=log, name='kumo-tracking', daemon=True) + worker.start() + worker.join(FINISH_SECONDS) + if worker.is_alive(): + logger.warning( + 'Experiment tracking did not finish within %ss; it goes on in ' + 'the background', FINISH_SECONDS) + return failed[0] if failed else None + + def _debug() -> bool: r"""Tracebacks of tracking failures only at DEBUG: the failure itself is reported once, by the call that tracked it.""" diff --git a/packages/kumo-engine/kumo/engine/connector/file.py b/packages/kumo-engine/kumo/engine/connector/file.py index be3839d2b..e5b161e41 100644 --- a/packages/kumo-engine/kumo/engine/connector/file.py +++ b/packages/kumo-engine/kumo/engine/connector/file.py @@ -11,12 +11,12 @@ from pathlib import Path from typing import Literal, Union -from kumo_core.artifact_export import ArtifactDestinationConflict from kumo_core.object_store_sink import ObjectStoreSinkBackend from kumo_core.sources import layout from kumo_core.sources.object_store import ObjectStoreSourceReader from kumo_core.sources.object_store_backend import ObjectStoreSourceBackend from kumo_core.types.enums import ConnectorType +from kumoapi.jobs import WriteMode from kumoapi.plugin import SourceConnector from kumoapi.source_table import FileType, LocalFileSourceTable from pyarrow.fs import LocalFileSystem @@ -168,7 +168,6 @@ def _destination(self, connector: Union[Connector, SourceConnector], else '') def preflight_write_files(self, connector, table_name, write_mode=None): - del write_mode _, destination = self._destination(connector, table_name) path = Path(destination) kind = layout.file_type(destination.rstrip('/')) or layout.PARQUET @@ -186,12 +185,9 @@ def preflight_write_files(self, connector, table_name, write_mode=None): try: if target is None and path.is_file(): raise _not_a_directory(destination) - if target is None and path.is_dir() and any(path.iterdir()): - # Each prediction run exports under a new identity, so - # publication never accepts an occupied directory. - raise ArtifactDestinationConflict( - f"Output directory '{destination}' already contains data " - 'from another export. Choose an empty destination path.') + if (target is None and write_mode is not WriteMode.OVERWRITE + and path.is_dir() and any(path.iterdir())): + raise self._conflict(destination) for root in roots: _probe_directory(root, destination) if target is not None and target.is_file(): diff --git a/packages/kumo-engine/kumo/engine/graph.py b/packages/kumo-engine/kumo/engine/graph.py index 2bc690eb3..df8bc8063 100644 --- a/packages/kumo-engine/kumo/engine/graph.py +++ b/packages/kumo-engine/kumo/engine/graph.py @@ -243,7 +243,8 @@ def schema_changes(before: Mapping[str, Any], ] for name in sorted(old.keys() & new.keys()): parts = sorted( - _SCHEMA_PARTS.get(key, key) + _SCHEMA_PARTS.get(key, key) + (_changed_columns( + old[name][key], new[name][key]) if key == 'cols' else '') for key in old[name].keys() | new[name].keys() if old[name].get(key) != new[name].get(key)) if parts: @@ -251,8 +252,21 @@ def schema_changes(before: Mapping[str, Any], return changes +def _changed_columns(before: List[str], after: List[str]) -> str: + old, new = ({ + column['name']: column + for column in map(json.loads, side or []) + } for side in (before, after)) + names = sorted(name for name in old.keys() | new.keys() + if old.get(name) != new.get(name)) + return f" ({', '.join(names)})" if names else '' + + def report_other_graph(message: str, before: Optional[Mapping[str, Any]], - after: Mapping[str, Any], stacklevel: int) -> None: + after: Mapping[str, Any], stacklevel: int, + strict: bool = False) -> None: + r"""Warn that a snapshot is of another graph, naming the schema + changes; ``strict`` (``on_error='raise'``) raises on a schema change.""" if before is None: warnings.warn( f"{message} The model does not record its training graph, so " @@ -261,6 +275,12 @@ def report_other_graph(message: str, before: Optional[Mapping[str, Any]], stacklevel=stacklevel + 1) return changes = schema_changes(before, after) + if changes and strict: + raise ValueError( + f"{message} Schema changes: {'; '.join(changes)}. With " + f"on_error='raise' a model scores only the graph it was trained " + f"on: declare the same tables, column types (Table.with_column(" + f"name, stype=...)) and edges.") if changes: warnings.warn(f"{message} Schema changes: {'; '.join(changes)}.", UserWarning, stacklevel=stacklevel + 1) diff --git a/packages/kumo-engine/kumo/engine/integrations/_mlflow_model.py b/packages/kumo-engine/kumo/engine/integrations/_mlflow_model.py index a1be11e94..bc85aa556 100644 --- a/packages/kumo-engine/kumo/engine/integrations/_mlflow_model.py +++ b/packages/kumo-engine/kumo/engine/integrations/_mlflow_model.py @@ -60,6 +60,7 @@ def _snapshot(self, uri: 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 @@ -89,7 +90,7 @@ def predict(trainer: Any, snapshot: GraphSnapshot, rows: pd.DataFrame, scored = trainer._predict( trainer._prediction_table(snapshot, entities, None), snapshot=snapshot, extra_fields=[_ANCHOR] if timed else [], - **options).frame + settings=resolve(), **options).frame return _in_input_order(rows, scored) diff --git a/packages/kumo-engine/kumo/engine/integrations/mlflow.py b/packages/kumo-engine/kumo/engine/integrations/mlflow.py index 96c616790..58d7a769f 100644 --- a/packages/kumo-engine/kumo/engine/integrations/mlflow.py +++ b/packages/kumo-engine/kumo/engine/integrations/mlflow.py @@ -1,8 +1,8 @@ # SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. # All rights reserved. # SPDX-License-Identifier: Apache-2.0 -r"""MLflow tracking and models for `Trainer.fit` -(``pip install 'kumo-engine[mlflow]'``).""" +r"""MLflow tracking and models for `Trainer.fit`, `Trainer.predict` and +monitoring (``pip install 'kumo-engine[mlflow]'``).""" import hashlib import json @@ -37,7 +37,12 @@ from kumo.engine._events import Event from kumo.engine._mlops import full_error -from kumo.engine._tracking import DROPPED_PARAMS, RegistrationError, RunRecord +from kumo.engine._tracking import ( + DROPPED_PARAMS, + RegistrationError, + RunRecord, + record_failure, +) from kumo.engine._version import __version__ from kumo.engine.connector._extras import import_provider from kumo.engine.integrations._mlflow_target import ( @@ -81,7 +86,9 @@ class MlflowTracking: tables (Unity Catalog then links them to a model registered from the run), file and object-store sources by URI. Params MLflow rejects are counted in ``kumo.dropped_params``. Concurrent fits in one process are - not supported. + not supported. As ``kumo.MLOps(tracking=...)``, it also records each + prediction run and each monitoring call as one run named after it, + placed like a fit run; ``run_name`` and ``tags`` apply to fit runs only. Args: experiment_name: experiment for the fit run, created if missing. @@ -112,8 +119,9 @@ class MlflowTracking: process-wide tracking URI for each call: Unity Catalog downloads and the SageMaker plugin route with it. - ``run_id`` is the MLflow id of the latest fit run, once it has started, - and ``run_url`` its link in the MLflow UI, where the server has one. + ``run_id`` is the MLflow id of the most recent tracked run, once it has + started, and ``run_url`` its link in the MLflow UI, where the server has + one. """ def __init__( @@ -319,6 +327,24 @@ def start(self, record: RunRecord) -> None: with self.target.scope(): _log_inputs(self._mlflow, client, self._parent, record.inputs) + def log_prediction_run(self, record: RunRecord) -> None: + self.start(record) + self.finish(record, record_failure(record)) + + def log_monitoring(self, record: RunRecord) -> None: + self.start(record) + experiment_id, run_id = self._experiment_id, self.run_id + client = self._client + self.finish(record, None) + with self.target.scope(): + for run in client.search_runs( + [experiment_id], filter_string=( + "tags.`kumo.kind` = 'prediction' AND " + "tags.`kumo.prediction_id` = " + f"'{record.tags['kumo.monitors']}'"), + max_results=1): + client.set_tag(run.info.run_id, 'kumo.monitored', run_id) + def on_trial(self, record: RunRecord, event: Event) -> None: key = (event.trial or 0, int(event.params['attempt'])) if event.name == 'trial_started': diff --git a/packages/kumo-engine/kumo/engine/integrations/snowflake.py b/packages/kumo-engine/kumo/engine/integrations/snowflake.py index 78b3e23b6..080df1d3c 100644 --- a/packages/kumo-engine/kumo/engine/integrations/snowflake.py +++ b/packages/kumo-engine/kumo/engine/integrations/snowflake.py @@ -1,8 +1,8 @@ # SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. # All rights reserved. # SPDX-License-Identifier: Apache-2.0 -r"""Snowflake Experiment Tracking for `Trainer.fit` -(``pip install 'kumo-engine[snowflake]'``).""" +r"""Snowflake Experiment Tracking for `Trainer.fit`, `Trainer.predict` and +monitoring (``pip install 'kumo-engine[snowflake]'``).""" import logging import re @@ -16,7 +16,12 @@ from kumo_core.sources.sql import resolve_snowflake_identifier from kumo.engine._events import Event -from kumo.engine._tracking import DROPPED_METRICS, DROPPED_PARAMS, RunRecord +from kumo.engine._tracking import ( + DROPPED_METRICS, + DROPPED_PARAMS, + RunRecord, + record_failure, +) from kumo.engine.connector._extras import import_provider logger = logging.getLogger(__name__) @@ -34,7 +39,9 @@ class SnowflakeTracking: trial attempts are logged in it under ``trial--attempt-/`` names and tags are logged as params. Snowflake has no dataset inputs, so a fit's source tables are not recorded as inputs. The Snowpark session's - parameters are not changed. + parameters are not changed. As ``kumo.MLOps(tracking=...)``, it also + records each prediction run and each monitoring call as one run named + after it; ``run_name`` and ``tags`` apply to fit runs only. Args: source: a ``SnowflakeConnector``, whose identity opens the session @@ -52,8 +59,8 @@ class SnowflakeTracking: native_model_lineage: keep the run open as a context for the whole fit, so models registered meanwhile are linked to it. - ``run_id`` is the latest fit run's name (Snowflake runs are identified - by name), once it has started. + ``run_id`` is the most recent tracked run's name (Snowflake runs are + identified by name), once it has started. """ platform = 'snowflake' @@ -144,7 +151,8 @@ def start(self, record: RunRecord) -> None: elif (not hasattr(tracking, '_run') or tracking._run is not None or getattr(tracking, '_session', None) is not self.session): logger.warning('Snowflake ExperimentTracking is in use in this ' - 'process; this fit is not tracked in Snowflake') + 'process; this %s run is not tracked in Snowflake', + record.tags.get('kumo.kind')) return self._tracking = tracking database = self.database or self.session.get_current_database() @@ -154,7 +162,9 @@ def start(self, record: RunRecord) -> None: database_name=database, schema_name=schema or 'PUBLIC') self._run = tracking.start_run( - self._new_run_name(tracking, experiment.name)) + self._new_run_name( + tracking, experiment.name, + _run_identifier(record.name) if record.name else self.run_name)) self.run_id = str(self._run.name) if self.native_model_lineage: self._run_context.enter_context(self._run) @@ -162,6 +172,14 @@ def start(self, record: RunRecord) -> None: self._log_params({**record.params, **record.tags}) self._log_artifacts(record.artifacts) + def log_prediction_run(self, record: RunRecord) -> None: + self.start(record) + self.finish(record, record_failure(record)) + + def log_monitoring(self, record: RunRecord) -> None: + self.start(record) + self.finish(record, None) + def on_trial(self, record: RunRecord, event: Event) -> None: if self._run is None: return @@ -209,9 +227,9 @@ def finish(self, record: Optional[RunRecord], self._run_context.close() self._reset() - def _new_run_name(self, tracking: Any, - experiment: Any) -> Optional[str]: - if self.run_name is None: + def _new_run_name(self, tracking: Any, experiment: Any, + run_name: Optional[str]) -> Optional[str]: + if run_name is None: return None # start_run(name) resumes, or raises on, any run matching name as a # SQL LIKE pattern and cannot create an existing run, so pick a name @@ -222,16 +240,16 @@ def _new_run_name(self, tracking: Any, experiment_name=experiment) ] short = uuid.uuid4().hex - for name in (self.run_name, _suffixed(self.run_name, short[:8]), - _suffixed(self.run_name, short)): + for name in (run_name, _suffixed(run_name, short[:8]), + _suffixed(run_name, short)): if not _taken(name, existing): break else: name = None - if name != self.run_name: + if name != run_name: logger.warning( 'Snowflake run %s already exists in experiment %s; recording ' - 'this fit as %s', self.run_name, self.experiment_name, + 'this run as %s', run_name, self.experiment_name, name or 'a generated run name') return name @@ -283,6 +301,13 @@ def _log_metrics(self, metrics: Mapping[str, float], limit: int, self._tracking.log_metrics(kept, step=step) +def _run_identifier(name: str) -> str: + r"""``name`` as an unquoted identifier: snowflake-ml puts a run's name + into its artifact URIs, where a quoted name is not found.""" + name = re.sub(r'[^A-Za-z0-9_]', '_', name) + return name if re.match(r'[A-Za-z_]', name) else f'_{name}' + + def _suffixed(name: str, suffix: str) -> str: if len(name) > 1 and name[0] == name[-1] == '"': return f'{name[:-1]}_{suffix}"' diff --git a/packages/kumo-engine/kumo/engine/pquery/prediction_table.py b/packages/kumo-engine/kumo/engine/pquery/prediction_table.py index f64912109..1c31ee593 100644 --- a/packages/kumo-engine/kumo/engine/pquery/prediction_table.py +++ b/packages/kumo-engine/kumo/engine/pquery/prediction_table.py @@ -3,10 +3,11 @@ # SPDX-License-Identifier: Apache-2.0 r"""Materialized prediction tables and their shared compile path.""" +import warnings from dataclasses import dataclass, replace from datetime import datetime from pathlib import Path -from typing import Any, Mapping, Optional, Tuple +from typing import Any, Dict, List, Literal, Mapping, Optional, Tuple from kumo_core.model_plan import \ TrainingTableGenerationPlan as CoreGenerationPlan @@ -70,6 +71,18 @@ class PredictionTable: compiler_options: Optional[Mapping[str, Any]] = None #: Effective temporal step used by the fitted forecast task. forecast_timeframe_step: Optional[int] = None + #: ``'pql'`` when generated from the query, ``'custom'`` from a source. + source: Literal['pql', 'custom'] = 'pql' + #: Supplied rows DiskGraph skipped (0 for a generated population), the + #: first `SKIPPED_SAMPLE` of their entity ids, and why. + skipped_rows: int = 0 + skipped_sample: Tuple[Any, ...] = () + skip_reason: Optional[str] = None + + @property + def requested_rows(self) -> int: + r"""Rows asked for: the scored rows plus the skipped ones.""" + return self.row_count + self.skipped_rows @classmethod def from_source_table( @@ -125,6 +138,7 @@ def from_source_table( request, PreparedCustomPredictionTable(str(rows_path)), ) + skipped = skipped_entities(generated.row_failures, rows_path) return cls( query=pq.query, table_path=Path(generated.table_path), @@ -135,6 +149,8 @@ def from_source_table( required_stages=request.task_dataset.required_stages, compiler_options=request.compiler_options, forecast_timeframe_step=request.forecast_timeframe_step, + source='custom', + **skipped, ) @@ -235,3 +251,70 @@ def _compile_prediction_table( ) generated = generate_prediction_table(request) return request, generated, graph_index_id + + +#: Skipped entity ids kept in a prediction record, at most; DiskGraph names +#: up to 5 rows per reason. +SKIPPED_SAMPLE = 20 +#: Why DiskGraph skips a supplied prediction row, by its condition name. +_SKIP_REASONS = { + 'unknown_entity': 'not in snapshot', + 'non_canonical_entity': 'not a canonical node in the snapshot', + 'null_entity': 'null entity id', + 'null_timestamp': 'null anchor time', +} + + +def skipped_entities(report: Optional[Mapping[str, Any]], + rows_path: Path) -> Dict[str, Any]: + r"""The skipped-row fields of a `PredictionTable` from DiskGraph's + ``row_failures`` report, warning once when rows were skipped. Sample ids + are read from the rows DiskGraph names (a few per condition).""" + skipped = int((report or {}).get('skipped_rows') or 0) + if not skipped: + return {} + conditions = report.get('by_condition') or [] + reasons = [(_SKIP_REASONS.get(entry.get('condition'), + str(entry.get('condition')).replace( + '_', ' ')), int(entry.get('count') or 0)) + for entry in conditions] + reason = (reasons[0][0] if len(reasons) == 1 else '; '.join( + f'{name} ({count:,})' for name, count in reasons)) or None + requested = int(report.get('rows_read') or 0) + warnings.warn( + f'Prediction skipped {skipped:,} of {requested:,} requested entities ' + f'({skipped / requested if requested else 1:.1%}): {reason}', + UserWarning, stacklevel=3) + return { + 'skipped_rows': skipped, + 'skipped_sample': _sample_ids(conditions, rows_path), + 'skip_reason': reason, + } + + +def _sample_ids(conditions: Any, rows_path: Path) -> Tuple[Any, ...]: + import re + + import pyarrow.parquet as pq + + ids: List[Any] = [] + try: + for entry in conditions: + for location in entry.get('samples') or (): + found = re.fullmatch(r'(.+) row group (\d+) row (\d+)', + str(location)) + if found is None or len(ids) >= SKIPPED_SAMPLE: + continue + part = Path(rows_path) + part = part / found.group(1) if part.is_dir() else part + reader = pq.ParquetFile(part) + group, row = int(found.group(2)), int(found.group(3)) + start = sum(reader.metadata.row_group(index).num_rows + for index in range(group)) + column = reader.read_row_group( + group, columns=[entry['column']]).column(0) + ids.append(column[row - start if row >= start else row] + .as_py()) + except Exception: + pass + return tuple(ids) diff --git a/packages/kumo-engine/kumo/engine/pquery/predictive_query.py b/packages/kumo-engine/kumo/engine/pquery/predictive_query.py index dab929730..5693daa38 100644 --- a/packages/kumo-engine/kumo/engine/pquery/predictive_query.py +++ b/packages/kumo-engine/kumo/engine/pquery/predictive_query.py @@ -43,12 +43,13 @@ """ import logging +from contextlib import contextmanager from copy import deepcopy from dataclasses import dataclass, replace from datetime import datetime from functools import cached_property from pathlib import Path -from typing import Any, Optional +from typing import Any, Iterator, Optional from kumo_core.model_plan import ModelPlan as CoreModelPlan from kumo_core.model_plan import ModelPlanValueError @@ -291,25 +292,10 @@ def generate_training_table( if errors: raise ModelPlanValueError(errors) - run_id = new_run_id() - run_dir = train_table_scratch_dir(run_id) + run_dir = train_table_scratch_dir(new_run_id()) run_dir.mkdir(parents=True, exist_ok=True) - - workflow = TrainingWorkflowRequest.from_spec( - graph=snapshot.graph_config, - predictive_query=self.query, - prepared_graph=PreparedGraph( - index_id=snapshot.id, - index_path=str(snapshot.index_path), - ), - run_id=run_id, - scratch_dir=str(run_dir), - artifact_dir=str(run_dir), - training_table_generation_plan=options, - ) - with no_encoder_inference_warnings(): - task_plan = build_training_task_plan(workflow) - generated = generate_train_table(task_plan.train_table) + workflow, task_plan, generated = self._generate( + snapshot, options, run_dir) generation_plan = options if workflow.query.query_type == 'temporal': @@ -342,6 +328,62 @@ def generate_training_table( ) return _finish(table, snapshot, task_plan.train_table, mlops) + def _generate(self, snapshot: GraphSnapshot, options: CoreGenerationPlan, + run_dir: Path, *, timestamps: tuple[int, ...] = (), + validate_targets: bool = True) -> tuple: + workflow = TrainingWorkflowRequest.from_spec( + graph=snapshot.graph_config, + predictive_query=self.query, + prepared_graph=PreparedGraph( + index_id=snapshot.id, + index_path=str(snapshot.index_path), + ), + run_id=run_dir.name, + scratch_dir=str(run_dir), + artifact_dir=str(run_dir), + training_table_generation_plan=options, + timestamps=timestamps, + ) + with no_encoder_inference_warnings(): + task_plan = build_training_task_plan(workflow) + return workflow, task_plan, generate_train_table( + task_plan.train_table, validate_targets=validate_targets) + + @contextmanager + def _label_parts( + self, snapshot: GraphSnapshot, timestamps: tuple[int, ...], + plan: Optional[TrainingTableGenerationPlan]) -> Iterator[list]: + r"""Generate this query's labels at ``timestamps`` on ``snapshot`` + and yield their Parquet parts, removed on exit.""" + import shutil + + from kumo_core.io.json import from_json_dict + from kumo_core.training.export import read_spec, stage_parts + + # The fit's split ranges would not cover later anchors. + options = CoreGenerationPlan() + if plan is not None: + options = replace(from_json_dict(to_json_dict(plan), + CoreGenerationPlan), + split=options.split) + run_dir = train_table_scratch_dir(new_run_id()) + run_dir.mkdir(parents=True, exist_ok=True) + staging = None + try: + _, _, generated = self._generate(snapshot, options, run_dir, + timestamps=timestamps, + validate_targets=False) + staging = Path(generated.staging_dir) + yield [ + str(staging / part) + for parts in stage_parts(read_spec(str(staging))).values() + for part in parts + ] + finally: + if staging is not None: + shutil.rmtree(staging, ignore_errors=True) + shutil.rmtree(run_dir, ignore_errors=True) + def generate_prediction_table( self, snapshot: GraphSnapshot, diff --git a/packages/kumo-engine/kumo/engine/trainer/_provenance.py b/packages/kumo-engine/kumo/engine/trainer/_provenance.py index 72c83129e..a8e5b3c98 100644 --- a/packages/kumo-engine/kumo/engine/trainer/_provenance.py +++ b/packages/kumo-engine/kumo/engine/trainer/_provenance.py @@ -3,14 +3,18 @@ # SPDX-License-Identifier: Apache-2.0 """A model bundle's ``_provenance.json``; excluded from the bundle id.""" +import warnings from dataclasses import replace from datetime import datetime, timezone from pathlib import Path from typing import List, Optional +from kumo_core.types import TaskType + from kumo.engine._manifest import ( Dependency, FitTiming, + HoldoutStats, ModelManifest, ModelParents, TrackingRef, @@ -20,6 +24,7 @@ stub, write_manifest, ) +from kumo.engine._score_stats import SCORE_COLUMNS from kumo.engine.pquery.training_table import TrainingTable from kumo.engine.trainer._spec import PROVENANCE_ROLE, listed_file @@ -31,6 +36,9 @@ def write_provenance( train_table: TrainingTable, warm_start: Optional[ModelManifest], started_at: datetime, + task_type: TaskType, + holdout_path: Optional[Path], + holdout_truncated: Optional[bool], fit_mode: Optional[str], tracking: Optional[TrackingRef] = None, dependencies: Optional[List[Dependency]] = None, @@ -39,6 +47,15 @@ def write_provenance( finished_at = datetime.now(timezone.utc) table = train_table.manifest + holdout: Optional[HoldoutStats] = None + if holdout_path is not None and task_type in SCORE_COLUMNS: + try: + holdout = holdout_baseline(holdout_path, task_type, + truncated=holdout_truncated) + except Exception as error: + warnings.warn( + f"Recording the model manifest without the holdout score " + f"statistics: {error}", UserWarning, stacklevel=3) write_manifest( bundle_dir / PROVENANCE_FILE, ModelManifest( @@ -56,6 +73,7 @@ def write_provenance( fit_mode=fit_mode, tracking=tracking, dependencies=dependencies, + holdout=holdout, **stamp(), )) add_bundle_file(bundle_dir, PROVENANCE_ROLE, PROVENANCE_FILE) @@ -112,6 +130,34 @@ def record_derivation(bundle_dir: Path, parent_dir: Path, derived_from=None if parent is None else embedded_model( parent)), dependencies=None if parent is None else parent.dependencies, + holdout=None if parent is None else parent.holdout, **stamp(), )) add_bundle_file(bundle_dir, PROVENANCE_ROLE, PROVENANCE_FILE) + + +def holdout_baseline(path: Path, task_type: TaskType, *, + truncated: Optional[bool]) -> HoldoutStats: + from kumo.engine._score_stats import ( + TOP_SCORE, + parquet_stats, + positive_rate, + top_score_stats, + ) + + if task_type is TaskType.BINARY_CLASSIFICATION: + stats = parquet_stats([path], ['True_PROB', 'TARGET']) + return HoldoutStats.from_dict({ + 'truncated': truncated, + 'rows': stats['TARGET']['count'], + 'columns': {'True_PROB': stats['True_PROB']}, + 'label_positive_rate': positive_rate(stats['TARGET']), + }) + if task_type.has_numerical_target: + columns = parquet_stats([path], ['TARGET_PRED', 'TARGET']) + return HoldoutStats(truncated=truncated, + rows=columns['TARGET_PRED']['count'], + columns=columns) + top = top_score_stats([path]) + return HoldoutStats(truncated=truncated, rows=top['count'], + columns={TOP_SCORE: top}) diff --git a/packages/kumo-engine/kumo/engine/trainer/_run_manifest.py b/packages/kumo-engine/kumo/engine/trainer/_run_manifest.py new file mode 100644 index 000000000..6a0951ae3 --- /dev/null +++ b/packages/kumo-engine/kumo/engine/trainer/_run_manifest.py @@ -0,0 +1,277 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +r"""Prediction-run manifest and output score statistics.""" + +import warnings +from dataclasses import replace +from datetime import datetime, timezone +from pathlib import Path +from typing import ( + Any, + Callable, + Dict, + List, + Mapping, + Optional, + Sequence, + Tuple, +) + +from kumo_core.types import TaskType +from kumoapi.jobs import WriteMode + +from kumo.engine._manifest import ( + ArtifactStats, + ModelManifest, + OutputDestination, + PredictionManifest, + PredictionOutput, + PredictionParents, + PredictionTableRecord, + RunScoring, + RunTimings, + SnapshotManifest, + embedded, + stamp, + stub, +) +from kumo.engine.trainer._provenance import embedded_model + +# Object stores locate a path themselves; other sinks' file paths are kept. +_OBJECT_STORES = frozenset({'LOCAL_FILE', 'S3', 'GCS'}) + + +def _now() -> str: + return datetime.now(timezone.utc).isoformat(timespec='seconds') + + +def start(*, result: Any, trainer: Any, snapshot: Any, prediction_table: Any, + task_type: TaskType, scoring: RunScoring, + started_at: str) -> Optional[PredictionManifest]: + r"""The run's manifest so far; None, with a warning, if it fails.""" + try: + return _manifest(result, trainer, snapshot, prediction_table, + task_type, scoring, started_at) + except Exception as error: + warnings.warn(f"Predicting without a prediction manifest: {error}", + UserWarning, stacklevel=3) + return None + + +def finish(result: Any, manifest: Optional[PredictionManifest], + output_config: Any, settings: Any) -> None: + r"""Export to ``output_config``, record the run (with the outputs + written, even if the export fails), store it in the artifact root and + track it; a storing or tracking failure raises last under + ``on_error='raise'``.""" + from kumo.engine._mlops import raise_on_error + from kumo.engine._tracking import track_prediction + + records = ([] if output_config is None else output_records( + output_config, result.id)) + written: List[PredictionOutput] = [] + error: Optional[BaseException] = None + untracked: Optional[BaseException] = None + unstored: Optional[BaseException] = None + try: + if output_config is not None: + result._export( + output_config, None if manifest is None else replace( + manifest, outputs=records), records, written) + except BaseException as raised: + error = raised + raise + finally: + if manifest is not None: + manifest = replace( + manifest, status='failed' if error else 'succeeded', + outputs=None if records is None else written, + timings=replace(manifest.timings, finished_at=_now())) + write(result, manifest) + if settings is not None: + if error is None and settings.wants('prediction'): + try: + store(result, manifest, settings) + except Exception as failure: + unstored = failure + untracked = track_prediction(manifest, settings.tracking, + error or unstored) + untracked = unstored or untracked + if untracked is not None: + raise_on_error(settings, untracked, + f'Trainer.predict (run {result.id}, outputs written)', + stacklevel=4) + + +def store(result: Any, manifest: Optional[PredictionManifest], + settings: Any) -> None: + r"""Copy the run's outputs and manifest to + ``/predictions///``.""" + from kumo_core.train.artifact_manifest import load_prediction_selection + + from kumo.engine.trainer.output_config import PredictionOutputConfig + + try: + selection = load_prediction_selection(result.selection, + validate_objects=False) + names = [item.artifact_type.value.lower() + for item in selection.artifacts] + base = f'predictions/{result.id}' + result._export( + PredictionOutputConfig(settings.root.connector, output_paths={ + name: settings.root.location(f'{base}/{name}/') + for name in names + }, write_mode=WriteMode.OVERWRITE), manifest) + except Exception as error: + raise RuntimeError(f"Could not store prediction run {result.id} in " + f"the artifact root: {error}") from error + + +def write(result: Any, manifest: PredictionManifest) -> None: + from kumo.engine._manifest import write_manifest + + try: + write_manifest(result._manifest_path, manifest) + except Exception as error: + warnings.warn(f"Could not write the prediction manifest: {error}", + UserWarning, stacklevel=3) + + +def _manifest(result: Any, trainer: Any, snapshot: Any, + prediction_table: Any, task_type: TaskType, scoring: RunScoring, + started_at: str) -> PredictionManifest: + from kumo_core.train.artifact_manifest import load_prediction_selection + + from kumo.engine._events import query_sha256 + + selection = load_prediction_selection(result.selection, + validate_objects=False) + model: Optional[ModelManifest] = trainer.result.manifest + scoring_snapshot = getattr(snapshot, 'manifest', None) + return _record_predictions( + result.id, { + item.artifact_type.value.lower(): + (item.num_rows, lambda kind=item.artifact_type: replace( + result, artifact_type=kind).files) + for item in selection.artifacts + }, + model=(stub('model', trainer.result.id) if model is None else model), + snapshot=(stub('snapshot', snapshot.id) + if scoring_snapshot is None else scoring_snapshot), + task_type=task_type, + table=PredictionTableRecord( + query_sha256=query_sha256(prediction_table.query), + source=prediction_table.source, + anchors=[ + datetime.fromtimestamp(ts / 1e9, timezone.utc).isoformat( + timespec='seconds') for ts in prediction_table.timestamps + ], + rows=prediction_table.row_count, + requested_rows=prediction_table.requested_rows, + scored_rows=prediction_table.row_count, + skipped_rows=prediction_table.skipped_rows, + skipped_sample=list(prediction_table.skipped_sample), + skip_reason=prediction_table.skip_reason, + ), scoring=scoring, started_at=started_at) + + +def _record_predictions( + run_id: str, + outputs: Mapping[str, Tuple[Optional[int], Callable[[], Sequence[str]]]], + *, + model: ModelManifest, + snapshot: SnapshotManifest, + task_type: TaskType, + table: PredictionTableRecord, + scoring: RunScoring, + started_at: str, +) -> PredictionManifest: + r"""The manifest of prediction run ``run_id`` of any ``model``: + ``outputs`` maps each artifact (``predictions``, ``embeddings``) to its + row count and a call listing its Parquet files, whose scores are + summarized for ``task_type`` (see the column contract in mlops.md); + ``table`` holds the requested, scored and skipped counts.""" + stats: Dict[str, Dict[str, Any]] = {} + for name, (rows, files) in outputs.items(): + stats[name] = {'rows': rows} + if name == 'predictions': + try: + stats[name].update(output_stats(files(), task_type)) + except Exception as error: + warnings.warn( + f"Recording the prediction manifest without output " + f"score statistics: {error}", UserWarning, stacklevel=5) + return PredictionManifest( + id=run_id, + parents=PredictionParents( + model=(model if model.reconstructed else embedded_model(model)), + snapshot=(snapshot if snapshot.reconstructed else + embedded(snapshot)), + ), + prediction_table=table, + scoring=scoring, + outputs=[], + status=None, + timings=RunTimings(started_at=started_at, finished_at=None), + stats={ + name: ArtifactStats.from_dict(record) + for name, record in stats.items() + }, + **stamp(), + ) + + +def output_records(output_config: Any, + run_id: Optional[str]) -> Optional[List[PredictionOutput]]: + r"""Records of the destinations run ``run_id`` writes to; None, with a + warning, if one cannot be described.""" + try: + return [ + PredictionOutput( + artifact=name, destination=destination_record( + output_config.connector, destination)) + for name, destination in output_config._targets(run_id).items() + ] + except Exception as error: + warnings.warn(f"Exporting without a prediction manifest: {error}", + UserWarning, stacklevel=4) + return None + + +def destination_record(connector: Any, destination: str) -> OutputDestination: + from kumo.engine._provenance import source_location + from kumo.engine.connector import is_file_destination + + kind = str(getattr(connector.source_type, 'value', connector.source_type)) + root = getattr(getattr(connector, 'config', None), 'root_dir', None) + if root is not None: + location = (f"{root.rstrip('/')}/{destination}" if '://' in root + else str(Path(root, destination))) + elif (kind not in _OBJECT_STORES + and is_file_destination(connector, destination)): + location = destination.rstrip('/') + else: + location = source_location(connector, destination) + return OutputDestination(connector=kind, location=location) + + +def output_stats(files: Sequence[str], + task_type: TaskType) -> Dict[str, Any]: + from kumo.engine._score_stats import ( + SCORE_COLUMNS, + TOP_SCORE, + parquet_stats, + positive_rate, + top_score_stats, + ) + + if not files or task_type not in SCORE_COLUMNS: + return {} + if task_type is TaskType.BINARY_CLASSIFICATION: + stats = parquet_stats(files, ['True_PROB', 'TARGET_PRED']) + return {'columns': {'True_PROB': stats['True_PROB']}, + 'predicted_positive_rate': positive_rate(stats['TARGET_PRED'])} + if task_type.has_numerical_target: + return {'columns': parquet_stats(files, ['TARGET_PRED'])} + return {'columns': {TOP_SCORE: top_score_stats(files)}} diff --git a/packages/kumo-engine/kumo/engine/trainer/output_config.py b/packages/kumo-engine/kumo/engine/trainer/output_config.py index 41c0c7ba5..160e9c834 100644 --- a/packages/kumo-engine/kumo/engine/trainer/output_config.py +++ b/packages/kumo-engine/kumo/engine/trainer/output_config.py @@ -6,6 +6,7 @@ from dataclasses import dataclass, field from typing import List, Mapping, Optional, cast +from kumo_core.sources import layout from kumo_core.sources.object_store_sessions import ObjectStoreSessions from kumoapi.jobs import MetadataField, PredictionArtifactType, WriteMode from kumoapi.plugin import ( @@ -15,7 +16,7 @@ ) from kumoapi.prediction_outputs import resolve_prediction_output_paths -from kumo.engine.connector import is_file_destination +from kumo.engine.connector import file_store, is_file_destination from kumo.engine.connector.base import Connector @@ -23,21 +24,42 @@ class PredictionOutputConfig: """Export via ``connector``. - Use ``output_paths`` for exact predictions/embeddings destinations, or + ``path`` names where predictions go: a directory or file on a file + destination (local, S3, GCS, a volume or stage), else a table. Use + ``output_paths`` for exact predictions/embeddings destinations, or ``artifact_type`` and ``table_name`` for one legacy destination. ``extra_fields`` applies to every selected artifact, only when passed to ``Trainer.predict``; ``PredictionResult.export`` rejects it. + A table appends by default and gets ``PREDICTION_ID`` and ``BUNDLE_ID`` + columns. A file destination gets a directory per run by default, + ``//``, or ``//`` when + ``path`` ends in a file name such as ``scores.csv``; + ``WriteMode.OVERWRITE`` replaces ``path`` itself. """ connector: Connector artifact_type: Optional[PredictionArtifactType] = None table_name: Optional[str] = None - write_mode: WriteMode = WriteMode.OVERWRITE + write_mode: Optional[WriteMode] = None extra_fields: List[MetadataField] = field(default_factory=list) output_paths: Optional[Mapping[str, str]] = None + path: Optional[str] = field(default=None, kw_only=True) def __post_init__(self) -> None: + if self.path is not None: + if (self.output_paths is not None + or self.artifact_type is not None + or self.table_name is not None): + raise ValueError('path replaces output_paths, artifact_type ' + 'and table_name; pass only one form.') + if is_file_destination(self.connector, self.path): + object.__setattr__(self, 'output_paths', + {'predictions': self.path}) + else: + object.__setattr__(self, 'artifact_type', + PredictionArtifactType.PREDICTIONS) + object.__setattr__(self, 'table_name', self.path) if self.output_paths is not None: paths = resolve_prediction_output_paths( self.output_paths, output_types=self.artifact_type, @@ -60,7 +82,11 @@ def __post_init__(self) -> None: object.__setattr__(self, 'output_paths', paths) elif self.artifact_type is None or self.table_name is None: raise ValueError( - 'Specify output_paths or artifact_type and table_name') + 'Specify path, output_paths, or artifact_type and table_name') + if self.write_mode is None and not all( + is_file_destination(self.connector, path) + for path in self.destinations().values()): + object.__setattr__(self, 'write_mode', WriteMode.APPEND) def destinations(self) -> dict[str, str]: if self.output_paths is not None: @@ -68,6 +94,43 @@ def destinations(self) -> dict[str, str]: assert self.artifact_type is not None and self.table_name is not None return {self.artifact_type.value.lower(): self.table_name} + def _targets(self, run_id: Optional[str]) -> dict[str, str]: + r"""The destinations a run ``run_id`` writes to.""" + if self.write_mode is not None or not run_id: + return self.destinations() + targets = {} + for name, path in self.destinations().items(): + head, sep, tail = path.rstrip('/').rpartition('/') + targets[name] = (f'{head}{sep}{run_id}/{path[len(head + sep):]}' + if layout.file_type(tail) else + f'{head}{sep}{tail}/{run_id}/') + return targets + + def _write(self, destination: str, files: List[str], **kwargs) -> str: + r"""Write ``files`` to ``destination``; ``OVERWRITE`` first empties + a directory destination unless it holds this export complete.""" + from kumo_core.object_store_sink import holds_export + + store = file_store(self.connector, destination) + export_id = kwargs.get('export_id') + if (self.write_mode is WriteMode.OVERWRITE and store is not None + and (destination.endswith('/') or len(files) > 1) + and not (export_id and holds_export(*store, export_id))): + store[0].delete_dir_contents(store[1], missing_dir_ok=True) + return self.connector.sink_backend.write_files( + self.connector, destination, files, + self.write_mode or WriteMode.OVERWRITE, **kwargs) + + def _id_columns(self, destination: str, prediction_id: Optional[str], + bundle_id: Optional[str]) -> dict[str, str]: + if is_file_destination(self.connector, destination): + return {} + return { + name: value + for name, value in (('PREDICTION_ID', prediction_id), + ('BUNDLE_ID', bundle_id)) if value + } + @dataclass(frozen=True) class TrainingTableOutputConfig: diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index e39170f02..91492ad45 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -16,6 +16,7 @@ from datetime import datetime, timezone from pathlib import Path from typing import ( + TYPE_CHECKING, Any, Callable, Dict, @@ -31,6 +32,7 @@ ) import pandas as pd +from kumo_core.artifacts import PROVENANCE_FILE from kumo_core.embedding_models import EmbeddingModelStore from kumo_core.instrumentation import format_error from kumo_core.model_plan import ModelPlan as CoreModelPlan @@ -66,7 +68,15 @@ with_text_encoders, ) from kumo.engine._events import Event, EventEmitter, query_sha256 -from kumo.engine._manifest import Dependency, ModelManifest, TrackingRef, stub +from kumo.engine._manifest import ( + Dependency, + ModelManifest, + PredictionManifest, + PredictionOutput, + RunScoring, + TrackingRef, + stub, +) from kumo.engine._mlops import full_error, warn_failure from kumo.engine._predictor import TRAINER_PREDICTOR from kumo.engine._scratch import new_run_id, trainer_scratch_dir @@ -76,6 +86,7 @@ from kumo.engine.graph import GraphSnapshot, report_other_graph from kumo.engine.pquery.prediction_table import PredictionTable from kumo.engine.pquery.training_table import TrainingTable +from kumo.engine.trainer import _run_manifest from kumo.engine.trainer._run_spec import FitContext, fit_spec, fitted_spec from kumo.engine.trainer._scoring import ( ScoringDefaults, @@ -98,6 +109,9 @@ check_model_location, ) +if TYPE_CHECKING: + from kumo.engine._monitoring import DriftReport, EvaluationResult + logger = logging.getLogger(__name__) @@ -201,6 +215,14 @@ def with_scoring_defaults(self, assert result is not None return result + def evaluate(self, *args: Any, **kwargs: Any) -> 'EvaluationResult': + """Evaluate a run of this model, as :meth:`Trainer.evaluate`.""" + return self._model('evaluate').evaluate(*args, **kwargs) + + def score_drift(self, *args: Any, **kwargs: Any) -> 'DriftReport': + """Score drift of a run of this model, as :meth:`Trainer.score_drift`.""" + return self._model('score_drift').score_drift(*args, **kwargs) + def _model(self, call: str) -> 'Trainer': trainer = self._trainer if trainer is None: @@ -239,13 +261,134 @@ class PredictionResult: ``num_rows`` reads bounded producer metadata; ``files`` resolves selected Parquet files. Accessing ``frame`` explicitly collects the whole result into memory and is intended only for outputs that fit in memory. + ``id`` names the run; ``manifest`` is its record (None if it could not + be written), from which ``bundle_id`` and ``snapshot_id`` are read. """ - selection: PredictionSelectionRef + selection: Optional[PredictionSelectionRef] artifact_type: PredictionArtifactType + _exported: Optional['_Export'] = field(default=None, compare=False, + repr=False) + + @classmethod + def load_from_connector(cls, connector: Connector, *, + path: str) -> 'PredictionResult': + r"""Load a Parquet directory written by :meth:`export` for monitoring. + + A `kumo.MLOps` ``connector`` reads ``path`` under its artifact root, + such as ``predictions/``; a run directory loads its + ``predictions/``. The export must be complete (its ``_SUCCESS`` + marker present). Only + its ``_provenance.json`` is fetched here; the Parquet parts are + downloaded, in parallel, when ``files`` is first read, into a local + directory removed when the result is garbage-collected. + """ + import tempfile + import weakref + + from pyarrow.fs import FileType + + from kumo.engine._manifest import read_manifest + from kumo.engine._mlops import on_root + from kumo.engine._scratch import get_scratch_dir + from kumo.engine.connector import file_store + + connector, path = on_root(connector, path) + store = file_store(connector, path) + if store is None: + raise TypeError(f'{path} is not a file destination of this ' + f'connector') + filesystem, root = store + root = root.rstrip('/') + if (filesystem.get_file_info(f'{root}/_SUCCESS').type + == FileType.NotFound and filesystem.get_file_info( + f'{root}/predictions/_SUCCESS').type == FileType.File): + root, path = f'{root}/predictions', f"{path.rstrip('/')}/predictions" + if filesystem.get_file_info( + f'{root}/_SUCCESS').type == FileType.NotFound: + raise FileNotFoundError( + f'{path} has no _SUCCESS marker: the export is incomplete or ' + f'still running') + scratch = get_scratch_dir() + scratch.mkdir(parents=True, exist_ok=True) + local = Path(tempfile.mkdtemp(prefix='fetched-predictions-', + dir=scratch)) + try: + with filesystem.open_input_stream( + f'{root}/{PROVENANCE_FILE}') as reader: + (local / PROVENANCE_FILE).write_bytes(reader.read()) + except (FileNotFoundError, OSError): + pass + manifest = read_manifest(local / PROVENANCE_FILE, PredictionManifest) + here = str(_run_manifest.destination_record(connector, + path).location).rstrip('/') + kind = next((PredictionArtifactType(record.artifact.upper()) + for record in (manifest and manifest.outputs or []) + if str(record.destination.location).rstrip('/') == here), + PredictionArtifactType.EMBEDDINGS + if root.rsplit('/', 1)[-1] == 'embeddings' else + PredictionArtifactType.PREDICTIONS) + exported = _Export(filesystem, root, local, kind) + weakref.finalize(exported, shutil.rmtree, str(local), True) + return cls(None, kind, _exported=exported) + + @property + def id(self) -> Optional[str]: + if self.selection is not None: + return self.selection.job_id + manifest = self.manifest + return None if manifest is None else manifest.id + + @property + def _manifest_path(self) -> Path: + if self._exported is not None: + return self._exported.local / PROVENANCE_FILE + return Path(self.selection.pointer_path).parent / PROVENANCE_FILE + + @property + def manifest(self) -> Optional[PredictionManifest]: + from kumo.engine._manifest import read_manifest + + return read_manifest(self._manifest_path, PredictionManifest) + + @property + def bundle_id(self) -> Optional[str]: + manifest = self.manifest + return None if manifest is None else manifest.parents.model.id + + @property + def snapshot_id(self) -> Optional[str]: + manifest = self.manifest + return None if manifest is None else manifest.parents.snapshot.id + + @property + def requested_rows(self) -> Optional[int]: + r"""Entities asked for, from the run's manifest.""" + return self._table_count('requested_rows') + + @property + def scored_rows(self) -> Optional[int]: + r"""Entities scored, from the run's manifest.""" + return self._table_count('scored_rows') + + @property + def skipped_rows(self) -> Optional[int]: + r"""Requested entities that were not scored (see the manifest's + ``prediction_table.skip_reason``), from the run's manifest.""" + return self._table_count('skipped_rows') + + def _table_count(self, name: str) -> Optional[int]: + manifest = self.manifest + return None if manifest is None else getattr( + manifest.prediction_table, name) @property def num_rows(self) -> int: + if self._exported is not None: + from pyarrow import parquet + return sum( + parquet.ParquetFile(path).metadata.num_rows + for path in self.files) from kumo_core.train import artifact_manifest selected = artifact_manifest.load_prediction_selection( self.selection, validate_objects=False) @@ -255,6 +398,11 @@ def num_rows(self) -> int: @property def files(self) -> tuple[str, ...]: from kumo_core.train import artifact_manifest + if self._exported is not None: + if self._exported.kind is not self.artifact_type: + raise ValueError( + f'This export holds {self._exported.kind.value.lower()}') + return self._exported.files() resolved = artifact_manifest.resolve_prediction_selection_artifact( self.selection, self.artifact_type) if resolved is None: @@ -275,23 +423,87 @@ def export(self, output_config: PredictionOutputConfig) -> None: 'extra_fields are chosen at prediction time and cannot be ' 'applied to an existing result; pass them to ' 'Trainer.predict(output_config=...)') - self._export(output_config) + self._export(output_config, self.manifest) - def _export(self, output_config: PredictionOutputConfig) -> None: + def _export(self, output_config: PredictionOutputConfig, + sidecar: Optional[PredictionManifest] = None, + records: Optional[List[PredictionOutput]] = None, + written: Optional[List[PredictionOutput]] = None) -> None: write_files = getattr(output_config.connector.sink_backend, 'write_files', None) if not callable(write_files): raise NotImplementedError( 'Prediction export requires FileArtifactSink.write_files') outputs = [] - for name, destination in output_config.destinations().items(): + for name, destination in output_config._targets(self.id).items(): kind = PredictionArtifactType(name.upper()) - files = PredictionResult(self.selection, kind).files + files = replace(self, artifact_type=kind).files outputs.append((kind, destination, files)) - for kind, destination, files in outputs: - write_files(output_config.connector, destination, list(files), - output_config.write_mode, - export_id=f'{self.selection.job_id}-{kind.value}') + payload = _sidecar_bytes(sidecar) + bundle_id = None if sidecar is None else sidecar.parents.model.id + for index, (kind, destination, files) in enumerate(outputs): + extra: Dict[str, Any] = ({'sidecars': {PROVENANCE_FILE: payload}} + if payload is not None and _holds_parquet( + output_config.connector, destination) + else {}) + columns = output_config._id_columns(destination, self.id, + bundle_id) + if columns: + extra['extra_columns'] = columns + output_config._write(destination, list(files), + export_id=f'{self.id}-{kind.value}', **extra) + if records is not None and written is not None: + written.append(records[index]) + + +@dataclass(frozen=True) +class _Export: + filesystem: Any + root: str + local: Path + kind: PredictionArtifactType + + def files(self) -> tuple[str, ...]: + parts = self.local / 'parts' + if not parts.is_dir(): + from kumo_core.sources.transfer import StagedFile, fetch_all + from pyarrow.fs import FileSelector, FileType + + def fetch(staged: StagedFile, target: str) -> None: + with self.filesystem.open_input_stream(staged.path) as reader, \ + open(target, 'wb') as writer: + shutil.copyfileobj(reader, writer, 1 << 20) + + partial = self.local / 'parts.partial' + shutil.rmtree(partial, ignore_errors=True) + listed = self.filesystem.get_file_info(FileSelector(self.root)) + fetch_all([ + StagedFile(info.path, info.size) for info in listed + if info.type == FileType.File and info.path.endswith( + '.parquet') + ], str(partial), fetch=fetch) + partial.rename(parts) + return tuple(sorted(map(str, parts.glob('*.parquet')))) + + +def _sidecar_bytes(manifest: Optional[PredictionManifest]) -> Optional[bytes]: + from kumo.engine._manifest import manifest_bytes + + try: + return None if manifest is None else manifest_bytes(manifest) + except Exception as error: + warnings.warn(f"Exporting without {PROVENANCE_FILE}: {error}", + UserWarning, stacklevel=4) + return None + + +def _holds_parquet(connector: Connector, destination: str) -> bool: + from kumo_core.sources import layout + + from kumo.engine.connector import is_file_destination + + return (is_file_destination(connector, destination) + and layout.file_type(destination.rstrip('/')) != layout.CSV) def _find_holdout(log_dir: str) -> tuple[Optional[Path], Optional[bool]]: @@ -430,6 +642,15 @@ def _published_holdout(trainer: Any) -> tuple[Optional[Path], Optional[bool]]: return _find_holdout(trainer.get_log_dir(fit_mode=published)) +def _option_source(argument: Any, requested: Any, value: Any, + fallback: str) -> str: + if value is None: + return 'none' + if argument is not None: + return 'argument' + return 'bundle' if requested is not None else fallback + + def _collect_garbage(stage: str) -> None: # Lightning links each trainer to its data module and model in reference # cycles; free a finished run now rather than at a later collection. @@ -965,7 +1186,10 @@ def _fit( warm_start._training_result.bundle_dir, warm_start_bundle_id), started_at=started_at, fit_mode=fit_mode, - tracking=tracking.pointer()), + tracking=tracking.pointer(), + task_type=TaskType(workflow.query.task_type), + holdout_path=holdout_path, + holdout_truncated=holdout_truncated), deps=deps, ) @@ -1063,6 +1287,7 @@ def predict( num_workers: int = 1, binary_classification_threshold: Optional[float] = None, prediction_table: Optional[PredictionTable] = None, + mlops: Optional[Any] = None, ) -> PredictionResult: r"""Score a snapshot and return durable file-backed results. @@ -1092,12 +1317,20 @@ def predict( thresholds change predicted labels without changing class scores. Omitted arguments fall back to the bundle's :attr:`scoring_defaults`. + + ``result.manifest`` records the run (``outputs`` lists where each + output was written). With MLOps on (``mlops`` or a + `kumo.mlops` block), the outputs and manifest are also stored under + ``/predictions//``. """ from kumo_core.prediction_options import ( validate_num_classes_to_return, validate_prediction_workers, ) + from kumo.engine._mlops import resolve + + settings = resolve(mlops) validate_prediction_workers(num_workers) validate_num_classes_to_return(num_classes_to_return) if source is not None and not isinstance(source, GraphSnapshot): @@ -1115,17 +1348,21 @@ def predict( 'and times; pass it or entities=/' 'prediction_time=, not both.') table = (prediction_table if prediction_table is not None else - self._prediction_table(snapshot, entities, prediction_time)) + self._prediction_table( + snapshot, entities, prediction_time, strict=( + settings is not None + and settings.on_error == 'raise'))) return self._predict( table, snapshot=snapshot, output_config=output_config, num_classes_to_return=num_classes_to_return, num_workers=num_workers, - binary_classification_threshold=binary_classification_threshold) + binary_classification_threshold=binary_classification_threshold, + settings=settings) def _prediction_table(self, snapshot: GraphSnapshot, entities: Optional[SourceTable], - prediction_time: Optional[datetime] - ) -> PredictionTable: + prediction_time: Optional[datetime], + strict: bool = False) -> PredictionTable: from kumo.engine.pquery.predictive_query import PredictiveQuery if self._training_query is None: @@ -1141,7 +1378,7 @@ def _prediction_table(self, snapshot: GraphSnapshot, f"Snapshot {snapshot.id} was taken of graph " f"{snapshot._graph_key}, but this model was trained on a " f"snapshot of graph {trained}.", self._graph_config, - snapshot.graph_config, stacklevel=3) + snapshot.graph_config, stacklevel=3, strict=strict) query = PredictiveQuery._from_snapshot(snapshot, self._training_query) plan = self._training_table_generation_plan return (query.generate_prediction_table( @@ -1158,9 +1395,11 @@ def _predict( num_classes_to_return: Optional[int] = None, num_workers: int = 1, binary_classification_threshold: Optional[float] = None, + settings: Optional[Any] = None, extra_fields: Sequence[str] = (), ) -> PredictionResult: training_result = self._fitted('predict') + arguments = binary_classification_threshold, num_classes_to_return if binary_classification_threshold is None: binary_classification_threshold = ( training_result.scoring_defaults.binary_classification_threshold) @@ -1199,6 +1438,7 @@ def _predict( parsed = parse_and_validate_query(prediction_table.query, native_graph) if num_classes_to_return is None and parsed.query.top_k is None: num_classes_to_return = training_result.scoring_defaults.num_classes_to_return + requested_classes = num_classes_to_return num_classes_to_return = resolve_num_classes_to_return( task_type, num_classes_to_return, @@ -1206,7 +1446,7 @@ def _predict( has_predictions=PredictionArtifactType.PREDICTIONS in artifacts, ) - write_files = None + run_id = f'{self._id}-predict-{new_run_id()}' if output_config is not None: write_files = getattr(output_config.connector.sink_backend, 'write_files', None) @@ -1217,7 +1457,7 @@ def _predict( preflight = getattr(output_config.connector.sink_backend, 'preflight_write_files', None) if callable(preflight): - for destination in destinations.values(): + for destination in output_config._targets(run_id).values(): preflight(output_config.connector, destination, write_mode=output_config.write_mode) @@ -1238,7 +1478,7 @@ def _predict( score_prediction, ) - run_id = f'{self._id}-predict-{new_run_id()}' + started_at = datetime.now(timezone.utc) run_dir = self.scratch_dir / 'runs' / run_id run_dir.mkdir(parents=True, exist_ok=True) @@ -1330,10 +1570,70 @@ def _predict( del encoders _collect_garbage('predict') result = PredictionResult(selection, artifact_type) - if output_config is not None: - result._export(output_config) + manifest = _run_manifest.start( + result=result, trainer=self, snapshot=snapshot, + prediction_table=prediction_table, + task_type=task_type, scoring=RunScoring( + binary_classification_threshold=threshold, + threshold_source=_option_source( + arguments[0], binary_classification_threshold, threshold, + 'default'), + num_classes_to_return=num_classes_to_return, + count_source=_option_source(arguments[1], requested_classes, + num_classes_to_return, 'query'), + num_workers=num_workers, + ), started_at=started_at.isoformat(timespec='seconds')) + _run_manifest.finish(result, manifest, output_config, settings) + return result + + def evaluate(self, run: PredictionResult, *, + labels_snapshot: GraphSnapshot, + mlops: Optional[Any] = None) -> 'EvaluationResult': + r"""Score a run against labels regenerated on ``labels_snapshot``, a + snapshot newer than the one the run scored, holding its outcomes. + + With MLOps tracking on (``mlops`` or a `kumo.mlops` block), the + result is logged as a new monitoring run of ``run``. + """ + from kumo.engine._mlops import resolve + from kumo.engine._monitoring import evaluate, track + + settings = resolve(mlops) + result = evaluate(self, run, labels_snapshot) + track(result, run, settings) return result + def score_drift( + self, + run: PredictionResult, + *, + previous: Optional[PredictionResult] = None, + snapshot: Optional[GraphSnapshot] = None, + training_snapshot: Optional[GraphSnapshot] = None, + mlops: Optional[Any] = None, + ) -> 'DriftReport': + r"""Compare a run's scores with the holdout and a previous run. + + Args: + previous: An earlier run, for run-over-run statistics. + snapshot: The snapshot the run scored. By default, the one the + run's manifest names, from the local snapshot store or, + with MLOps on, the artifact root. + training_snapshot: The snapshot this model trained on; found + like ``snapshot`` by default. + mlops: Logs the report as a new monitoring run of ``run`` when + tracking is on, like `evaluate`. + """ + from kumo.engine._mlops import resolve + from kumo.engine._monitoring import score_drift, track + + settings = resolve(mlops) + report = score_drift(self, run, previous=previous, snapshot=snapshot, + training_snapshot=training_snapshot, + root=None if settings is None else settings.root) + track(report, run, settings) + return report + @classmethod def load( cls, diff --git a/packages/kumo-engine/test/_snowflake_experiment_fake.py b/packages/kumo-engine/test/_snowflake_experiment_fake.py index 27b1cc778..a9715caa6 100644 --- a/packages/kumo-engine/test/_snowflake_experiment_fake.py +++ b/packages/kumo-engine/test/_snowflake_experiment_fake.py @@ -2,7 +2,8 @@ # All rights reserved. # SPDX-License-Identifier: Apache-2.0 """A recording stand-in for `snowflake.ml.experiment.ExperimentTracking` that -keeps its observable rules (singleton, one active run, LIKE-matched resume).""" +keeps its observable rules (singleton, one active run, LIKE-matched resume, +no artifacts for a run whose name needs quoting).""" import os import re @@ -158,6 +159,11 @@ def log_metrics(self, metrics: Dict[str, float], step: int = 0) -> None: def log_artifact(self, local_path: str, artifact_path: Optional[str] = None) -> None: self._record('log_artifact', local_path, artifact_path) + # The artifact URI quotes no identifier, so a run whose name needs + # quoting is not found there. + if not re.fullmatch(r'[A-Z_][A-Z0-9_$]*', self._run.name): + raise RuntimeError(f'Run "{self._run.name}" does not exist in ' + f'Experiment {self.experiment}.') with open(local_path, encoding='utf-8') as file: self._active()['artifacts'][os.path.basename(local_path)] = ( file.read()) diff --git a/packages/kumo-engine/test/test_bundle_provenance.py b/packages/kumo-engine/test/test_bundle_provenance.py index f07b93a10..182cd7c42 100644 --- a/packages/kumo-engine/test/test_bundle_provenance.py +++ b/packages/kumo-engine/test/test_bundle_provenance.py @@ -11,6 +11,7 @@ import pandas as pd import pytest from kumo_core.artifacts import add_bundle_file +from kumo_core.types import TaskType from kumoapi.model_plan import ModelPlan, TrainingTableGenerationPlan import kumo.engine as kumo @@ -97,7 +98,8 @@ def test_each_fit_returns_a_new_result(fitted): assert second is not first and trainer.result is second -def _finish_stub_fit(bundle, trainer, warm_start=None, settings=None): +def _finish_stub_fit(bundle, trainer, warm_start=None, settings=None, + holdout_path=None): (bundle / 'model').mkdir(parents=True) (bundle / 'model' / 'config.json').write_text('{}') (bundle / MANIFEST_FILE).write_text(json.dumps({'artifacts': []})) @@ -108,7 +110,9 @@ def _finish_stub_fit(bundle, trainer, warm_start=None, settings=None): snapshot_id='snap', export_config=None, training_table_generation=TrainingTableGenerationPlan(), provenance=dict(train_table=table, warm_start=warm_start, - started_at=datetime.now(timezone.utc), fit_mode=None)) + started_at=datetime.now(timezone.utc), fit_mode=None, + task_type=TaskType.BINARY_CLASSIFICATION, + holdout_path=holdout_path, holdout_truncated=False)) _store_and_register(result.bundle_dir, result.manifest, settings) return result @@ -166,3 +170,17 @@ def fail(*args, **kwargs): assert result.manifest is None assert re.fullmatch('[0-9a-f]{64}', result.bundle_id) assert trainer.result is result + + +def test_holdout_stats_failure_keeps_the_manifest(tmp_path, monkeypatch): + def fail(*args, **kwargs): + raise ValueError('unreadable holdout') + + monkeypatch.setattr('kumo.engine.trainer._provenance.holdout_baseline', + fail) + with pytest.warns(UserWarning, match='unreadable holdout'): + result = _finish_stub_fit(tmp_path / 'bundle', + kumo.Trainer(ModelPlan()), + holdout_path=tmp_path / 'test_df.parquet') + assert result.manifest.parents.training_table.id == 'tt-1' + assert result.manifest.holdout is None diff --git a/packages/kumo-engine/test/test_filestore_sink.py b/packages/kumo-engine/test/test_filestore_sink.py index c58fc0bd6..54d0a5344 100644 --- a/packages/kumo-engine/test/test_filestore_sink.py +++ b/packages/kumo-engine/test/test_filestore_sink.py @@ -8,9 +8,10 @@ import pandas as pd import pytest +from kumo_core.artifact_export import ArtifactDestinationConflict from kumo_core.chunked_fs import ChunkedFileHandler from kumo_core.object_store_sink import FilesystemSink -from kumoapi.jobs import PredictionArtifactType +from kumoapi.jobs import PredictionArtifactType, WriteMode from kumoapi.plugin import FileStoreSink from pyarrow.fs import PyFileSystem @@ -178,3 +179,17 @@ def test_publishing_an_identical_bundle_again_reuses_the_copy( lambda *a: uploads.append(a) or upload(*a)) files.publish_model(connector, str(bundle), path='@stage/models/m') assert len(uploads) == 1 + + +def test_an_occupied_directory_is_refused_by_its_full_location( + tmp_path, connector): + files = FilesystemSink(connector.sink_backend) + part = tmp_path / 'part-0.parquet' + pd.DataFrame({'x': range(3)}).to_parquet(part) + files.write_files(connector, '@stage/out/', [str(part)]) + + files.preflight_write_files(connector, '@stage/new/') + files.preflight_write_files(connector, '@stage/out/', + write_mode=WriteMode.OVERWRITE) + with pytest.raises(ArtifactDestinationConflict, match="'@stage/out/'"): + files.preflight_write_files(connector, '@stage/out/') diff --git a/packages/kumo-engine/test/test_graph_mismatch.py b/packages/kumo-engine/test/test_graph_mismatch.py index 6714c5ba9..892532d98 100644 --- a/packages/kumo-engine/test/test_graph_mismatch.py +++ b/packages/kumo-engine/test/test_graph_mismatch.py @@ -31,13 +31,14 @@ def _moved(tables): table['source_table_name'] = 'prod.' + table['source_table_name'] -def _predict(trainer, snapshot, monkeypatch): +def _predict(trainer, snapshot, monkeypatch, **kwargs): def stop(*args, **kwargs): raise _Stop monkeypatch.setattr(PredictiveQuery, '_from_snapshot', stop) with pytest.raises(_Stop): - trainer.predict(snapshot, binary_classification_threshold=.5) + trainer.predict(snapshot, binary_classification_threshold=.5, + **kwargs) def test_predict_says_what_differs_from_the_training_graph( @@ -60,6 +61,15 @@ def dropped(tables): r'\w+\. Schema changes: events edges\.$'): _predict(trainer, _other(snapshot, dropped), monkeypatch) + strict = kumo.MLOps(str(tmp_path / 'root'), on_error='raise') + with pytest.raises(ValueError, match=r"Schema changes: events edges\. " + r"With on_error='raise'"): + trainer.predict(_other(snapshot, dropped), mlops=strict, + binary_classification_threshold=.5) + with warnings.catch_warnings(): + warnings.filterwarnings('error', message='Snapshot .* was taken') + _predict(trainer, _other(snapshot, _moved), monkeypatch, mlops=strict) + monkeypatch.setattr(trainer, '_graph_config', None) with pytest.warns(UserWarning, match='does not record its training graph'): _predict(trainer, _other(snapshot, _moved), monkeypatch) @@ -86,4 +96,4 @@ def config(location, columns, tables=('users', )): assert schema_changes( config('dev', [age]), config('dev', [{**age, 'dtype': 'float64'}], tables=('users', 'x')) - ) == ['table x added', 'users columns'] + ) == ['table x added', 'users columns (age)'] diff --git a/packages/kumo-engine/test/test_holdout_df.py b/packages/kumo-engine/test/test_holdout_df.py index 290cce9cb..c4b74af09 100644 --- a/packages/kumo-engine/test/test_holdout_df.py +++ b/packages/kumo-engine/test/test_holdout_df.py @@ -15,6 +15,7 @@ import kumo.engine as kumo from kumo.engine._scratch import get_scratch_dir +from kumo.engine.trainer._provenance import embedded_model from kumo.engine.trainer.trainer import _find_holdout @@ -93,23 +94,31 @@ def test_holdout_df_after_fit(graph): assert len(frame) > 0 assert {'ENTITY', 'TARGET_PRED'} <= set(frame.columns) assert not (result.bundle_dir / 'test_df.parquet').exists() + holdout = result.manifest.holdout + assert holdout.rows == len(frame) + assert holdout.label_positive_rate == frame['TARGET'].mean() + assert holdout.truncated is False + assert embedded_model(result.manifest).holdout is None -@pytest.mark.parametrize('refit_trainval,refit_full,explain,suffix', [ - (False, False, False, ''), - (True, False, True, '_refit_trainval'), - (True, True, False, None), +@pytest.mark.parametrize('refit_trainval,refit_full,explain,model', [ + (False, False, False, 'fit'), + (True, False, True, 'refit_trainval'), + (True, True, True, None), ]) def test_holdout_is_the_published_models(graph, refit_trainval, refit_full, - explain, suffix): + explain, model): result = _fit(graph, QUERY, refit_trainval, refit_full, explain) - if suffix is None: + assert result.manifest.fit_mode == (model or 'refit_full') + if model is None: assert result.holdout_path is None + assert result.manifest.holdout is None else: + suffix = '' if model == 'fit' else '_refit_trainval' assert re.fullmatch(rf'version_\d+{suffix}', result.holdout_path.parent.name) - assert len(result.holdout_df()) > 0 + assert result.manifest.holdout.rows == len(result.holdout_df()) def test_no_holdout(graph): @@ -119,6 +128,7 @@ def test_no_holdout(graph): assert result.holdout_path is None assert not result.holdout_truncated + assert result.manifest.holdout is None with pytest.raises(FileNotFoundError, match='No holdout'): result.holdout_df() @@ -138,6 +148,7 @@ def test_a_loaded_or_fetched_model_keeps_its_holdout(graph, tmp_path, assert loaded.bundle_id == result.bundle_id assert loaded.holdout_truncated is True pd.testing.assert_frame_equal(loaded.holdout_df(), expected) + assert loaded.manifest.holdout == result.manifest.holdout def test_a_bundle_without_a_holdout_says_so(graph, monkeypatch): @@ -174,6 +185,7 @@ def test_fit_reports_truncation(graph, monkeypatch, caplog): result = _fit(graph, QUERY) assert result.holdout_truncated is True + assert result.manifest.holdout.truncated is True with caplog.at_level('WARNING'): assert len(result.holdout_df()) > 0 assert 'truncated' in caplog.text diff --git a/packages/kumo-engine/test/test_job_templates.py b/packages/kumo-engine/test/test_job_templates.py new file mode 100644 index 000000000..a3cb78017 --- /dev/null +++ b/packages/kumo-engine/test/test_job_templates.py @@ -0,0 +1,206 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import importlib.util +import json +import py_compile +import sys +from pathlib import Path +from types import SimpleNamespace + +import numpy as np +import pandas as pd +import pytest +from kumoapi.model_plan import ModelPlan + +from kumo.engine.pquery import PredictiveQuery + + +def _examples() -> Path: + for parent in Path(__file__).resolve().parents: + candidate = parent / 'examples' / 'mlops' + if (candidate / 'jobs.py').is_file(): + return candidate + pytest.skip('examples/mlops is not in this checkout') + + +def _module(name: str): + spec = importlib.util.spec_from_file_location( + f'template_{name}', _examples() / f'{name}.py') + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +@pytest.mark.parametrize('name', [ + 'jobs.py', 'fit_job.py', 'score_job.py', 'sagemaker/pipeline.py', + 'snowflake/register_jobs.py' +]) +def test_python_compiles(name, tmp_path): + py_compile.compile(str(_examples() / name), cfile=str(tmp_path / 'out'), + doraise=True) + + +def _data(root: Path, prefix: str = '') -> None: + root.mkdir() + rng = np.random.default_rng(0) + pd.DataFrame({'USER_ID': np.arange(50) + }).to_parquet(root / f'{prefix}users.parquet') + pd.DataFrame({'ITEM_ID': np.arange(20) + }).to_parquet(root / f'{prefix}items.parquet') + pd.DataFrame({ + 'ORDER_ID': np.arange(3000), + 'USER_ID': rng.integers(0, 50, 3000), + 'ITEM_ID': rng.integers(0, 20, 3000), + 'ORDER_TIME': pd.Timestamp('2024-01-01') + + pd.to_timedelta(rng.integers(0, 120, 3000), unit='D'), + }).to_parquet(root / f'{prefix}orders.parquet') + + +def test_fit_then_score_locally(tmp_path, monkeypatch, capsys): + _data(tmp_path / 'data', prefix='demo_') + plan = ModelPlan() + plan.training_job.num_experiments = 1 + plan.training_job.disable_explain = True + plan.optimization.max_epochs = 1 + plan.optimization.max_steps_per_epoch = 2 + monkeypatch.setattr(PredictiveQuery, 'suggest_model_plan', + lambda self, snapshot: plan) + monkeypatch.syspath_prepend(str(_examples())) + monkeypatch.delitem(sys.modules, 'jobs', raising=False) + root = tmp_path / 'artifacts' + flags = [ + '--source', + str(tmp_path / 'data'), '--artifact-root', + str(root), '--scratch-dir', + str(tmp_path / 'scratch'), '--table-prefix=demo_' + ] + + fit = _module('fit_job').main([ + *flags, '--query', 'PREDICT COUNT(orders.*, 0, 7, days) > 0 ' + 'FOR EACH users.USER_ID' + ]) + assert capsys.readouterr().out.splitlines()[-1] == f'model: models/{fit.id}' + assert (root / 'models' / fit.id / '_provenance.json').is_file() + assert any((root / 'snapshots').iterdir()) + assert any((root / 'training_tables').iterdir()) + + score = _module('score_job').main( + [*flags, '--model-path', f'models/{fit.id}']) + assert score.num_rows == 50 + stored = root / 'predictions' / score.id / 'predictions' + assert any(stored.glob('*.parquet')) + assert json.loads((stored / '_provenance.json').read_text())[ + 'parents']['model']['id'] == fit.id + + +def test_the_shared_graph_declares_text_columns(tmp_path, monkeypatch): + import kumo.engine as kumo + + data = tmp_path / 'data' + _data(data) + users = pd.read_parquet(data / 'users.parquet') + users.assign(BIO='likes red shoes').to_parquet(data / 'users.parquet') + jobs = _module('jobs') + monkeypatch.setattr(jobs, 'TEXT_COLUMNS', ('users.BIO', )) + graph = jobs.build_graph(kumo.FileConnector(root_dir=str(data))) + [bio] = [column for column in graph.tables['users'].definition.cols + if column.name == 'BIO'] + assert bio.stype == 'text' + + +def test_jobs_fail_on_tracking_errors_unless_told_to_warn(tmp_path): + jobs = _module('jobs') + flags = ['--artifact-root', str(tmp_path), '--source', str(tmp_path)] + for extra, expected in (([], 'raise'), (['--on-error', 'warn'], 'warn')): + args = jobs.parse_args([*flags, *extra], score=False) + assert jobs.mlops(args, None).on_error == expected + + +@pytest.mark.parametrize('name', ['fit_job', 'score_job']) +def test_scripts_import_without_dunder_file(name, monkeypatch): + path = _examples() / f'{name}.py' + monkeypatch.setattr(sys, 'argv', [str(path)]) + monkeypatch.setattr(sys, 'path', list(sys.path)) + monkeypatch.delitem(sys.modules, 'jobs', raising=False) + namespace = {'__name__': 'spark_python_task'} + exec(compile(path.read_text(), str(path), 'exec'), namespace) + assert callable(namespace['main']) + + +def test_sagemaker_steps_take_their_own_instance_types_and_the_alias( + monkeypatch): + made = {} + + def fake(name): + class Fake: + def __init__(self, *args, **kwargs): + made[name] = kwargs + + def fit(self, **kwargs): + return kwargs + + def run(self, **kwargs): + return kwargs + + def upsert(self, **kwargs): + pass + + return Fake + + for module, names in { + 'sagemaker.estimator': ['Estimator'], + 'sagemaker.inputs': ['TrainingInput'], + 'sagemaker.processing': ['ProcessingInput', 'Processor'], + 'sagemaker.workflow.pipeline': ['Pipeline'], + 'sagemaker.workflow.pipeline_context': ['PipelineSession'], + 'sagemaker.workflow.steps': ['ProcessingStep', 'TrainingStep'], + }.items(): + monkeypatch.setitem(sys.modules, module, + SimpleNamespace(**{n: fake(n) + for n in names})) + flags = [ + f'--{flag}=x' for flag in ('role', 'image', 'code', 'data-root', + 'artifact-root', 'tracking-arn') + ] + _module('sagemaker/pipeline').main([ + *flags, '--fit-instance-type', 'ml.g5.xlarge', + '--score-instance-type', 'ml.m5.xlarge', '--alias', 'candidate' + ]) + assert made['Estimator']['instance_type'] == 'ml.g5.xlarge' + assert made['Processor']['instance_type'] == 'ml.m5.xlarge' + for arguments in (made['Estimator']['container_arguments'], + made['ProcessingStep']['step_args']['arguments']): + assert arguments[arguments.index('--alias') + 1] == 'candidate' + assert arguments[arguments.index('--on-error') + 1] == 'raise' + assert arguments[arguments.index('--model-name') + 1] == 'users-churn' + + +def test_snowflake_tasks_pass_names_and_strict_mode(monkeypatch): + calls = [] + + class Definition: + @classmethod + def register(cls, *args, **kwargs): + return cls() + + def to_sql(self, job_args, use_async): + calls.append(job_args) + return 'CALL x()' + + monkeypatch.setitem(sys.modules, 'snowflake.ml.jobs', + SimpleNamespace(MLJobDefinition=Definition)) + session = SimpleNamespace( + sql=lambda sql: SimpleNamespace(collect=lambda: None)) + _module('snowflake/register_jobs').register( + session, compute_pool='p', stage='s', database='D', schema='S', + warehouse='W', artifact_root='s3://b/k/', requirements=[], + external_access_integrations=[], model_name='CHURN', + predictions_table='CHURN_OUT', table_prefix='DEMO_') + train, score = calls + for arguments in (train, score): + assert arguments[arguments.index('--model-name') + 1] == 'CHURN' + assert arguments[arguments.index('--on-error') + 1] == 'raise' + assert '--table-prefix=DEMO_' in arguments + assert score[score.index('--predictions-table') + 1] == 'CHURN_OUT' diff --git a/packages/kumo-engine/test/test_mlflow_model.py b/packages/kumo-engine/test/test_mlflow_model.py index a62f90694..ee3aab4c0 100644 --- a/packages/kumo-engine/test/test_mlflow_model.py +++ b/packages/kumo-engine/test/test_mlflow_model.py @@ -93,6 +93,12 @@ def test_fit_logs_registers_and_serves_the_model(table_and_snapshot, params={'snapshot': str(tmp_path / 'exports' / location)}) assert list(served) == ['ENTITY', 'TARGET_PRED'] assert len(served) == len(range(0, 400, 7)) + with pytest.warns(UserWarning, match='skipped 1 of 2 requested ' + r'entities \(50\.0%\): not in snapshot'): + partial = mlflow.pyfunc.load_model('models:/m@champion').predict( + pd.DataFrame({'ENTITY': [7, 10**9]}), + params={'snapshot': str(tmp_path / 'exports' / location)}) + assert list(partial['ENTITY']) == [7] assert load_model('models:/m@champion').result.id == result.id champion = kumo.Trainer.load_from_connector( cfg, location=kumo.ModelRegistryLocation('m', alias='champion')) diff --git a/packages/kumo-engine/test/test_mlflow_tracking.py b/packages/kumo-engine/test/test_mlflow_tracking.py index 920e92272..9b8d1dae6 100644 --- a/packages/kumo-engine/test/test_mlflow_tracking.py +++ b/packages/kumo-engine/test/test_mlflow_tracking.py @@ -385,12 +385,67 @@ def test_fit_inputs_come_from_the_embedded_parents(): assert fit_inputs(None) == {} +def test_a_monitoring_run_is_new_and_tags_its_prediction_run(tracking_uri): + from kumo.engine._tracking import RunRecord, log_run + from kumo.engine.integrations.mlflow import MlflowTracking + + client = mlflow.MlflowClient() + experiment = client.create_experiment('scores') + predicted = client.create_run(experiment, tags={ + 'kumo.kind': 'prediction', 'kumo.prediction_id': 'pred-1' + }).info.run_id + client.set_terminated(predicted) + + log_run( + MlflowTracking(experiment_name='scores', log_system_metrics=False), + 'log_monitoring', + RunRecord(name='monitor-pred-1', metrics={'psi': 0.1, 'ks': 0.2}, + tags={'kumo.kind': 'monitoring', 'kumo.monitors': 'pred-1', + 'kumo.bundle_id': 'b'})) + + monitor, = [run for run in _runs(client) + if run.data.tags['kumo.kind'] == 'monitoring'] + assert monitor.info.run_name == 'monitor-pred-1' + assert monitor.info.status == 'FINISHED' + assert monitor.data.metrics == {'psi': 0.1, 'ks': 0.2} + assert monitor.data.tags['kumo.monitors'] == 'pred-1' + assert client.get_run(predicted).data.tags['kumo.monitored'] == ( + monitor.info.run_id) + + +def test_prediction_run_nests_under_the_active_run(tracking_uri): + from dataclasses import replace + + from test_tracking_contract import PREDICTION_MANIFEST + + from kumo.engine._tracking import track_prediction + from kumo.engine.integrations.mlflow import MlflowTracking + + client = mlflow.MlflowClient() + tracking = MlflowTracking(experiment_name='scores', + log_system_metrics=False) + with mlflow.start_run() as outer: + track_prediction(PREDICTION_MANIFEST, tracking) + assert mlflow.active_run().info.run_id == outer.info.run_id + track_prediction(replace(PREDICTION_MANIFEST, id='trainer-predict-2'), + tracking) + + nested, top = sorted( + (run for run in _runs(client) if 'kumo.kind' in run.data.tags), + key=lambda run: 'mlflow.parentRunId' not in run.data.tags) + assert nested.data.tags['mlflow.parentRunId'] == outer.info.run_id + assert 'mlflow.parentRunId' not in top.data.tags + assert client.get_experiment(top.info.experiment_id).name == 'scores' + assert (top.info.run_name, top.info.status) == ('trainer-predict-2', + 'FINISHED') + + def test_link_prediction_metric_names_are_logged(tracking_uri): from kumo.engine._tracking import RunRecord from kumo.engine.integrations.mlflow import MlflowTracking tracker = MlflowTracking(experiment_name='lp') - tags = {'kumo.kind': 'fit', 'kumo.run_id': 'r'} + tags = {'kumo.kind': 'fit'} tracker.start(RunRecord(name='r', tags=tags)) tracker.finish( RunRecord(name='r', tags=tags, metrics={ @@ -404,3 +459,23 @@ def test_link_prediction_metric_names_are_logged(tracking_uri): 'test_map_at_10': 0.5, 'val_precision_at_10': 0.25 } + + +def test_link_prediction_monitoring_metrics_are_logged(tracking_uri): + from kumo.engine._tracking import RunRecord + from kumo.engine.integrations.mlflow import MlflowTracking + + tracker = MlflowTracking(experiment_name='lp') + tracker.log_monitoring( + RunRecord( + name='m', tags={ + 'kumo.kind': 'monitoring', + 'kumo.monitors': 'p' + }, metrics={ + 'map@10': 0.5, + 'precision@10': 0.25 + })) + + run = mlflow.get_run(tracker.run_id) + assert run.info.status == 'FINISHED' + assert run.data.metrics == {'map_at_10': 0.5, 'precision_at_10': 0.25} diff --git a/packages/kumo-engine/test/test_mlops.py b/packages/kumo-engine/test/test_mlops.py index 397ddddd9..7fd7a4eac 100644 --- a/packages/kumo-engine/test/test_mlops.py +++ b/packages/kumo-engine/test/test_mlops.py @@ -65,7 +65,7 @@ def test_a_root_is_a_uri_or_a_connector_and_path(tmp_path): assert uri.root.location('models/m') == 'models/m' assert uri.root.uri('models/m') == f'{tmp_path}/models/m' assert isinstance(uri.root.connector, kumo.FileConnector) - assert uri.kinds == {'snapshot', 'training_table', 'model'} + assert uri.kinds == {'snapshot', 'training_table', 'model', 'prediction'} connector = kumo.FileConnector(root_dir=str(tmp_path)) pair = kumo.MLOps(artifact_root=(connector, 'volume/kumo/'), diff --git a/packages/kumo-engine/test/test_model_agnostic_predictions.py b/packages/kumo-engine/test/test_model_agnostic_predictions.py new file mode 100644 index 000000000..ccd61f4ca --- /dev/null +++ b/packages/kumo-engine/test/test_model_agnostic_predictions.py @@ -0,0 +1,64 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Prediction runs and monitoring work for any model's output files, with a +model manifest that records no fit.""" + +from types import SimpleNamespace + +import numpy as np +import pandas as pd +from kumo_core.types import TaskType + +from kumo.engine._manifest import ( + ModelManifest, + PredictionTableRecord, + RunScoring, + stub, +) +from kumo.engine._monitoring import score_drift +from kumo.engine._tracking import prediction_record +from kumo.engine.trainer._run_manifest import _record_predictions + +MODEL = ModelManifest(id='toy-model') + + +def _run(tmp_path, name, seed): + scores = np.random.default_rng(seed).uniform(size=200) + path = tmp_path / f'{name}.parquet' + pd.DataFrame({'ENTITY': range(200), 'True_PROB': scores, + 'TARGET_PRED': scores > .5}).to_parquet(path) + manifest = _record_predictions( + name, {'predictions': (200, lambda: [str(path)])}, model=MODEL, + snapshot=stub('snapshot', 'g@d'), + task_type=TaskType.BINARY_CLASSIFICATION, + table=PredictionTableRecord(query_sha256='q', source='entities', + anchors=[], rows=200, + requested_rows=205, scored_rows=200, + skipped_rows=5), + scoring=RunScoring(binary_classification_threshold=.5, + threshold_source='argument', + num_classes_to_return=None, count_source='none', + num_workers=1), + started_at='2026-10-28T00:00:00+00:00') + return SimpleNamespace(id=name, manifest=manifest, bundle_id=MODEL.id) + + +def test_any_models_prediction_run_is_recorded_and_monitored(tmp_path): + run = _run(tmp_path, 'run-2', 1) + stats = run.manifest.stats['predictions'] + assert stats.rows == 200 and stats.columns['True_PROB']['count'] == 200 + assert run.manifest.parents.model == MODEL + + record = prediction_record(run.manifest) + assert record.metrics['predictions.skipped_rows'] == 5 + assert record.tags['kumo.bundle_id'] == 'toy-model' + + report = score_drift(SimpleNamespace(result=SimpleNamespace( + manifest=MODEL, id=MODEL.id)), run, + previous=_run(tmp_path, 'run-1', 2), snapshot=None, + training_snapshot=None) + assert report.reasons['holdout'] == ('the model records no holdout ' + 'score statistics') + assert report.psi is None and report.previous_psi is not None + assert report.bundle_matches is True diff --git a/packages/kumo-engine/test/test_monitoring.py b/packages/kumo-engine/test/test_monitoring.py new file mode 100644 index 000000000..076cec208 --- /dev/null +++ b/packages/kumo-engine/test/test_monitoring.py @@ -0,0 +1,731 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import gc +import math +import os +import time +from dataclasses import dataclass, replace +from types import SimpleNamespace +from typing import Any, Optional + +import numpy as np +import pandas as pd +import pytest +from kumoapi.graph import Edge +from kumoapi.model_plan import ModelPlan +from kumoapi.typing import Stype + +import kumo.engine as kumo +from kumo.engine._manifest import ArtifactStats +from kumo.engine._monitoring import _ks, _psi, _wasserstein, evaluate +from kumo.engine._scratch import get_scratch_dir +from kumo.engine.trainer.trainer import PredictionResult, TrainingResult + +CUTOFF = pd.Timestamp('2024-03-31') +ANCHOR = CUTOFF - pd.Timedelta(days=1) +PREDICTED_AT = pd.Timestamp('2024-05-01', tz='UTC') + + +@pytest.fixture(scope='module') +def snapshots(tmp_path_factory): + root = tmp_path_factory.mktemp('monitoring') + previous = get_scratch_dir() + kumo.configure(scratch_dir=root / 'scratch') + rng = np.random.default_rng(0) + data = root / 'data' + data.mkdir() + flag = pd.Series(rng.integers(0, 2, 60).astype(bool), dtype='boolean') + orders = pd.DataFrame({ + 'EID': np.arange(1500), + 'USER_ID': rng.integers(0, 60, 1500), + 'AMT': rng.gamma(2.0, 10.0, 1500), + 'TS': pd.Timestamp('2024-01-01') + + pd.to_timedelta(rng.integers(0, 120, 1500), unit='D'), + }) + + def write(users_flag, rows): + pd.DataFrame({ + 'ID': np.arange(60), + 'FLAG': users_flag + }).to_parquet(data / 'users.parquet') + rows.to_parquet(data / 'orders.parquet') + + hidden = flag.copy() + hidden[:12] = pd.NA + write(hidden, orders[orders.TS < CUTOFF]) + # File versions are modification times in seconds. + for path in data.iterdir(): + os.utime(path, (time.time() - 3600, ) * 2) + connector = kumo.FileConnector(root_dir=str(data)) + users = kumo.Table.from_source_table(connector['users.parquet'], + primary_key='ID').infer_metadata() + events = kumo.Table.from_source_table( + connector['orders.parquet'], primary_key='EID', + time_column='TS').infer_metadata().with_column('USER_ID', + stype=Stype.ID) + graph = kumo.Graph( + tables={ + 'users': users, + 'orders': events + }, edges=[Edge(src_table='orders', fkey='USER_ID', dst_table='users')]) + before = graph.snapshot() + write(flag, orders) + after = graph.snapshot() + yield SimpleNamespace(graph=graph, before=before, after=after, + orders=orders, flag=flag) + kumo.configure(scratch_dir=previous) + + +@pytest.fixture(scope='module') +def classifier(snapshots): + pquery = kumo.PredictiveQuery( + graph=snapshots.graph, query='PREDICT COUNT(orders.*, 0, 7, days) > 0 ' + 'FOR EACH users.ID') + plan = ModelPlan() + plan.training_job.disable_explain = True + plan.optimization.max_epochs = 1 + plan.optimization.max_steps_per_epoch = 2 + trainer = kumo.Trainer(model_plan=plan) + trainer.fit(pquery.generate_training_table(snapshots.before), + snapshot=snapshots.before) + table = pquery.generate_prediction_table(snapshots.before) + return trainer, trainer.predict(table, snapshot=snapshots.before, + binary_classification_threshold=0.5) + + +@dataclass +class _Run: + frame: pd.DataFrame + artifact_type: Any = None + manifest: Any = None + id: str = 'run' + bundle_id: Optional[str] = None + + @property + def files(self): + path = get_scratch_dir() / f'{id(self.frame)}.parquet' + self.frame.to_parquet(path) + return [path] + + +def _evaluate(snapshots, query, frame, anchors=()): + trainer = SimpleNamespace(_training_query=query, result=None, + training_table_generation_plan=None) + manifest = SimpleNamespace( + prediction_table=SimpleNamespace(anchors=[str(a) for a in anchors]), + timings=SimpleNamespace(started_at=PREDICTED_AT.isoformat())) + return evaluate(trainer, _Run(frame, manifest=manifest), snapshots.after) + + +def _window(orders, anchor): + rows = orders[(orders.TS > anchor) + & (orders.TS <= anchor + pd.Timedelta(days=7))] + return rows.groupby('USER_ID') + + +def test_evaluate_matured_run(snapshots, classifier): + trainer, run = classifier + started = pd.Timestamp.now(tz='UTC') + result = trainer.evaluate(run, labels_snapshot=snapshots.after) + + assert pd.Timestamp(result.anchors[0]).tz_convert(None) == ANCHOR + assert (result.matured, result.bundle_matches, + result.labelled_fraction) == (True, True, 1.0) + assert result.rows == result.labelled == 60 and result.reasons == {} + active = set(_window(snapshots.orders, ANCHOR).groups) + rows = pd.read_parquet(result.actuals).set_index('ENTITY') + labels = rows['TARGET'] + assert labels.to_dict() == {user: user in active for user in range(60)} + expected = run.frame.set_index('ENTITY').loc[rows.index] + assert (rows['True_PROB'] == expected['True_PROB']).all() + assert (rows['TARGET_PRED'] == expected['TARGET_PRED']).all() + assert rows['EVALUATED_AT'].nunique() == 1 + assert started <= rows['EVALUATED_AT'].iloc[0] <= pd.Timestamp.now( + tz='UTC') + + from sklearn.metrics import roc_auc_score + joined = run.frame.set_index('ENTITY').join(labels) + assert result.metrics['auroc'] == pytest.approx( + roc_auc_score(joined['TARGET'], joined['True_PROB']), abs=1e-6) + + with pytest.warns(UserWarning, match='labels are not mature yet'): + immature = trainer.evaluate(run, labels_snapshot=snapshots.before) + assert immature.matured is False + assert immature.metrics == {'auroc': None, 'ap': None} + assert 'undefined' in immature.reasons['metrics'] + + +def test_evaluate_immature_run_returns_metrics_and_warns( + snapshots, classifier, monkeypatch): + import kumo_core.pquery + + trainer, run = classifier + monkeypatch.setattr(kumo_core.pquery, '_target_maximum', + lambda *args: ANCHOR + pd.Timedelta(days=1)) + with pytest.warns(UserWarning, match='labels are not mature yet'): + result = trainer.evaluate(run, labels_snapshot=snapshots.after) + assert result.matured is False and result.labelled == 60 + assert result.metrics['auroc'] is not None and result.reasons == {} + + +def test_evaluate_joins_several_anchors(snapshots, classifier, monkeypatch): + trainer, run = classifier + anchors = [ANCHOR - pd.Timedelta(days=14), ANCHOR] + stacked = _Run(pd.concat([ + run.frame.assign(ANCHOR_TIMESTAMP=anchor.tz_localize('UTC')) + for anchor in anchors + ], ignore_index=True)) + recorded = run.manifest + monkeypatch.setattr(PredictionResult, 'files', + property(lambda _: stacked.files)) + monkeypatch.setattr( + PredictionResult, 'manifest', property(lambda _: replace( + recorded, prediction_table=replace( + recorded.prediction_table, anchors=[ + anchor.tz_localize('UTC').isoformat() + for anchor in anchors + ])))) + result = trainer.evaluate(run, labels_snapshot=snapshots.after) + + assert (result.rows, result.labelled, result.matured) == (120, 120, True) + labels = pd.read_parquet(result.actuals).set_index( + ['ANCHOR_TIMESTAMP', 'ENTITY'])['TARGET'] + for anchor in anchors: + active = set(_window(snapshots.orders, anchor).groups) + assert labels[anchor].to_dict() == { + user: user in active + for user in range(60) + } + + monkeypatch.undo() + monkeypatch.setattr(PredictionResult, 'manifest', property(lambda _: None)) + result = trainer.evaluate(run, labels_snapshot=snapshots.after) + assert result.reasons['metrics'] == 'the run records no anchor timestamps' + assert (result.metrics, result.actuals) == ({}, None) + + +def test_evaluate_regression_and_static_queries(snapshots): + frame = pd.DataFrame({'ENTITY': np.arange(60), 'TARGET_PRED': 10.0}) + result = _evaluate(snapshots, + 'PREDICT SUM(orders.AMT, 0, 7, days) FOR EACH users.ID', + frame, [ANCHOR]) + sums = _window(snapshots.orders, ANCHOR)['AMT'].sum() + error = 10.0 - pd.Series({user: sums.get(user, 0.0) for user in range(60)}) + assert result.matured and result.labelled_fraction == 1.0 + assert result.metrics['mae'] == pytest.approx(error.abs().mean(), rel=1e-5) + assert result.metrics['rmse'] == pytest.approx( + math.sqrt((error**2).mean()), rel=1e-5) + + frame = pd.DataFrame({'ENTITY': np.arange(12), 'True_PROB': .5}) + result = _evaluate(snapshots, 'PREDICT users.FLAG FOR EACH users.ID', + frame) + assert (result.matured, result.anchors) == (None, ()) + assert result.rows == 12 and result.labelled_fraction == 1.0 + rows = pd.read_parquet(result.actuals).set_index('ENTITY') + assert all(rows['TARGET'][user] == snapshots.flag[user] + for user in range(12)) + assert (rows['ANCHOR_TIMESTAMP'] == PREDICTED_AT).all() + assert (rows['True_PROB'] == .5).all() + + unknown = pd.DataFrame({'ENTITY': np.arange(1000, 1012), 'True_PROB': .5}) + result = _evaluate(snapshots, 'PREDICT users.FLAG FOR EACH users.ID', + unknown) + assert (result.rows, result.labelled, result.actuals) == (12, 0, None) + assert result.reasons['metrics'] == 'no prediction has a label' + with pytest.raises(ValueError, match='No evaluated rows'): + result.export(None) + + +def test_evaluation_exports_its_rows(snapshots, classifier, tmp_path): + from kumo.engine._monitoring import EvaluationResult + from kumo.engine.trainer.output_config import PredictionOutputConfig + + trainer, run = classifier + result = trainer.evaluate(run, labels_snapshot=snapshots.after) + connector = kumo.FileConnector(root_dir=str(tmp_path)) + output = PredictionOutputConfig(connector, + output_paths={'predictions': 'evaluated/'}) + result.export(output) + pd.testing.assert_frame_equal( + pd.read_parquet(tmp_path / 'evaluated' / run.id), + pd.read_parquet(result.actuals)) + + loaded = EvaluationResult.load_from_connector( + connector, path=f'evaluated/{run.id}/') + assert replace(loaded, actuals=None) == replace(result, actuals=None) + pd.testing.assert_frame_equal(pd.read_parquet(loaded.actuals), + pd.read_parquet(result.actuals)) + + written = (tmp_path / 'evaluated' / run.id / '_SUCCESS').stat().st_mtime_ns + result.export(output) + assert (tmp_path / 'evaluated' / run.id / + '_SUCCESS').stat().st_mtime_ns == written + + +def test_evaluation_appends_to_a_table_with_its_ids(snapshots, classifier): + from unittest.mock import Mock + + from kumoapi.jobs import PredictionArtifactType, WriteMode + + from kumo.engine.trainer.output_config import PredictionOutputConfig + + trainer, run = classifier + result = trainer.evaluate(run, labels_snapshot=snapshots.after) + sink = SimpleNamespace(write_files=Mock()) + output = PredictionOutputConfig( + SimpleNamespace(sink_backend=sink), PredictionArtifactType.PREDICTIONS, + 'EVALUATED') + result.export(output) + (_, table, _, mode), kwargs = sink.write_files.call_args + assert (table, mode) == ('EVALUATED', WriteMode.APPEND) + assert kwargs['extra_columns'] == {'PREDICTION_ID': run.id, + 'BUNDLE_ID': trainer.result.id} + + +def test_a_snowflake_table_gets_the_types_model_monitor_needs( + snapshots, classifier, monkeypatch): + import pyarrow as pa + import pyarrow.parquet as parquet + + from kumo.engine.trainer.output_config import PredictionOutputConfig + + trainer, run = classifier + result = trainer.evaluate(run, labels_snapshot=snapshots.after) + written = [] + sink = SimpleNamespace(write_files=lambda connector, table, files, mode, + **_: written.extend( + parquet.read_table(f) for f in files)) + monkeypatch.setattr(kumo.SnowflakeConnector, 'sink_backend', sink) + connector = kumo.SnowflakeConnector(warehouse='WH', database='ML', + schema_name='S') + result.export(PredictionOutputConfig(connector, path='EVALUATED')) + schema = pa.concat_tables(written).schema + assert schema.field('TARGET').type == pa.int8() + assert schema.field('EVALUATED_AT').type == pa.timestamp('ns') + original = pd.read_parquet(result.actuals) + assert original.TARGET.dtype == bool + assert pa.concat_tables(written).column('TARGET').to_pylist() == ( + original.TARGET.astype(int).tolist()) + + +def _join(tmp_path, task, predictions, labels): + import json + + import diskgraph + + predictions.to_parquet(tmp_path / 'predictions.parquet', + row_group_size=97) + labels.to_parquet(tmp_path / 'labels.parquet', row_group_size=89) + spec = {'task': task, 'prediction_keys': ['ENTITY'], + 'label_keys': ['ENTITY1'], 'label_target': 'TARGET1', + 'spill_dir': str(tmp_path / 'spill'), 'partition_rows': 500} + return json.loads(diskgraph.evaluate_prediction_parquet( + [str(tmp_path / 'predictions.parquet')], + [str(tmp_path / 'labels.parquet')], json.dumps(spec))) + + +def test_realized_metrics_match_fit_time_definitions(tmp_path): + import torch + from kumo_core.evaluation.classification import ( + BinaryAUROC, + BinaryAveragePrecision, + ) + from torch_geometric.metrics import LinkPredMAP, LinkPredPrecision + + rng = np.random.default_rng(0) + n = 2000 + score = np.round(rng.random(n), 2).astype(np.float32) + target = rng.random(n) < score + joined = _join(tmp_path, 'binary', + pd.DataFrame({'ENTITY': np.arange(n), 'True_PROB': score}), + pd.DataFrame({'ENTITY1': np.arange(n)[::-1], + 'TARGET1': target[::-1]})) + for name, metric in (('auroc', BinaryAUROC()), + ('ap', BinaryAveragePrecision())): + metric.update(torch.tensor(score), torch.tensor(target).long()) + assert joined['metrics'][name] == pytest.approx( + float(metric.compute()), abs=1e-6) + + entities, k, items = 300, 5, 40 + ranked = -np.sort(-rng.random((entities, k)), axis=1) + classes = np.stack([rng.choice(items, k, replace=False) + for _ in range(entities)]) + truth = [rng.choice(items, rng.integers(0, 4), replace=False) + for _ in range(entities)] + joined = _join( + tmp_path, 'link', + pd.DataFrame({'ENTITY': np.repeat(np.arange(entities), k), + 'CLASS': classes.ravel(), 'SCORE': ranked.ravel()}), + pd.DataFrame({'ENTITY1': np.arange(entities), + 'TARGET1': [list(t) for t in truth]})) + index = torch.tensor(np.stack([ + np.repeat(np.arange(entities), [len(t) for t in truth]), + np.concatenate(truth)])) + for name, metric in ((f'map@{k}', LinkPredMAP(k)), + (f'precision@{k}', LinkPredPrecision(k))): + metric.update(torch.tensor(classes), index) + assert joined['metrics'][name] == pytest.approx( + float(metric.compute()), abs=1e-6) + assert joined['labelled'] == entities + + +def _summary(values): + return {'count': len(values), 'min': float(np.min(values)), + 'max': float(np.max(values)), + 'percentiles': list(np.quantile(values, np.arange(1, 100) / 100))} + + +def test_psi_reads_percentile_curves(): + rng = np.random.default_rng(0) + baseline = _summary(rng.beta(2, 5, 50_000)) + + assert _psi(baseline, baseline) == 0 + assert _psi(_summary(rng.beta(2, 5, 50_000)), baseline) < .01 + assert _psi(_summary(rng.beta(5, 2, 50_000)), baseline) > .2 + tied = _summary(np.r_[np.zeros(40_000), rng.random(10_000)]) + assert _psi(tied, tied) == 0 + assert _psi({'quantiles': {'q50': .5}}, baseline) is None + + +def test_ks_and_wasserstein_from_percentile_curves(): + rng = np.random.default_rng(0) + a, b = rng.normal(0, 1, 100_000), rng.normal(0.5, 2, 100_000) + ks = max(abs(np.searchsorted(np.sort(a), x, 'right') / len(a) - + np.searchsorted(np.sort(b), x, 'right') / len(b)) + for x in np.linspace(-6, 6, 4001)) + levels = (np.arange(10_000) + .5) / 10_000 + w1 = np.mean(np.abs(np.quantile(a, levels) - np.quantile(b, levels))) + + assert _ks(_summary(a), _summary(b)) == pytest.approx(ks, abs=0.02) + assert _wasserstein(_summary(a), _summary(b)) == pytest.approx(w1, rel=.02) + assert _ks(_summary(a), _summary(a)) == _wasserstein( + _summary(a), _summary(a)) == 0 + # Point masses: 30% more ones moves the CDF jump at 0 by 0.3. + zeros = _summary(np.r_[np.zeros(500), np.ones(500)]) + ones = _summary(np.r_[np.zeros(200), np.ones(800)]) + assert _ks(zeros, ones) == pytest.approx(0.3, abs=0.02) + assert _wasserstein(zeros, ones) == pytest.approx(0.3, abs=0.02) + assert _ks(zeros, None) is None and _wasserstein({}, zeros) is None + + +def test_score_drift(snapshots, classifier, monkeypatch): + trainer, run = classifier + report = trainer.score_drift(run, previous=run, snapshot=snapshots.after, + training_snapshot=snapshots.before) + + assert report.score == 'True_PROB' and math.isfinite(report.psi) + assert math.isfinite(report.ks) and math.isfinite(report.wasserstein) + assert report.holdout_truncated is False and report.bundle_matches + assert (report.previous_psi, report.previous_ks, + report.previous_wasserstein) == (0, 0, 0) + assert report.input_changes['tables.orders.rows'] == (len( + snapshots.orders[snapshots.orders.TS < CUTOFF]), len(snapshots.orders)) + assert not any(key.endswith('_reverse') for key in report.input_changes) + assert report.reasons == {} + amount = report.features['orders']['AMT'] + assert amount['type'] == 'numerical' + assert amount['delta']['num_unique'] > 0 + + found = trainer.score_drift(run) + assert found.features == trainer.score_drift( + run, snapshot=snapshots.before, + training_snapshot=snapshots.before).features + assert (report.prediction_id, report.previous_prediction_id) == (run.id, + run.id) + + shifted = _summary(np.full(500, 0.999)) + report = trainer.score_drift( + SimpleNamespace(id='r', bundle_id=None, manifest=SimpleNamespace( + stats={'predictions': ArtifactStats( + rows=500, columns={'True_PROB': shifted})}))) + assert report.psi > 0.2 and report.ks > 0.5 + assert report.bundle_matches is None + assert 'input_changes' in report.reasons + + report = trainer.score_drift(SimpleNamespace(manifest=None)) + assert (report.prediction_id, report.psi) == (None, None) + assert report.reasons + monkeypatch.setattr(TrainingResult, 'manifest', property(lambda self: None)) + report = trainer.score_drift(run) + assert (report.psi, report.ks) == (None, None) + assert 'holdout' in report.reasons['holdout'] + + +def test_monitoring_a_run_loaded_from_its_export(snapshots, classifier, + tmp_path): + from kumo.engine.trainer.output_config import PredictionOutputConfig + + trainer, run = classifier + connector = kumo.FileConnector(root_dir=str(tmp_path)) + output = PredictionOutputConfig(connector, path='p/') + run.export(output) + success = tmp_path / 'p' / run.id / '_SUCCESS' + written = success.stat().st_mtime_ns + run.export(output) + assert success.stat().st_mtime_ns == written + + path = f'p/{run.id}/' + loaded = PredictionResult.load_from_connector(connector, path=path) + parts = loaded._exported.local / 'parts' + assert loaded.manifest.id == run.id and not parts.exists() + assert trainer.score_drift(loaded).psi is not None + assert not parts.exists() + + assert (loaded.id, loaded.bundle_id, loaded.snapshot_id, + loaded.num_rows) == (run.id, run.bundle_id, run.snapshot_id, + run.num_rows) + assert trainer.evaluate(loaded, labels_snapshot=snapshots.after).metrics == \ + trainer.evaluate(run, labels_snapshot=snapshots.after).metrics + assert trainer.score_drift(loaded).psi == trainer.score_drift(run).psi + + local = loaded._exported.local + assert local.is_dir() + del loaded + gc.collect() + assert not local.exists() + + (tmp_path / path / '_SUCCESS').unlink() + with pytest.raises(FileNotFoundError, match='_SUCCESS'): + PredictionResult.load_from_connector(connector, path=path) + + from kumo.engine.trainer import _run_manifest + + cfg = kumo.MLOps(str(tmp_path / 'root')) + _run_manifest.store(run, run.manifest, cfg) + stored = PredictionResult.load_from_connector(cfg, + path=f'predictions/{run.id}') + assert (stored.id, stored.num_rows) == (run.id, run.num_rows) + + +def test_tracked_monitoring_logs_a_new_run(snapshots, classifier, tmp_path): + trainer, run = classifier + + class Tracker: + records: list = [] + + def start(self, record): + pass + + def on_trial(self, record, event): + pass + + def finish(self, record, error): + pass + + def log_monitoring(self, record): + self.records.append(record) + + cfg = kumo.MLOps(artifact_root=str(tmp_path / 'root'), tracking=Tracker()) + report = trainer.score_drift(run, mlops=cfg) + with kumo.mlops(cfg): + trainer.evaluate(run, labels_snapshot=snapshots.after) + + drift, evaluation = Tracker.records + assert (drift.name, evaluation.name) == (f'drift-{run.id}', + f'evaluate-{run.id}') + for record in (drift, evaluation): + assert record.tags['kumo.kind'] == 'monitoring' + assert record.tags['kumo.monitors'] == run.id + assert record.tags['kumo.bundle_id'] == run.bundle_id + assert drift.metrics['psi'] == report.psi + assert drift.metrics['wasserstein'] == report.wasserstein + assert 'auroc' in evaluation.metrics and 'rows' in evaluation.metrics + + +def test_a_monitoring_tracking_failure_warns_or_raises(classifier, tmp_path, + caplog): + trainer, run = classifier + + class Tracker: + def log_monitoring(self, record): + raise ConnectionError('tracking server down') + + for on_error in ('warn', 'raise'): + cfg = kumo.MLOps(str(tmp_path / 'root'), tracking=Tracker(), + on_error=on_error) + if on_error == 'warn': + assert trainer.score_drift(run, mlops=cfg).psi is not None + else: + with pytest.raises(RuntimeError, match='tracking server down' + ) as raised: + trainer.score_drift(run, mlops=cfg) + assert isinstance(raised.value.__cause__, ConnectionError) + assert ('Experiment tracking failed: ConnectionError: tracking server ' + 'down; the monitoring run is not tracked') in caplog.messages + + +def test_drift_reads_snapshot_statistics_from_the_artifact_root( + snapshots, classifier, tmp_path): + trainer, _ = classifier + pquery = kumo.PredictiveQuery( + graph=snapshots.graph, query='PREDICT COUNT(orders.*, 0, 7, days) > 0 ' + 'FOR EACH users.ID') + run = trainer.predict(pquery.generate_prediction_table(snapshots.after), + snapshot=snapshots.after, + binary_classification_threshold=0.5) + passed = trainer.score_drift(run, snapshot=snapshots.after, + training_snapshot=snapshots.before) + root = kumo.FileConnector(root_dir=str(tmp_path / 'root')) + for snapshot in (snapshots.before, snapshots.after): + snapshot.export(kumo.SnapshotOutputConfig(root, 'snapshots')) + + previous = get_scratch_dir() + kumo.configure(scratch_dir=tmp_path / 'fresh') + try: + alone = trainer.score_drift(run) + found = trainer.score_drift( + run, mlops=kumo.MLOps(artifact_root=str(tmp_path / 'root'))) + finally: + kumo.configure(scratch_dir=previous) + + assert (alone.features, alone.input_changes) == ({}, None) + for name in ('features', 'input_changes'): + assert alone.reasons[name].endswith( + 'training and scoring snapshot: pass training_snapshot=' + 'GraphSnapshot.load_from_connector(store, path=...), snapshot=' + 'GraphSnapshot.load_from_connector(store, path=...), or turn MLOps ' + 'on with the artifact root that stores them') + assert found.reasons == passed.reasons == {} + assert found.features == passed.features + assert found.features['orders']['AMT']['delta']['num_unique'] > 0 + assert found.input_changes == passed.input_changes + assert found.input_changes['tables.orders.rows'] == (len( + snapshots.orders[snapshots.orders.TS < CUTOFF]), len(snapshots.orders)) + assert not (tmp_path / 'fresh' / 'snapshots').exists() + + +@pytest.fixture(scope='module') +def shifted(tmp_path_factory): + root = tmp_path_factory.mktemp('features') + previous = get_scratch_dir() + kumo.configure(scratch_dir=root / 'scratch') + rng = np.random.default_rng(1) + data = root / 'data' + data.mkdir() + + def users(mean, plans): + pd.DataFrame({ + 'ID': np.arange(400), + 'AGE': rng.normal(mean, 5.0, 400), + 'PLAN': rng.choice(['free', 'pro', 'team'], 400, p=plans), + 'BIO': rng.choice(['likes red shoes', 'reads books'], 400), + }).to_parquet(data / 'users.parquet') + + connector = kumo.FileConnector(root_dir=str(data)) + users(40.0, [0.6, 0.3, 0.1]) + graph = kumo.Graph(tables={ + 'users': kumo.Table.from_source_table( + connector['users.parquet'], primary_key='ID').infer_metadata() + .with_column('PLAN', stype=Stype.categorical) + .with_column('BIO', stype=Stype.text) + }, edges=[]) + before = graph.snapshot() + time.sleep(1.1) + users(55.0, [0.1, 0.3, 0.6]) + after = graph.snapshot() + yield before, after, root + kumo.configure(scratch_dir=previous) + + +def _compare(trainer, snapshot, training_snapshot, reasons): + from kumo.engine._monitoring import _features, _snapshot_stats + + return _features(trainer, _snapshot_stats( + trainer, SimpleNamespace(), snapshot, training_snapshot, None), + reasons) + + +def _encoding(root, columns): + from kumo_core.encoders.config import EncodersConfig + from kumo_core.io.json import save_json + from kumoapi.encoder import Index, Numerical + + path = root / f"encoders-{'-'.join(columns)}.json" + save_json(EncodersConfig(encoders={'users': { + column: Index() if column == 'PLAN' else Numerical() + for column in columns + }}), str(path)) + return SimpleNamespace(result=SimpleNamespace( + manifest=None, artifact_paths={'encoders': str(path)})) + + +def test_feature_drift_detects_a_numeric_and_a_categorical_shift(shifted): + before, after, root = shifted + reasons = {} + features = _compare(_encoding(root, ['AGE', 'PLAN']), after, before, + reasons) + + assert reasons == {} + age, plan = features['users']['AGE'], features['users']['PLAN'] + assert age['type'] == 'numerical' + assert age['standardized_mean_shift'] > 2 + assert age['delta']['mean'] == pytest.approx(15.0, abs=2.0) + assert set(age['delta']) >= {'mean', 'std', 'min', 'max', 'q25', 'q50', + 'q75', 'skewness', 'kurtosis', 'null_rate', + 'num_unique'} + assert plan['type'] == 'categorical' and plan['psi'] > 0.5 + assert 0 < plan['js'] <= 1 and plan['category_counts_capped'] is False + + same = _compare(_encoding(root, ['AGE', 'PLAN']), before, before, {}) + assert same['users']['AGE']['standardized_mean_shift'] == 0 + assert same['users']['PLAN']['psi'] == pytest.approx(0, abs=1e-9) + + +def test_feature_drift_explains_what_it_cannot_compare(shifted): + before, after, root = shifted + reasons = {} + features = _compare(_encoding(root, ['AGE', 'BIO', 'GONE']), after, + before, reasons) + assert list(features['users']) == ['AGE'] + assert reasons['features'] == ( + 'column statistics missing on a side for users.GONE; no drift ' + 'statistics for the type of users.BIO (text)') + + reasons = {} + assert _compare(_encoding(root, ['AGE']), None, None, reasons) == {} + assert 'no column statistics' in reasons['features'] + + +def test_feature_metrics_stay_far_below_a_run_limit(): + from kumo.engine._monitoring import ( + FEATURE_METRICS_TOP, + DriftReport, + feature_metrics, + track, + ) + + features = { + f'table_{t}': { + f'col {c}@x': {'standardized_mean_shift': float(c), 'psi': float(t)} + for c in range(40) + } + for t in range(30) + } + metrics = feature_metrics(features) + assert len(metrics) <= 4 * FEATURE_METRICS_TOP + assert all(set(name) <= set('abcdefghijklmnopqrstuvwxyz' + 'ABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789_.-') + for name in metrics) + + class Tracker: + records: list = [] + + def log_monitoring(self, record): + self.records.append(record) + + report = DriftReport( + prediction_id='p', bundle_id='b', score='True_PROB', psi=0.1, ks=0.1, + wasserstein=0.1, holdout_truncated=False, bundle_matches=True, + previous_prediction_id=None, previous_psi=None, previous_ks=None, + previous_wasserstein=None, input_changes={'tables.t.rows': (1, 2)}, + reasons={}, features=features) + track(report, SimpleNamespace(id='p', bundle_id='b'), + SimpleNamespace(tracking=Tracker())) + deadline = time.monotonic() + 10 + while not Tracker.records and time.monotonic() < deadline: + time.sleep(0.05) + [record] = Tracker.records + assert len(record.metrics) < 100 + assert 'feature_drift.json' in record.artifacts diff --git a/packages/kumo-engine/test/test_prediction_completion.py b/packages/kumo-engine/test/test_prediction_completion.py index 255ba9a7a..849f99b91 100644 --- a/packages/kumo-engine/test/test_prediction_completion.py +++ b/packages/kumo-engine/test/test_prediction_completion.py @@ -184,6 +184,7 @@ def _destination(self, connector, coordinate): output_config=output, binary_classification_threshold=.5) assert result.num_rows == 8 + target /= result.id assert (target / '_SUCCESS').exists() exported = sorted(target.glob('*.parquet')) assert len(exported) == len(result.files) == workers @@ -238,6 +239,7 @@ def _destination(self, connector, coordinate): num_workers=workers, binary_classification_threshold=.5) for kind, path in paths.items(): + path = f'{path}{result.id}/' files = sorted(Path(path).glob('*.parquet')) assert len(files) == workers assert sum(pq.ParquetFile(p).metadata.num_rows for p in files) == 8 @@ -248,9 +250,6 @@ def _destination(self, connector, coordinate): assert (Path(path) / '_SUCCESS').exists() assert not Path(path.rstrip('/') + '_' + kind).exists() result.export(config) - with pytest.raises(ValueError, match='another export'): - trainer.predict(table, snapshot=snapshot, output_config=config, - binary_classification_threshold=.5) def test_export_missing_artifact_does_not_publish_partial_outputs( @@ -339,8 +338,8 @@ def interrupt(storage, source, destination, **kwargs): with pytest.raises(RuntimeError, match='embedding transfer interrupted'): result.export(retry) - prediction_dir = tmp_path / 'retry/predictions' - embedding_dir = tmp_path / 'retry/embeddings' + prediction_dir = tmp_path / 'retry/predictions' / result.id + embedding_dir = tmp_path / 'retry/embeddings' / result.id assert (prediction_dir / '_SUCCESS').exists() assert not (embedding_dir / '_SUCCESS').exists() assert len(list(embedding_dir.glob('*.parquet'))) == 1 diff --git a/packages/kumo-engine/test/test_prediction_manifest.py b/packages/kumo-engine/test/test_prediction_manifest.py new file mode 100644 index 000000000..f629fce60 --- /dev/null +++ b/packages/kumo-engine/test/test_prediction_manifest.py @@ -0,0 +1,436 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import hashlib +import json +from pathlib import Path +from types import SimpleNamespace + +import numpy as np +import pandas as pd +import pyarrow as pa +import pyarrow.parquet as pq +import pytest +from kumoapi.graph import Edge +from kumoapi.jobs import PredictionArtifactType +from kumoapi.model_plan import ModelPlan +from kumoapi.typing import Stype + +import kumo.engine as kumo +from kumo.engine._manifest import ( + ArtifactStats, + OutputDestination, + PredictionManifest, + PredictionParents, + PredictionTableRecord, + RunScoring, + embedded, + read_manifest, + stub, + write_manifest, +) +from kumo.engine._scratch import get_scratch_dir +from kumo.engine.trainer._provenance import embedded_model +from kumo.engine.trainer.output_config import PredictionOutputConfig + + +@pytest.fixture(scope='module') +def graph(tmp_path_factory): + root = tmp_path_factory.mktemp('graph') + previous = get_scratch_dir() + kumo.configure(scratch_dir=root / 'scratch') + rng = np.random.default_rng(0) + data = root / 'data' + data.mkdir() + pd.DataFrame({'ID': np.arange(50)}).to_parquet(data / 'users.parquet') + pd.DataFrame({ + 'EID': np.arange(2000), + 'USER_ID': rng.integers(0, 50, 2000), + 'TS': pd.Timestamp('2024-01-01') + + pd.to_timedelta(rng.integers(0, 120, 2000), unit='D'), + }).to_parquet(data / 'orders.parquet') + connector = kumo.FileConnector(root_dir=str(data)) + users = kumo.Table.from_source_table(connector['users.parquet'], + primary_key='ID').infer_metadata() + orders = kumo.Table.from_source_table( + connector['orders.parquet'], + primary_key='EID', time_column='TS').infer_metadata().with_column( + 'USER_ID', stype=Stype.ID) + graph = kumo.Graph( + tables={'users': users, 'orders': orders}, + edges=[Edge(src_table='orders', fkey='USER_ID', dst_table='users')]) + yield graph, graph.snapshot() + kumo.configure(scratch_dir=previous) + + +@pytest.fixture(scope='module') +def classifier(graph): + graph, snapshot = graph + pquery = kumo.PredictiveQuery( + graph=graph, query='PREDICT COUNT(orders.*, 0, 7, days) > 0 ' + 'FOR EACH users.ID') + plan = ModelPlan() + plan.training_job.disable_explain = True + plan.optimization.max_epochs = 1 + plan.optimization.max_steps_per_epoch = 2 + trainer = kumo.Trainer(model_plan=plan) + trainer.fit(pquery.generate_training_table(snapshot), snapshot=snapshot) + return trainer, pquery.generate_prediction_table(snapshot) + + +def _quantiles_match(summary, values): + values = np.sort(values) + levels = np.arange(1, 100) / 100 + low, high = (np.quantile(values, np.clip(levels + d, 0, 1)) + for d in (-.03, .03)) + assert ((low <= summary['percentiles']) + & (summary['percentiles'] <= high)).all() + assert summary['count'] == len(values) + assert (summary['min'], summary['max']) == (values[0], values[-1]) + + +def test_manifest_records_run_and_directory_export(graph, classifier, + tmp_path, recwarn): + _, snapshot = graph + trainer, table = classifier + output = PredictionOutputConfig( + connector=kumo.FileConnector(root_dir=str(tmp_path)), + output_paths={'predictions': 'preds/', 'embeddings': 'embs/'}) + result = trainer.predict(table, snapshot=snapshot, output_config=output, + binary_classification_threshold=0.4) + assert not [w for w in recwarn if 'location' in str(w.message)] + + manifest = result.manifest + assert manifest.kind == 'prediction' + assert manifest.id == result.id == result.selection.job_id + assert manifest.parents == PredictionParents( + model=embedded_model(trainer.result.manifest), + snapshot=embedded(snapshot.manifest)) + assert manifest.parents.model.holdout is None + assert (result.bundle_id, result.snapshot_id) == (trainer.result.id, + snapshot.id) + anchor = pd.Timestamp(table.timestamps[0], unit='ns', tz='UTC') + assert manifest.prediction_table == PredictionTableRecord( + query_sha256=hashlib.sha256(table.query.encode()).hexdigest(), + source='pql', anchors=[anchor.isoformat()], rows=table.row_count, + requested_rows=table.row_count, scored_rows=table.row_count, + skipped_rows=0, skipped_sample=[]) + assert manifest.scoring == RunScoring( + binary_classification_threshold=0.4, threshold_source='argument', + num_classes_to_return=None, count_source='none', num_workers=1) + assert [o.destination for o in manifest.outputs] == [ + OutputDestination(connector='LOCAL_FILE', + location=str(tmp_path / name / result.id)) + for name in ('preds', 'embs') + ] + assert manifest.timings.started_at <= manifest.timings.finished_at + assert manifest.status == 'succeeded' + assert manifest.stats['embeddings'] == ArtifactStats(rows=table.row_count) + + frame = result.frame + stats = manifest.stats['predictions'] + assert stats.rows == len(frame) + _quantiles_match(stats.columns['True_PROB'], frame['True_PROB']) + assert stats.predicted_positive_rate == pytest.approx( + (frame['True_PROB'] > 0.4).mean()) + + for name in ('preds', 'embs'): + sidecar = json.loads( + (tmp_path / name / result.id / '_provenance.json').read_text()) + assert sidecar['id'] == manifest.id + assert sidecar['outputs'] == manifest.to_dict()['outputs'] + + +def test_a_model_without_a_manifest_is_a_stub_parent(graph, classifier, + monkeypatch): + _, snapshot = graph + trainer, table = classifier + monkeypatch.setattr(type(trainer.result), 'manifest', + property(lambda _: None)) + result = trainer.predict(table, snapshot=snapshot, + binary_classification_threshold=0.5) + assert result.manifest.parents.model == stub('model', trainer.result.id) + assert result.bundle_id == trainer.result.id + + +def test_mlops_stores_the_run_after_its_export(graph, classifier, tmp_path): + _, snapshot = graph + trainer, table = classifier + output = PredictionOutputConfig( + connector=kumo.FileConnector(root_dir=str(tmp_path / 'out')), + output_paths={'predictions': 'preds/'}) + with kumo.mlops(kumo.MLOps(artifact_root=str(tmp_path / 'root'))): + result = trainer.predict(table, snapshot=snapshot, + output_config=output, + binary_classification_threshold=0.5) + stored = tmp_path / 'root' / 'predictions' / result.id / 'predictions' + assert sorted(p.name for p in stored.glob('*.parquet')) == sorted( + Path(p).name for p in result.files) + copy = json.loads((stored / '_provenance.json').read_text()) + assert copy == result.manifest.to_dict() + assert [o['destination']['location'] for o in copy['outputs']] == [ + str(tmp_path / 'out' / 'preds' / result.id) + ] + + off = kumo.MLOps(artifact_root=str(tmp_path / 'off'), store={'model'}) + trainer.predict(table, snapshot=snapshot, mlops=off, + binary_classification_threshold=0.5) + assert not (tmp_path / 'off' / 'predictions').exists() + + +@pytest.mark.parametrize('on_error', ['warn', 'raise']) +def test_a_failed_store_of_a_run_follows_on_error(graph, classifier, + tmp_path, monkeypatch, + on_error): + _, snapshot = graph + trainer, table = classifier + monkeypatch.setattr( + kumo.PredictionResult, '_export', + lambda self, *args: (_ for _ in ()).throw(OSError('ACCESS_DENIED'))) + records = [] + tracking = SimpleNamespace(log_prediction_run=records.append) + cfg = kumo.MLOps(str(tmp_path / 'root'), tracking=tracking, + on_error=on_error) + if on_error == 'raise': + with pytest.raises(RuntimeError, match='Could not store prediction'): + trainer.predict(table, snapshot=snapshot, mlops=cfg, + binary_classification_threshold=0.5) + else: + with pytest.warns(UserWarning, match='ACCESS_DENIED'): + trainer.predict(table, snapshot=snapshot, mlops=cfg, + binary_classification_threshold=0.5) + [record] = records + assert 'ACCESS_DENIED' in record.tags['kumo.error'] + + +def test_prediction_manifest_round_trips_unchanged(graph, classifier, + tmp_path): + _, snapshot = graph + trainer, table = classifier + result = trainer.predict(table, snapshot=snapshot, + binary_classification_threshold=0.4) + newer = json.loads(result._manifest_path.read_text()) + newer['future'] = {'a': 1} + newer['scoring']['calibration'] = 'none' + newer['stats']['predictions']['predicted_positive_rate'] = None + source = tmp_path / 'newer.json' + source.write_text(json.dumps(newer, indent=2, sort_keys=True) + '\n') + + manifest = read_manifest(source, PredictionManifest) + write_manifest(tmp_path / 'out.json', manifest) + + assert (tmp_path / 'out.json').read_bytes() == source.read_bytes() + assert manifest.extra == {'future': {'a': 1}} + assert manifest.scoring.extra == {'calibration': 'none'} + assert manifest.stats['predictions'].predicted_positive_rate is None + + +@pytest.mark.parametrize('argument, source', [(None, 'bundle'), + (0.4, 'argument')]) +def test_manifest_records_bundle_scoring_defaults_as_bundle( + graph, classifier, argument, source): + _, snapshot = graph + trainer, table = classifier + child = trainer.result.with_scoring_defaults( + kumo.ScoringDefaults(binary_classification_threshold=0.3)) + scoring = child.predict( + table, snapshot=snapshot, + binary_classification_threshold=argument).manifest.scoring + assert (scoring.binary_classification_threshold, + scoring.threshold_source) == (argument or 0.3, source) + + +def test_top_scores_per_entity_and_anchor_across_files(tmp_path): + from kumo_core.types import TaskType + + from kumo.engine.trainer._run_manifest import output_stats + + rng = np.random.default_rng(1) + entity = np.repeat(np.arange(200), 6) + frame = pd.DataFrame({ + 'ENTITY': entity, + 'ANCHOR_TIMESTAMP': pd.to_datetime( + np.tile(np.repeat(['2024-01-01', '2024-02-01'], 3), 200)), + 'CLASS': rng.integers(0, 50, entity.size), + 'SCORE': rng.random(entity.size), + }) + paths = [] + for index, part in enumerate((frame.iloc[:601], frame.iloc[601:])): + paths.append(str(tmp_path / f'part-{index}.parquet')) + pq.write_table(pa.Table.from_pandas(part, preserve_index=False), + paths[-1], row_group_size=100) + + stats = output_stats(paths, TaskType.MULTICLASS_CLASSIFICATION) + link = output_stats(paths, TaskType.TEMPORAL_LINK_PREDICTION) + + tops = frame.groupby(['ENTITY', 'ANCHOR_TIMESTAMP'])['SCORE'].max() + _quantiles_match(stats['columns']['top_score'], tops) + assert link == stats + + +def test_plain_export_records_nothing_and_skips_non_parquet_sidecars( + graph, classifier, tmp_path): + _, snapshot = graph + trainer, table = classifier + result = trainer.predict(table, snapshot=snapshot, + binary_classification_threshold=0.5) + before = result.manifest + calls = [] + connector = SimpleNamespace( + source_type='SNOWFLAKE', config=SimpleNamespace(), + source_backend=SimpleNamespace( + table_location=lambda _, name: f'DB.SCHEMA.{name}'), + sink_backend=SimpleNamespace( + write_files=lambda *args, **kwargs: calls.append(kwargs))) + result.export(PredictionOutputConfig(connector, + PredictionArtifactType.PREDICTIONS, + 'PREDS')) + assert len(calls) == 1 and 'sidecars' not in calls[0] + + files = kumo.FileConnector(root_dir=str(tmp_path)) + result.export(PredictionOutputConfig( + connector=files, output_paths={'predictions': 'preds/'})) + result.export(PredictionOutputConfig( + connector=files, output_paths={'predictions': 'preds.csv/'})) + sidecar = json.loads( + (tmp_path / 'preds' / result.id / '_provenance.json').read_text()) + assert sidecar == before.to_dict() and sidecar['outputs'] == [] + assert list((tmp_path / result.id / 'preds.csv').glob('*.csv')) + assert not (tmp_path / result.id / 'preds.csv' / '_provenance.json' + ).exists() + assert result.manifest == before + + +def test_manifest_failures_warn_and_keep_outputs(graph, classifier, tmp_path, + monkeypatch): + from kumo.engine.trainer import _run_manifest + + _, snapshot = graph + trainer, table = classifier + + def run(name): + output = PredictionOutputConfig( + connector=kumo.FileConnector(root_dir=str(tmp_path)), + output_paths={'predictions': f'{name}/'}) + result = trainer.predict(table, snapshot=snapshot, + output_config=output, + binary_classification_threshold=0.5) + written = pd.concat( + pd.read_parquet(p) + for p in sorted((tmp_path / name).glob('*/*.parquet'))) + return result, written.sort_values('ENTITY').reset_index(drop=True) + + def fail(*args, **kwargs): + raise RuntimeError('unavailable') + + _, exported = run('baseline') + monkeypatch.setattr(_run_manifest, 'output_stats', fail) + with pytest.warns(UserWarning, match='without output score statistics'): + partial, partial_exported = run('partial') + assert partial.manifest.stats == { + 'predictions': ArtifactStats(rows=table.row_count) + } + monkeypatch.setattr(_run_manifest, 'destination_record', fail) + with pytest.warns(UserWarning, match='Exporting without a prediction'): + unknown, unknown_exported = run('unknown') + assert unknown.manifest.outputs is None + monkeypatch.setattr(_run_manifest, '_manifest', fail) + with pytest.warns(UserWarning, match='without a prediction manifest'): + bare, bare_exported = run('bare') + assert bare.manifest is None and bare.bundle_id is None + assert not (tmp_path / 'bare' / '_provenance.json').exists() + # Scores agree to float tolerance: GPU scoring is not bit-reproducible, + # so a label may differ only where its probability sits at the threshold. + clear = (exported['True_PROB'] - 0.5).abs() > 1e-3 + for other in (partial_exported, unknown_exported, bare_exported): + pd.testing.assert_frame_equal( + other.drop(columns='TARGET_PRED'), + exported.drop(columns='TARGET_PRED'), check_exact=False, + atol=1e-4) + assert other['TARGET_PRED'][clear].equals( + exported['TARGET_PRED'][clear]) + + +class _Recorder: + def __init__(self, raises=False): + self.records, self.raises = [], raises + + def start(self, record): + pass + + def on_trial(self, record, event): + pass + + def finish(self, record, error): + pass + + def log_prediction_run(self, record): + if self.raises: + raise RuntimeError('tracker down') + self.records.append(record) + + +def test_failed_export_still_writes_and_tracks_the_run(graph, classifier, + tmp_path, monkeypatch): + _, snapshot = graph + trainer, table = classifier + connector = kumo.FileConnector(root_dir=str(tmp_path)) + sink = type(connector.sink_backend) + write_files = sink.write_files + + def fail_second(self, connector, name, *args, **kwargs): + if name.startswith('embs'): + raise OSError('disk full') + return write_files(self, connector, name, *args, **kwargs) + + monkeypatch.setattr(sink, 'write_files', fail_second) + output = PredictionOutputConfig( + connector=connector, + output_paths={'predictions': 'preds/', 'embeddings': 'embs/'}) + recorder = _Recorder() + cfg = kumo.MLOps(artifact_root=str(tmp_path / 'root'), tracking=recorder) + with pytest.raises(OSError, match='disk full'): + trainer.predict(table, snapshot=snapshot, output_config=output, + binary_classification_threshold=0.5, mlops=cfg) + + manifest, = [ + json.loads(path.read_text()) for path in trainer.scratch_dir.glob( + 'predict_output/*/_provenance.json') + if json.loads(path.read_text())['status'] == 'failed' + ] + assert manifest['outputs'] == [{ + 'artifact': 'predictions', + 'destination': {'connector': 'LOCAL_FILE', + 'location': str(tmp_path / 'preds' / manifest['id'])}, + }] + assert manifest['timings']['finished_at'] is not None + assert not (tmp_path / 'root' / 'predictions').exists() + record, = recorder.records + assert record.tags['kumo.error'] == 'OSError: disk full' + assert record.tags['kumo.prediction_id'] == manifest['id'] + + +def test_tracking_records_the_run_and_never_changes_it(graph, classifier, + tmp_path, caplog): + from kumo.engine._tracking import prediction_record + + _, snapshot = graph + trainer, table = classifier + recorder = _Recorder() + + def predict(tracking): + cfg = kumo.MLOps(artifact_root=str(tmp_path / 'root'), + tracking=tracking, store=set()) + return trainer.predict(table, snapshot=snapshot, mlops=cfg, + binary_classification_threshold=0.4) + + untracked = trainer.predict(table, snapshot=snapshot, + binary_classification_threshold=0.4) + tracked = predict(recorder) + failed = predict(_Recorder(raises=True)) + + for result in (tracked, failed): + pd.testing.assert_frame_equal(result.frame, untracked.frame) + assert recorder.records == [prediction_record(tracked.manifest)] + assert 'tracker down' in caplog.text diff --git a/packages/kumo-engine/test/test_prediction_output_columns.py b/packages/kumo-engine/test/test_prediction_output_columns.py index 84eeb2631..a384742ab 100644 --- a/packages/kumo-engine/test/test_prediction_output_columns.py +++ b/packages/kumo-engine/test/test_prediction_output_columns.py @@ -121,7 +121,7 @@ def _predict(graph, fitted, fields, kinds=('predictions', ), **options): kind: pd.concat( pd.read_parquet(path) - for path in sorted((root / 'out' / paths[kind]).glob('*.parquet'))) + for path in sorted((root / 'out' / paths[kind]).glob('*/*.parquet'))) for kind in kinds } @@ -202,7 +202,8 @@ def test_link_prediction_respects_explicit_count( assert frame.groupby('ENTITY').size().eq(expected).all() if export: exported = pd.concat( - pd.read_parquet(path) for path in tmp_path.glob('ranking/*.parquet')) + pd.read_parquet(path) + for path in tmp_path.glob('ranking/*/*.parquet')) assert exported.groupby('ENTITY').size().eq(expected).all() @@ -242,7 +243,7 @@ def test_link_embeddings_do_not_require_count(graph, link, tmp_path): trainer.predict(table, snapshot=snapshot, output_config=output) - assert list(tmp_path.glob('embeddings/*.parquet')) + assert list(tmp_path.glob('embeddings/*/*.parquet')) def test_later_export_rejects_extra_fields(graph, classifier, tmp_path): @@ -262,7 +263,7 @@ def test_later_export_rejects_extra_fields(graph, classifier, tmp_path): result.export( PredictionOutputConfig(connector=connector, output_paths={'predictions': 'good/'})) - copied = sorted((tmp_path / 'good').glob('*.parquet')) + copied = sorted((tmp_path / 'good' / result.id).glob('*.parquet')) written = [p.read_bytes() for p in copied] expected = [Path(f).read_bytes() for f in sorted(result.files)] assert written == expected diff --git a/packages/kumo-engine/test/test_prediction_output_preflight.py b/packages/kumo-engine/test/test_prediction_output_preflight.py index 2592c4ec8..7fb9cce07 100644 --- a/packages/kumo-engine/test/test_prediction_output_preflight.py +++ b/packages/kumo-engine/test/test_prediction_output_preflight.py @@ -2,14 +2,17 @@ # All rights reserved. # SPDX-License-Identifier: Apache-2.0 +import json import os +import pandas as pd import pytest from kumo_core.artifact_export import ArtifactDestinationConflict -from kumoapi.jobs import PredictionArtifactType +from kumoapi.jobs import PredictionArtifactType, WriteMode import kumo.engine as kumo from kumo.engine.trainer.output_config import PredictionOutputConfig +from kumo.engine.trainer.trainer import PredictionResult _NON_ROOT = pytest.mark.skipif(os.geteuid() == 0, reason='root bypasses mode permissions') @@ -48,7 +51,7 @@ def test_unwritable_output_fails_before_prediction(fitted, tmp_path, outputs, binary_classification_threshold=.5) finally: (tmp_path / locked).chmod(0o700) - assert error.value.filename == f'{tmp_path}/{denied}' + assert error.value.filename.startswith(f"{tmp_path}/{denied.rstrip('/')}/") assert set(runs.glob('*-predict-*')) == before assert [path.name for path in tmp_path.iterdir()] == ['locked'] @@ -58,18 +61,64 @@ def test_unwritable_output_fails_before_prediction(fitted, tmp_path, outputs, assert list(tmp_path.glob(destination.rstrip('/') + '*')) -def test_occupied_output_directory_fails_before_prediction(fitted, tmp_path): +def _predict(trainer, table, snapshot, output): + return trainer.predict(table, snapshot=snapshot, output_config=output, + binary_classification_threshold=.5) + + +def test_runs_write_to_their_own_directories(fitted, tmp_path): + trainer, table, snapshot = fitted + connector = kumo.FileConnector(root_dir=str(tmp_path)) + output = PredictionOutputConfig(connector, + output_paths={'predictions': 'scores/'}) + runs = [_predict(trainer, table, snapshot, output) for _ in range(2)] + + assert sorted(p.name for p in (tmp_path / 'scores').iterdir()) == sorted( + run.id for run in runs) + for run in runs: + [record] = run.manifest.outputs + assert record.destination.location == f'{tmp_path}/scores/{run.id}' + loaded = PredictionResult.load_from_connector( + connector, path=f'scores/{run.id}/') + assert (loaded.id, loaded.num_rows) == (run.id, run.num_rows) + + +def test_overwrite_replaces_the_destination(fitted, tmp_path): trainer, table, snapshot = fitted (tmp_path / 'results').mkdir() (tmp_path / 'results' / 'old.txt').write_text('other') output = PredictionOutputConfig( kumo.FileConnector(root_dir=str(tmp_path)), - artifact_type=PredictionArtifactType.PREDICTIONS, table_name='results/') - runs = trainer.scratch_dir / 'runs' - before = set(runs.glob('*-predict-*')) + output_paths={'predictions': 'results/'}, + write_mode=WriteMode.OVERWRITE) - with pytest.raises(ArtifactDestinationConflict, match='another export'): - trainer.predict(table, snapshot=snapshot, output_config=output, - binary_classification_threshold=.5) - assert set(runs.glob('*-predict-*')) == before - assert [p.name for p in (tmp_path / 'results').iterdir()] == ['old.txt'] + for _ in range(2): + run = _predict(trainer, table, snapshot, output) + assert not (tmp_path / 'results' / 'old.txt').exists() + assert run.manifest.outputs[0].destination.location == \ + f'{tmp_path}/results' + assert json.loads((tmp_path / 'results' / '_provenance.json' + ).read_text())['id'] == run.id + assert pd.read_parquet(tmp_path / 'results').shape[0] == run.num_rows + + +def test_occupied_output_directory_fails_before_scoring(fitted, tmp_path, + monkeypatch): + trainer, table, snapshot = fitted + monkeypatch.setattr('kumo.engine.trainer.trainer.new_run_id', + lambda: 'fixed') + taken = tmp_path / 'results' / f'{trainer._id}-predict-fixed' + taken.mkdir(parents=True) + (taken / 'old.txt').write_text('other') + + def score(**_): + raise AssertionError('scored before the destination was checked') + + monkeypatch.setattr('kumo_core.prediction.score_prediction', score) + output = PredictionOutputConfig( + kumo.FileConnector(root_dir=str(tmp_path)), + artifact_type=PredictionArtifactType.PREDICTIONS, table_name='results/') + with pytest.raises(ArtifactDestinationConflict, + match=f"'file://{taken}/'.*write_mode=WriteMode"): + _predict(trainer, table, snapshot, output) + assert [p.name for p in taken.iterdir()] == ['old.txt'] diff --git a/packages/kumo-engine/test/test_prediction_table_output.py b/packages/kumo-engine/test/test_prediction_table_output.py new file mode 100644 index 000000000..2b24cc5c5 --- /dev/null +++ b/packages/kumo-engine/test/test_prediction_table_output.py @@ -0,0 +1,90 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from types import SimpleNamespace + +import pandas as pd +import pytest +from kumoapi.jobs import PredictionArtifactType, WriteMode + +import kumo.engine as kumo +from kumo.engine.trainer.output_config import PredictionOutputConfig + + +class TableSink: + def __init__(self): + self.tables = {} + + def write_files(self, connector, table_name, files, + write_mode=WriteMode.OVERWRITE, project_id=None, *, + export_id=None, extra_columns=None): + frame = pd.concat(pd.read_parquet(f) for f in files).assign( + **(extra_columns or {})) + if write_mode is WriteMode.APPEND and table_name in self.tables: + frame = pd.concat([self.tables[table_name], frame]) + self.tables[table_name] = frame.reset_index(drop=True) + return table_name + + +def _output(sink, **options): + return PredictionOutputConfig( + SimpleNamespace(source_type='SNOWFLAKE', sink_backend=sink), + PredictionArtifactType.PREDICTIONS, 'PREDS', **options) + + +def test_table_outputs_append_runs_with_their_ids(fitted): + trainer, table, snapshot = fitted + sink = TableSink() + runs = [ + trainer.predict(table, snapshot=snapshot, output_config=_output(sink), + binary_classification_threshold=.5) for _ in range(2) + ] + + written = sink.tables['PREDS'] + assert written.groupby('PREDICTION_ID').size().to_dict() == { + run.id: run.num_rows + for run in runs + } + assert set(written['BUNDLE_ID']) == {trainer.result.id} + + last = trainer.predict( + table, snapshot=snapshot, binary_classification_threshold=.5, + output_config=_output(sink, write_mode=WriteMode.OVERWRITE)) + assert set(sink.tables['PREDS']['PREDICTION_ID']) == {last.id} + assert len(sink.tables['PREDS']) == last.num_rows + + +def test_file_outputs_default_to_a_run_directory_without_id_columns( + fitted, tmp_path): + trainer, table, snapshot = fitted + connector = kumo.FileConnector(root_dir=str(tmp_path)) + output = PredictionOutputConfig(connector, output_paths={'predictions': + 'preds/'}) + assert output.write_mode is None + assert _output(TableSink()).write_mode is WriteMode.APPEND + + run = trainer.predict(table, snapshot=snapshot, output_config=output, + binary_classification_threshold=.5) + columns = set(pd.read_parquet(tmp_path / 'preds' / run.id).columns) + assert not columns & {'PREDICTION_ID', 'BUNDLE_ID'} + + appended = PredictionOutputConfig(connector, output_paths={ + 'predictions': 'other/' + }, write_mode=WriteMode.APPEND) + with pytest.raises(NotImplementedError, match='APPEND'): + trainer.predict(table, snapshot=snapshot, output_config=appended, + binary_classification_threshold=.5) + + +def test_path_names_a_file_destination_or_a_table(): + files = kumo.FileConnector(root_dir='.') + assert PredictionOutputConfig(files, path='scores/').destinations() == { + 'predictions': 'scores/'} + table = SimpleNamespace(source_type='SNOWFLAKE', sink_backend=None) + config = PredictionOutputConfig(table, path='CHURN_SCORES') + assert config.destinations() == {'predictions': 'CHURN_SCORES'} + assert config.write_mode is WriteMode.APPEND + with pytest.raises(ValueError, match='only one form'): + PredictionOutputConfig(files, path='scores/', + output_paths={'predictions': 'x/'}) diff --git a/packages/kumo-engine/test/test_predictive_query.py b/packages/kumo-engine/test/test_predictive_query.py index 2f77622a3..de12f3eb9 100644 --- a/packages/kumo-engine/test/test_predictive_query.py +++ b/packages/kumo-engine/test/test_predictive_query.py @@ -309,7 +309,7 @@ def test_generated_weight_rejected_before_artifacts(connector_kind, tmp_path, ) monkeypatch.setattr( 'kumo.engine.pquery.predictive_query.generate_train_table', - lambda request: generated) + lambda request, **kwargs: generated) table = query.generate_training_table(snapshot) @@ -395,7 +395,7 @@ def test_forecast_generation_retains_effective_plan(graph_snapshot, tmp_path, ) monkeypatch.setattr( 'kumo.engine.pquery.predictive_query.generate_train_table', - lambda plan: generated) + lambda plan, **kwargs: generated) table = query.generate_training_table( snapshot, @@ -661,6 +661,7 @@ def generate(actual_request, custom): staging_dir=str(tmp_path / 'staging'), row_count=3, tt_id='custom-prediction', + row_failures={'rows_read': 3, 'skipped_rows': 0}, ) monkeypatch.setattr( diff --git a/packages/kumo-engine/test/test_query_from_snapshot.py b/packages/kumo-engine/test/test_query_from_snapshot.py index c22709da8..95968b794 100644 --- a/packages/kumo-engine/test/test_query_from_snapshot.py +++ b/packages/kumo-engine/test/test_query_from_snapshot.py @@ -121,3 +121,35 @@ def test_each_fit_result_scores_with_its_own_model(fitted): own = kumo.Trainer.load(result.bundle_dir) pd.testing.assert_frame_equal( score, _scores(own, prediction_table, snapshot)) + + +def test_entities_missing_from_the_snapshot_are_recorded_as_skipped( + fitted, tmp_path): + trainer, graph_table, snapshot = fitted + anchor = pd.Timestamp(graph_table.timestamps[0]) + pd.DataFrame({ + 'ENTITY1': [*range(10), *range(1000, 1004)], + 'TIMESTAMP': [anchor] * 14, + }).to_parquet(tmp_path / 'rows.parquet') + rows = kumo.FileConnector(root_dir=str(tmp_path))['rows.parquet'] + + with pytest.warns(UserWarning, match=r'^Prediction skipped 4 of 14 ' + r'requested entities \(28\.6%\): not in snapshot$'): + result = trainer.predict(snapshot, entities=rows, + binary_classification_threshold=0.5) + assert sorted(result.frame['ENTITY']) == list(range(10)) + assert (result.requested_rows, result.scored_rows, + result.skipped_rows) == (14, 10, 4) + table = result.manifest.prediction_table + assert sorted(table.skipped_sample) == [1000, 1001, 1002, 1003] + assert table.skip_reason == 'not in snapshot' + + from kumo.engine._tracking import prediction_record + metrics = prediction_record(result.manifest).metrics + assert metrics['predictions.skipped_rows'] == 4 + assert metrics['predictions.skipped_fraction'] == pytest.approx(4 / 14) + + full = trainer.predict(snapshot, binary_classification_threshold=0.5) + assert full.skipped_rows == 0 + assert full.requested_rows == full.scored_rows == 50 + assert full.manifest.prediction_table.skipped_sample == [] diff --git a/packages/kumo-engine/test/test_score_stats.py b/packages/kumo-engine/test/test_score_stats.py new file mode 100644 index 000000000..f4052e22d --- /dev/null +++ b/packages/kumo-engine/test/test_score_stats.py @@ -0,0 +1,86 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import numpy as np +import pandas as pd +import pytest +from kumo_core.types import TaskType + +from kumo.engine._score_stats import parquet_stats, top_score_stats +from kumo.engine.trainer._provenance import holdout_baseline + + +def _close_to_ranks(summary, values): + values = np.sort(values) + assert summary['count'] == len(values) + assert summary['min'] == values[0] and summary['max'] == values[-1] + assert summary['mean'] == pytest.approx(values.mean()) + ranks = np.searchsorted(values, summary['percentiles']) / len(values) + assert np.abs(ranks - np.arange(1, 100) / 100).max() < .02 + + +def test_long_format_holdouts_take_each_predictions_top_score(tmp_path): + rng = np.random.default_rng(0) + entities, k = 400, 3 + score = rng.random((entities, k)) + frame = pd.DataFrame({ + 'ENTITY': np.repeat(np.arange(entities), k), + 'CLASS': [f'c{i}' for i in range(entities * k)], + 'SCORE': score.ravel(), + 'LABEL': np.tile([True, False, False], entities), + 'PREDICTED': np.tile([True, False, False], entities), + }) + frame['TARGET'] = np.repeat(frame['CLASS'].to_numpy()[::k], k) + path = tmp_path / 'test_df.parquet' + frame.to_parquet(path, row_group_size=7) + parts = [tmp_path / 'part-0.parquet', tmp_path / 'part-1.parquet'] + frame.iloc[:601].to_parquet(parts[0]) + frame.iloc[601:].to_parquet(parts[1]) + + for task_type in (TaskType.MULTILABEL_CLASSIFICATION, + TaskType.MULTICLASS_CLASSIFICATION): + holdout = holdout_baseline(path, task_type, truncated=None) + assert holdout.rows == entities + _close_to_ranks(holdout.columns['top_score'], score.max(axis=1)) + assert holdout.truncated is None + _close_to_ranks(top_score_stats(parts), score.max(axis=1)) + + +def test_non_finite_scores_are_not_counted(tmp_path): + pd.DataFrame({ + 'True_PROB': [.2, np.nan, np.inf, .4, None, -np.inf], + 'TARGET': [1, 0, 1, 0, 1, 0], + }).to_parquet(tmp_path / 'binary.parquet') + + summary = parquet_stats([tmp_path / 'binary.parquet'], + ['True_PROB'])['True_PROB'] + + assert summary['count'] == 2 + assert summary['min'] == .2 and summary['max'] == .4 + + +def test_binary_and_regression_holdouts(tmp_path): + rng = np.random.default_rng(1) + prob = rng.random(1000) + pd.DataFrame({ + 'TARGET': (rng.random(1000) < .3).astype(int), + 'True_PROB': prob, + }).to_parquet(tmp_path / 'binary.parquet') + pd.DataFrame({ + 'TARGET': [1., 2., 5.], + 'TARGET_PRED': [1., 3., 4.] + }).to_parquet(tmp_path / 'regression.parquet') + + binary = holdout_baseline(tmp_path / 'binary.parquet', + TaskType.BINARY_CLASSIFICATION, truncated=False) + regression = holdout_baseline(tmp_path / 'regression.parquet', + TaskType.REGRESSION, truncated=False) + + assert binary.rows == 1000 + assert binary.label_positive_rate == pytest.approx( + pd.read_parquet(tmp_path / 'binary.parquet')['TARGET'].mean()) + _close_to_ranks(binary.columns['True_PROB'], prob) + assert regression.rows == 3 + assert regression.columns['TARGET']['max'] == 5. + assert regression.columns['TARGET_PRED']['mean'] == pytest.approx(8 / 3) diff --git a/packages/kumo-engine/test/test_scoring_defaults.py b/packages/kumo-engine/test/test_scoring_defaults.py index fd8812113..7ed3ad4a4 100644 --- a/packages/kumo-engine/test/test_scoring_defaults.py +++ b/packages/kumo-engine/test/test_scoring_defaults.py @@ -59,6 +59,7 @@ def test_derived_model_thresholds_unless_explicit(fitted): assert child.id != parent.id assert child.manifest.id == child.id assert child.manifest.parents.training_table is None + assert child.manifest.holdout == parent.manifest.holdout is not None derived_from = child.manifest.parents.derived_from assert derived_from.id == parent.id assert derived_from.parents.training_table == ( diff --git a/packages/kumo-engine/test/test_semantic_carry_forward.py b/packages/kumo-engine/test/test_semantic_carry_forward.py index 5968ef805..e0c349f84 100644 --- a/packages/kumo-engine/test/test_semantic_carry_forward.py +++ b/packages/kumo-engine/test/test_semantic_carry_forward.py @@ -213,7 +213,8 @@ def load(*args, **kwargs): assert calls[0]['job_start_time'].tzinfo is timezone.utc sink.write_files.assert_called_once_with( output.connector, output.table_name, list(result.files), - output.write_mode, export_id=f'{result.selection.job_id}-PREDICTIONS') + output.write_mode, export_id=f'{result.selection.job_id}-PREDICTIONS', + extra_columns={'PREDICTION_ID': result.id}) def test_predict_reuses_generated_table_without_recompiling( diff --git a/packages/kumo-engine/test/test_snowflake_tracking.py b/packages/kumo-engine/test/test_snowflake_tracking.py index 7cbbacec9..e7b0032d2 100644 --- a/packages/kumo-engine/test/test_snowflake_tracking.py +++ b/packages/kumo-engine/test/test_snowflake_tracking.py @@ -14,11 +14,12 @@ from test_tracking_contract import ( CONTEXT, EVENTS, + PREDICTION_MANIFEST, RESULT, join_tracking_threads, ) -from kumo.engine._tracking import TrackingSession +from kumo.engine._tracking import TrackingSession, prediction_record from kumo.engine.trainer._run_spec import fit_spec, fitted_spec FAKED = ('__init__', 'set_experiment', 'start_run', 'end_run', 'log_params', @@ -123,10 +124,12 @@ def test_an_instance_in_use_is_left_alone(fake, busy, caplog): calls = len(user.calls) _replay(_tracking(session)) + _tracking(session).log_prediction_run( + prediction_record(PREDICTION_MANIFEST)) assert user.calls[calls:] == [] assert user.warnings_as_errors is True - assert caplog.text.count('not tracked in Snowflake') == 1 + assert caplog.text.count('not tracked in Snowflake') == 2 def test_location_defaults_to_the_session_not_the_last_fit(fake): @@ -160,7 +163,7 @@ def test_an_existing_run_is_never_resumed(fake, taken, caplog): assert set(user.runs) == {taken, new} assert user.runs[taken]['params'] == {} assert user.runs[new]['status'] == 'FINISHED' - assert 'recording this fit as' in caplog.text + assert 'recording this run as' in caplog.text def test_metric_and_param_limits(fake, monkeypatch): @@ -222,6 +225,20 @@ def test_a_hung_snowflake_call_does_not_stall_the_fit(fake, monkeypatch, assert run['status'] == 'FINISHED' +def test_prediction_runs_are_never_resumed(fake, caplog): + tracking = _tracking() + record = prediction_record(PREDICTION_MANIFEST) + tracking.log_prediction_run(record) + tracking.log_prediction_run(record) + + first, second = fake._instance.runs + assert first == 'TRAINER_PREDICT_1' + assert re.fullmatch('TRAINER_PREDICT_1_[0-9A-F]{8}', second) + assert {run['status'] for run in fake._instance.runs.values()} == { + 'FINISHED'} + assert 'already exists' in caplog.text + + class _SchemaSession: def __init__(self, schemas): self.schemas = schemas diff --git a/packages/kumo-engine/test/test_tracking_contract.py b/packages/kumo-engine/test/test_tracking_contract.py index 0d7bedbf1..9774ed0d7 100644 --- a/packages/kumo-engine/test/test_tracking_contract.py +++ b/packages/kumo-engine/test/test_tracking_contract.py @@ -1,7 +1,8 @@ # SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. # All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Every tracking adapter records the same run records as a recorded fit.""" +"""Every tracking adapter records the same run records as a recorded fit +and a recorded prediction run.""" import json import threading @@ -14,11 +15,19 @@ import pytest from kumo_core.artifacts import PROVENANCE_FILE +from kumo.engine._manifest import ( + ArtifactStats, + PredictionManifest, + PredictionParents, + stub, +) from kumo.engine._tracking import ( DROPPED_METRICS, DROPPED_PARAMS, TrackingSession, + prediction_record, run_record, + track_prediction, ) from kumo.engine.tracking import Event, RunRecord from kumo.engine.trainer._run_spec import FitContext, fit_spec, fitted_spec @@ -81,6 +90,53 @@ def _epochs(trial, *values): id=BUNDLE, bundle_dir=Path('no-bundle')) +PREDICTION_RUN = 'trainer-predict-1' +PREDICTION_MANIFEST = PredictionManifest.from_dict({ + 'kind': 'prediction', + 'schema_version': 1, + 'id': PREDICTION_RUN, + 'parents': { + 'model': {'kind': 'model', 'schema_version': 1, 'id': BUNDLE, + 'reconstructed': True}, + 'snapshot': { + 'kind': 'snapshot', 'schema_version': 1, 'id': 'g@d', + 'parents': {'source_tables': INPUTS}, + }, + }, + 'created_at': '2026-10-21T00:00:00+00:00', + 'environment': None, + 'prediction_table': { + 'query_sha256': '69dab700', + 'source': 'pql', + 'anchors': ['2024-04-29T00:00:00+00:00'], + 'rows': 50, + }, + 'scoring': { + 'binary_classification_threshold': 0.4, + 'threshold_source': 'argument', + 'num_classes_to_return': None, + 'count_source': 'none', + 'num_workers': None, + }, + 'outputs': [], + 'status': 'succeeded', + 'timings': {'started_at': '2026-10-21T00:00:00+00:00', + 'finished_at': '2026-10-21T00:01:00+00:00'}, + 'stats': { + 'predictions': { + 'rows': 50, + 'columns': { + 'True_PROB': { + 'count': 50, + 'percentiles': [q / 100 for q in range(1, 100)], + } + }, + 'predicted_positive_rate': 0.62, + }, + }, +}) + + class _Recorder: def __init__(self, tracker=None): self.calls: List[Tuple[str, Any]] = [] @@ -156,6 +212,7 @@ class Adapter: nested_runs: bool = False step_metrics: bool = False artifacts: bool = False + run_name: Callable[[str], str] = str def gaps(expected: Dict[str, RunRecord], observed: Dict[str, Observed], @@ -178,7 +235,8 @@ def differ(what, have, want): found.append(f'{run_id} {what}: {have!r} != {want!r}') if record.name is not None: - differ('name', seen.name, record.name) + differ('name', adapter.run_name(seen.name), + adapter.run_name(record.name)) differ('params', dict(seen.params), dict(kept_params(record.params))) differ('metrics', {k: seen.metrics.get(k) for k in record.metrics}, dict(record.metrics)) @@ -259,11 +317,14 @@ def snowflake_adapter(monkeypatch): import _snowflake_experiment_fake fake = _snowflake_experiment_fake.install(monkeypatch) - from kumo.engine.integrations.snowflake import SnowflakeTracking + from kumo.engine.integrations.snowflake import ( + SnowflakeTracking, + _run_identifier, + ) def observe(): observed = {} - for run in fake._instance.runs.values(): + for run_name, run in fake._instance.runs.items(): names: Dict[Optional[str], Dict[str, Dict[str, Any]]] = {} def bucket(key): @@ -285,8 +346,9 @@ def bucket(key): runs['steps'][name] = (*runs['steps'].get(name, ()), (step, value)) for prefix, runs in names.items(): - observed[prefix or FIT] = Observed( - name=prefix, + key = prefix or runs['tags'].get('kumo.prediction_id', FIT) + observed[key] = Observed( + name=prefix or run_name, artifacts=run['artifacts'] if prefix is None else {}, **runs) return observed @@ -294,7 +356,8 @@ def bucket(key): return Adapter( make=lambda: SnowflakeTracking(_snowflake_experiment_fake.FakeSession(), experiment_name='ents', tags=USER_TAGS), - observe=observe, step_metrics=True, artifacts=True) + observe=observe, step_metrics=True, artifacts=True, + run_name=lambda name: _run_identifier(name).upper()) @pytest.fixture(params=['mlflow', 'snowflake']) @@ -407,3 +470,69 @@ def start(self, record): assert elapsed < 1.5 assert 'did not finish within' in caplog.text + + +def test_prediction_record_of_the_recorded_run(): + record = prediction_record(PREDICTION_MANIFEST) + + assert record.name == PREDICTION_RUN + assert record.params == { + 'anchors': '["2024-04-29T00:00:00+00:00"]', + 'binary_classification_threshold': '0.4', + 'threshold_source': 'argument', + 'count_source': 'none', + } + assert record.metrics == { + 'prediction_table.rows': 50.0, + 'predictions.rows': 50.0, + 'predictions.True_PROB.count': 50.0, + **{ + f'predictions.True_PROB.q{q:02d}': q / 100 + for q in (1, 5, 25, 50, 75, 95, 99) + }, + 'predictions.predicted_positive_rate': 0.62, + } + assert record.tags == {'kumo.kind': 'prediction', + 'kumo.prediction_id': PREDICTION_RUN, + 'kumo.bundle_id': BUNDLE} + assert record.inputs == INPUTS + assert json.loads(record.artifacts[PROVENANCE_FILE]) == ( + PREDICTION_MANIFEST.to_dict()) + + +def test_prediction_record_of_a_partial_manifest(): + record = prediction_record( + replace(PREDICTION_MANIFEST, + parents=PredictionParents(model=stub('model', BUNDLE), + snapshot=stub('snapshot', 'g@d')), + scoring=replace(PREDICTION_MANIFEST.scoring, + num_classes_to_return=3), + stats={'predictions': ArtifactStats(rows=150, columns={ + 'top_score': {'count': 150, 'percentiles': []}}), + 'embeddings': ArtifactStats(rows=150)}), + RuntimeError('disk full')) + + assert record.params['num_classes_to_return'] == '3' + assert record.metrics == {'prediction_table.rows': 50.0, + 'predictions.rows': 150.0, + 'predictions.top_score.count': 150.0, + 'embeddings.rows': 150.0} + assert record.inputs == {} + assert record.tags['kumo.error'] == 'RuntimeError: disk full' + + +def test_adapter_records_the_same_prediction_run(adapter): + track_prediction(PREDICTION_MANIFEST, adapter.make()) + + expected = {PREDICTION_RUN: prediction_record(PREDICTION_MANIFEST)} + assert gaps(expected, adapter.observe(), adapter) == [] + + +def test_prediction_run_without_a_manifest_or_hook_is_skipped(caplog): + recorder = _Recorder() + recorder.log_prediction_run = recorder.calls.append + track_prediction(None, recorder) + track_prediction(PREDICTION_MANIFEST, _Recorder()) + + assert recorder.calls == [] + assert 'not tracked' in caplog.text diff --git a/packages/kumo-snowflake/tests/test_sink.py b/packages/kumo-snowflake/tests/test_sink.py index 1c348eba3..7f95d0fc1 100644 --- a/packages/kumo-snowflake/tests/test_sink.py +++ b/packages/kumo-snowflake/tests/test_sink.py @@ -141,3 +141,38 @@ def open_session(_connector): pass else: raise AssertionError("append must not be retried after token expiry") + + +def test_an_append_loads_files_stamped_with_extra_columns(monkeypatch, + tmp_path): + import pandas as pd + import pyarrow.parquet as pq + + backend = SnowflakeSinkBackend() + artifact = tmp_path / "part-0.parquet" + pd.DataFrame({"ENTITY": ["a"]}).to_parquet(artifact) + + @contextmanager + def open_session(_connector): + yield _Session() + + monkeypatch.setattr(backend, "_session", open_session) + monkeypatch.setattr(sink, "_require_snowpark", lambda _session: None) + loads = [] + + def load(target, table_name, path, replace, *, artifact_files, + upload_local): + prepare = upload_local.keywords["prepare_file"] + with prepare(path) as staged: + loads.append((table_name, replace, + pq.read_table(staged).to_pydict())) + + monkeypatch.setattr("kumo_snowflake.write.write_artifacts_to_table", load) + monkeypatch.setattr(backend, "_qualified", lambda _c, name: name) + backend.write_files(SimpleNamespace(), "PREDICTIONS", [str(artifact)], + WriteMode.APPEND, extra_columns={"PREDICTION_ID": "p1"}) + + assert loads == [("PREDICTIONS", False, { + "ENTITY": ["a"], + "PREDICTION_ID": ["p1"] + })] diff --git a/tests/public/Dockerfile.dockerignore b/tests/public/Dockerfile.dockerignore index 4bedfe32e..ddae86d2c 100644 --- a/tests/public/Dockerfile.dockerignore +++ b/tests/public/Dockerfile.dockerignore @@ -12,6 +12,8 @@ !packages/kumo-core/src/kumo_core/csv_rules.py !packages/kumo-core/src/kumo_core/custom_sql.py !packages/kumo-core/src/kumo_core/artifact_staging.py +# Table sinks stamp run ids on exported prediction files with it. +!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 !packages/kumo-core/src/kumo_core/config.py