diff --git a/3rdparty/notices/inventory.json b/3rdparty/notices/inventory.json index d0dc3d7bc..c79d73931 100644 --- a/3rdparty/notices/inventory.json +++ b/3rdparty/notices/inventory.json @@ -515,6 +515,7 @@ "rust:cpp-demangle@0.4.5", "rust:cpufeatures@0.2.17", "rust:cpufeatures@0.3.0", + "rust:crc32fast@1.5.0", "rust:crossbeam-deque@0.8.6", "rust:crossbeam-epoch@0.9.21", "rust:crossbeam-utils@0.8.21", @@ -3642,7 +3643,8 @@ }, { "artifact_ids": [ - "diskgraph-extension" + "diskgraph-extension", + "diskgraph-ingest" ], "copyright": [ "Copyright (c) 2018 Sam Rijs, Alex Crichton and contributors" diff --git a/THIRD-PARTY.txt b/THIRD-PARTY.txt index 1394a494e..df85abc32 100644 --- a/THIRD-PARTY.txt +++ b/THIRD-PARTY.txt @@ -5152,7 +5152,7 @@ cpufeatures 0.3.0 crc32fast 1.5.0 ------------------------------------------------------------------------------ License : MIT OR Apache-2.0 - Distributed in: diskgraph._c.abi3.so + Distributed in: diskgraph._c.abi3.so, diskgraph/bin/ingest Copyright : Copyright (c) 2018 Sam Rijs, Alex Crichton and contributors Legal text c6596eb7be85: LICENSE-APACHE diff --git a/examples/mlops/README.md b/examples/mlops/README.md new file mode 100644 index 000000000..d534573e8 --- /dev/null +++ b/examples/mlops/README.md @@ -0,0 +1,61 @@ +# Scheduled kumo-engine jobs + +Two plain Python job scripts that use the normal `kumo.engine` API inside +`kumo.job()`, and the wiring that runs them on a schedule: + +- `fit_job.py` builds a snapshot (exported under `KUMO_ARTIFACT_ROOT` when + set), trains the users-churn model, publishes it and moves the `champion` + alias to it. +- `score_job.py` loads `champion`, scores a fresh snapshot (or the exported + one named by `KUMO_SNAPSHOT`) with `ANCHOR_TIMESTAMP` on every row, and + writes the predictions. It passes a 0.5 threshold unless the model carries + scoring defaults. +- `jobs.py` holds what both share: the graph (users, items, orders), the + connectors and the registry choice, all read from environment variables + (listed in its docstring). Replace it with your own module. + +| 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 | + +## Run locally + +With `users.parquet`, `items.parquet` and `orders.parquet` in a directory: + +```bash +cd examples/mlops +export KUMO_SOURCE=file KUMO_DATA_ROOT= KUMO_ARTIFACT_ROOT= +python fit_job.py --job-id local --job-run-id run-1 +KUMO_MODEL_PATH=runs/run-1/model python score_job.py --job-id local --job-run-id run-2 +``` + +Without `KUMO_REGISTRY` the fit publishes the model to +`/runs//model` (its summary's `outputs.model.uri`; the +`version` there is null, as no registry assigned one), +and the score job loads the path `KUMO_MODEL_PATH` names under +`KUMO_ARTIFACT_ROOT`; scheduled score runs have their own run ids, so set it to +the fit run you want to score with. The query needs +several 30-day windows of order history; for a small sample set `KUMO_QUERY` +to a shorter window and `KUMO_NUM_EXPERIMENTS=1`. + +## What every scheduler sees + +- **Exit codes.** `0` ok, `1` failure, `143` SIGTERM (timeout or cancel). + The summary's `error.kind` says whether a failure is `transient` (worth a + retry), `permanent` or `unknown`. +- **Output.** Each script prints one JSON line on stdout when it ends: + `status`, `exit_code`, `scheduler`, `bundle_id`, `snapshot_id`, `run_id`, + `outputs` (full URIs of the exported snapshot and published model, or of + the prediction destinations), `warnings` and `error`. Logs go to stderr. +- **Job ids.** `--job-id`/`--job-run-id` (or `KUMO_JOB_ID`/`KUMO_JOB_RUN_ID`) + are recorded as `scheduler: {platform, job_id, job_run_id}` in the snapshot + and training-table manifests and the bundle's `provenance.json`, and as + `kumo.scheduler`/`kumo.job_id`/`kumo.job_run_id` on a tracked fit run. + Without them the engine detects SageMaker and Snowflake ML Jobs itself. + +## Limitations + +`score_job.py` does not yet tag output rows with lineage columns or record a +tracked prediction run; both arrive with the prediction-run PRs. diff --git a/examples/mlops/databricks/README.md b/examples/mlops/databricks/README.md new file mode 100644 index 000000000..ae917d2ae --- /dev/null +++ b/examples/mlops/databricks/README.md @@ -0,0 +1,85 @@ +# Databricks Jobs + +`databricks.yml` is an Asset Bundle with two jobs on a single-node CPU job +cluster. Each task is a `spark_python_task` running a job script with +`--job-id {{job.id}} --job-run-id {{job.run_id}}`; the cluster's +`spark_env_vars` configure `jobs.py`. + +- **`kumo_train`** runs `fit_job.py` on the `users`, `items` and `orders` + tables in `.`. `KUMO_REGISTRY=mlflow` logs the fit to + `MLFLOW_EXPERIMENT_NAME`, registers the model as a new version of + `..users_churn` and moves the `champion` alias to it. The + snapshot is exported to the volume. +- **`kumo_score`** runs `score_job.py` daily at 03:00 UTC (created paused). + It loads `models:/..users_churn@champion`, builds a fresh + snapshot and writes `users_churn_predictions` in the connector's schema. Set + `KUMO_SNAPSHOT=snapshots/` to score an exported 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 3) on a + volume and set `kumo_engine_wheel` and `requirements`. Start the + requirements file with `--prefer-binary`: DBR 16.4's setuptools cannot + build the SQLAlchemy sdist MLflow pulls. Each task installs the + requirements, then the wheel. + Use a standard (non-ML) DBR runtime: the ML runtime's preinstalled + torch, torchvision and pyOpenSSL conflict with kumo-engine's pins. + Databricks GPU nodes require a GPU ML runtime, so use a CPU node + type with local NVMe disks (for example m5d.xlarge). +2. Replace the `<...>` placeholders in `databricks.yml` (or pass + `--var catalog=...`). The parent folder of the MLflow experiment + path must already exist in the workspace. +3. Adjust the `dev` target (workspace `root_path`) 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 + +- **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__`). Databricks sets no job-id + environment variable, so ids come from `{{job.id}}`/`{{job.run_id}}`. +- **Dedicated access mode** (`SINGLE_USER`): standard compute runs as a + low-privilege user and has no GPUs. `KUMO_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. +- **One registry format.** A `DatabricksConnector` publish and an MLflow + pyfunc model are different formats; use one per registered model. +- **Job clusters.** Both jobs repeat the cluster block; YAML anchors would + remove the copy but are untested with `databricks bundle validate`. + +## Verified + +The bundle was deployed and both jobs were run on a single-node cluster against +synthetic `users`, `items` and `orders` tables: both exit 0 with the JSON +summary and the Databricks job and run ids, the fit registers a new +`users_churn` version and moves `champion`, and the score job loads it and +writes the predictions table. + +## 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..d0acfeece --- /dev/null +++ b/examples/mlops/databricks/databricks.yml @@ -0,0 +1,136 @@ +# 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, starting with --prefer-binary + default: /Volumes////wheels/requirements.txt + spark_version: + description: a standard (non-ML) runtime, for example 16.4.x-scala2.12 + default: + node_type_id: + description: a CPU node type with local NVMe disks, for example m5d.xlarge + default: + experiment: + description: MLflow experiment (workspace path) for train and score runs + default: /Shared/kumo-engine/users-churn + +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: + KUMO_SOURCE: databricks + KUMO_DATABRICKS_CATALOG: ${var.catalog} + KUMO_DATABRICKS_SCHEMA: ${var.schema} + KUMO_DATABRICKS_WAREHOUSE_ID: ${var.warehouse_id} + KUMO_ARTIFACT_ROOT: /Volumes/${var.catalog}/${var.schema}/${var.volume}/kumo + MLFLOW_TRACKING_URI: databricks + MLFLOW_REGISTRY_URI: databricks-uc + MLFLOW_EXPERIMENT_NAME: ${var.experiment} + KUMO_REGISTRY: mlflow + KUMO_SCRATCH_DIR: /local_disk0/kumo + KUMO_MODEL_NAME: ${var.catalog}.${var.schema}.users_churn + 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: + - --job-id + - "{{job.id}}" + - --job-run-id + - "{{job.run_id}}" + + 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: + KUMO_SOURCE: databricks + KUMO_DATABRICKS_CATALOG: ${var.catalog} + KUMO_DATABRICKS_SCHEMA: ${var.schema} + KUMO_DATABRICKS_WAREHOUSE_ID: ${var.warehouse_id} + KUMO_ARTIFACT_ROOT: /Volumes/${var.catalog}/${var.schema}/${var.volume}/kumo + MLFLOW_TRACKING_URI: databricks + MLFLOW_REGISTRY_URI: databricks-uc + MLFLOW_EXPERIMENT_NAME: ${var.experiment} + KUMO_REGISTRY: mlflow + KUMO_SCRATCH_DIR: /local_disk0/kumo + KUMO_MODEL_NAME: ${var.catalog}.${var.schema}.users_churn + KUMO_PREDICTIONS_TABLE: users_churn_predictions + 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: + - --job-id + - "{{job.id}}" + - --job-run-id + - "{{job.run_id}}" + +targets: + dev: + default: true + workspace: + root_path: /Workspace/Shared/kumo-engine/.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..bc8cd4238 --- /dev/null +++ b/examples/mlops/fit_job.py @@ -0,0 +1,59 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Train and publish the users-churn model: ``python fit_job.py [--job-id ID +--job-run-id ID]``; see README.md.""" + +import os +import sys +from dataclasses import asdict +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 + +# 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' + + +def main(argv: Optional[Sequence[str]] = None) -> None: + job_id, job_run_id = jobs.start(argv) + with kumo.job(job_id=job_id, job_run_id=job_run_id) as run: + connector = jobs.source() + graph = jobs.build_graph(connector) + snapshot = graph.snapshot() + run.record(snapshot_id=snapshot.id) + if os.environ.get('KUMO_ARTIFACT_ROOT'): + run.record(outputs={ + 'snapshot': + jobs.artifact_uri(snapshot.export(jobs.artifacts(), 'snapshots')) + }) + + query = kumo.PredictiveQuery(graph=graph, + query=os.environ.get('KUMO_QUERY', QUERY)) + table = query.generate_training_table(snapshot) + plan = query.suggest_model_plan(snapshot) + if os.environ.get('KUMO_NUM_EXPERIMENTS'): + plan.training_job.num_experiments = int( + os.environ['KUMO_NUM_EXPERIMENTS']) + tracking, export_config = jobs.model_destination(connector) + result = kumo.Trainer(plan).fit(table, snapshot=snapshot, + tracking=tracking, + export_config=export_config) + + registered = getattr(tracking, 'registered', None) + published = registered or result.published_model + model = published and asdict(published) + if model and not os.environ.get('KUMO_REGISTRY'): + model['version'] = None + run.record(bundle_id=result.bundle_id, + run_id=getattr(tracking, 'run_id', None), + outputs={'model': model}) + + +if __name__ == '__main__': + main() diff --git a/examples/mlops/jobs.py b/examples/mlops/jobs.py new file mode 100644 index 000000000..1c2888d25 --- /dev/null +++ b/examples/mlops/jobs.py @@ -0,0 +1,189 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The graph, connectors and model registry the job scripts share. + +Configured by environment variables the job definition sets; credentials come +from the platform identity. + +- ``KUMO_SOURCE``: ``databricks`` (``KUMO_DATABRICKS_CATALOG``/``_SCHEMA``/ + ``_WAREHOUSE_ID``), ``snowflake`` (``KUMO_SNOWFLAKE_WAREHOUSE``/ + ``_DATABASE``/``_SCHEMA``), or ``s3``/``file`` (``KUMO_DATA_ROOT`` holding + ``users``, ``items`` and ``orders`` parquet files). +- ``KUMO_ARTIFACT_ROOT``: ``s3://``, ``gs://`` or a local/volume path that + receives snapshot exports, path-published models and predictions + (optional with a registry and a predictions table). +- ``KUMO_REGISTRY``: ``mlflow`` (register ``KUMO_MODEL_NAME`` through + ``MlflowTracking``), ``snowflake`` (the Snowflake Model Registry, tracked in + Snowflake Experiment Tracking) or unset (fit publishes to + ``KUMO_ARTIFACT_ROOT/runs//model``; scoring loads ``KUMO_MODEL_PATH``, + that path relative to ``KUMO_ARTIFACT_ROOT``). With a registry, fit moves + the ``champion`` alias to the new version and scoring loads ``champion``. +- ``KUMO_SCRATCH_DIR``: local disk for snapshots, tables and bundles. +- ``KUMO_PREDICTIONS_TABLE``: write predictions to this table of the source + connector instead of ``KUMO_ARTIFACT_ROOT/runs//predictions/``. +""" + +import argparse +import os +from typing import Any, Optional, Sequence, Tuple + +from kumoapi.jobs import MetadataField, PredictionArtifactType +from kumoapi.typing import Stype + +from kumo.engine import ( + Edge, + Graph, + ModelOutputConfig, + ModelRegistryLocation, + PredictionOutputConfig, + Table, + Trainer, +) +from kumo.engine.connector import ( + Connector, + DatabricksConnector, + FileConnector, + S3Connector, + SnowflakeConnector, + connector_for_root, +) + +CHAMPION = 'champion' + + +def _env(name: str) -> str: + value = os.environ.get(name) + if not value: + raise ValueError(f'Set the {name} environment variable') + return value + + +def start(argv: Optional[Sequence[str]] = None) -> Tuple[Any, Any]: + """Applies ``KUMO_SCRATCH_DIR``; returns ``--job-id``/``--job-run-id``.""" + from kumo.engine import configure + + parser = argparse.ArgumentParser() + parser.add_argument('--job-id') + parser.add_argument('--job-run-id') + args = parser.parse_args(argv) + if os.environ.get('KUMO_SCRATCH_DIR'): + configure(scratch_dir=os.environ['KUMO_SCRATCH_DIR']) + return args.job_id, args.job_run_id + + +def source() -> Connector: + kind = _env('KUMO_SOURCE') + if kind == 'databricks': + return DatabricksConnector.from_notebook( + catalog=_env('KUMO_DATABRICKS_CATALOG'), + schema_name=_env('KUMO_DATABRICKS_SCHEMA'), + warehouse_id=_env('KUMO_DATABRICKS_WAREHOUSE_ID')) + if kind == 'snowflake': + return SnowflakeConnector.from_spcs_env( + warehouse=_env('KUMO_SNOWFLAKE_WAREHOUSE'), + database=_env('KUMO_SNOWFLAKE_DATABASE'), + schema_name=_env('KUMO_SNOWFLAKE_SCHEMA')) + if kind in ('s3', 'file'): + return connector_for_root(_env('KUMO_DATA_ROOT')) + raise ValueError(f'Unknown KUMO_SOURCE {kind!r}') + + +def artifacts() -> Connector: + return connector_for_root(_env('KUMO_ARTIFACT_ROOT')) + + +def artifact_uri(path: str) -> str: + """``path`` under ``KUMO_ARTIFACT_ROOT`` as a full URI or path.""" + return f"{_env('KUMO_ARTIFACT_ROOT').rstrip('/')}/{path}" + + +def _run_path() -> str: + run = _env('KUMO_JOB_RUN_ID').rstrip('/').rsplit('/', 1)[-1] + return f'runs/{run}' + + +def run_artifacts() -> Connector: + """``KUMO_ARTIFACT_ROOT/runs/``.""" + return connector_for_root(artifact_uri(_run_path())) + + +def build_graph(connector: Connector) -> Graph: + + def table(alias: str, **kwargs: str) -> Table: + 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) + return Graph( + tables={ + 'users': table('users', primary_key='USER_ID'), + 'items': table('items', primary_key='ITEM_ID'), + 'orders': orders, + }, edges=[ + Edge(src_table='orders', fkey='USER_ID', dst_table='users'), + Edge(src_table='orders', fkey='ITEM_ID', dst_table='items'), + ]) + + +def model_destination( + connector: Connector) -> Tuple[Any, Optional[ModelOutputConfig]]: + """The ``tracking`` and ``export_config`` that publish a fit.""" + registry = os.environ.get('KUMO_REGISTRY') + if registry == 'mlflow': + from kumo.engine.integrations.mlflow import MlflowTracking + + return MlflowTracking(registered_model_name=_env('KUMO_MODEL_NAME'), + aliases=[CHAMPION]), None + if registry == 'snowflake': + from kumo.engine.integrations.snowflake import SnowflakeTracking + + assert isinstance(connector, SnowflakeConnector) + name = _env('KUMO_MODEL_NAME') + return (SnowflakeTracking(connector.session(), experiment_name=name), + ModelOutputConfig(connector, location=ModelRegistryLocation( + name, alias=CHAMPION))) + return None, ModelOutputConfig(run_artifacts(), path='model') + + +def load_champion(connector: Connector) -> Trainer: + registry = os.environ.get('KUMO_REGISTRY') + if registry == 'mlflow': + from kumo.engine.integrations.mlflow import load_model + + return load_model(f"models:/{_env('KUMO_MODEL_NAME')}@{CHAMPION}") + if registry == 'snowflake': + return Trainer.load_from_connector( + connector, location=ModelRegistryLocation( + _env('KUMO_MODEL_NAME'), alias=CHAMPION)) + return Trainer.load_from_connector(artifacts(), + path=_env('KUMO_MODEL_PATH')) + + +def prediction_output(connector: Connector) -> PredictionOutputConfig: + fields = [MetadataField.ANCHOR_TIMESTAMP] + table = os.environ.get('KUMO_PREDICTIONS_TABLE') + if table: + return PredictionOutputConfig( + connector, artifact_type=PredictionArtifactType.PREDICTIONS, + table_name=table, extra_fields=fields) + return PredictionOutputConfig(run_artifacts(), + output_paths={'predictions': 'predictions/'}, + extra_fields=fields) + + +def prediction_destinations(output: PredictionOutputConfig) -> dict: + """Where ``output`` writes: full URIs, or the table name.""" + if output.output_paths is None: + return output.destinations() + return { + name: artifact_uri(f'{_run_path()}/{path}') + for name, path in output.output_paths.items() + } diff --git a/examples/mlops/sagemaker/README.md b/examples/mlops/sagemaker/README.md new file mode 100644 index 000000000..7ec1d7a40 --- /dev/null +++ b/examples/mlops/sagemaker/README.md @@ -0,0 +1,79 @@ +# 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 `KUMO_DATA_ROOT`. + `KUMO_REGISTRY=mlflow` logs the fit to the managed MLflow tracking server, + registers the model as a new version of `users-churn` and moves the + `champion` alias to it. +2. **`Predict`** (processing job, after `Fit`) runs `python score_job.py`, + loading `models:/users-churn@champion` and writing + `/runs//predictions/`. + +To score on a schedule without retraining, put `Predict` in its own pipeline +with a `PipelineSchedule` trigger (`pipeline.put_triggers`). + +## Before running + +1. Build an ECR image with `kumo-engine[mlflow,s3]`, its dependency wheels + and `sagemaker-mlflow` (plus the CUDA stack for GPU instance types). +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 \ + --instance-type --start + ``` + 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 + +- **Job ids.** The training step passes the pipeline name and execution ARN + as `KUMO_JOB_ID`/`KUMO_JOB_RUN_ID` (`container_arguments` takes only plain + strings); the processing step passes them as `--job-id`/`--job-run-id`. +- **MLflow.** `MLFLOW_TRACKING_URI` is the tracking server or MLflow App ARN. + Managed MLflow servers run MLflow 3.0; the engine needs a newer client + (3.11 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. +- **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 + +The 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. An MLflow App ARN as tracking URI is unverified. + +## 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..992520274 --- /dev/null +++ b/examples/mlops/sagemaker/pipeline.py @@ -0,0 +1,111 @@ +# 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, Dict, 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, + instance_type: str, +) -> Any: + from sagemaker.estimator import Estimator + from sagemaker.inputs import TrainingInput + from sagemaker.processing import ProcessingInput, Processor + from sagemaker.workflow.execution_variables import ExecutionVariables + from sagemaker.workflow.pipeline import Pipeline + from sagemaker.workflow.pipeline_context import PipelineSession + from sagemaker.workflow.steps import ProcessingStep, TrainingStep + + session = PipelineSession() + env: Dict[str, Any] = { + 'KUMO_SOURCE': 's3', + 'KUMO_DATA_ROOT': data_root, + 'KUMO_ARTIFACT_ROOT': artifact_root, + 'MLFLOW_TRACKING_URI': tracking_arn, + 'MLFLOW_EXPERIMENT_NAME': experiment, + 'KUMO_REGISTRY': 'mlflow', + 'KUMO_MODEL_NAME': 'users-churn', + } + + estimator = Estimator( + image_uri=image, + role=role, + instance_count=1, + instance_type=instance_type, + container_entry_point=['python', f'{TRAIN_CODE}/fit_job.py'], + environment={ + **env, + 'KUMO_SCRATCH_DIR': '/tmp/kumo', + 'KUMO_JOB_ID': ExecutionVariables.PIPELINE_NAME, + 'KUMO_JOB_RUN_ID': ExecutionVariables.PIPELINE_EXECUTION_ARN, + }, + 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=instance_type, + entrypoint=['python', f'{PROCESS_CODE}/score_job.py'], + env={ + **env, 'KUMO_SCRATCH_DIR': '/opt/ml/processing/scratch' + }, + sagemaker_session=session, + ) + predict = ProcessingStep( + name='Predict', + step_args=processor.run( + inputs=[ProcessingInput(source=code, destination=PROCESS_CODE)], + arguments=[ + '--job-id', + ExecutionVariables.PIPELINE_NAME, + '--job-run-id', + ExecutionVariables.PIPELINE_EXECUTION_ARN, + ], + ), + 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('--instance-type', required=True) + 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..bcb8c5311 --- /dev/null +++ b/examples/mlops/score_job.py @@ -0,0 +1,44 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Score with the champion model: ``python score_job.py [--job-id ID +--job-run-id ID]``; see README.md. ``KUMO_SNAPSHOT`` (a path under +``KUMO_ARTIFACT_ROOT``) scores an exported snapshot instead of fresh data.""" + +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) -> None: + job_id, job_run_id = jobs.start(argv) + with kumo.job(job_id=job_id, job_run_id=job_run_id) as run: + connector = jobs.source() + trainer = jobs.load_champion(connector) + run.record(bundle_id=trainer.bundle_id) + graph = jobs.build_graph(connector) + exported = os.environ.get('KUMO_SNAPSHOT') + snapshot = (kumo.GraphSnapshot.load_from_connector( + jobs.artifacts(), path=exported, graph=graph) + if exported else graph.snapshot()) + run.record(snapshot_id=snapshot.id) + + output = jobs.prediction_output(connector) + defaults = trainer.scoring_defaults.binary_classification_threshold + result = trainer.predict_snapshot( + snapshot, output_config=output, + binary_classification_threshold=None + if defaults is not None else 0.5) + run.record(run_id=result.selection.job_id, + outputs=jobs.prediction_destinations(output)) + + +if __name__ == '__main__': + main() diff --git a/examples/mlops/snowflake/README.md b/examples/mlops/snowflake/README.md new file mode 100644 index 000000000..07c2f1e03 --- /dev/null +++ b/examples/mlops/snowflake/README.md @@ -0,0 +1,89 @@ +# 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 from `task.sql`. + +- **`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 (`KUMO_REGISTRY=snowflake`). 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, publish without the alias in +`jobs.model_destination` and move `champion` yourself with +`kumo.engine.set_model_alias`. +If moving the alias fails, the fit still publishes, warns, and scoring keeps +the previous `champion`. + +Each Task body reads `SYSTEM$TASK_RUNTIME_INFO` for the task name and graph +run id, passes them as `--job-id`/`--job-run-id`, and runs the job +synchronously with the statement `MLJobDefinition.to_sql(use_async=False)` +generated. `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 ``. +3. Run: + ```bash + python examples/mlops/snowflake/register_jobs.py --connection \ + --compute-pool --stage --database \ + --schema --warehouse \ + --requirement 'kumo-engine[snowflake]' --external-access-integration + ``` +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). +- Without job-id flags the engine records `SNOWFLAKE_SERVICE_NAME` as + `job_run_id` and leaves `job_id` null. +- The Task's synchronous `CALL` holds the task warehouse while the job runs; + a serverless task avoids that. +- `KUMO_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, the `$$`-quoted Task body through +`session.sql`, 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 +- `SYSTEM$TASK_RUNTIME_INFO`: + https://docs.snowflake.com/en/sql-reference/functions/system_task_runtime_info +- `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..32f6f59c8 --- /dev/null +++ b/examples/mlops/snowflake/register_jobs.py @@ -0,0 +1,113 @@ +# 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 + +HERE = Path(__file__).resolve().parent +PAYLOAD = HERE.parent +JOB_IDS = ['--job-id', '__KUMO_JOB_ID__', '--job-run-id', '__KUMO_JOB_RUN_ID__'] +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: + text = (HERE / 'task.sql').read_text() + text = text[text.index('CREATE OR REPLACE TASK'):] + if '$$' in execute_ml_job_sql: + raise ValueError('the job SQL cannot be embedded in a $$ string') + for slot, value in (('', task_name), ('', warehouse), + ('', schedule), + ('', str(timeout_ms)), + ('', execute_ml_job_sql)): + text = text.replace(slot, value) + return text + + +def register(session: Any, *, compute_pool: str, stage: str, database: str, + schema: str, warehouse: str, requirements: List[str], + external_access_integrations: List[str]) -> Dict[str, str]: + from snowflake.ml.jobs import MLJobDefinition + + env = { + 'KUMO_SOURCE': 'snowflake', + 'KUMO_SNOWFLAKE_WAREHOUSE': warehouse, + 'KUMO_SNOWFLAKE_DATABASE': database, + 'KUMO_SNOWFLAKE_SCHEMA': schema, + 'KUMO_REGISTRY': 'snowflake', + 'KUMO_MODEL_NAME': 'USERS_CHURN', + 'KUMO_PREDICTIONS_TABLE': 'USERS_CHURN_PREDICTIONS', + 'KUMO_SCRATCH_DIR': '/tmp/kumo', + } + 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, + env_vars=env, + 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_IDS, + 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'): + parser.add_argument(flag, required=True) + parser.add_argument('--requirement', action='append', required=True) + parser.add_argument('--external-access-integration', action='append', + default=[]) + 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, + requirements=args.requirement, + external_access_integrations=args.external_access_integration) + for kind, name in tasks.items(): + print(f'{kind}: {name} (suspended; ALTER TASK {name} RESUME)') + + +if __name__ == '__main__': + main() diff --git a/examples/mlops/snowflake/task.sql b/examples/mlops/snowflake/task.sql new file mode 100644 index 000000000..48e261caf --- /dev/null +++ b/examples/mlops/snowflake/task.sql @@ -0,0 +1,25 @@ +-- SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +-- All rights reserved. +-- SPDX-License-Identifier: Apache-2.0 + +-- register_jobs.py fills the <...> slots once per job; the block swaps this +-- run's task name and graph run id into the job arguments. +CREATE OR REPLACE TASK + WAREHOUSE = + SCHEDULE = '' + USER_TASK_TIMEOUT_MS = + SUSPEND_TASK_AFTER_NUM_FAILURES = 3 +AS +DECLARE + job_sql STRING DEFAULT $$$$; + task_name STRING; + run_id STRING; +BEGIN + SELECT SYSTEM$TASK_RUNTIME_INFO('CURRENT_TASK_NAME') INTO :task_name; + SELECT SYSTEM$TASK_RUNTIME_INFO('CURRENT_TASK_GRAPH_RUN_GROUP_ID') + INTO :run_id; + job_sql := REPLACE(REPLACE(job_sql, '__KUMO_JOB_ID__', task_name), + '__KUMO_JOB_RUN_ID__', run_id); + EXECUTE IMMEDIATE :job_sql; + RETURN run_id; +END; diff --git a/packages/diskgraph/Cargo.lock b/packages/diskgraph/Cargo.lock index be75fbfdd..233fec7c1 100644 --- a/packages/diskgraph/Cargo.lock +++ b/packages/diskgraph/Cargo.lock @@ -1739,6 +1739,7 @@ dependencies = [ "bytes", "chrono", "clap", + "crc32fast", "crossterm", "dashmap", "dirs", diff --git a/packages/diskgraph/diskgraph-ingest/Cargo.toml b/packages/diskgraph/diskgraph-ingest/Cargo.toml index 1f29515e3..d068650d2 100644 --- a/packages/diskgraph/diskgraph-ingest/Cargo.toml +++ b/packages/diskgraph/diskgraph-ingest/Cargo.toml @@ -55,6 +55,7 @@ epserde = { workspace = true } memmap2 = "0.9" libc = "0.2" rayon = "1" +crc32fast = "1.5" jemallocator = { version = "0.5", optional = true, features = ["profiling"] } jemalloc-ctl = { version = "0.5", optional = true } sysinfo = "0.30" diff --git a/packages/diskgraph/diskgraph-ingest/src/snapshot_mirror.rs b/packages/diskgraph/diskgraph-ingest/src/snapshot_mirror.rs index 7ea4edbcc..bfba23cfb 100644 --- a/packages/diskgraph/diskgraph-ingest/src/snapshot_mirror.rs +++ b/packages/diskgraph/diskgraph-ingest/src/snapshot_mirror.rs @@ -36,6 +36,7 @@ use anyhow::{Context, Result}; use object_store::path::Path as ObjectPath; use parking_lot::Mutex; use serde::{Deserialize, Serialize}; +use std::collections::BTreeMap; use diskgraph_fs::{CloudTransfer, CloudTransferTicket}; @@ -488,21 +489,9 @@ pub(crate) fn validate_manifest_rel_path(rel: &str) -> Result<(), &'static str> Ok(()) } -/// Verify the snapshot rooted at `output_dir` against its `_MANIFEST.json`. -/// Used by the consumer side after a Python-driven dataplane→local copy -/// completes, to catch lifecycle-mid-delete races and partial copies -/// before the snapshot is handed to a workload. -/// -/// Pass on success; [`CorruptionError`] on any deviation from the -/// manifest contract. Errors are precise about what failed so the -/// caller's alerting carries actionable detail (which file is missing, -/// expected vs actual size, etc.). -pub fn verify_snapshot_integrity(output_dir: &std::path::Path) -> std::result::Result<(), CorruptionError> { - let success_path = output_dir.join(SUCCESS_MARKER_FILENAME); - if !success_path.exists() { - return Err(CorruptionError::MissingSuccess(success_path)); - } - +/// Read and validate `_MANIFEST.json` at `output_dir`: schema version and +/// every rel_path's confinement to the snapshot root. +fn read_manifest(output_dir: &std::path::Path) -> std::result::Result { let manifest_path = output_dir.join(MANIFEST_FILENAME); let manifest_bytes = match std::fs::read(&manifest_path) { Ok(b) => b, @@ -548,6 +537,25 @@ pub fn verify_snapshot_integrity(output_dir: &std::path::Path) -> std::result::R }); } } + Ok(manifest) +} + +/// Verify the snapshot rooted at `output_dir` against its `_MANIFEST.json`. +/// Used by the consumer side after a Python-driven dataplane→local copy +/// completes, to catch lifecycle-mid-delete races and partial copies +/// before the snapshot is handed to a workload. +/// +/// Pass on success; [`CorruptionError`] on any deviation from the +/// manifest contract. Errors are precise about what failed so the +/// caller's alerting carries actionable detail (which file is missing, +/// expected vs actual size, etc.). +pub fn verify_snapshot_integrity(output_dir: &std::path::Path) -> std::result::Result<(), CorruptionError> { + let success_path = output_dir.join(SUCCESS_MARKER_FILENAME); + if !success_path.exists() { + return Err(CorruptionError::MissingSuccess(success_path)); + } + + let manifest = read_manifest(output_dir)?; let mut missing: Vec = Vec::new(); let mut size_mismatches: Vec<(String, u64, u64)> = Vec::new(); @@ -584,6 +592,71 @@ pub fn verify_snapshot_integrity(output_dir: &std::path::Path) -> std::result::R Ok(()) } +/// CRC-32 of every `range_bytes` span of each file `_MANIFEST.json` lists, +/// and of `_MANIFEST.json` itself, keyed by rel_path (an empty file has one +/// zero entry). Lets a copy of the snapshot be checked for content as well +/// as size; the manifest format is unchanged. +pub fn snapshot_range_checksums( + output_dir: &std::path::Path, + range_bytes: u64, +) -> std::result::Result>, CorruptionError> { + use rayon::prelude::*; + use std::os::unix::fs::FileExt; + + const CHUNK: u64 = 1 << 20; + if range_bytes == 0 { + return Err(CorruptionError::Io { + path: output_dir.to_path_buf(), + detail: "range_bytes must be positive".into(), + }); + } + let mut files = read_manifest(output_dir)?.files; + let manifest_size = std::fs::metadata(output_dir.join(MANIFEST_FILENAME)) + .map(|meta| meta.len()) + .unwrap_or(0); + files.push(ManifestEntry { + rel_path: MANIFEST_FILENAME.into(), + size_bytes: manifest_size, + }); + let ranges: Vec<(usize, u64)> = files + .iter() + .enumerate() + .flat_map(|(index, entry)| { + let count = entry.size_bytes.div_ceil(range_bytes).max(1); + (0..count).map(move |range| (index, range)) + }) + .collect(); + let crcs = ranges + .par_iter() + .map(|&(index, range)| { + let entry = &files[index]; + let path = output_dir.join(&entry.rel_path); + let io = |e: std::io::Error| CorruptionError::Io { + path: path.clone(), + detail: e.to_string(), + }; + let file = std::fs::File::open(&path).map_err(io)?; + let start = range * range_bytes; + let end = entry.size_bytes.min(start + range_bytes); + let mut hasher = crc32fast::Hasher::new(); + let mut buffer = vec![0u8; CHUNK.min(end - start) as usize]; + let mut offset = start; + while offset < end { + let length = CHUNK.min(end - offset) as usize; + file.read_exact_at(&mut buffer[..length], offset).map_err(io)?; + hasher.update(&buffer[..length]); + offset += length as u64; + } + Ok(hasher.finalize()) + }) + .collect::, CorruptionError>>()?; + let mut out: BTreeMap> = BTreeMap::new(); + for (&(index, _), crc) in ranges.iter().zip(crcs) { + out.entry(files[index].rel_path.clone()).or_default().push(crc); + } + Ok(out) +} + #[cfg(test)] mod tests { use super::*; @@ -933,4 +1006,35 @@ mod tests { verify_snapshot_integrity(local_root.path())?; Ok(()) } + + #[test] + fn range_checksums_cover_every_listed_file_and_the_manifest() { + let tmp = tempfile::tempdir().unwrap(); + let big: Vec = (0..10u32).flat_map(|i| i.to_le_bytes()).collect(); + write_marker_files(tmp.path(), &[("a.bin", &big), ("nested/empty", b"")]); + write_manifest( + tmp.path(), + vec![ + ManifestEntry { + rel_path: "a.bin".into(), + size_bytes: 40, + }, + ManifestEntry { + rel_path: "nested/empty".into(), + size_bytes: 0, + }, + ], + ); + let crcs = snapshot_range_checksums(tmp.path(), 16).unwrap(); + let manifest = std::fs::read(tmp.path().join(MANIFEST_FILENAME)).unwrap(); + assert_eq!(crcs["a.bin"], [&big[..16], &big[16..32], &big[32..]].map(crc32fast::hash)); + assert_eq!(crcs["nested/empty"], vec![0]); + assert_eq!(crcs[MANIFEST_FILENAME], manifest.chunks(16).map(crc32fast::hash).collect::>()); + + std::fs::write(tmp.path().join("a.bin"), &big[..20]).unwrap(); + assert!(matches!( + snapshot_range_checksums(tmp.path(), 16).unwrap_err(), + CorruptionError::Io { .. } + )); + } } diff --git a/packages/diskgraph/diskgraph-pybind/build.rs b/packages/diskgraph/diskgraph-pybind/build.rs index 9bd833888..9b3c22cf0 100644 --- a/packages/diskgraph/diskgraph-pybind/build.rs +++ b/packages/diskgraph/diskgraph-pybind/build.rs @@ -9,8 +9,37 @@ mod native_cpu_guard { )); } +fn build_commit() -> String { + println!("cargo:rerun-if-env-changed=DISKGRAPH_BUILD_COMMIT"); + if let Ok(commit) = std::env::var("DISKGRAPH_BUILD_COMMIT") { + return commit; + } + let git = |args: &[&str]| { + std::process::Command::new("git") + .args(args) + .output() + .ok() + .filter(|out| out.status.success()) + .map(|out| String::from_utf8_lossy(&out.stdout).trim().to_owned()) + .filter(|out| !out.is_empty()) + }; + for name in ["HEAD", "packed-refs"] { + if let Some(path) = git(&["rev-parse", "--git-path", name]) { + println!("cargo:rerun-if-changed={path}"); + } + } + if let Some(head) = git(&["symbolic-ref", "-q", "HEAD"]) { + if let Some(path) = git(&["rev-parse", "--git-path", &head]) { + println!("cargo:rerun-if-changed={path}"); + } + } + git(&["rev-parse", "--short=12", "HEAD"]).unwrap_or_else(|| "unknown".into()) +} + fn main() { pyo3_build_config::use_pyo3_cfgs(); + println!("cargo:rerun-if-changed=build.rs"); + println!("cargo:rustc-env=DISKGRAPH_BUILD_COMMIT={}", build_commit()); if std::env::var_os("CARGO_FEATURE_EXTENSION_MODULE").is_some() { native_cpu_guard::reject_native_cpu_build("diskgraph Python extension"); diff --git a/packages/diskgraph/diskgraph-pybind/python/diskgraph/__init__.py b/packages/diskgraph/diskgraph-pybind/python/diskgraph/__init__.py index d157a5b0c..5b9f5d0b6 100644 --- a/packages/diskgraph/diskgraph-pybind/python/diskgraph/__init__.py +++ b/packages/diskgraph/diskgraph-pybind/python/diskgraph/__init__.py @@ -3,6 +3,8 @@ # SPDX-License-Identifier: Apache-2.0 from ._c import ( + __version__, + BUILD_COMMIT, export_tt_artifacts, convert_kumo_config, ingest, @@ -38,6 +40,7 @@ write_success_marker, touch_success_marker, verify_snapshot_integrity, + snapshot_range_checksums, has_resumable_ingest, validated_index_resume_progress, compute_col_stats_for_parquet, @@ -139,6 +142,7 @@ def __getattr__(name): "write_success_marker", "touch_success_marker", "verify_snapshot_integrity", + "snapshot_range_checksums", "has_resumable_ingest", "validated_index_resume_progress", "compute_col_stats_for_parquet", diff --git a/packages/diskgraph/diskgraph-pybind/python/tests/test_exception_exports.py b/packages/diskgraph/diskgraph-pybind/python/tests/test_exception_exports.py index 7744976d2..d4b29806e 100644 --- a/packages/diskgraph/diskgraph-pybind/python/tests/test_exception_exports.py +++ b/packages/diskgraph/diskgraph-pybind/python/tests/test_exception_exports.py @@ -11,3 +11,27 @@ def test_public_corruption_error_catches_snapshot_verifier(tmp_path): with pytest.raises(diskgraph.DiskgraphCorruptionError): diskgraph.verify_snapshot_integrity(str(tmp_path)) + + +def test_build_identity_is_exported(): + assert diskgraph.__version__ + assert diskgraph.BUILD_COMMIT + + +def test_snapshot_range_checksums_match_zlib(tmp_path): + import json + import zlib + + payload = bytes(range(256)) * 5 + (tmp_path / "a.bin").write_bytes(payload) + (tmp_path / "_MANIFEST.json").write_text(json.dumps({ + "schema_version": 1, + "files": [{"rel_path": "a.bin", "size_bytes": len(payload)}], + })) + crcs = diskgraph.snapshot_range_checksums(str(tmp_path), 512) + assert crcs["a.bin"] == [zlib.crc32(payload[i:i + 512]) + for i in range(0, len(payload), 512)] + assert crcs["_MANIFEST.json"] == [ + zlib.crc32((tmp_path / "_MANIFEST.json").read_bytes())] + with pytest.raises(diskgraph.DiskgraphCorruptionError): + diskgraph.snapshot_range_checksums(str(tmp_path / "missing"), 512) diff --git a/packages/diskgraph/diskgraph-pybind/src/lib.rs b/packages/diskgraph/diskgraph-pybind/src/lib.rs index c907055d1..a47daef4d 100644 --- a/packages/diskgraph/diskgraph-pybind/src/lib.rs +++ b/packages/diskgraph/diskgraph-pybind/src/lib.rs @@ -1553,6 +1553,21 @@ fn verify_snapshot_integrity(py: Python<'_>, output_dir: &str) -> PyResult<()> { .map_err(|e| DiskgraphCorruptionError::new_err(e.to_string())) } +/// CRC-32 of every `range_bytes` span of each file in the snapshot's +/// `_MANIFEST.json` and of the manifest itself, keyed by relative path. +#[pyfunction] +fn snapshot_range_checksums( + py: Python<'_>, + output_dir: &str, + range_bytes: u64, +) -> PyResult>> { + let output_dir = std::path::PathBuf::from(output_dir); + py.detach(|| { + diskgraph_ingest::snapshot_mirror::snapshot_range_checksums(&output_dir, range_bytes) + }) + .map_err(|e| DiskgraphCorruptionError::new_err(e.to_string())) +} + #[pyfunction] fn has_resumable_ingest(output_dir: &str) -> PyResult { diskgraph_ingest::initial_resume_phase(std::path::Path::new(output_dir)) @@ -1593,6 +1608,8 @@ fn diskgraph_pybind(m: &Bound<'_, PyModule>) -> PyResult<()> { "DiskgraphCorruptionError", m.py().get_type::(), )?; + m.add("__version__", env!("CARGO_PKG_VERSION"))?; + m.add("BUILD_COMMIT", env!("DISKGRAPH_BUILD_COMMIT"))?; m.add_function(wrap_pyfunction!(export_tt_artifacts, m)?)?; m.add_function(wrap_pyfunction!(convert_kumo_config, m)?)?; m.add_function(wrap_pyfunction!(ingest, m)?)?; @@ -1606,6 +1623,7 @@ fn diskgraph_pybind(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_function(wrap_pyfunction!(success_marker::write_success_marker, m)?)?; m.add_function(wrap_pyfunction!(success_marker::touch_success_marker, m)?)?; m.add_function(wrap_pyfunction!(verify_snapshot_integrity, m)?)?; + m.add_function(wrap_pyfunction!(snapshot_range_checksums, m)?)?; 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)?)?; diff --git a/packages/diskgraph/third_party/artifact-manifest.json b/packages/diskgraph/third_party/artifact-manifest.json index 1f62d43ef..d22f0bc33 100644 --- a/packages/diskgraph/third_party/artifact-manifest.json +++ b/packages/diskgraph/third_party/artifact-manifest.json @@ -3241,6 +3241,13 @@ "source_kind": "registry-or-git", "version": "0.3.0" }, + { + "license": "MIT OR Apache-2.0", + "name": "crc32fast", + "source": "registry+https://github.com/rust-lang/crates.io-index", + "source_kind": "registry-or-git", + "version": "1.5.0" + }, { "license": "MIT OR Apache-2.0", "name": "crossbeam-deque", @@ -5202,6 +5209,6 @@ ] } ], - "cargo_lock_sha256": "297851eb08303a13a97f091c68cf3ce5ccac633ad3d7fe09724c1506e486cdaf", + "cargo_lock_sha256": "a1c319a5eb83c1a60f056c213932c12e385f0cff468da4db0021947f11b7aa07", "schema": 1 } diff --git a/packages/kumo-api/kumoapi/plugin.py b/packages/kumo-api/kumoapi/plugin.py index 2bb9dd998..013734e6f 100644 --- a/packages/kumo-api/kumoapi/plugin.py +++ b/packages/kumo-api/kumoapi/plugin.py @@ -21,10 +21,12 @@ :class:`~kumoai_studio.backend.LocalBackend`. """ +import re from dataclasses import dataclass, field from typing import ( Any, Callable, + Dict, List, Mapping, Optional, @@ -41,9 +43,13 @@ from kumoapi.typing import Dtype __all__ = [ + 'BUNDLE_ID_KEY', + 'ENGINE_MARKER', + 'REGISTRY_KEY_PREFIX', 'ArtifactLocation', 'ArtifactSink', 'IngestResult', + 'ModelAliasingSink', 'ModelFetchingSink', 'ModelListingSink', 'ModelPublishingSink', @@ -54,8 +60,20 @@ 'SourceBackend', 'SourceConnector', 'SourceCredentials', + 'check_model_alias', + 'registry_lineage', ] +#: Prefix of the keys an engine export records in registry metadata. +REGISTRY_KEY_PREFIX = 'kumo_' +#: Registry metadata key marking a version as a Kumo engine bundle. +ENGINE_MARKER = 'kumo_engine' +#: Registry metadata key holding the version's bundle id; the bundle's own +#: provenance.json holds its lineage. +BUNDLE_ID_KEY = 'kumo_bundle_id' +_ALIAS = re.compile(r'[A-Za-z_][A-Za-z0-9_]{0,254}') +_RESERVED_ALIAS = re.compile(r'DEFAULT|FIRST|LAST|LATEST|V[0-9]+', re.I) + @dataclass(frozen=True) class SourceCredentials: @@ -147,13 +165,16 @@ class ModelRegistryLocation: Args: model_name: Registry-local model name. version: Exact version to write or read. ``None`` lets the registry - assign the next one on export, and selects the latest on import. - An assigned version comes back on the export result. + assign the next one on export; on import it selects Snowflake's + default or Databricks' latest Kumo version. comment: Optional description recorded on the version. + alias: Set on the new version after export; on import, the version + it points at. Exclusive with :obj:`version`. """ model_name: str version: Optional[str] = None comment: Optional[str] = None + alias: Optional[str] = None def __post_init__(self) -> None: if (not isinstance(self.model_name, str) or not self.model_name @@ -164,6 +185,10 @@ def __post_init__(self) -> None: raise ValueError( 'version must be a non-empty version, or None to let the ' 'registry choose') + if self.alias is not None and (not self.alias + or self.version is not None): + raise ValueError('alias must be a non-empty alias, and is ' + 'mutually exclusive with version') @dataclass(frozen=True) @@ -687,7 +712,8 @@ def fetch_model( local_dir: Directory to place the bundle in. Created if absent. path: Source for a path-addressed target. location: A registry location. A version of ``None`` selects the - latest, which is resolved here and reported back. + registry's choice (see :class:`ModelRegistryLocation`), which + is resolved here and reported back. project_id: Authorizing project, where one applies. Returns: @@ -697,6 +723,44 @@ def fetch_model( ... +@runtime_checkable +class ModelAliasingSink(ArtifactSink, Protocol): + r"""A registry sink that sets aliases and fetches by + :attr:`ModelRegistryLocation.alias`.""" + def set_model_alias( + self, + connector: SourceConnector, + location: ModelRegistryLocation, + alias: str, + ) -> None: + r"""Points :obj:`alias` at ``location.version``, moving it if held.""" + ... + + +def check_model_alias(alias: str) -> None: + r"""Raises :class:`ValueError` unless every model registry can set + :obj:`alias`: a letter or underscore, then letters, digits or underscores + (at most 255), and not ``DEFAULT``, ``FIRST``, ``LAST``, ``LATEST`` or + ``v`` in any case.""" + if (not isinstance(alias, str) or not _ALIAS.fullmatch(alias) + or _RESERVED_ALIAS.fullmatch(alias)): + raise ValueError( + f'{alias!r} cannot be set as a model alias: use letters, digits ' + f'and underscores, starting with a letter or underscore, and ' + f'avoid the reserved DEFAULT, FIRST, LAST, LATEST and v ' + f'(for DEFAULT use ALTER MODEL ... SET DEFAULT_VERSION).') + + +def registry_lineage(recorded: Mapping[str, Any]) -> Dict[str, str]: + r"""The ``kumo_*`` entries of a version's registry metadata, such as + :data:`BUNDLE_ID_KEY`, as strings and without :data:`ENGINE_MARKER`.""" + return { + key: str(value) + for key, value in recorded.items() + if key.startswith(REGISTRY_KEY_PREFIX) and key != ENGINE_MARKER + } + + @runtime_checkable class ModelListingSink(ArtifactSink, Protocol): r"""A sink that can enumerate the models published through it. diff --git a/packages/kumo-api/test/test_plugin_registry.py b/packages/kumo-api/test/test_plugin_registry.py new file mode 100644 index 000000000..52ca7a69f --- /dev/null +++ b/packages/kumo-api/test/test_plugin_registry.py @@ -0,0 +1,33 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +import pytest + +from kumoapi.plugin import check_model_alias, registry_lineage + + +@pytest.mark.parametrize('alias', [ + '', ' ', 'default', 'LAST', 'Latest', 'v3', 'my-alias', '2nd', 'a' * 256, + None +]) +def test_check_model_alias_refuses_what_a_registry_reserves(alias): + with pytest.raises(ValueError): + check_model_alias(alias) + + +def test_check_model_alias_accepts_a_plain_name(): + check_model_alias('champion') + check_model_alias('_v2_prod') + + +def test_registry_lineage_keeps_kumo_keys_as_strings(): + recorded = { + 'kumo_engine': '1', + 'kumo_bundle_id': 'b', + 'auc': 0.9, + 'kumo_seconds': 3 + } + assert registry_lineage(recorded) == { + 'kumo_bundle_id': 'b', + 'kumo_seconds': '3' + } diff --git a/packages/kumo-bigquery/src/kumo_bigquery/identity.py b/packages/kumo-bigquery/src/kumo_bigquery/identity.py index dede046ed..d22d8aada 100644 --- a/packages/kumo-bigquery/src/kumo_bigquery/identity.py +++ b/packages/kumo-bigquery/src/kumo_bigquery/identity.py @@ -6,7 +6,11 @@ from typing import Any -from kumo_core.sources.identity import epoch_second, hour_version +from kumo_core.sources.identity import ( + SourceVersion, + epoch_second, + hour_version, +) LOGICAL_VIEW = "VIEW" UNDATED_VERSION = "-1" @@ -23,8 +27,8 @@ def stamp(table: Any) -> str: if is_logical_view(table): return hour_version() if table.modified is None: - return UNDATED_VERSION - return epoch_second(table.modified) + return SourceVersion(UNDATED_VERSION) + return SourceVersion(epoch_second(table.modified), "timestamp") def data_version(client: Any, fqn: str) -> str: diff --git a/packages/kumo-bigquery/src/kumo_bigquery/session.py b/packages/kumo-bigquery/src/kumo_bigquery/session.py index 36dada9c8..9e5283ab6 100644 --- a/packages/kumo-bigquery/src/kumo_bigquery/session.py +++ b/packages/kumo-bigquery/src/kumo_bigquery/session.py @@ -278,6 +278,11 @@ def _credentials(self, connector: SourceConnector) -> Credentials: _require_credentials(credentials) return credentials + def table_location(self, connector: SourceConnector, + table_name: str) -> str: + """Fully qualified name of ``table_name`` in the connector's target.""" + return _name(connector, table_name) + def _client(self, connector: SourceConnector) -> Any: return _clients.query_client( self._principal(connector), project=_billing_project(connector), diff --git a/packages/kumo-bigquery/tests/test_identity.py b/packages/kumo-bigquery/tests/test_identity.py index 28d98fc1d..82f74bbc2 100644 --- a/packages/kumo-bigquery/tests/test_identity.py +++ b/packages/kumo-bigquery/tests/test_identity.py @@ -20,6 +20,7 @@ def stamp(**overrides): def test_a_physical_table_is_dated_by_the_time_it_was_modified(): """A version carrying `.0` is a different path segment and a different key.""" assert stamp() == MODIFIED_VERSION + assert stamp().kind == "timestamp" assert stamp(modified=MODIFIED.replace( microsecond=900000)) == (MODIFIED_VERSION) assert stamp( @@ -28,6 +29,7 @@ def test_a_physical_table_is_dated_by_the_time_it_was_modified(): def test_a_logical_view_falls_back_to_the_hour_window(): assert stamp(table_type="VIEW") == hour_version() + assert stamp(table_type="VIEW").kind == "hour_bucket" @pytest.mark.parametrize("table_type", @@ -39,6 +41,7 @@ def test_only_the_exact_view_type_falls_back(table_type): def test_a_table_without_a_modification_time_is_reported_undated(): assert stamp(modified=None) == "-1" + assert stamp(modified=None).kind == "unknown" def test_the_stamp_reads_the_catalog_and_scans_nothing(): diff --git a/packages/kumo-core/src/kumo_core/artifact_export.py b/packages/kumo-core/src/kumo_core/artifact_export.py index 0d73620ca..311b972bc 100644 --- a/packages/kumo-core/src/kumo_core/artifact_export.py +++ b/packages/kumo-core/src/kumo_core/artifact_export.py @@ -199,18 +199,22 @@ def publish_export_destination(output_path, *, fs): def export_selected_files(files, output_path, *, fs, identity, file_type='parquet', convert_csv=None, - copy_parquet=None, temporary_directory=None): + copy_parquet=None, temporary_directory=None, + sidecars={}): """Publish the hosted selected-file layout and success marker last. Filesystem/authentication and retry wrappers stay with the caller. The directory ownership, output names and CSV retry decisions are shared. + ``sidecars`` (name to bytes) land in a directory destination before + ``_SUCCESS``. """ import os.path as osp import tempfile + output_is_dir = output_path.endswith('/') or len(files) > 1 output_path, _ = prepare_export_destination( - output_path, output_is_dir=output_path.endswith('/') or len(files) > 1, - file_type=file_type, fs=fs, identity=identity) + output_path, output_is_dir=output_is_dir, file_type=file_type, fs=fs, + identity=identity) if file_type == 'parquet': if copy_parquet is None: @@ -230,6 +234,14 @@ def export_selected_files(files, output_path, *, fs, identity, else: raise ValueError(f'Unsupported artifact file type: {file_type}') + if output_is_dir: + for name, payload in sidecars.items(): + try: + fs.write_bytes(osp.join(output_path, name), payload) + except Exception as error: + import warnings + warnings.warn(f'Exporting without {name}: {error}', + UserWarning, stacklevel=2) return publish_export_destination(output_path, fs=fs) diff --git a/packages/kumo-core/src/kumo_core/artifacts.py b/packages/kumo-core/src/kumo_core/artifacts.py index b6e7a49c4..a2910deaf 100644 --- a/packages/kumo-core/src/kumo_core/artifacts.py +++ b/packages/kumo-core/src/kumo_core/artifacts.py @@ -13,12 +13,23 @@ import tempfile from dataclasses import dataclass, field from pathlib import Path -from typing import Any, ClassVar, Mapping, Protocol, runtime_checkable - -from .io.filesystem import local_path +from typing import ( + TYPE_CHECKING, + Any, + ClassVar, + Iterable, + Mapping, + Protocol, + runtime_checkable, +) + +from .io.filesystem import atomic_write_bytes, fsync_directory, local_path from .io.json import load_json, save_json from .manifest import RunManifest, RunPaths +if TYPE_CHECKING: + from .types import TaskType + @runtime_checkable class ArtifactStore(Protocol): @@ -62,36 +73,7 @@ def write_manifest(self, paths: RunPaths, manifest: RunManifest) -> None: ensure_ascii=False, allow_nan=False, ) + "\n" - temporary_path: Path | None = None - try: - with tempfile.NamedTemporaryFile( - mode="w", - encoding="utf-8", - dir=paths.manifest_path.parent, - prefix=f".{paths.manifest_path.name}.", - suffix=".tmp", - delete=False, - ) as temporary: - temporary_path = Path(temporary.name) - temporary.write(payload) - temporary.flush() - os.fsync(temporary.fileno()) - os.replace(temporary_path, paths.manifest_path) - temporary_path = None - self._fsync_directory(paths.manifest_path.parent) - finally: - if temporary_path is not None: - temporary_path.unlink(missing_ok=True) - - @staticmethod - def _fsync_directory(directory: Path) -> None: - if not hasattr(os, "O_DIRECTORY"): - return - descriptor = os.open(directory, os.O_RDONLY | os.O_DIRECTORY) - try: - os.fsync(descriptor) - finally: - os.close(descriptor) + atomic_write_bytes(paths.manifest_path, payload.encode("utf-8")) @dataclass(frozen=True) @@ -258,7 +240,7 @@ def publish( encoding="utf-8", ) os.replace(staging, destination) - FilesystemArtifactStore._fsync_directory(destination.parent) + fsync_directory(destination.parent) return publication finally: if staging.exists(): @@ -273,14 +255,16 @@ def _sha256(path: Path) -> str: return digest.hexdigest() +PROVENANCE_FILE = 'provenance.json' + + def verify_model_bundle(bundle: Path) -> str: """Verify fetched artifacts before parsing or staging a model bundle. - Preserves the local SDK import verifier's error and content-identity - contract. Paths must stay inside the bundle after symlink resolution; - declared sizes and SHA-256 digests must match. Older manifests may omit - sizes or leave training-spec.json unlisted; its digest still participates - in the returned content identity. + Returns the bundle id, a SHA-256 over every verified file except + ``provenance.json``. Paths must stay inside the bundle after symlink + resolution; declared sizes and SHA-256 digests must match. Older + manifests may omit sizes or leave training-spec.json unlisted. """ from .sources import layout from .sources.errors import PermanentSourceError @@ -298,7 +282,7 @@ def verify_model_bundle(bundle: Path) -> str: artifacts = manifest.get('artifacts') problems: list[str] = [] - verified: list[tuple[str, str, str]] = [] + verified: list[tuple[str, str]] = [] seen = set() claimed = set() if not artifacts: @@ -350,7 +334,7 @@ def verify_model_bundle(bundle: Path) -> str: if VitalArtifactPublicationPolicy._sha256(path) != expected: problems.append(f'{relative} does not match its digest') continue - verified.append((role, safe, expected)) + verified.append((safe, expected)) if problems: raise ValueError( f"The model bundle at {bundle} did not survive the transfer: " @@ -360,11 +344,73 @@ def verify_model_bundle(bundle: Path) -> str: if not spec_path.resolve().is_relative_to(bundle): raise ValueError('Training specification is not below the bundle') if 'training-spec.json' not in seen and spec_path.is_file(): - verified.append(('training_spec', 'training-spec.json', + verified.append(('training-spec.json', VitalArtifactPublicationPolicy._sha256(spec_path))) - payload = json.dumps(sorted(verified), sort_keys=True, - separators=(',', ':'), default=str) - return hashlib.sha256(payload.encode('utf-8')).hexdigest()[:16] + return _bundle_id(verified) + + +def bundle_id_from_manifest(bundle: Path) -> str: + """The bundle id from the digests a freshly published bundle lists.""" + manifest_path = Path(bundle) / VitalArtifactPublicationPolicy.MANIFEST_NAME + manifest = json.loads(manifest_path.read_text(encoding='utf-8')) + return _bundle_id((str(entry['relative_path']), str(entry['sha256'])) + for entry in manifest['artifacts']) + + +def _bundle_id(pairs: Iterable[tuple[str, str]]) -> str: + payload = json.dumps( + sorted(pair for pair in pairs if pair[0] != PROVENANCE_FILE), + separators=(',', ':')) + return hashlib.sha256(payload.encode('utf-8')).hexdigest() + + +def add_bundle_file(bundle_dir: Path, role: str, relative_path: str) -> None: + """List one file already written in a bundle in its publication manifest.""" + bundle_dir = Path(bundle_dir) + path = bundle_dir / relative_path + manifest_path = bundle_dir / VitalArtifactPublicationPolicy.MANIFEST_NAME + manifest = json.loads(manifest_path.read_text(encoding='utf-8')) + manifest['artifacts'] = [ + entry for entry in manifest.get('artifacts') or [] + if entry.get('relative_path') != relative_path + and entry.get('role') != role + ] + [ + PublishedArtifact( + role=role, + relative_path=relative_path, + size_bytes=path.stat().st_size, + sha256=VitalArtifactPublicationPolicy._sha256(path), + ).to_dict() + ] + _write_bundle_manifest(bundle_dir, manifest) + + +def remove_bundle_file(bundle_dir: Path, role: str) -> None: + """Delete the file listed under ``role`` and its manifest entry.""" + from .sources import layout + + bundle_dir = Path(bundle_dir) + manifest_path = bundle_dir / VitalArtifactPublicationPolicy.MANIFEST_NAME + manifest = json.loads(manifest_path.read_text(encoding='utf-8')) + entries = manifest.get('artifacts') or [] + removed = [entry for entry in entries if entry.get('role') == role] + if not removed: + return + manifest['artifacts'] = [ + entry for entry in entries if entry.get('role') != role + ] + _write_bundle_manifest(bundle_dir, manifest) + for entry in removed: + (bundle_dir / layout.confined(entry['relative_path'])).unlink( + missing_ok=True) + + +def _write_bundle_manifest(bundle_dir: Path, + manifest: Mapping[str, Any]) -> None: + payload = json.dumps(manifest, indent=2, sort_keys=True) + '\n' + atomic_write_bytes( + bundle_dir / VitalArtifactPublicationPolicy.MANIFEST_NAME, + payload.encode('utf-8')) @runtime_checkable @@ -456,6 +502,14 @@ def resolve_vital_artifact_sources( } +def bundle_task_type(artifacts: Mapping[str, str]) -> TaskType: + """The task type a bundle's prediction schema records.""" + from .types import TaskType + + schema = load_json(artifacts['prediction_schema'], dict) + return TaskType(schema['task_type']) + + def stage_model_bundle(artifacts: Mapping[str, str], dest: Path) -> Path: r"""Copies a trained run's artifacts into the layout scoring expects. diff --git a/packages/kumo-core/src/kumo_core/automl/auto_trainer.py b/packages/kumo-core/src/kumo_core/automl/auto_trainer.py index 5dd818966..3abe36170 100644 --- a/packages/kumo-core/src/kumo_core/automl/auto_trainer.py +++ b/packages/kumo-core/src/kumo_core/automl/auto_trainer.py @@ -12,6 +12,7 @@ import math import os import os.path as osp +import time import traceback import warnings from dataclasses import asdict, dataclass, fields, is_dataclass @@ -39,7 +40,7 @@ AutoTrainerProgress, AutoTrainingProgressLogger, ) -from kumo_core.instrumentation import TrainingInstrumentation +from kumo_core.instrumentation import TrainingInstrumentation, format_error from kumo_core.io import filesystem from kumo_core.io.json import load_json, save_json from kumo_core.job_config import JobConfig, TrainStatus @@ -458,6 +459,7 @@ def _train( metrics_client_job_name: str, pruning: Callback | None = None, status: TrainStatus | None = None, + trial: int | None = None, ) -> list[float]: r"""Config-driven training function. @@ -713,7 +715,7 @@ def _train( trainer.pl_trainer.callbacks.append(pruning) # Add progress logging callback: - self.training_progress_logger.set_current_trainer(trainer) + self.training_progress_logger.set_current_trainer(trainer, trial) trainer.pl_trainer.callbacks.append(self.training_progress_logger) # Save AutoTrainer YAML (Trainer YAML will be saved within @@ -847,7 +849,7 @@ def _run( logger.info(f"Initialized '{self.searcher}'") self.training_progress_logger = AutoTrainingProgressLogger( - self.log_dir, self.num_experiments) + self.log_dir, self.num_experiments, on_event=self._emit) tune_mode = MetricStore.tune_mode(self.tune_metric) metrics_job_name = self.runtime.metrics_job_name() @@ -872,12 +874,17 @@ def _run( logger.info(f"Starting AutoML experiment " f"{i + 1}/{self.num_experiments}") + self._emit('trial_started', trial=i + 1, params={ + **self.searcher.sampled_params, + 'attempt': j + 1, + }) batch_size = cfg.optimization.batch_size[0] with self.instrumentation.timing( "auto_trainer.run.status_add_trial", trial=i + 1): status.add_trial(batch_size) status.save(self.home_dir) + trial_start = time.monotonic() with self.instrumentation.timing( "auto_trainer.run.train_trial", trial=i + 1): val_performance = self._train( @@ -885,7 +892,9 @@ def _run( metrics_client_job_name=metrics_job_name, pruning=pruning, status=status, + trial=i + 1, ) + trial_seconds = time.monotonic() - trial_start best_performance = _eval_fn[tune_mode](val_performance) logger.info(f"AutoML experiment " f"{i + 1}/{self.num_experiments} finished " @@ -902,9 +911,21 @@ def _run( "auto_trainer.run.status_save_success", trial=i + 1): status.save(self.home_dir) + self._emit('trial_ended', trial=i + 1, metrics={ + self.val_tune_metric: best_performance, + 'seconds': trial_seconds, + }, params={ + 'status': 'succeeded', + 'attempt': j + 1, + }) break except RuntimeError as e: + self._emit('trial_ended', trial=i + 1, params={ + 'status': 'failed', + 'attempt': j + 1, + 'error': format_error(e), + }) stacktrace = ''.join( traceback.format_exception(e, e, e.__traceback__)) logger.warning(stacktrace) @@ -957,6 +978,16 @@ def _run( should_init=init_best_trainer): return self.get_best_trainer(should_init=init_best_trainer) + def _emit(self, name: str, **fields: Any) -> None: + event = getattr(self.instrumentation, 'event', None) + if event is None: + return + try: + event(name, **fields) + except Exception: + logger.warning(f"Could not report progress event '{name}'", + exc_info=True) + def _handle_oom_retry( self, cfg: ModelPlan, diff --git a/packages/kumo-core/src/kumo_core/automl/optuna_searcher.py b/packages/kumo-core/src/kumo_core/automl/optuna_searcher.py index db5a03bda..fad0adb28 100644 --- a/packages/kumo-core/src/kumo_core/automl/optuna_searcher.py +++ b/packages/kumo-core/src/kumo_core/automl/optuna_searcher.py @@ -60,6 +60,7 @@ def suggest(self) -> PlanT: for _ in range(self.fallback_sampler.max_attempts): self.current_trial = self.study.ask() cfg = self.fallback_sampler.sample_instance() + self.sampled_params = self.fallback_sampler.sampled_params if cfg not in [prev_cfg for prev_cfg, _ in self.state]: return cfg return cfg diff --git a/packages/kumo-core/src/kumo_core/automl/runtime.py b/packages/kumo-core/src/kumo_core/automl/runtime.py index 209680edf..880246456 100644 --- a/packages/kumo-core/src/kumo_core/automl/runtime.py +++ b/packages/kumo-core/src/kumo_core/automl/runtime.py @@ -196,9 +196,10 @@ def resolve_warm_start( encoders; they do not retain hosted AutoML directory layouts. """ if artifacts is not None: - source_schema = load_json(artifacts["prediction_schema"], dict) + from kumo_core.artifacts import bundle_task_type + validate_warm_start_task_type(task.task_type, - TaskType(source_schema["task_type"])) + bundle_task_type(artifacts)) checkpoint = artifacts["selected_checkpoint"] plan_path = artifacts["resolved_model_plan"] encoder_dir = posixpath.dirname(artifacts["encoders"]) diff --git a/packages/kumo-core/src/kumo_core/automl/searcher.py b/packages/kumo-core/src/kumo_core/automl/searcher.py index f8810e7b6..9116bf389 100644 --- a/packages/kumo-core/src/kumo_core/automl/searcher.py +++ b/packages/kumo-core/src/kumo_core/automl/searcher.py @@ -39,6 +39,7 @@ def __init__(self, space: PlanT, max_attempts: int = 50) -> None: self.space = space self.max_attempts = max_attempts self.state: list[tuple[PlanT, list[float]]] = [] + self.sampled_params: dict[str, Any] = {} @property def tune_metric(self) -> str: @@ -68,6 +69,7 @@ def sample_value(self, key: str, value: list[Any]) -> Any: def sample_instance(self) -> PlanT: instance = copy.deepcopy(self.space) + self.sampled_params = {} for field in fields(instance): section = getattr(instance, field.name) @@ -76,6 +78,7 @@ def sample_instance(self) -> PlanT: continue assert isinstance(value, list) + searched = len(value) > 1 if isinstance(value[0], PlanMixin): index_of_plan = self.sample_value(key, list(range(len(value)))) @@ -85,17 +88,25 @@ def sample_instance(self) -> PlanT: subvalue, MissingType): continue assert isinstance(subvalue, list) + searched = searched or len(subvalue) > 1 unique_key = f"{key}{index_of_plan}_{subkey}" setattr( plan, subkey, [self.sample_value(unique_key, subvalue)], ) - setattr(section, key, [plan]) + chosen = plan else: - setattr(section, key, [self.sample_value(key, value)]) + chosen = self.sample_value(key, value) + setattr(section, key, [chosen]) + if searched: + self.sampled_params[f"{field.name}.{key}"] = ( + chosen if chosen is None + or isinstance(chosen, (bool, int, float, str)) else + str(chosen)) return instance def __repr__(self) -> str: return f"{self.__class__.__name__}()" + diff --git a/packages/kumo-core/src/kumo_core/callbacks/train_progress_logger.py b/packages/kumo-core/src/kumo_core/callbacks/train_progress_logger.py index 45f124db5..abdbd7803 100644 --- a/packages/kumo-core/src/kumo_core/callbacks/train_progress_logger.py +++ b/packages/kumo-core/src/kumo_core/callbacks/train_progress_logger.py @@ -5,7 +5,7 @@ import logging import os.path as osp from datetime import datetime, timezone -from typing import Any +from typing import Any, Callable import kumoapi.jobs as jobs_api from kumo_core.io.json import load_json, save_json @@ -140,17 +140,22 @@ class AutoTrainingProgressLogger(MonitoringCallback): # pragma: no cover # Constant to account for overhead not captured in trial estimation: OVERHEAD_CONSTANT = 1.30 - def __init__(self, log_dir: str, num_experiments: int): + def __init__(self, log_dir: str, num_experiments: int, + on_event: Callable[..., None] | None = None): super().__init__() self.path = osp.join(log_dir, self.FILE_NAME) + self.on_event = on_event + self.trial: int | None = None self.progress = load_json(self.path, AutoTrainerProgress, raise_on_missing=False) if self.progress is None: self.progress = AutoTrainerProgress(total_trials=num_experiments) - def set_current_trainer(self, trainer: Any | None): + def set_current_trainer(self, trainer: Any | None, + trial: int | None = None): self.trainer = trainer + self.trial = trial def on_train_start(self, pl_trainer, pl_module): if self.progress.start_time is None: @@ -184,6 +189,20 @@ def on_train_epoch_end(self, pl_trainer, pl_module): save_json(self.progress, self.path) + if self.on_event is not None and self.trial is not None: + try: + progress = self.trainer.training_progress_logger.progress + epoch = progress.num_epochs_completed + metrics = progress.metrics[epoch - 1] + self.on_event('epoch_ended', trial=self.trial, epoch=epoch, + metrics={ + **metrics.train_metrics, + **metrics.validation_metrics + }) + except Exception: + logger.warning('Could not report epoch_ended', + exc_info=True) + def on_validation_end(self, pl_trainer, pl_module): if pl_trainer.state.fn == TrainerFn.VALIDATING: save_json(self.progress, self.path) diff --git a/packages/kumo-core/src/kumo_core/instrumentation.py b/packages/kumo-core/src/kumo_core/instrumentation.py index ab887f8aa..141bcd748 100644 --- a/packages/kumo-core/src/kumo_core/instrumentation.py +++ b/packages/kumo-core/src/kumo_core/instrumentation.py @@ -24,6 +24,9 @@ def span(self, name: str, **fields: Any) -> ContextManager[None]: def timing(self, name: str, **fields: Any) -> ContextManager[None]: ... + def event(self, name: str, **fields: Any) -> None: + return None + class NoOpTrainingInstrumentation: def record_metrics_job_name(self, job_name: str) -> None: @@ -37,6 +40,9 @@ def span(self, name: str, **fields: Any) -> Iterator[None]: def timing(self, name: str, **fields: Any) -> Iterator[None]: yield + def event(self, name: str, **fields: Any) -> None: + return None + _instrumentation: TrainingInstrumentation = NoOpTrainingInstrumentation() @@ -54,6 +60,10 @@ def timing(name: str, **fields: Any) -> ContextManager[None]: return _instrumentation.timing(name, **fields) +def format_error(error: BaseException) -> str: + return f'{type(error).__name__}: {error}'[:200] + + def trace() -> Callable[[Callable[..., T]], Callable[..., T]]: def decorate(function: Callable[..., T]) -> Callable[..., T]: @wraps(function) diff --git a/packages/kumo-core/src/kumo_core/io/filesystem.py b/packages/kumo-core/src/kumo_core/io/filesystem.py index 7ac2de5be..c8c227148 100644 --- a/packages/kumo-core/src/kumo_core/io/filesystem.py +++ b/packages/kumo-core/src/kumo_core/io/filesystem.py @@ -7,6 +7,7 @@ import os import shutil +import tempfile from collections.abc import Callable from pathlib import Path, PurePosixPath from typing import Any @@ -85,6 +86,37 @@ def local_path(path: str) -> Path: return Path(fs_path) +def fsync_directory(directory: str | Path) -> None: + if not hasattr(os, "O_DIRECTORY"): + return + descriptor = os.open(directory, os.O_RDONLY | os.O_DIRECTORY) + try: + os.fsync(descriptor) + finally: + os.close(descriptor) + + +def atomic_write_bytes(path: str | Path, payload: bytes) -> None: + """Replace a local file so readers see the old or the new bytes.""" + path = Path(path) + temporary: Path | None = None + try: + with tempfile.NamedTemporaryFile(dir=path.parent, + prefix=f".{path.name}.", + suffix=".tmp", + delete=False) as out: + temporary = Path(out.name) + out.write(payload) + out.flush() + os.fsync(out.fileno()) + os.replace(temporary, path) + temporary = None + fsync_directory(path.parent) + finally: + if temporary is not None: + temporary.unlink(missing_ok=True) + + def listdir(path: str) -> list[str]: fs, fs_path = _resolve(path) return [fs.unstrip_protocol(item) for item in fs.ls(fs_path, detail=False)] 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 ec34f89f9..3d673a524 100644 --- a/packages/kumo-core/src/kumo_core/object_store_sink.py +++ b/packages/kumo-core/src/kumo_core/object_store_sink.py @@ -31,6 +31,12 @@ def _destination(self, connector: SourceConnector, """Return the authenticated filesystem and validated bucket/key path.""" raise NotImplementedError + def filesystem(self, connector: SourceConnector, + path: str) -> Tuple[FileSystem, str]: + """The authenticated filesystem and resolved ``path``.""" + filesystem, resolved = self._destination(connector, path) + return filesystem, resolved.rstrip('/') + def write_table( self, connector: SourceConnector, @@ -55,7 +61,8 @@ def write_table( def write_files(self, connector, table_name, files, write_mode=WriteMode.OVERWRITE, project_id=None, *, - export_id=None, extra_columns=None): + export_id=None, extra_columns=None, + sidecars: Mapping[str, bytes] = {}): """Export selected Parquet without collecting their rows.""" del project_id if write_mode is not WriteMode.OVERWRITE: @@ -69,7 +76,8 @@ def write_files(self, connector, table_name, files, _, output = export_selected_files(files, path, fs=storage, identity=export_id or str(files), file_type=kind, - convert_csv=storage.convert_csv) + convert_csv=storage.convert_csv, + sidecars=sidecars) return f'{self.scheme}://{output}' def publish_model( @@ -282,8 +290,11 @@ def read_text(self, path): return reader.readall().decode('utf-8') def write_text(self, path, value): + self.write_bytes(path, value.encode('utf-8')) + + def write_bytes(self, path, value): with self.filesystem.open_output_stream(path) as writer: - writer.write(value.encode('utf-8')) + writer.write(value) def makedirs(self, path, exist_ok=True): _prepared(self.filesystem, path) diff --git a/packages/kumo-core/src/kumo_core/prediction.py b/packages/kumo-core/src/kumo_core/prediction.py index 885e3ae89..8bb56d2e3 100644 --- a/packages/kumo-core/src/kumo_core/prediction.py +++ b/packages/kumo-core/src/kumo_core/prediction.py @@ -507,6 +507,7 @@ class _PredictionWorkerConfig(NamedTuple): forecast_options: Optional[Dict[str, int]] embedding_models: Optional[EmbeddingModelStore] = None ignore_time_col: bool = False + constant_columns: Optional[Dict[str, str]] = None def _score_prediction_shard( @@ -570,6 +571,7 @@ def _score_prediction_shard( in config.artifacts), top_k=config.top_k, output_extra_fields=config.extra_fields or None, + output_constant_columns=config.constant_columns, job_start_time=config.job_start_time, pred_time_col_name=config.pred_time_col_name, ignore_time_col=config.ignore_time_col, @@ -606,6 +608,7 @@ def score_prediction( forecast_options: Mapping[str, int] | None = None, embedding_models: EmbeddingModelStore | None = None, encoders: Any = None, + constant_columns: Mapping[str, str] | None = None, ) -> Any: """Score partitions and publish bounded metadata using the hosted protocol. @@ -678,6 +681,7 @@ def score_prediction( forecast_options=(dict(forecast_options) if forecast_options is not None else None), embedding_models=embedding_models, + constant_columns=dict(constant_columns or {}), ) if worker_count == 1: references = [ diff --git a/packages/kumo-core/src/kumo_core/sources/identity.py b/packages/kumo-core/src/kumo_core/sources/identity.py index ac69d6b92..0e3d8e664 100644 --- a/packages/kumo-core/src/kumo_core/sources/identity.py +++ b/packages/kumo-core/src/kumo_core/sources/identity.py @@ -20,6 +20,26 @@ CUSTOM_SQL_MAX_STALENESS_HOURS = 24 +class SourceVersion(str): + """A version string that also says what kind of evidence it is: + ``exact``, ``timestamp``, ``listing``, ``hour_bucket`` or ``unknown``.""" + + kind: str + dependencies: Mapping[str, str] | None + + def __new__(cls, value: str, kind: str = "unknown", + dependencies: Mapping[str, str] | None = None + ) -> "SourceVersion": + version = super().__new__(cls, value) + version.kind = kind + version.dependencies = dependencies + return version + + +def version_kind(version: object) -> str: + return getattr(version, "kind", "unknown") + + def epoch_second(moment: datetime) -> str: """Truncated by string rather than through ``int``: the two disagree just below the epoch, and this value is a raw path segment.""" @@ -29,8 +49,10 @@ def epoch_second(moment: datetime) -> str: def hour_version() -> str: """Fallback for a source with no readable version; it expires itself rather than needing invalidation.""" - return epoch_second( - datetime.now(timezone.utc).replace(minute=0, second=0, microsecond=0)) + return SourceVersion( + epoch_second( + datetime.now(timezone.utc).replace(minute=0, second=0, + microsecond=0)), "hour_bucket") def sql_data_version( @@ -83,10 +105,11 @@ def sql_data_version( versions[leaf] = _indeterminate() else: versions[leaf] = str(stamped) - return custom_sql_data_version( - sql, connector_type, versions, cols or {}, - max_staleness_bucket=int(time()) // - (CUSTOM_SQL_MAX_STALENESS_HOURS * 3600)) + return SourceVersion( + custom_sql_data_version( + sql, connector_type, versions, cols or {}, + max_staleness_bucket=int(time()) // + (CUSTOM_SQL_MAX_STALENESS_HOURS * 3600)), "hour_bucket", versions) #: The prefix every sentinel carries, so one can be recognized without being diff --git a/packages/kumo-core/src/kumo_core/sources/object_store.py b/packages/kumo-core/src/kumo_core/sources/object_store.py index 4b7ac0a37..bf1058d07 100644 --- a/packages/kumo-core/src/kumo_core/sources/object_store.py +++ b/packages/kumo-core/src/kumo_core/sources/object_store.py @@ -52,7 +52,7 @@ TransientSourceError, UnsupportedReadModeError, ) -from .identity import epoch_second +from .identity import SourceVersion, epoch_second from .object_store_access import ObjectFetch, ObjectIdentity from .parallel import run_in_parallel from .prepare import ReadOutput, run_prepare @@ -101,14 +101,16 @@ def _source_version(path: str, entries: Sequence[Any]) -> str: if len(entries) == 1 and entries[0].path.rstrip("/") == path.rstrip("/"): latest = entries[0].mtime - return epoch_second(latest) if latest is not None else "-1" + return (SourceVersion(epoch_second(latest), "timestamp") + if latest is not None else SourceVersion("-1")) # A maximum mtime misses changes to older files and added/deleted parts. metadata = [ (entry.path, entry.size, getattr(entry, "mtime_ns", None), entry.mtime.isoformat() if entry.mtime is not None else None) for entry in sorted(entries, key=lambda entry: entry.path) ] - return hashlib.sha256(json.dumps(metadata).encode()).hexdigest() + return SourceVersion( + hashlib.sha256(json.dumps(metadata).encode()).hexdigest(), "listing") class _CsvSource(NamedTuple): @@ -489,7 +491,7 @@ def _copy(self, path: str, found: Sequence[Any], dest: str, if self.fetch is not None: return self._download_objects(path, found, dest, worker_ceiling, parquet=True) - with _copies(found, worker_ceiling) as workers: + with copies(found, worker_ceiling) as workers: return fetch_all([ StagedFile( entry.path, entry.size @@ -519,7 +521,7 @@ def download(entry: Any) -> str: return target try: - with _copies(found, worker_ceiling) as workers, ThreadPoolExecutor( + with copies(found, worker_ceiling) as workers, ThreadPoolExecutor( max_workers=workers) as pool: futures = [pool.submit(download, entry) for entry in found] try: @@ -814,8 +816,8 @@ def _parquet_batches(fragment: Any, @contextmanager -def _copies(found: Sequence[Any], - worker_ceiling: int | None) -> Iterator[int]: +def copies(found: Sequence[Any], + worker_ceiling: int | None = None) -> Iterator[int]: """Arrow-backed downloaders copy on Arrow's IO pool.""" with memory_share( COPY_WORKER_BYTES, lambda ceiling: min( diff --git a/packages/kumo-core/src/kumo_core/sources/object_store_backend.py b/packages/kumo-core/src/kumo_core/sources/object_store_backend.py index 69f3025f7..f5f903d0a 100644 --- a/packages/kumo-core/src/kumo_core/sources/object_store_backend.py +++ b/packages/kumo-core/src/kumo_core/sources/object_store_backend.py @@ -130,7 +130,7 @@ def get_data_version( ) -> str: del project_id reader, name = self._reader(connector, table_name) - return str(reader.data_version(name)) + return reader.data_version(name) def prepare_ingest( self, diff --git a/packages/kumo-core/src/kumo_core/sources/object_store_sessions.py b/packages/kumo-core/src/kumo_core/sources/object_store_sessions.py index 1e9dd307a..951b780bb 100644 --- a/packages/kumo-core/src/kumo_core/sources/object_store_sessions.py +++ b/packages/kumo-core/src/kumo_core/sources/object_store_sessions.py @@ -6,7 +6,7 @@ from abc import ABC, abstractmethod from dataclasses import dataclass from typing import Any, Callable, Optional, Tuple -from urllib.parse import unquote, urlsplit +from urllib.parse import quote, unquote, urlsplit from kumoapi.plugin import SourceConnector from pyarrow.fs import FileSystem @@ -16,6 +16,8 @@ from .object_store_access import ObjectStoreAccess from .principals import IdentityRef, Principal, resolve +_KEY_SAFE = "/-._~!$&'()*+,;=:@" + @dataclass(frozen=True) class ObjectStoreRoot: @@ -104,6 +106,15 @@ def _over(self, connector: SourceConnector, bucket: str, root: str, fetch=access.fetch, object_identity=access.object_identity) + def table_location(self, connector: SourceConnector, + table_name: str) -> str: + """Full ``scheme://bucket/key`` of ``table_name``, without credentials.""" + root = self._root_for(connector) + if root is None: + return self.paths.source_uri(table_name).rstrip('/') + key = quote(posixpath.join(root.prefix, table_name), safe=_KEY_SAFE) + return f'{self.scheme}://{root.bucket}/{key}'.rstrip('/') + def _placement(self, connector: SourceConnector, table_name: str) -> Tuple[str, str, str, Optional[str]]: root = self._root_for(connector) diff --git a/packages/kumo-core/src/kumo_core/train/_artifact_storage.py b/packages/kumo-core/src/kumo_core/train/_artifact_storage.py index 6e6cc5ebd..06153d6f7 100644 --- a/packages/kumo-core/src/kumo_core/train/_artifact_storage.py +++ b/packages/kumo-core/src/kumo_core/train/_artifact_storage.py @@ -16,6 +16,8 @@ from pathlib import Path from typing import Any, Callable +from kumo_core.io.filesystem import atomic_write_bytes, fsync_directory + class _LocalFiles: @staticmethod @@ -54,7 +56,7 @@ def write_text_if_absent(path, text): os.link(temporary, path) except FileExistsError: return False - _fsync_directory(Path(path).parent) + fsync_directory(Path(path).parent) return True @@ -74,20 +76,9 @@ def _temporary_text(path, text): os.unlink(temporary) -def _fsync_directory(directory): - if not hasattr(os, "O_DIRECTORY"): - return - descriptor = os.open(directory, os.O_RDONLY | os.O_DIRECTORY) - try: - os.fsync(descriptor) - finally: - os.close(descriptor) - - def _save_json_atomic(value, path): - with _temporary_text(path, json.dumps(value, indent=2)) as temporary: - os.replace(temporary, path) - _fsync_directory(Path(path).parent) + Path(path).parent.mkdir(parents=True, exist_ok=True) + atomic_write_bytes(path, json.dumps(value, indent=2).encode('utf-8')) @dataclass(frozen=True) diff --git a/packages/kumo-core/src/kumo_core/train/diskgraph_nested_trainer.py b/packages/kumo-core/src/kumo_core/train/diskgraph_nested_trainer.py index a8d197382..218db9962 100644 --- a/packages/kumo-core/src/kumo_core/train/diskgraph_nested_trainer.py +++ b/packages/kumo-core/src/kumo_core/train/diskgraph_nested_trainer.py @@ -1346,6 +1346,7 @@ def predict( ignore_time_col: bool = False, pred_time_col_name: str | None = None, output_extra_fields: dict[str, list[str]] | None = None, + output_constant_columns: Mapping[str, str] | None = None, max_output_file_size: int = 1024**3, progress_log_dir: str | None = None, progress_log_id: int = 0, @@ -1403,6 +1404,7 @@ def predict( ignore_time_col=ignore_time_col, pred_time_col_name=pred_time_col_name, output_extra_fields=output_extra_fields, + output_constant_columns=output_constant_columns, max_output_file_size=max_output_file_size, progress_log_dir=progress_log_dir, progress_log_id=progress_log_id, partition_id=partition_id, @@ -1585,6 +1587,7 @@ def write_prediction(self, prediction: dict[str, Any]) -> None: output_extra_fields=output_extra_fields, job_start_time=job_start_time, write_predictions=emit_predictions, + constant_columns=output_constant_columns, artifact_stats_sample_rows=({ artifact_type: artifact_stats_collector.sample_rows_needed( @@ -1698,6 +1701,7 @@ def _predict_lp( ignore_time_col: bool = False, pred_time_col_name: str | None = None, output_extra_fields: dict[str, list[str]] | None = None, + output_constant_columns: Mapping[str, str] | None = None, max_output_file_size: int = 1024**3, progress_log_dir: str | None = None, progress_log_id: int = 0, @@ -1982,6 +1986,7 @@ def write_prediction(self, prediction: dict[str, Any], output_extra_fields=output_extra_fields, job_start_time=job_start_time, write_predictions=emit_predictions, + constant_columns=output_constant_columns, artifact_stats_sample_rows=({ artifact_type: artifact_stats_collector.sample_rows_needed( diff --git a/packages/kumo-core/src/kumo_core/train/formatting.py b/packages/kumo-core/src/kumo_core/train/formatting.py index 53ce12c3a..0d5585d73 100644 --- a/packages/kumo-core/src/kumo_core/train/formatting.py +++ b/packages/kumo-core/src/kumo_core/train/formatting.py @@ -31,6 +31,15 @@ class LinkPredictionTable(Protocol): _ANCHOR_TIME_COL = "ANCHOR_TIMESTAMP" +_CLASS_TASKS = frozenset({ + TaskType.MULTICLASS_CLASSIFICATION, + TaskType.MULTILABEL_CLASSIFICATION, +}) +_LONG_FORMAT_TASKS = _CLASS_TASKS | { + TaskType.MULTILABEL_RANKING, + TaskType.TEMPORAL_LINK_PREDICTION, + TaskType.STATIC_LINK_PREDICTION, +} def _get_col_names( @@ -67,6 +76,23 @@ def _get_col_names( return result +def prediction_output_columns(task_type: TaskType) -> list[str]: + r"""The prediction columns a scored row carries after its entity column, + for a target named ``TARGET`` and a time column named ``TIMESTAMP``.""" + if task_type in _CLASS_TASKS: + return ['PREDICTED', 'CLASS', 'SCORE'] + if task_type.is_ranking: + return ['CLASS', 'SCORE'] + if task_type == TaskType.FORECASTING: + return ['TIMESTAMP', 'TARGET_PRED'] + if task_type == TaskType.BINARY_CLASSIFICATION: + return [ + 'TARGET_PRED', + *_get_col_names(pd.Index([]), task_type, is_wide_format=True) + ] + return ['TARGET_PRED'] + + def _format_prediction_sanity_checks( target_encoder: TargetEncoder, task_type: TaskType, @@ -284,13 +310,7 @@ def format_prediction( 'metadata. Rename the time column before forecasting.') format: Literal['wide', 'long'] - if task_type in { - TaskType.MULTICLASS_CLASSIFICATION, - TaskType.MULTILABEL_CLASSIFICATION, - TaskType.MULTILABEL_RANKING, - TaskType.TEMPORAL_LINK_PREDICTION, - TaskType.STATIC_LINK_PREDICTION, - }: + if task_type in _LONG_FORMAT_TASKS: format = 'long' else: assert top_k is None, f"`top_k` is not supported for {task_type}." diff --git a/packages/kumo-core/src/kumo_core/train/io.py b/packages/kumo-core/src/kumo_core/train/io.py index f47dbf2a8..c25131926 100644 --- a/packages/kumo-core/src/kumo_core/train/io.py +++ b/packages/kumo-core/src/kumo_core/train/io.py @@ -16,7 +16,7 @@ import threading from concurrent.futures import Future, ThreadPoolExecutor from pathlib import Path -from typing import Any, TypeAlias +from typing import Any, Mapping, TypeAlias import pandas as pd import pyarrow as pa @@ -161,6 +161,7 @@ def __init__( spool_root: str | os.PathLike[str] | None = None, max_row_group_rows: int = _DEFAULT_ROW_GROUP_ROWS, stats_sample_rows: int | None = None, + constant_columns: Mapping[str, str] | None = None, ) -> None: if max_file_size <= 0 or max_buffer_size <= 0: raise ValueError("Parquet stream size limits must be positive") @@ -190,6 +191,7 @@ def __init__( ) self._peak_row_group_bytes = 0 self._schema: pa.Schema | None = None + self._constant_columns = dict(constant_columns or {}) self._index_offset = 0 self._closed = False self._aborted = False @@ -296,6 +298,18 @@ def _to_table(self, frame: Frame) -> pa.Table: table = table.cast(self._schema.remove_metadata(), safe=False) return table.replace_schema_metadata(self._schema.metadata) + def _file_schema(self) -> pa.Schema: + schema = self._schema or pa.schema([]) + for name in self._constant_columns: + schema = schema.append(pa.field(name, pa.string())) + return schema + + def _stamped(self, table: pa.Table) -> pa.Table: + for name, value in self._constant_columns.items(): + table = table.append_column( + name, pa.repeat(pa.scalar(value, pa.string()), table.num_rows)) + return table + def _slices(self, table: pa.Table, max_size: int) -> list[pa.Table]: if table.num_rows == 0: return [] @@ -486,9 +500,8 @@ def flush_group() -> None: slot_held = True self._check_error() local_path = self._temp_dir / f"part-{part_index:05d}.parquet" - assert self._schema is not None - writer = pq.ParquetWriter(local_path, self._schema) - table = pa.concat_tables(group) + writer = pq.ParquetWriter(local_path, self._file_schema()) + table = self._stamped(pa.concat_tables(group)) writer.write_table(table, row_group_size=table.num_rows) current_size += group_size current_rows += group_rows @@ -594,7 +607,7 @@ def close(self) -> None: except Exception as error: close_error = close_error or error if close_error is None and self.collect_stats: - schema = self._schema or pa.schema([]) + schema = self._file_schema() sample_rows = len(self._sample_prefix_logical_bytes) - 1 self._final_stats = FrameWriterStats( path_prefix=self.path_prefix, diff --git a/packages/kumo-core/src/kumo_core/train/output.py b/packages/kumo-core/src/kumo_core/train/output.py index ce24b5b0b..40e7f8c7e 100644 --- a/packages/kumo-core/src/kumo_core/train/output.py +++ b/packages/kumo-core/src/kumo_core/train/output.py @@ -739,6 +739,7 @@ def __init__( write_predictions: bool = True, artifact_stats_sample_rows: Mapping[PredictionArtifactType, int] | None = None, + constant_columns: Mapping[str, str] | None = None, ) -> None: makedirs(folder, exist_ok=True) prediction_dir = osp.join( @@ -762,6 +763,7 @@ def __init__( stats_sample_rows=(None if artifact_stats_sample_rows is None else artifact_stats_sample_rows.get( PredictionArtifactType.PREDICTIONS, 0)), + constant_columns=constant_columns, ) self._embedding_dir = embedding_dir self._embedding_prefix = osp.join(embedding_dir, @@ -770,6 +772,7 @@ def __init__( self._max_buffer_size = max_buffer_size self._output_extra_fields = output_extra_fields self._job_start_time = job_start_time + self._constant_columns = constant_columns self._upload_workers = upload_workers self._spool_root = spool_root self._collect_artifact_stats = artifact_stats_sample_rows is not None @@ -939,6 +942,7 @@ def write_embedding( stats_sample_rows=(None if not self._collect_artifact_stats else self._artifact_stats_sample_rows.get( PredictionArtifactType.EMBEDDINGS, 0)), + constant_columns=self._constant_columns, ) self._embedding_writer.write(embedding) diff --git a/packages/kumo-core/src/kumo_core/train/trainer.py b/packages/kumo-core/src/kumo_core/train/trainer.py index cc129a815..458a6536f 100644 --- a/packages/kumo-core/src/kumo_core/train/trainer.py +++ b/packages/kumo-core/src/kumo_core/train/trainer.py @@ -1500,6 +1500,7 @@ def get_log_dir( self, local: bool = False, original: bool = False, + fit_mode: FitMode | None = None, ) -> str: r"""Switch the log directory based on `self.fit_mode`. `version_xx/`: fit mode. @@ -1509,16 +1510,17 @@ def get_log_dir( If `original=True`, then we get the original (non-refitting) log directory. This way, we can update `config.json` in the original - directory""" + directory. `fit_mode` overrides `self.fit_mode`.""" + fit_mode = self.fit_mode if fit_mode is None else fit_mode path = self.root_dir if not local else self.local_root_dir # In case the `root_dir` contains a trainer config, it is a log dir # for a previous model version. if fsutil.exists(osp.join(path, 'config.json')): - if not original and self.fit_mode.is_refit: + if not original and fit_mode.is_refit: if path[-1] == '/': path = path[:-1] - path = f"{path}_{self.fit_mode.value}" + path = f"{path}_{fit_mode.value}" fsutil.makedirs(path, exist_ok=True) return path else: @@ -1528,10 +1530,10 @@ def get_log_dir( path = osp.join(path, self.name) if self.track_version: path = osp.join(path, self.version_dir_name) - if not original and self.fit_mode.is_refit: + if not original and fit_mode.is_refit: if path[-1] == '/': path = path[:-1] - path = f"{path}_{self.fit_mode.value}" + path = f"{path}_{fit_mode.value}" fsutil.makedirs(path, exist_ok=True) return path @@ -2507,16 +2509,21 @@ def _preprocess_top_k(self, top_k: int | None): # Batch prediction ######################################################## + def published_fit_mode(self) -> FitMode | None: + r"""The fit mode predictions use by default: the model trained on the + most data, or `None` before fitting.""" + if self.refit_cfg.refitted_full: + return FitMode.refit_full + if self.refit_cfg.refitted_trainval: + return FitMode.refit_trainval + if self.refit_cfg.fitted: + return FitMode.fit + return None + def _set_fit_mode_to_use(self, fit_mode_to_use: FitMode | None = None): - # Infer `fit_mode_to_use` - Use model that's trained on most data: if fit_mode_to_use is None: - if self.refit_cfg.refitted_full: - fit_mode_to_use = FitMode.refit_full - elif self.refit_cfg.refitted_trainval: - fit_mode_to_use = FitMode.refit_trainval - elif self.refit_cfg.fitted: - fit_mode_to_use = FitMode.fit - else: + fit_mode_to_use = self.published_fit_mode() + if fit_mode_to_use is None: raise RuntimeError("Trainer is not fitted yet.") else: self._check_readiness_for_prediction(fit_mode_to_use) diff --git a/packages/kumo-core/src/kumo_core/training/runtime.py b/packages/kumo-core/src/kumo_core/training/runtime.py index 096584f72..14ada8ff5 100644 --- a/packages/kumo-core/src/kumo_core/training/runtime.py +++ b/packages/kumo-core/src/kumo_core/training/runtime.py @@ -20,6 +20,7 @@ if TYPE_CHECKING: from kumo_core.embedding_models import EmbeddingModelStore + from kumo_core.instrumentation import TrainingInstrumentation class TrainingDataRuntime(Protocol): @@ -202,6 +203,7 @@ def local_training_bindings( training_input: Any, run_name: Optional[str] = None, embedding_models: Optional[EmbeddingModelStore] = None, + instrumentation: Optional[TrainingInstrumentation] = None, ) -> Any: r"""Binds one resolved training input to this process. @@ -219,7 +221,7 @@ def local_training_bindings( resolved_task=training_input.task, run_name=run_name, embedding_models=embedding_models), - instrumentation=NoOpTrainingInstrumentation(), + instrumentation=instrumentation or NoOpTrainingInstrumentation(), resolve_vital_artifacts=resolve_vital_artifact_sources, # The vital bundle lands beside the run. A platform backend would then # upload it to a registry; locally the path is the identity. diff --git a/packages/kumo-core/tests/sources/parquet/test_object_store.py b/packages/kumo-core/tests/sources/parquet/test_object_store.py index 4c2001d9a..5989a09f5 100644 --- a/packages/kumo-core/tests/sources/parquet/test_object_store.py +++ b/packages/kumo-core/tests/sources/parquet/test_object_store.py @@ -144,9 +144,9 @@ def get_file_info(self, target): reader = ObjectStoreSourceReader(MetadataOnlyFilesystem(), root=str(root)) version = reader.data_version(table) if table in ("events", "hive"): - assert len(version) == 64 + assert (len(version), version.kind) == (64, "listing") else: - assert int(version) > 0 + assert int(version) > 0 and version.kind == "timestamp" def test_version_of_exact_root_does_not_probe_other_suffixes(root): diff --git a/packages/kumo-core/tests/sources/unit/test_sql_identity.py b/packages/kumo-core/tests/sources/unit/test_sql_identity.py index 7f3ad3993..2d18419d9 100644 --- a/packages/kumo-core/tests/sources/unit/test_sql_identity.py +++ b/packages/kumo-core/tests/sources/unit/test_sql_identity.py @@ -36,6 +36,12 @@ def test_a_referenced_table_that_moved_gives_a_new_version(): assert stamped(version="v1") != stamped(version="v2") +def test_the_version_keeps_the_dependency_versions_it_folds(): + version = stamped(version="v1") + assert version.kind == "hour_bucket" + assert set(version.dependencies.values()) == {"v1"} + + def test_a_formatting_only_edit_does_not_move_the_version(): assert stamped( sql="select ID,\n AMOUNT from ORDERS where AMOUNT > 0") == (stamped()) diff --git a/packages/kumo-core/tests/test_artifact_export.py b/packages/kumo-core/tests/test_artifact_export.py index 39aba03b2..b3701f00f 100644 --- a/packages/kumo-core/tests/test_artifact_export.py +++ b/packages/kumo-core/tests/test_artifact_export.py @@ -222,3 +222,35 @@ def interrupted(source, destination, **kwargs): pd.testing.assert_frame_equal(pd.read_csv(target / 'part-00000.csv'), expected) assert (target / '_SUCCESS').exists() + + + +def test_sidecars_are_best_effort_and_precede_success(tmp_path): + import pyarrow.parquet as pq + from kumo_core.artifact_export import export_selected_files + from kumo_core.object_store_sink import _ArrowExportFiles + from pyarrow.fs import LocalFileSystem + + source = tmp_path / 'part-0.parquet' + pq.write_table(pa.table({'id': [1]}), source) + target = tmp_path / 'output' + fs = _ArrowExportFiles(LocalFileSystem(), None) + write_bytes, touch = fs.write_bytes, fs.touch + + def denied(path, value): + if path.endswith('bad.json'): + raise PermissionError('denied') + write_bytes(path, value) + + def checked_touch(path): + assert (target / 'meta.json').read_bytes() == b'{}' + touch(path) + + fs.write_bytes, fs.touch = denied, checked_touch + with pytest.warns(UserWarning, match='denied'): + export_selected_files([str(source)], f'{target}/', fs=fs, + identity='job-1', sidecars={ + 'meta.json': b'{}', + 'bad.json': b'{}' + }) + assert (target / '_SUCCESS').exists() diff --git a/packages/kumo-core/tests/test_artifacts.py b/packages/kumo-core/tests/test_artifacts.py index 35c2cef65..8d407c713 100644 --- a/packages/kumo-core/tests/test_artifacts.py +++ b/packages/kumo-core/tests/test_artifacts.py @@ -15,7 +15,11 @@ FilesystemVitalArtifactPublisher, VitalArtifactPublicationPolicy, VitalArtifactSources, + add_bundle_file, + bundle_id_from_manifest, + bundle_task_type, resolve_vital_artifact_sources, + verify_model_bundle, ) from kumo_core.config import NativeRuntimeConfig from kumo_core.manifest import RunManifest, RunPaths @@ -199,6 +203,7 @@ def test_vital_source_resolution_is_core_owned(tmp_path: Path, refit) -> None: class SourceTrainer(SimpleNamespace): _set_fit_mode_to_use = Trainer._set_fit_mode_to_use + published_fit_mode = Trainer.published_fit_mode @property def best_ckpt(self): @@ -236,3 +241,33 @@ def best_ckpt(self): assert sources.selected_checkpoint == checkpoints[selected_mode] assert sources.graph_manifest == bundle / "graph.json" assert json.loads(sources.metrics.read_text()) == {"test_accuracy": 0.8} + + + +def test_bundle_id_ignores_provenance_only(tmp_path: Path) -> None: + sources, _ = _vital_sources(tmp_path) + bundle = tmp_path / "bundle" + VitalArtifactPublicationPolicy().publish(sources, bundle) + bundle_id = verify_model_bundle(bundle) + + for payload in ('{"a": 1}', '{"a": 2}'): + (bundle / "provenance.json").write_text(payload) + add_bundle_file(bundle, "provenance", "provenance.json") + assert verify_model_bundle(bundle) == bundle_id + assert bundle_id_from_manifest(bundle) == bundle_id + roles = [ + entry["role"] for entry in json.loads( + (bundle / "publication-manifest.json").read_text())["artifacts"] + ] + assert roles.count("provenance") == 1 + + (bundle / "metrics.json").write_text('{"auc": 0}') + add_bundle_file(bundle, "metrics", "metrics.json") + assert verify_model_bundle(bundle) != bundle_id + + +def test_bundle_task_type_reads_the_prediction_schema(tmp_path: Path) -> None: + schema = tmp_path / "prediction-schema.json" + schema.write_text('{"task_type": "regression"}') + assert bundle_task_type({"prediction_schema": str(schema)}).value == ( + "regression") diff --git a/packages/kumo-core/tests/test_prediction_output_columns.py b/packages/kumo-core/tests/test_prediction_output_columns.py new file mode 100644 index 000000000..eee36d837 --- /dev/null +++ b/packages/kumo-core/tests/test_prediction_output_columns.py @@ -0,0 +1,23 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +from __future__ import annotations + +import pytest +from kumo_core.train.formatting import prediction_output_columns +from kumo_core.types import TaskType + + +@pytest.mark.parametrize('task_type, columns', [ + (TaskType.BINARY_CLASSIFICATION, + ['TARGET_PRED', 'False_PROB', 'True_PROB']), + (TaskType.REGRESSION, ['TARGET_PRED']), + (TaskType.FORECASTING, ['TIMESTAMP', 'TARGET_PRED']), + (TaskType.MULTICLASS_CLASSIFICATION, ['PREDICTED', 'CLASS', 'SCORE']), + (TaskType.MULTILABEL_CLASSIFICATION, ['PREDICTED', 'CLASS', 'SCORE']), + (TaskType.MULTILABEL_RANKING, ['CLASS', 'SCORE']), + (TaskType.TEMPORAL_LINK_PREDICTION, ['CLASS', 'SCORE']), + (TaskType.STATIC_LINK_PREDICTION, ['CLASS', 'SCORE']), +]) +def test_prediction_output_columns(task_type, columns) -> None: + assert prediction_output_columns(task_type) == columns diff --git a/packages/kumo-core/tests/test_searched_params.py b/packages/kumo-core/tests/test_searched_params.py new file mode 100644 index 000000000..c59c8390a --- /dev/null +++ b/packages/kumo-core/tests/test_searched_params.py @@ -0,0 +1,38 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import json + +from kumo_core.automl.optuna_searcher import TPESearcher +from kumo_core.model_plan import ModelPlan + + +def test_searcher_records_only_varied_fields_as_json(): + space = ModelPlan() + space.training_job.tune_metric = 'mae' + space.optimization.base_lr = [0.1, 0.01] + space.optimization.batch_size = [512] + space.optimization.optimizer = ['Adam', 'AdamW'] + space.neighbor_sampling.num_neighbors = [[16, 8], [8, 4]] + space.model_architecture.channels = [64, 128] + searcher = TPESearcher(space, seed=0) + + def chosen(cfg): + return { + 'optimization.base_lr': cfg.optimization.base_lr[0], + 'optimization.optimizer': cfg.optimization.optimizer[0], + 'model_architecture.channels': cfg.model_architecture.channels[0], + 'neighbor_sampling.num_neighbors': + str(cfg.neighbor_sampling.num_neighbors[0]), + } + + cfg = searcher.suggest() + assert searcher.sampled_params == chosen(cfg) + json.dumps(searcher.sampled_params) + + searcher.max_attempts = 0 + cfg = searcher.suggest() + assert searcher.sampled_params == chosen(cfg) + assert (searcher.sampled_params + is searcher.fallback_sampler.sampled_params) diff --git a/packages/kumo-core/tests/test_train_output.py b/packages/kumo-core/tests/test_train_output.py index ca14616a7..d5a74ae9e 100644 --- a/packages/kumo-core/tests/test_train_output.py +++ b/packages/kumo-core/tests/test_train_output.py @@ -171,6 +171,35 @@ def test_streaming_writer_stats_match_post_parquet_prefixes(tmp_path, table): assert stats.file_count == len(paths) +def test_constant_columns_are_appended_to_every_part(tmp_path): + from kumo_core.train.artifact_stats import schema_fingerprint + stamp = {"KUMO_RUN_ID": "run-1", "KUMO_BUNDLE_ID": "b1"} + writer = StreamingFrameWriter( + str(tmp_path / "part"), + max_file_size=256, + max_buffer_size=128, + stats_sample_rows=10, + index=False, + constant_columns=stamp, + ) + source = pa.table({"ID": np.arange(100, dtype=np.int64), + "COUNT": pa.array([None, 3] * 50, pa.int64())}) + writer.write(source) + writer.close() + + paths = sorted(tmp_path.glob("part-*.parquet")) + assert len(paths) > 1 + out = _read_parts(tmp_path) + assert out.schema.names == ["ID", "COUNT", *stamp] + assert out.select(["ID", "COUNT"]).equals(source) + for name, value in stamp.items(): + assert out.column(name).type == pa.string() + assert set(out.column(name).to_pylist()) == {value} + assert writer.stats.num_cols == 4 + assert writer.stats.schema_fingerprint == schema_fingerprint( + pq.read_schema(paths[0])) + + def test_stats_enabled_writer_enforces_row_group_memory_ceiling(tmp_path): writer = StreamingFrameWriter( str(tmp_path / "part"), diff --git a/packages/kumo-databricks/src/kumo_databricks/identity.py b/packages/kumo-databricks/src/kumo_databricks/identity.py index ae6023a93..c4ea11088 100644 --- a/packages/kumo-databricks/src/kumo_databricks/identity.py +++ b/packages/kumo-databricks/src/kumo_databricks/identity.py @@ -8,7 +8,11 @@ from databricks.sql.client import Cursor from kumo_core.sources.errors import SourceNotFoundError -from kumo_core.sources.identity import epoch_second, hour_version +from kumo_core.sources.identity import ( + SourceVersion, + epoch_second, + hour_version, +) from kumo_databricks.errors import is_not_found DESCRIBE = "DESCRIBE EXTENDED {source}" @@ -65,5 +69,5 @@ def table_version(cursor: Cursor, source: str) -> str: rows = cursor.fetchall() if rows: commit = epoch_second(pd.Timestamp(rows[0][1])) - return f"{rows[0][0]}_{commit}" + return SourceVersion(f"{rows[0][0]}_{commit}", "exact") return hour_version() diff --git a/packages/kumo-databricks/src/kumo_databricks/registry.py b/packages/kumo-databricks/src/kumo_databricks/registry.py index d2d88de09..05a50f4a7 100644 --- a/packages/kumo-databricks/src/kumo_databricks/registry.py +++ b/packages/kumo-databricks/src/kumo_databricks/registry.py @@ -10,7 +10,6 @@ import os import shutil import tempfile -from dataclasses import replace from pathlib import Path from typing import Any, Iterator, Mapping, Optional, cast from urllib.parse import quote @@ -20,12 +19,15 @@ from kumo_databricks.session import _DatabricksSessions, _target_config from kumo_databricks.sql import normalize_table_coordinate from kumoapi.plugin import ( + BUNDLE_ID_KEY, + ENGINE_MARKER, ModelRegistryLocation, PublishedModel, SourceConnector, + check_model_alias, + registry_lineage, ) -_ENGINE_MARKER = 'kumo_engine' _MODELS = '/api/2.1/unity-catalog/models' _LOG = logging.getLogger(__name__) @@ -56,6 +58,17 @@ def _tags(version: Mapping[str, Any]) -> dict[str, str]: return {tag['key']: tag['value'] for tag in version.get('tags', [])} +def _alias_path(name: str, alias: str) -> str: + return f'{_MODELS}/{quote(name, safe="")}/aliases/{quote(alias, safe="")}' + + +def _version_number(location: ModelRegistryLocation) -> Optional[int]: + version = location.version + if version is not None and (not version.isdecimal() or int(version) < 1): + raise ValueError('Databricks model versions are positive integers.') + return None if version is None else int(version) + + def _published( name: str, version: Mapping[str, Any], @@ -133,7 +146,9 @@ def publish_bundle( verify_model_bundle(root) _validate_catalog(dict(metadata=dict(metadata), metrics=dict(metrics))) catalog_name, schema, model = name.split('.') - tags = {_ENGINE_MARKER: '1'} + tags = {ENGINE_MARKER: '1'} + if metadata.get(BUNDLE_ID_KEY): + tags[BUNDLE_ID_KEY] = str(metadata[BUNDLE_ID_KEY]) try: client.registered_models.create(catalog_name=catalog_name, schema_name=schema, name=model) @@ -200,7 +215,7 @@ def _bundle_metadata(bundle: Path) -> tuple[dict[str, str], dict[str, float]]: def _download(client: WorkspaceClient, remote: str, destination: Path) -> None: for entry in client.files.list_directory_contents(remote): # A server-provided path must stay an immediate child of this directory. - path = entry.path or '' + path = (entry.path or '').rstrip('/') prefix = remote + '/' child = path[len(prefix):] if path.startswith(prefix) else '' if not child or '/' in child or child in ('.', '..') or '\\' in child: @@ -219,13 +234,18 @@ def _download(client: WorkspaceClient, remote: str, destination: Path) -> None: def fetch_bundle(client: WorkspaceClient, name: str, local_dir: Path, *, - version: Optional[str] = None) -> PublishedModel: + version: Optional[str] = None, + alias: Optional[str] = None) -> PublishedModel: from kumo_core.artifacts import verify_model_bundle - if version is None: + if alias is not None: + aliased = cast(dict[str, Any], + client.api_client.do('GET', _alias_path(name, alias))) + model_version = _get_version(client, name, aliased['version']) + elif version is None: versions = [ v for v in _versions(client, name) if v.get('status') == 'READY' - and _tags(v).get(_ENGINE_MARKER) == '1' + and _tags(v).get(ENGINE_MARKER) == '1' ] if not versions: raise ValueError('No ready Kumo model version in this registry.') @@ -233,7 +253,7 @@ def fetch_bundle(client: WorkspaceClient, name: str, local_dir: Path, *, else: model_version = _get_version(client, name, version) if model_version.get('status') != 'READY' or _tags(model_version).get( - _ENGINE_MARKER) != '1': + ENGINE_MARKER) != '1': raise ValueError('This version is not a ready Kumo model bundle.') destination = local_dir destination.parent.mkdir(parents=True, exist_ok=True) @@ -251,9 +271,12 @@ def fetch_bundle(client: WorkspaceClient, name: str, local_dir: Path, *, 'The registered model must contain one Kumo bundle.') bundle = bundles[0] verify_model_bundle(bundle) - published = _published(name, model_version) metadata, metrics = _bundle_metadata(bundle) - published = replace(published, metadata=metadata, metrics=metrics) + published = _published( + name, model_version, metadata={ + **metadata, + **registry_lineage(_tags(model_version)) + }, metrics=metrics) backup = Path(tmp) / 'previous' if destination.exists(): os.replace(destination, backup) @@ -277,8 +300,10 @@ def list_models(client: WorkspaceClient, catalog: str, continue for version in _versions(client, name): if version.get('status') == 'READY' and _tags(version).get( - _ENGINE_MARKER) == '1': - found.append(_published(name, version)) + ENGINE_MARKER) == '1': + found.append( + _published(name, version, + metadata=registry_lineage(_tags(version)))) return sorted(found, key=lambda item: (item.name, int(item.version))) @@ -321,15 +346,26 @@ def fetch_model( location: Optional[ModelRegistryLocation] = None, project_id: Optional[str] = None, ) -> PublishedModel: - """Download the named version, or the latest ready Kumo version.""" + """Download the named or aliased version, else the latest Kumo one.""" del project_id name = _address(connector, path, location) assert location is not None - if location.version is not None and (not location.version.isdecimal() - or int(location.version) < 1): - raise ValueError('Databricks model versions are positive integers.') + _version_number(location) return fetch_bundle(self._client(connector), name, - Path(local_dir).resolve(), version=location.version) + Path(local_dir).resolve(), version=location.version, + alias=getattr(location, 'alias', None)) + + def set_model_alias(self, connector: SourceConnector, + location: ModelRegistryLocation, alias: str) -> None: + """Point an alias at a version; UC moves it from any other.""" + name = _address(connector, None, location) + number = _version_number(location) + if number is None: + raise ValueError('Setting an alias needs a ModelRegistryLocation ' + 'naming a version.') + check_model_alias(alias) + self._client(connector).api_client.do('PUT', _alias_path(name, alias), + body=dict(version_num=number)) def list_models(self, connector: SourceConnector, *, project_id: Optional[str] = None) -> list[PublishedModel]: diff --git a/packages/kumo-databricks/src/kumo_databricks/session.py b/packages/kumo-databricks/src/kumo_databricks/session.py index af5a9209e..b274cd234 100644 --- a/packages/kumo-databricks/src/kumo_databricks/session.py +++ b/packages/kumo-databricks/src/kumo_databricks/session.py @@ -194,6 +194,11 @@ def _principal(self, connector: SourceConnector) -> Principal: f'accepted, as {self._default_hint}.'), ) + def table_location(self, connector: SourceConnector, + table_name: str) -> str: + """Fully qualified name of ``table_name`` in the connector's target.""" + return self._qualified(connector, self._bare(connector, table_name)) + @staticmethod def _qualified(connector: SourceConnector, table_name: Optional[str] = None) -> str: diff --git a/packages/kumo-databricks/tests/test_identity.py b/packages/kumo-databricks/tests/test_identity.py index 28bd35f48..03aeee1ad 100644 --- a/packages/kumo-databricks/tests/test_identity.py +++ b/packages/kumo-databricks/tests/test_identity.py @@ -49,6 +49,7 @@ def stamp(recorder=None, source=SOURCE, **overrides): def test_a_delta_table_is_stamped_with_its_version_and_commit_time(): """The first two history columns are the stamp; a fraction is dropped.""" assert stamp() == COMMIT_VERSION + assert stamp().kind == "exact" assert stamp(history=[(947, f"{COMMITTED}.9", "WRITE", "extra")]) == (COMMIT_VERSION) @@ -60,6 +61,7 @@ def test_the_stamp_changes_when_the_same_version_is_recommitted(): def test_a_table_with_no_history_falls_back_to_the_hour_window(): assert stamp(history=[]) == hour_version() + assert stamp(history=[]).kind == "hour_bucket" @pytest.mark.parametrize("refusal", [ diff --git a/packages/kumo-databricks/tests/test_registry.py b/packages/kumo-databricks/tests/test_registry.py index 050fcf18e..4e4c8cee4 100644 --- a/packages/kumo-databricks/tests/test_registry.py +++ b/packages/kumo-databricks/tests/test_registry.py @@ -15,6 +15,8 @@ from kumo_databricks.sink import DatabricksSinkBackend from kumoapi.data_source import DatabricksConnectorResourceConfig from kumoapi.plugin import ( + ENGINE_MARKER, + ModelAliasingSink, ModelFetchingSink, ModelListingSink, ModelPublishingSink, @@ -61,6 +63,7 @@ def test_sink_implements_registry_protocols(): assert isinstance(sink, ModelPublishingSink) assert isinstance(sink, ModelFetchingSink) assert isinstance(sink, ModelListingSink) + assert isinstance(sink, ModelAliasingSink) def test_publish_uses_connector_client_and_reports_assigned_version( @@ -141,7 +144,7 @@ def test_invalid_fetch_version_fails_before_authentication( def version(number, kumo=True, status='READY'): - tags = {registry._ENGINE_MARKER: '1'} if kumo else {} + tags = {ENGINE_MARKER: '1'} if kumo else {} return dict(version=number, status=status, tags=[dict(key=k, value=v) for k, v in tags.items()]) @@ -155,6 +158,7 @@ def __init__(self): upload=self.upload, list_directory_contents=self.list_directory, download=self.download) self.objects = {} + self.aliases = {} self.calls = [] self.deleted = [] self.versions = [ @@ -173,6 +177,12 @@ def delete(self, name, number): def do(self, method, path, query=None, body=None): self.calls.append((method, path, query, body)) + if '/aliases/' in path: + alias = path.rsplit('/', 1)[-1] + if method == 'PUT': + self.aliases[alias] = body['version_num'] + return {} + return dict(version=self.aliases[alias]) if method == 'POST': assert all(len(tag['value']) <= 256 for tag in body['tags']) self.versions.append( @@ -207,8 +217,11 @@ def list_directory(self, path): if key.startswith(path + '/'): relative = key[len(path) + 1:] child = relative.split('/')[0] - children[child] = SimpleNamespace(path=path + '/' + child, - is_directory='/' in relative) + directory = '/' in relative + # The Files API lists directories with a trailing slash. + children[child] = SimpleNamespace( + path=f"{path}/{child}{'/' if directory else ''}", + is_directory=directory) return children.values() def download(self, path): @@ -485,3 +498,49 @@ def test_invalid_bundle_metadata_keeps_destination(client, bundle, tmp_path, registry.fetch_bundle(client, name='main.analytics.churn', local_dir=destination, version='7') assert (destination / 'weights.bin').read_bytes() == b'previous' + + +def test_bundle_id_is_tagged_listed_and_fetched_by_version_and_alias( + connector, client, bundle, tmp_path, monkeypatch): + lineage = dict(kumo_bundle_id='ab' * 32) + sink = DatabricksSinkBackend() + monkeypatch.setattr(sink, '_client', lambda _: client) + published = sink.publish_model( + connector, str(bundle), location=ModelRegistryLocation('churn'), + metadata=dict(spec_version='1', kumo_engine='0', kumo_run_id='run-1', + **lineage)) + assert registry._tags(client.versions[-1]) == { + 'kumo_engine': '1', + **lineage + } + listed = registry.list_models(client, catalog='main', schema='analytics') + assert next(m for m in listed if m.version == '7').metadata == lineage + + sink.set_model_alias( + connector, ModelRegistryLocation('churn', version=published.version), + 'prod') + assert client.calls[-1] == ( + 'PUT', + '/api/2.1/unity-catalog/models/main.analytics.churn/aliases/prod', None, + dict(version_num=7)) + for location in (ModelRegistryLocation('churn', version='7'), + ModelRegistryLocation('churn', alias='prod')): + fetched = sink.fetch_model(connector, str(tmp_path / 'restored'), + location=location) + assert fetched.version == '7' + assert fetched.metadata == { + 'spec_version': '1', + 'predictive_query': 'PREDICT COUNT(*)', + 'snapshot_id': 'snapshot', + **lineage + } + + +def test_set_alias_needs_a_positive_version(connector, monkeypatch): + sink = DatabricksSinkBackend() + monkeypatch.setattr(sink, '_client', + lambda _: pytest.fail('resolved identity')) + with pytest.raises(ValueError, match='positive integer'): + sink.set_model_alias(connector, + ModelRegistryLocation('churn', version='0'), + 'prod') diff --git a/packages/kumo-databricks/tests/test_uc_registry.py b/packages/kumo-databricks/tests/test_uc_registry.py index 8866cf80f..8acca5647 100644 --- a/packages/kumo-databricks/tests/test_uc_registry.py +++ b/packages/kumo-databricks/tests/test_uc_registry.py @@ -76,7 +76,7 @@ def handle_request(self): return self.send(dict(registered_models=[model])) if path.endswith('/models/versions') and self.command == 'POST': assert payload['source'].startswith('file:') - assert payload['tags'] == [{'key': 'kumo_engine', 'value': '1'}] + assert {'key': 'kumo_engine', 'value': '1'} in payload['tags'] versions[identity] = dict(version=1, status='PENDING_REGISTRATION', tags=payload['tags']) diff --git a/packages/kumo-engine/README.md b/packages/kumo-engine/README.md index 3847ae5bd..af0a54fb3 100644 --- a/packages/kumo-engine/README.md +++ b/packages/kumo-engine/README.md @@ -135,6 +135,40 @@ Snapshot indexes are reused by graph definition and source version; changed versions get separate directories. Builds for the same graph use a filesystem lock. See the [implementation](https://github.com/NVIDIA/kumo-engine/blob/main/packages/kumo-engine/kumo/engine/graph.py). +To build an index on one machine and train on another, export it through an +S3, GCS or file connector and load it there: + +```python +location = graph.snapshot().export(S3Connector(), "s3://bucket/snapshots") +# On the training machine: +snapshot = GraphSnapshot.load_from_connector(S3Connector(), path=location) +``` + +`export` writes `///` and returns it. It first +removes `kumo_snapshot.json` and anything else at that location, uploads every +file, and writes `kumo_snapshot.json` (with per-range CRC32s) last, so an +interrupted export is never loadable and re-running starts over. A load either +completes or leaves nothing; a load that overlaps a re-export of the same +location can fail and should be retried. To keep several builds of one +snapshot (for example from two engine versions), export them to different +paths. +Snowflake, Databricks and BigQuery connectors cannot hold an export: use a +`FileConnector` on a mounted volume (e.g. `/Volumes///`) +or the S3/GCS connector on the bucket behind it. Load only from paths whose +writers you trust: checksums detect corruption, not tampering. + +`snapshot.source_changes(graph)` reports, per table, whether the source moved +on since the snapshot was built, from metadata only: + +```python +{"users": {"status": "unchanged", "recorded": "1792539971", "current": "1792539971"}, + "events": {"status": "changed", "recorded": "1792539971", "current": "1792539974"}} +``` + +The status is `unknown` when there is nothing reliable to compare: no recorded +version, a different connector or location, Custom SQL, or a version whose +`version_kind` is `hour_bucket` or `unknown` (views, undated sources). `load_from_connector(..., graph=graph)` warns when tables `changed`. + The return values expose the retained files directly: ```python @@ -150,6 +184,21 @@ print(prediction_result.files) # selected Parquet files predictions = prediction_result.frame # collect them into a pandas DataFrame ``` +### Progress events + +`fit(callbacks=[...])` calls each callback with every progress `Event`: +`fit_started`, then per search trial `trial_started`, `epoch_ended` and +`trial_ended`, then `fit_ended`. Ended events carry `params["status"]` and, on +failure, a short `params["error"]`. A callback that raises is logged and not +called again; it never fails the fit. For JSON lines on stdout: + +```python +import json + +trainer.fit(training_table, snapshot=snapshot, + callbacks=[lambda e: print(json.dumps(e.to_dict()))]) +``` + ### Holdout predictions `fit_result.holdout_df()` reads test predictions from `holdout_path` into a pandas @@ -194,6 +243,26 @@ for evaluation. Older bundles may contain an AutoML search space instead of the fitted architecture. They support prediction but are unreliable warm-start sources. +### Scoring defaults + +A bundle can carry the prediction arguments `predict` uses when a call omits them: + +```python +from kumo.engine import ModelOutputConfig, ModelRegistryLocation, ScoringDefaults + +tuned = trainer.with_scoring_defaults( + ScoringDefaults(binary_classification_threshold=0.37)) +tuned.predict(prediction_table, snapshot=snapshot) # thresholds at 0.37 +tuned.export(ModelOutputConfig(registry, location=ModelRegistryLocation("CHURN_GNN"))) +``` + +`with_scoring_defaults` returns a trainer on a copy with its own `bundle_id`; the copy's +`provenance.parents.derived_from` names the original. The threshold applies to +binary and multilabel classification, `num_classes_to_return` to multiclass, multilabel +and link prediction. Explicit `predict` arguments and a query's `TOP k` win; +`ScoringDefaults()` removes the defaults. `Trainer.export(config)` publishes a fitted or +loaded bundle like `fit(export_config=...)`. + ### Caller-supplied training and prediction rows A connector table can supply labeled training rows directly. The table uses @@ -300,6 +369,171 @@ Custom forecasting rows must include precomputed `future_labels` for training, year-over-year features when enabled. Constructors preserve these sequences; they do not derive them from entity, target, and timestamp columns. +## Experiment tracking + +`fit(tracking=...)` records a fit in an experiment tracker: one run per fit, +with each hyperparameter-search trial attempt as a child run (MLflow) or +inside the fit run (Snowflake). It is opt-in; tracking errors are logged as warnings and never fail the fit. Only global rank 0 tracks. +Tracker calls run on a background thread, so a slow or unreachable tracking +server does not hold up training: at the end of the fit, `fit` waits at most +30 s for tracking to finish (the rest continues in the background). During +the fit, MLflow HTTP calls use 2 retries and a 30 s timeout unless +`MLFLOW_HTTP_REQUEST_MAX_RETRIES` / `MLFLOW_HTTP_REQUEST_TIMEOUT` are set. + +```python +from kumo.engine.integrations.mlflow import MlflowTracking # kumo-engine[mlflow] + +result = trainer.fit(training_table, snapshot=snapshot, + tracking=MlflowTracking(experiment_name="users-spend")) +``` + +MLflow runs go to `MLFLOW_TRACKING_URI`: on Databricks this is automatic; for +SageMaker managed MLflow set it to the tracking server ARN with the +`sagemaker-mlflow` plugin installed. With nothing configured MLflow writes a +local store in the working directory. Per-epoch metrics use step = epoch. +After the fit, `tracking.run_id` is the fit run's MLflow run id. + +The fit's source tables (from the snapshot manifest, plus a custom training +table's source) are logged as MLflow run inputs: Databricks tables as Delta +tables at the snapshot's Delta version, so Unity Catalog links each table to a +model registered from the run in a Databricks notebook or job; file and +object-store sources by URI. Other warehouse tables and Custom SQL sources are +not logged as inputs; they stay in the bundle's `provenance.json`. + +Snowflake records the fit in Snowflake Experiment Tracking (Snowsight: AI & ML +» Experiments), using a Snowpark session: + +```python +from kumo.engine.integrations.snowflake import SnowflakeTracking # kumo-engine[snowflake] + +result = trainer.fit(training_table, snapshot=snapshot, + tracking=SnowflakeTracking(session, experiment_name="users_spend")) +``` + +Snowflake has no child runs or tags, so trial attempts are logged in the fit +run as `trial-1-attempt-1/` params and metrics, and tags are params of +the same name. `native_model_lineage=True` links a model the fit publishes +through Snowflake lineage. Names beyond Snowflake's 200 metrics and 1000 params +per run are counted in `kumo.dropped_metrics` / `kumo.dropped_params`. If a +Snowflake run is already active in the process, the fit is not tracked there; +an existing run is never resumed (a taken `run_name` gets a suffix). Snowflake +Experiment Tracking has no dataset inputs, so source tables are not recorded +there. The Snowpark session's parameters are left as they are. After the fit, +`tracking.run_id` is the fit run's name. + +Other systems plug in through the `kumo.engine.tracking.Tracking` protocol, +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's run after each of its events, and `finish(record, error)` with the +fit run's final record (`None` when tracking stopped early). A tracker that +raises is not called again for that fit and its run is ended as failed. +Params are the resolved model plan (fit) or searched hyperparameters and +`attempt` (trial). Runs carry pointers, not lineage: `kumo.kind` and +`kumo.run_id` (plus `kumo.parent_run`, `kumo.status` and `kumo.error` for +outcomes); lineage stays in the model bundle, whose `provenance.json` is logged +on the fit run when the bundle has one, with `query.pql` and `model_plan.json`. +`RunRecord.inputs` maps each source table to its `connector`, `location`, +`version` and `version_kind`. Your `tags=` are added as given, except `kumo.*` +names and param names. + +The fit run also carries `kumo.bundle_id` once the fit has a bundle, and +under a scheduler `kumo.scheduler`, `kumo.job_id` and `kumo.job_run_id`. In the +other direction, a tracked fit's bundle records `"tracking": {"run_id": ...}` +in `provenance.json`, the `kumo.run_id` every tracker tags its fit run with. + +MLflow models: `log_model(trainer, snapshot=...)` logs a bundle as a pyfunc +model, with the bundle's source tables as inputs of the run it is logged in, +that scores where `kumo-engine[s3,gcs]` of the same version is installed (its +pip requirement); `load_model(uri)` restores the `Trainer`. +`MlflowTracking(registered_model_name=..., aliases=...)` logs, registers and +aliases the fit's model on the fit run (a failed fit logs none). `fit()` waits +for the upload and registration before it returns, for up to an hour, with a +600 s MLflow HTTP timeout during the upload (unless +`MLFLOW_HTTP_REQUEST_TIMEOUT` is set); if they fail, the tracking thread stops +or the hour passes, `fit()` warns, still returns its result, the run ends +`FAILED`, and a model that finishes uploading later is not registered. The model's metadata is `kumo_bundle_id`; its `provenance.json` +artifact is the bundle's. A tracker can do the same through the optional +`log_model(result, snapshot)` protocol member. + +```python +from kumo.engine.integrations.mlflow import MlflowTracking + +result = trainer.fit(training_table, snapshot=snapshot, tracking=MlflowTracking( + experiment_name="/Shared/users_spend", + registered_model_name="main.ml.users_spend", aliases=["champion"])) +``` + +```python +model = mlflow.pyfunc.load_model("models:/users-spend@champion") +scores = model.predict(entities, params={"snapshot": "s3:///snapshots/"}) +``` + +The pyfunc model loads a snapshot on its first `predict` with that URI and +keeps it for the life of the process: each serving replica, and each Spark +executor of `mlflow.pyfunc.spark_udf`, downloads the snapshot once and keeps +every snapshot it has loaded (no eviction), so its disk must hold them. With +`serving_snapshot=` (on `log_model` or +`MlflowTracking`), each replica downloads that snapshot when it starts, and +`predict` uses it when no `snapshot` param is given. +`params={"population": True}` scores every entity and returns the whole +frame in memory; for large populations use `Trainer.predict` with a +`PredictionOutputConfig`. Each `predict` call works in its own scratch +directory and removes it before returning. + +The `mlflow` extra needs an MLflow 3.13 or later client: DiskGraph pins +pyarrow 24, which older MLflow releases exclude. SageMaker managed MLflow +servers run MLflow 3.0; a 3.17 client against a 3.0 server handles logging, +registration, aliases and `load_model`. + +On Databricks, the parent folder of an MLflow experiment path must exist. +Registering to Unity Catalog from outside Databricks also needs `boto3`, which +the `mlflow` extra does not install. + +#### Which registry + +There are two ways to register a model; pick one per registered model, since +the two store different formats. Both record only the `kumo_bundle_id` +pointer; lineage stays in the bundle's `provenance.json`, which an MLflow +model also carries as its `provenance.json` artifact. An MLflow model logs +each source table (and a custom training table) as a run input, Unity +Catalog tables as Delta datasets. Both accept the same aliases +(`kumoapi.plugin.check_model_alias`). + +| | Connector registry | MLflow model | +|---|---|---| +| Publish | `fit(export_config=ModelOutputConfig(connector, location=ModelRegistryLocation(name, alias=...)))` or `Trainer.export` | `MlflowTracking(registered_model_name=..., aliases=[...])` or `log_model(...)` | +| Load | `Trainer.load_from_connector(connector, location=...)` | `mlflow.pyfunc.load_model(uri)` or `load_model(uri)` | +| Promote | `set_model_alias(connector, location, alias)` | MLflow's `set_registered_model_alias` | +| Databricks | `DatabricksConnector` (Unity Catalog) | Unity Catalog through MLflow | +| SageMaker | `S3Connector` path | SageMaker managed MLflow (`MLFLOW_TRACKING_URI` = tracking server ARN) | +| Snowflake | `SnowflakeConnector` (Snowflake Model Registry) | not available | +| Plain VM | `FileConnector`, `S3Connector`, `GCSConnector` path (no aliases) | any MLflow tracking server | + +## Scheduled jobs + +Wrap a job script's body in `kumo.job()` so a scheduler gets one JSON summary +line on stdout and an exit code it can act on: + +```python +import kumo.engine as kumo + +with kumo.job() as run: + snapshot = graph.snapshot() + result = kumo.Trainer(plan).fit(training_table, snapshot=snapshot) + run.record(bundle_id=result.bundle_id, snapshot_id=snapshot.id) +``` + +The summary carries `status`, `exit_code`, the scheduler context, the recorded +ids and outputs, the first 20 distinct warnings other than deprecations (300 +characters each), and `error` on failure. The line is printed +however the body ends, including Ctrl-C and `sys.exit()`. The job exits 0 on +success, 1 on failure and 143 on SIGTERM; `error.kind` says whether a failure +is `transient` (a source, connection or timeout error worth retrying), +`permanent` (a source refused the request) or `unknown`. Scheduler ids come +from `job(job_id=..., job_run_id=...)`, `KUMO_JOB_ID`/`KUMO_JOB_RUN_ID`, or the +detected platform, and are recorded in every manifest and tracking run the job +writes. Logs go to stderr. + ## BigQuery staged ingestion BigQuery streams through the Storage Read API by default. For large tables, @@ -376,7 +610,18 @@ cols=..., source_table=events.sql(query).to_source_table(), ...))`. ## Model-plan suggestions Given `pquery = PredictiveQuery(graph, query)` and `snapshot = graph.snapshot()`, -suggest and customize a plan before training: +suggest and customize a plan before training. A scoring job without `Graph` +code loads the snapshot the training job exported +(`GraphSnapshot.load_from_connector(store, path=location)`, where `location` is +what `graph.snapshot().export(store, 'snapshots')` returned) and scores it with +`trainer.predict_snapshot(snapshot)`, which reuses the trained query and +table-generation plan (`entities=` scores given rows instead; +`prediction_time=`, `output_config=` and the thresholds are as in `predict`; +it warns when the snapshot was taken of a different graph than the training +snapshot). +`PredictiveQuery.from_snapshot(snapshot, trainer.training_query)` binds the +query directly. Its graph check only means the snapshot belongs to the query's +graph, not that it is compatible with the trained model. ```python from kumo.engine import Trainer @@ -416,13 +661,14 @@ restored = Trainer.load(result.bundle_dir) # pass the recorded path For an exported model, use `Trainer.load_from_connector(connector, path="model/version")`, or -`location=ModelRegistryLocation(...)` for Snowflake Model Registry. Rebuild -the live graph and snapshot before generating prediction inputs. Bundles -contain trained artifacts, not source data or graph storage. +`location=ModelRegistryLocation(...)` for Snowflake or Databricks. A scoring +job with `Graph` code rebuilds the live graph and snapshot; one without loads an +exported snapshot with `GraphSnapshot.load_from_connector(...)` and calls +`trainer.predict_snapshot(snapshot)`. Bundles contain trained artifacts, not +source data or graph storage. With `export_config`, `result.published_model` records the assigned name, version, -and URI; otherwise it is `None`. Snowflake also records training metrics and -metadata for model discovery and comparison. Retain the assigned version for reload: +and URI; otherwise it is `None`. Retain the assigned version for reload: ```python result = trainer.fit( @@ -438,6 +684,41 @@ restored = Trainer.load_from_connector(snowflake, location=ModelRegistryLocation S3, GCS, and file connectors expose the published path or URI in the same result. +### Publish and promote + +Snowflake Model Registry and Databricks Unity Catalog connectors support aliases: + +```python +from kumo.engine import ModelOutputConfig, ModelRegistryLocation, set_model_alias + +result = trainer.fit(training_table, snapshot=snapshot, export_config=ModelOutputConfig( + registry, location=ModelRegistryLocation("CHURN_GNN", alias="challenger"))) +published = result.published_model + +# Promote: the alias moves from any version that held it. +set_model_alias(registry, ModelRegistryLocation( + published.name, version=published.version), "prod") + +# Scoring jobs load whichever version "prod" points at. +trainer = Trainer.load_from_connector(registry, location=ModelRegistryLocation( + "CHURN_GNN", alias="prod")) +``` + +An alias is a letter or underscore followed by letters, digits or underscores; +`DEFAULT`, `FIRST`, `LAST`, `LATEST` and `v` are reserved, and `fit()` +rejects an invalid export alias before training. The alias is set after +publishing; if that fails, `fit()` warns and still returns the published model. +Snowflake allows one alias per version; if moving an alias fails, it is put +back on the version that held it. Snowflake's `DEFAULT`, `FIRST` and `LAST` can +be loaded but not set. File, S3 and GCS connectors have no aliases. + +Each published version records only pointers: the `kumo_engine` marker and +`kumo_bundle_id` (Snowflake version metrics, Databricks version tags). Training +metrics stay in the bundle's `metrics.json` and on the tracked fit run. The bundle's own `provenance.json` holds its +lineage. `list_models` and `fetch_model` return `kumo_bundle_id` in +`PublishedModel.metadata`; `load_from_connector` +warns if the fetched bundle differs from the recorded `kumo_bundle_id`. + ## Prediction and embedding destinations Request both artifacts with separate destinations. Mapping keys select which @@ -469,6 +750,14 @@ Choose `extra_fields` during `predict`; `export()` rejects them. Directory retri can reuse their own output; different exports need an empty destination. Explicit `.parquet`/`.csv` paths retain single-file/multipart behavior. +`PredictionOutputConfig(..., lineage_columns=True)` passed to `predict` writes +constant `KUMO_RUN_ID`, `KUMO_BUNDLE_ID` and `KUMO_SNAPSHOT_ID` string columns +into every output row (`result.files`, `result.frame` and every export); unknown +ids are omitted. It changes the destination's column set, so enable it on a new +table or directory: on a table created without it, Snowflake `OVERWRITE`, +BigQuery `APPEND` and Databricks `APPEND` fail, and Snowflake `APPEND` silently +drops the new columns. `export()` copies whatever columns the result has. + Destinations export independently; failures may leave completed outputs and partial directories. Retry `result.export()` with the same destinations. Directory retries delete `_SUCCESS` before copying and restore it on completion, so delete permission diff --git a/packages/kumo-engine/kumo/engine/__init__.py b/packages/kumo-engine/kumo/engine/__init__.py index fe68f71be..766f5875a 100644 --- a/packages/kumo-engine/kumo/engine/__init__.py +++ b/packages/kumo-engine/kumo/engine/__init__.py @@ -25,6 +25,8 @@ from kumoapi.plugin import ModelRegistryLocation from kumoapi.table import Column +from kumo.engine._events import Event +from kumo.engine._job import JobRun, job from kumo.engine._version import __version__ from kumo.engine.config import configure from kumo.engine.connector import ( @@ -51,14 +53,19 @@ ModelOutputConfig, PredictionOutputConfig, PredictionResult, + ScoringDefaults, Trainer, TrainingResult, TrainingTableOutputConfig, + set_model_alias, ) __all__ = [ "__version__", "configure", + "Event", + "job", + "JobRun", "EmbeddingModelStore", "GloVeVectors", "TokenEmbeddingModel", @@ -89,9 +96,11 @@ "Trainer", "TrainingResult", "PredictionResult", + "ScoringDefaults", "ModelOutputConfig", "PredictionOutputConfig", "TrainingTableOutputConfig", + "set_model_alias", ] ensure_nofile() diff --git a/packages/kumo-engine/kumo/engine/_events.py b/packages/kumo-engine/kumo/engine/_events.py new file mode 100644 index 000000000..7309aab19 --- /dev/null +++ b/packages/kumo-engine/kumo/engine/_events.py @@ -0,0 +1,86 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +r"""Progress events reported by `Trainer.fit`.""" + +import hashlib +import logging +import math +import time +from dataclasses import asdict, dataclass, field +from typing import ( + Any, + Callable, + Dict, + Literal, + Mapping, + Optional, + Sequence, + cast, +) + +from kumo_core.instrumentation import NoOpTrainingInstrumentation + +logger = logging.getLogger(__name__) + +EventName = Literal['fit_started', 'trial_started', 'epoch_ended', + 'trial_ended', 'fit_ended'] + + +@dataclass(frozen=True) +class Event: + r"""One progress event. ``trial`` and ``epoch`` count from 1.""" + + name: EventName + run_id: str + time: float + trial: Optional[int] = None + epoch: Optional[int] = None + metrics: Mapping[str, float] = field(default_factory=dict) + params: Mapping[str, Any] = field(default_factory=dict) + + def to_dict(self) -> Dict[str, Any]: + return asdict(self) + + +def query_sha256(query: str) -> str: + return hashlib.sha256(query.encode()).hexdigest() + + +def _finite_floats(metrics: Mapping[str, Any]) -> Dict[str, float]: + result = {} + for key, value in metrics.items(): + try: + number = float(value) + except (TypeError, ValueError): + continue + if math.isfinite(number): + result[key] = number + return result + + +class EventEmitter(NoOpTrainingInstrumentation): + def __init__(self, run_id: str, + callbacks: Sequence[Callable[[Event], None]]) -> None: + self.run_id = run_id + self._callbacks = list(callbacks) + + def emit(self, name: EventName, *, trial: Optional[int] = None, + epoch: Optional[int] = None, + metrics: Optional[Mapping[str, float]] = None, + params: Optional[Mapping[str, Any]] = None) -> None: + event = Event(name=name, run_id=self.run_id, time=time.time(), + trial=trial, epoch=epoch, + metrics=_finite_floats(metrics or {}), + params=dict(params or {})) + for callback in list(self._callbacks): + try: + callback(event) + except Exception: + logger.warning( + 'Progress callback %r failed; it is disabled for the ' + 'rest of run %s', callback, self.run_id, exc_info=True) + self._callbacks.remove(callback) + + def event(self, name: str, **fields: Any) -> None: + self.emit(cast(EventName, name), **fields) diff --git a/packages/kumo-engine/kumo/engine/_job.py b/packages/kumo-engine/kumo/engine/_job.py new file mode 100644 index 000000000..88dd40859 --- /dev/null +++ b/packages/kumo-engine/kumo/engine/_job.py @@ -0,0 +1,184 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +r"""Run a Python script as a scheduled job: :func:`job`.""" + +import contextlib +import json +import logging +import os +import signal +import sys +import threading +import traceback +import warnings +from typing import Any, Dict, Iterator, Mapping, NoReturn, Optional + +_MAX_MESSAGE = 500 +_MAX_WARNING = 300 +_MAX_WARNINGS = 20 +_DEPRECATIONS = (DeprecationWarning, PendingDeprecationWarning, FutureWarning) + + +class Terminated(BaseException): + r"""The process received SIGTERM.""" + + +class JobRun: + r"""What a job records for its summary line; see :func:`job`.""" + def __init__(self) -> None: + self.summary: Dict[str, Any] = { + 'bundle_id': None, + 'snapshot_id': None, + 'run_id': None, + 'outputs': {}, + 'warnings': [], + } + + def record(self, *, bundle_id: Optional[str] = None, + snapshot_id: Optional[str] = None, run_id: Optional[str] = None, + outputs: Optional[Mapping[str, Any]] = None) -> None: + r"""Adds ids and outputs (name to absolute URI or table) to the + summary; ``None`` leaves a field.""" + for key, value in (('bundle_id', bundle_id), + ('snapshot_id', snapshot_id), ('run_id', run_id)): + if value is not None: + self.summary[key] = value + self.summary['outputs'].update(outputs or {}) + + +def _error_kind(error: BaseException) -> str: + from kumo_core.sources.errors import ( + PermanentSourceError, + TransientSourceError, + ) + + cause: Optional[BaseException] = error + while cause is not None: + if isinstance(cause, PermanentSourceError): + return 'permanent' + if isinstance(cause, + (TransientSourceError, ConnectionError, TimeoutError)): + return 'transient' + cause = cause.__cause__ + return 'unknown' + + +def _exit_code(error: Optional[BaseException]) -> int: + if error is None: + return 0 + if isinstance(error, SystemExit): + code = error.code + return code if isinstance(code, int) else 0 if code is None else 1 + if isinstance(error, KeyboardInterrupt): + return 130 + if isinstance(error, Terminated): + return 143 + return 1 + + +@contextlib.contextmanager +def job(*, job_id: Optional[str] = None, + job_run_id: Optional[str] = None) -> Iterator[JobRun]: + r"""Runs the body as one scheduled job. + + ``job_id``/``job_run_id`` (else ``KUMO_JOB_ID``/``KUMO_JOB_RUN_ID`` or the + detected scheduler) are recorded in the manifests and tracking tags the + body writes. Logs go to stderr. On exit, however the body ends, one JSON + line goes to stdout: ``status``, ``exit_code``, the scheduler context, what + the body passed to :meth:`JobRun.record`, the first 20 distinct warnings + other than deprecations (300 characters each), and ``error`` + (``type``, ``kind``: ``transient``/``permanent``/``unknown``, ``message``) + on failure. A failed body exits the process with 1; SIGTERM with 143. + + Example:: + + with kumo.job() as run: + result = Trainer(plan).fit(table, snapshot=snapshot) + run.record(bundle_id=result.bundle_id, snapshot_id=snapshot.id) + """ + logging.basicConfig(level=logging.INFO, stream=sys.stderr) + ids = {'KUMO_JOB_ID': job_id, 'KUMO_JOB_RUN_ID': job_run_id} + previous = {name: os.environ.get(name) for name in ids} + for name, value in ids.items(): + if value is not None: + os.environ[name] = value + run = JobRun() + shown = warnings.showwarning + showing = threading.local() + + def show(message: Any, category: Any, *args: Any, **kwargs: Any) -> None: + if getattr(showing, 'active', False): + warnings._showwarnmsg_impl( # type: ignore[attr-defined] + warnings.WarningMessage(message, category, *args, **kwargs)) + return + text = str(message)[:_MAX_WARNING] + recorded = run.summary['warnings'] + if (not issubclass(category, _DEPRECATIONS) and text not in recorded + and len(recorded) < _MAX_WARNINGS): + recorded.append(text) + showing.active = True + try: + shown(message, category, *args, **kwargs) + finally: + showing.active = False + + warnings.showwarning = show + error: Optional[BaseException] = None + try: + with _sigterm_raises(): + yield run + except BaseException as raised: + error = raised + if isinstance(raised, (Exception, Terminated)): + traceback.print_exception(raised, file=sys.stderr) + raise SystemExit(_exit_code(raised)) from raised + raise + finally: + warnings.showwarning = shown + try: + _print_summary(run, error) + finally: + for name, value in previous.items(): + if value is None: + os.environ.pop(name, None) + else: + os.environ[name] = value + + +def _print_summary(run: JobRun, error: Optional[BaseException]) -> None: + from kumo.engine._provenance import scheduler_context + + code = _exit_code(error) + scheduler = scheduler_context() + line = { + 'event': 'kumo.engine.job', + 'status': 'succeeded' if code == 0 else 'failed', + 'exit_code': code, + 'scheduler': None if scheduler is None else scheduler.to_dict(), + **run.summary, + } + if code != 0 and error is not None: + line['error'] = { + 'type': type(error).__name__, + 'kind': _error_kind(error), + 'message': str(error)[:_MAX_MESSAGE], + } + print(json.dumps(line, default=str), flush=True) + + +def _terminate(signum: int, frame: Any) -> NoReturn: + raise Terminated('Terminated by SIGTERM') + + +@contextlib.contextmanager +def _sigterm_raises() -> Iterator[None]: + try: + previous = signal.signal(signal.SIGTERM, _terminate) + except ValueError: # not the main thread + yield + return + try: + yield + finally: + signal.signal(signal.SIGTERM, previous) diff --git a/packages/kumo-engine/kumo/engine/_manifest.py b/packages/kumo-engine/kumo/engine/_manifest.py new file mode 100644 index 000000000..fd066e57e --- /dev/null +++ b/packages/kumo-engine/kumo/engine/_manifest.py @@ -0,0 +1,286 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Shared JSON manifest format for engine artifacts.""" + +import json +import types +import warnings +from collections.abc import Mapping as AbcMapping +from dataclasses import dataclass, field, fields +from datetime import datetime, timezone +from pathlib import Path +from typing import ( + Any, + Dict, + List, + Mapping, + Optional, + Self, + Type, + TypeVar, + Union, + get_args, + get_origin, + get_type_hints, +) + +from kumo_core.io.filesystem import atomic_write_bytes + +SCHEMA_VERSION = 1 +_OMIT_NONE = 'omit_none' + +R = TypeVar('R', bound='Record') +M = TypeVar('M', bound='Manifest') + + +def omit_none() -> Any: + """A field written only when it is not None.""" + return field(default=None, metadata={_OMIT_NONE: True}) + + +@dataclass(frozen=True, kw_only=True) +class Record: + """A JSON object; keys this version does not know are kept in ``extra``.""" + + extra: Mapping[str, Any] = field(default_factory=dict, repr=False, + compare=False) + + def to_dict(self) -> Dict[str, Any]: + data = dict(self.extra) + for item in fields(self): + value = getattr(self, item.name) + if item.name != 'extra' and not (value is None + and item.metadata.get(_OMIT_NONE)): + data[item.name] = _plain(value) + return data + + @classmethod + def from_dict(cls: Type[R], data: Mapping[str, Any]) -> R: + hints = get_type_hints(cls) + names = {item.name for item in fields(cls)} - {'extra'} + return cls( + **{ + name: _typed(hints[name], data[name]) + for name in names if name in data + }, extra={ + key: value + for key, value in data.items() if key not in names + }) + + +def _plain(value: Any) -> Any: + if isinstance(value, Record): + return value.to_dict() + if isinstance(value, AbcMapping): + return {key: _plain(item) for key, item in value.items()} + if isinstance(value, (list, tuple)): + return [_plain(item) for item in value] + return value + + +def _typed(hint: Any, value: Any) -> Any: + origin, args = get_origin(hint), get_args(hint) + if value is None: + return None + if origin in (Union, types.UnionType): + return _typed(next(arg for arg in args if arg is not type(None)), value) + if isinstance(hint, type) and issubclass(hint, Record): + return hint.from_dict(value) + if origin in (dict, AbcMapping) and args: + return {key: _typed(args[1], item) for key, item in value.items()} + if origin is list and args: + return [_typed(args[0], item) for item in value] + return value + + +@dataclass(frozen=True, kw_only=True) +class Environment(Record): + kumo_engine: Optional[str] + build_commit: Optional[str] + kumo_core: Optional[str] + diskgraph: Optional[str] + diskgraph_build_commit: Optional[str] + torch: Optional[str] + python: Optional[str] + + +@dataclass(frozen=True, kw_only=True) +class SchedulerContext(Record): + platform: str + job_id: Optional[str] + job_run_id: Optional[str] + + +@dataclass(frozen=True, kw_only=True) +class SourceTableEntry(Record): + connector: Optional[str] + location: Optional[str] + version: Optional[str] + version_kind: str + dependencies: Optional[Dict[str, str]] = omit_none() + normalized_sql_fingerprint: Optional[str] = omit_none() + sql: Optional[str] = omit_none() + dialect: Optional[str] = omit_none() + namespace: Optional[List[Any]] = omit_none() + + +@dataclass(frozen=True, kw_only=True) +class Manifest(Record): + kind: str + schema_version: int = SCHEMA_VERSION + id: Optional[str] + created_at: Optional[str] + environment: Optional[Environment] + scheduler: Optional[SchedulerContext] + reconstructed: Optional[bool] = omit_none() + + +def stamp() -> Dict[str, Any]: + """Envelope fields for a manifest written now.""" + from kumo.engine._provenance import environment, scheduler_context + + return dict( + created_at=datetime.now(timezone.utc).isoformat(timespec='seconds'), + environment=environment(), scheduler=scheduler_context()) + + +#: Envelope fields for a manifest rebuilt after the fact. +RECONSTRUCTED: Dict[str, Any] = dict(created_at=None, environment=None, + scheduler=None, reconstructed=True) + + +@dataclass(frozen=True, kw_only=True) +class SnapshotParents(Record): + source_tables: Dict[str, SourceTableEntry] + + +@dataclass(frozen=True, kw_only=True) +class SnapshotManifest(Manifest): + kind: str = 'graph_snapshot' + parents: SnapshotParents + record_sources: Optional[bool] = omit_none() + + +@dataclass(frozen=True, kw_only=True) +class TrainingTableParents(Record): + snapshot: str + source_table: Optional[SourceTableEntry] = omit_none() + + +@dataclass(frozen=True, kw_only=True) +class TrainingTableManifest(Manifest): + kind: str = 'training_table' + parents: TrainingTableParents + query: str + generation_plan: Dict[str, Any] + task_type: str + split_seed: Optional[int] + + +@dataclass(frozen=True, kw_only=True) +class BundleParents(Record): + training_table: Optional[str] = None + warm_start: Optional[str] = None + derived_from: Optional[str] = omit_none() + + +@dataclass(frozen=True, kw_only=True) +class FitTiming(Record): + started_at: str + finished_at: str + seconds: float + + +@dataclass(frozen=True, kw_only=True) +class BundleTrainingTable(Record): + """Written with ``source_table``'s keys inline.""" + + source: Optional[str] + snapshot: Optional[str] + source_table: Optional[SourceTableEntry] = None + + def to_dict(self) -> Dict[str, Any]: + data = {**self.extra, 'source': self.source, 'snapshot': self.snapshot} + if self.source_table is not None: + data.update(self.source_table.to_dict()) + return data + + @classmethod + def from_dict(cls, data: Mapping[str, Any]) -> Self: + rest = { + key: value + for key, value in data.items() if key not in ('source', 'snapshot') + } + if 'version_kind' not in rest: + return cls(source=data.get('source'), snapshot=data.get('snapshot'), + extra=rest) + return cls(source=data.get('source'), snapshot=data.get('snapshot'), + source_table=SourceTableEntry.from_dict(rest)) + + +@dataclass(frozen=True, kw_only=True) +class TrackingRef(Record): + run_id: str + + +@dataclass(frozen=True, kw_only=True) +class BundleProvenance(Manifest): + """A reconstructed record (a derived bundle whose parent had none) writes + only the fields it knows.""" + + kind: str = 'model_bundle' + parents: BundleParents + fit: Optional[FitTiming] = None + training_table: Optional[BundleTrainingTable] = None + source_tables: Optional[Dict[str, SourceTableEntry]] = None + fit_mode: Optional[str] = None + tracking: Optional[TrackingRef] = omit_none() + + def to_dict(self) -> Dict[str, Any]: + data = super().to_dict() + if not self.reconstructed: + return data + envelope = {item.name for item in fields(Manifest)} + parents = {k: v for k, v in data['parents'].items() if v is not None} + return { + **{ + k: v + for k, v in data.items() if v is not None or k in envelope + }, 'parents': parents + } + + +def write_manifest(path: Path, + manifest: Union[Record, Mapping[str, Any]]) -> None: + if isinstance(manifest, Record): + manifest = manifest.to_dict() + payload = json.dumps(manifest, indent=2, sort_keys=True) + '\n' + atomic_write_bytes(path, payload.encode('utf-8')) + + +def read_manifest(path: Path, kind: Type[M]) -> Optional[M]: + try: + manifest = json.loads(path.read_text(encoding='utf-8')) + except FileNotFoundError: + return None + except (OSError, ValueError) as error: + warnings.warn(f"Ignoring unreadable manifest {path}: {error}", + UserWarning, stacklevel=3) + return None + version = (manifest.get('schema_version') + if isinstance(manifest, dict) else None) + if type(version) is not int or not 1 <= version <= SCHEMA_VERSION: + warnings.warn( + f"Ignoring {path}: manifest schema_version {version!r} is not " + f"supported by this kumo-engine (up to {SCHEMA_VERSION}).", + UserWarning, + stacklevel=3, + ) + return None + try: + return kind.from_dict(manifest) + except (AttributeError, TypeError, ValueError) as error: + warnings.warn(f"Ignoring unreadable manifest {path}: {error}", + UserWarning, stacklevel=3) + return None diff --git a/packages/kumo-engine/kumo/engine/_provenance.py b/packages/kumo-engine/kumo/engine/_provenance.py new file mode 100644 index 000000000..e463b68e6 --- /dev/null +++ b/packages/kumo-engine/kumo/engine/_provenance.py @@ -0,0 +1,174 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Runtime versions and scheduler job identity recorded in manifests.""" + +import hashlib +import json +import os +import platform +import subprocess +import warnings +from functools import cache +from pathlib import Path +from typing import Any, Dict, Optional + +from kumo.engine._manifest import ( + Environment, + SchedulerContext, + SourceTableEntry, +) + +PROCESSING_JOB_CONFIG = Path('/opt/ml/config/processingjobconfig.json') +_LOCATED_CONNECTORS = frozenset( + {'LOCAL_FILE', 'S3', 'GCS', 'SNOWFLAKE', 'DATABRICKS', 'BIGQUERY'}) + + +@cache +def _build_commit(version: str) -> Optional[str]: + if '+' in version: + return version.split('+', 1)[1] + from kumo.engine import _version + stamped = getattr(_version, 'BUILD_COMMIT', None) + if stamped: + return stamped + root = Path(__file__).resolve().parents[2] + if not (root / 'pyproject.toml').is_file(): + return None + command = ['git', '-C', str(root), 'rev-parse', '--short', 'HEAD'] + try: + return subprocess.run(command, capture_output=True, text=True, + check=True, timeout=10).stdout.strip() or None + except (OSError, subprocess.SubprocessError): + return None + + +def _diskgraph() -> Any: + try: + import diskgraph + return diskgraph + except ImportError: + return None + + +def environment() -> Environment: + from kumo_core.version import ( + get_kumo_version, + installed_distribution_version, + ) + + from kumo.engine._version import __version__ + + diskgraph = _diskgraph() + return Environment( + kumo_engine=get_kumo_version(), + build_commit=_build_commit(__version__), + kumo_core=installed_distribution_version('kumo-core'), + diskgraph=getattr(diskgraph, '__version__', None), + diskgraph_build_commit=getattr(diskgraph, 'BUILD_COMMIT', None), + torch=installed_distribution_version('torch'), + python=platform.python_version(), + ) + + +def _context(platform_name: str, job_id: Optional[str], + job_run_id: Optional[str]) -> SchedulerContext: + return SchedulerContext(platform=platform_name, job_id=job_id, + job_run_id=job_run_id) + + +def _detect() -> Optional[SchedulerContext]: + env = os.environ + if env.get('TRAINING_JOB_NAME'): + return _context('sagemaker', env['TRAINING_JOB_NAME'], + env.get('TRAINING_JOB_ARN')) + try: + processing = json.loads(PROCESSING_JOB_CONFIG.read_text()) + except (OSError, ValueError): + processing = None + if isinstance(processing, dict) and processing.get('ProcessingJobName'): + return _context('sagemaker', processing['ProcessingJobName'], + processing.get('ProcessingJobArn')) + service = env.get('SNOWFLAKE_SERVICE_NAME') + if service: + return _context('snowflake', None, service) + if env.get('DATABRICKS_RUNTIME_VERSION'): + return _context('databricks', None, None) + return None + + +def scheduler_context() -> Optional[SchedulerContext]: + env = os.environ + scheduler = env.get('KUMO_SCHEDULER') + job_id = env.get('KUMO_JOB_ID') + job_run_id = env.get('KUMO_JOB_RUN_ID') + detected = _detect() + if scheduler or job_id or job_run_id: + if detected is not None and scheduler in (None, '', + detected.platform): + scheduler = detected.platform + job_id = job_id or detected.job_id + job_run_id = job_run_id or detected.job_run_id + return _context(scheduler or 'other', job_id, job_run_id) + return detected + + +def source_location(connector: Any, table_name: str, + sql: Optional[str] = None) -> str: + if sql is not None: + return 'sql:' + hashlib.sha256(sql.encode()).hexdigest() + if _connector_kind(connector) not in _LOCATED_CONNECTORS: + return table_name + return connector.source_backend.table_location(connector, table_name) + + +def _connector_kind(connector: Any) -> str: + kind = connector.source_type + return str(getattr(kind, 'value', kind)) + + +def source_table(connector: Any, table_name: str, sql: Optional[str], + version: Any, *, + record_sources: bool = True) -> SourceTableEntry: + from kumo_core.sources.identity import version_kind + + entry: Dict[str, Any] = {'connector': None, 'location': None} + dependencies = getattr(version, 'dependencies', None) + try: + entry['connector'] = _connector_kind(connector) + entry['location'] = (source_location(connector, table_name, sql) + if record_sources or sql is not None else + table_name) + if sql is not None: + _record_sql(entry, connector, sql, record_sources) + except Exception as error: + warnings.warn(f"Could not record a source table's location: {error}", + UserWarning, stacklevel=3) + return SourceTableEntry( + version=None if version is None else str(version), + version_kind=version_kind(version), + dependencies=None if dependencies is None else { + name: str(value) + for name, value in dependencies.items() + }, + **entry, + ) + + +def _record_sql(entry: Dict[str, Any], connector: Any, sql: str, + record_sources: bool) -> None: + from kumo_core.custom_sql import ( + CUSTOM_SQL_NAMESPACE_ATTRS, + custom_sql_synthetic_name, + ) + from kumo_core.types import ConnectorType + + kind = ConnectorType[entry['connector']] + entry['normalized_sql_fingerprint'] = custom_sql_synthetic_name(sql, kind) + if record_sources: + entry['sql'] = sql + entry['dialect'] = entry['connector'].lower() + entry['namespace'] = [ + getattr(connector.config, attr) + for attr in CUSTOM_SQL_NAMESPACE_ATTRS[kind] + ] diff --git a/packages/kumo-engine/kumo/engine/_scratch.py b/packages/kumo-engine/kumo/engine/_scratch.py index fb0c62a3f..144011342 100644 --- a/packages/kumo-engine/kumo/engine/_scratch.py +++ b/packages/kumo-engine/kumo/engine/_scratch.py @@ -36,11 +36,15 @@ import hashlib import uuid +from contextlib import contextmanager +from contextvars import ContextVar from pathlib import Path -from typing import Union +from typing import Iterator, Optional, Union _DEFAULT_ROOT = Path.home() / ".cache" / "kumo-engine" _root = _DEFAULT_ROOT +_scoped_root: ContextVar[Optional[Path]] = ContextVar('kumo_scratch_root', + default=None) def set_scratch_dir(path: Union[str, Path]) -> None: @@ -66,7 +70,18 @@ def set_scratch_dir(path: Union[str, Path]) -> None: def get_scratch_dir() -> Path: - return _root + return _scoped_root.get() or _root + + +@contextmanager +def scoped_scratch_dir(path: Path) -> Iterator[None]: + r"""Nests this thread's (context's) new scratch directories under + ``path`` instead of the configured root.""" + token = _scoped_root.set(path) + try: + yield + finally: + _scoped_root.reset(token) def new_run_id() -> str: diff --git a/packages/kumo-engine/kumo/engine/_snapshot_io.py b/packages/kumo-engine/kumo/engine/_snapshot_io.py new file mode 100644 index 000000000..33823d18d --- /dev/null +++ b/packages/kumo-engine/kumo/engine/_snapshot_io.py @@ -0,0 +1,446 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import hashlib +import json +import logging +import os +import re +import threading +import time +from concurrent.futures import ThreadPoolExecutor +from functools import partial +from pathlib import Path +from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple + +from pyarrow.fs import FileSystem + +_FORMAT = 1 +_MANIFEST = 'kumo_snapshot.json' +_INDEX_MANIFEST = '_MANIFEST.json' +_RANGE_BYTES = 64 << 20 +_CHUNK = 8 << 20 +_ATTEMPTS = 3 +_KEY = re.compile('[0-9a-f]{64}') +_PROGRESS_SECONDS = 30 + +logger = logging.getLogger(__name__) + + +class ConcurrentExportError(TimeoutError): + pass + + +def _location(connector: Any, path: str) -> Tuple[FileSystem, str]: + from kumo_core.object_store_sink import ObjectStoreSinkBackend + + try: + sink = connector.sink_backend + except Exception: + sink = None + if not isinstance(sink, ObjectStoreSinkBackend): + raise NotImplementedError( + f'{type(connector).__name__} cannot hold a snapshot export; use ' + 'an S3Connector, GCSConnector or a FileConnector on a volume.') + return sink.filesystem(connector, path) + + +def _files(filesystem: FileSystem) -> Any: + from kumo_core.object_store_sink import _ArrowExportFiles + + return _ArrowExportFiles(filesystem, None) + + +def _check_key(graph_key: Any, data_version: Any, snapshot_id: Any) -> None: + from kumo.engine.graph import _snapshot_id + + if not (isinstance(graph_key, str) and isinstance(data_version, str) + and _KEY.fullmatch(graph_key) and _KEY.fullmatch(data_version) + and snapshot_id == _snapshot_id(graph_key, data_version)): + raise ValueError(f'{snapshot_id!r} is not a snapshot id.') + + +class _Client: + def __init__(self, connector: Any, path: str) -> None: + self._make = lambda: _location(connector, path)[0] + self.current = self._make() + self._lock = threading.Lock() + + def renew(self, stale: FileSystem) -> None: + with self._lock: + if self.current is stale: + self.current = self._make() + + +class _Progress: + def __init__(self, action: str, path: str, total: int) -> None: + self.action, self.path, self.total = action, path, total + self.moved = 0 + self.started = self.logged = time.monotonic() + self._lock = threading.Lock() + + def add(self, size: int) -> None: + with self._lock: + self.moved += size + now = time.monotonic() + if (now - self.logged < _PROGRESS_SECONDS + and self.moved < self.total): + return + self.logged = now + logger.info('Snapshot %s of %s: %d of %d bytes, %.2f GB/s', + self.action, self.path, self.moved, self.total, + self.moved / max(now - self.started, 1e-9) / 1e9) + + +def _transfer(connector: Any, path: str, items: Sequence[Any], + sizes: Sequence[int], work: Callable[..., Any], + action: str) -> List[Any]: + from kumo_core.sources.object_store import copies + from kumo_core.sources.parallel import run_in_parallel + + if not items: + return [] + stop = threading.Event() + with copies(items) as workers: + # After `copies`: Arrow's IO pool size sets the connection pool. + client = _Client(connector, path) + logger.info('Snapshot %s of %s: %d files, %d bytes, %d workers', + action, path, len(items), sum(sizes), workers) + progress = _Progress(action, path, sum(sizes)) + + def run(item: Any, size: int) -> Any: + result = _retried(client, work, item, stop) + progress.add(size) + return result + + return run_in_parallel( + [partial(run, item, size) for item, size in zip(items, sizes)], + workers=workers, stop=stop) + + +def _retried(client: _Client, work: Callable[..., Any], item: Any, + stop: threading.Event) -> Any: + attempt = 0 + while True: + _stopped(stop) + current = client.current + try: + return work(current, item, stop) + except (FileNotFoundError, PermissionError): + raise + except OSError as error: + if attempt + 1 == _ATTEMPTS: + raise + logger.warning('Retrying a snapshot transfer after: %s', error) + # Arrow keeps the credentials a client was built with. + client.renew(current) + stop.wait(2**attempt) + attempt += 1 + + +def _stopped(stop: threading.Event) -> None: + from kumo_core.sources.parallel import WorkCancelled + + if stop.is_set(): + raise WorkCancelled('abandoned after another transfer failed') + + +def _read_json(filesystem: FileSystem, path: str) -> Optional[Any]: + try: + return json.loads(_files(filesystem).read_text(path)) + except FileNotFoundError: + return None + + +def _open_output(filesystem: FileSystem, path: str) -> Any: + _files(filesystem).makedirs(path.rsplit('/', 1)[0]) + return filesystem.open_output_stream(path) + + +def _write(filesystem: FileSystem, path: str, payload: bytes) -> None: + with _open_output(filesystem, path) as writer: + writer.write(payload) + + +def _fingerprint(index_manifest: bytes) -> str: + return hashlib.sha256(index_manifest).hexdigest() + + +def _snapshot_manifest(snapshot: Any) -> bytes: + from kumo.engine._manifest import ( + RECONSTRUCTED, + SnapshotManifest, + SnapshotParents, + ) + from kumo.engine.graph import SNAPSHOT_MANIFEST_FILE + + try: + return (Path(snapshot.index_path) / SNAPSHOT_MANIFEST_FILE).read_bytes() + except OSError: + manifest = SnapshotManifest( + id=snapshot.id, parents=SnapshotParents(source_tables={}), + **RECONSTRUCTED) + return (json.dumps(manifest.to_dict(), indent=2, sort_keys=True) + + '\n').encode() + + +def export(snapshot: Any, connector: Any, path: str) -> str: + from diskgraph import snapshot_range_checksums + + from kumo.engine._provenance import environment + from kumo.engine.graph import SNAPSHOT_MANIFEST_FILE + + source = Path(snapshot.index_path) + graph_key, data_version = snapshot.graph_key, snapshot.data_version + _check_key(graph_key, data_version, snapshot.id) + location = f"{path.rstrip('/')}/{graph_key}/{data_version[:16]}" + index_manifest = (source / _INDEX_MANIFEST).read_bytes() + sizes = { + entry['rel_path']: int(entry['size_bytes']) + for entry in json.loads(index_manifest)['files'] + } + snapshot_manifest = _snapshot_manifest(snapshot) + filesystem, root = _location(connector, location) + _delete_marker(filesystem, root) + filesystem.delete_dir_contents(root, missing_dir_ok=True) + # An export that committed while this one wiped no longer holds its files. + _delete_marker(filesystem, root) + + def upload(client: FileSystem, name: str, stop: threading.Event) -> None: + size = sizes[name] + with open(source / name, 'rb', buffering=0) as reader, _open_output( + client, f'{root}/index/{name}') as writer: + for offset in range(0, size, _CHUNK): + _stopped(stop) + chunk = reader.read(min(_CHUNK, size - offset)) + if len(chunk) != min(_CHUNK, size - offset): + raise OSError(f'{source / name} changed while it was ' + 'exported') + writer.write(chunk) + + def send(names: List[str], action: str) -> None: + names = sorted(names, key=lambda name: -sizes[name]) + _transfer(connector, location, names, [sizes[name] for name in names], + upload, action) + + with ThreadPoolExecutor(1) as pool: + checksums = pool.submit(snapshot_range_checksums, str(source), + _RANGE_BYTES) + send(list(sizes), 'export') + crc32 = checksums.result() + marker = json.dumps({ + 'format': _FORMAT, + 'snapshot_id': snapshot.id, + 'graph_key': graph_key, + 'data_version': data_version, + 'index_manifest': _fingerprint(index_manifest), + 'graph_config': snapshot.graph_config, + 'environment': environment().to_dict(), + 'range_bytes': _RANGE_BYTES, + 'crc32': crc32, + }).encode() + for _ in range(_ATTEMPTS): + filesystem, root = _location(connector, location) + _write(filesystem, f'{root}/index/{_INDEX_MANIFEST}', index_manifest) + _write(filesystem, f'{root}/{SNAPSHOT_MANIFEST_FILE}', + snapshot_manifest) + _write(filesystem, f'{root}/{_MANIFEST}', marker) + if not _files(filesystem).exists(f'{root}/{_MANIFEST}'): + logger.warning('A later export to %s took it over.', location) + return location + missing = _missing(filesystem, root, + {**sizes, _INDEX_MANIFEST: len(index_manifest)}) + if not missing: + return location + logger.warning('Another export to %s removed %d files; uploading ' + 'them again.', location, len(missing)) + send([name for name in missing if name in sizes], 'repair') + raise ConcurrentExportError( + f'Another export to {location} keeps replacing it; run ' + 'GraphSnapshot.export again once it finishes.') + + +def _delete_marker(filesystem: FileSystem, root: str) -> None: + try: + filesystem.delete_file(f'{root}/{_MANIFEST}') + except FileNotFoundError: + pass + + +def _missing(filesystem: FileSystem, root: str, + sizes: Dict[str, int]) -> List[str]: + from pyarrow.fs import FileSelector, FileType + + prefix = f'{root}/index/' + found = { + info.path[len(prefix):]: info.size + for info in filesystem.get_file_info( + FileSelector(prefix, recursive=True, allow_not_found=True)) + if info.type is FileType.File + } + return [name for name, size in sizes.items() if found.get(name) != size] + + +def _corrupt(path: str, detail: str) -> ValueError: + return ValueError(f'The snapshot export at {path} is corrupt ({detail}); ' + 'delete it and run GraphSnapshot.export again.') + + +def read(connector: Any, path: str) -> Tuple[str, dict]: + filesystem, root = _location(connector, path) + manifest = _read_json(filesystem, f'{root}/{_MANIFEST}') + if manifest is None: + raise FileNotFoundError(f'No complete snapshot export at {path}; ' + 're-run export.') + if not isinstance(manifest, dict): + raise _corrupt(path, 'manifest is not a JSON object') + if manifest.get('format') != _FORMAT: + raise ValueError( + f"Snapshot export at {path} has format " + f"{manifest.get('format')!r}; this kumo-engine reads format " + f"{_FORMAT}. Load it with the kumo-engine that exported it.") + _check_key(manifest.get('graph_key'), manifest.get('data_version'), + manifest.get('snapshot_id')) + span, crcs = manifest.get('range_bytes'), manifest.get('crc32') + if (type(span) is not int or span <= 0 or not isinstance(crcs, dict) + or not isinstance(manifest.get('index_manifest'), str) + or not isinstance(manifest.get('graph_config'), dict)): + raise _corrupt(path, 'malformed manifest') + return root, manifest + + +def _download(root: str, staging: Path, client: FileSystem, + item: Tuple[str, int, int], stop: threading.Event) -> None: + name, offset, length = item + # Each open is a HEAD: pyarrow cannot open a file from a known size. + with client.open_input_file(f'{root}/index/{name}') as reader: + fd = os.open(staging / name, os.O_WRONLY) + try: + for start in range(offset, offset + length, _CHUNK): + _stopped(stop) + want = min(_CHUNK, offset + length - start) + chunk = reader.read_at(want, start) + if len(chunk) != want: + raise ValueError(f'{name} is truncated') + os.pwrite(fd, chunk, start) + getattr(os, 'fdatasync', os.fsync)(fd) + finally: + os.close(fd) + + +def _index_files( + path: str, root: str, manifest: dict, connector: Any +) -> Tuple[bytes, List[Tuple[str, int]], Optional[bytes]]: + from kumo_core.sources.layout import confined + + from kumo.engine.graph import SNAPSHOT_MANIFEST_FILE + + filesystem, _ = _location(connector, path) + with filesystem.open_input_stream( + f'{root}/index/{_INDEX_MANIFEST}') as reader: + index_manifest = reader.readall() + if _fingerprint(index_manifest) != manifest['index_manifest']: + raise _corrupt(path, f'{_INDEX_MANIFEST} is not the exported one') + try: + files = [(confined(entry['rel_path']), int(entry['size_bytes'])) + for entry in json.loads(index_manifest)['files']] + except Exception as error: + raise _corrupt(path, f'bad {_INDEX_MANIFEST}: {error}') from error + if len({name for name, _ in files}) != len(files): + raise _corrupt(path, f'{_INDEX_MANIFEST} lists a file twice') + try: + with filesystem.open_input_stream( + f'{root}/{SNAPSHOT_MANIFEST_FILE}') as reader: + snapshot_manifest = reader.readall() + except FileNotFoundError: + snapshot_manifest = None + return index_manifest, files, snapshot_manifest + + +def load_into(connector: Any, path: str, root: str, manifest: dict, + target: Path) -> None: + from diskgraph import ( + DiskgraphCorruptionError, + snapshot_range_checksums, + verify_snapshot_integrity, + ) + from kumo_core.diskgraph.snapshot import ( + DiskgraphIngestResult, + write_diskgraph_success_marker, + ) + from kumo_core.io.filesystem import fsync_directory + from kumo_core.sources.transfer import discard_destination + + from kumo.engine.graph import SNAPSHOT_MANIFEST_FILE + + span = manifest['range_bytes'] + index_manifest, files, snapshot_manifest = _index_files( + path, root, manifest, connector) + staging = target.with_name(f'{target.name}.loading') + discard_destination(str(target)) + discard_destination(str(staging)) + try: + staging.mkdir(parents=True) + needed = sum(size for _, size in files) + disk = os.statvfs(staging) + if disk.f_bavail * disk.f_frsize < needed: + raise OSError(f'Loading this snapshot needs {needed} free bytes ' + f'in {staging}.') + for name, size in files: + (staging / name).parent.mkdir(parents=True, exist_ok=True) + fd = os.open(staging / name, os.O_WRONLY | os.O_CREAT | os.O_EXCL) + try: + if size: + os.posix_fallocate(fd, 0, size) + finally: + os.close(fd) + ranges = sorted(((name, offset, min(span, size - offset)) + for name, size in files + for offset in range(0, size, span)), + key=lambda item: -item[2]) + try: + _transfer(connector, path, ranges, [item[2] for item in ranges], + partial(_download, root, staging), 'load') + except ValueError as error: + raise _corrupt(path, str(error)) from error + (staging / _INDEX_MANIFEST).write_bytes(index_manifest) + write_diskgraph_success_marker(str(staging), + DiskgraphIngestResult.for_index(target)) + try: + verify_snapshot_integrity(str(staging)) + loaded = snapshot_range_checksums(str(staging), span) + except DiskgraphCorruptionError as error: + raise _corrupt(path, str(error)) from error + for name, crcs in sorted(loaded.items()): + expected = manifest['crc32'].get(name) + if crcs != expected: + offset = next((index * span for index, crc in enumerate(crcs) + if not expected or index >= len(expected) + or crc != expected[index]), 0) + raise _corrupt( + path, f'{name} failed its checksum at offset {offset}') + _check_readable(staging, manifest) + if snapshot_manifest is not None: + (staging / SNAPSHOT_MANIFEST_FILE).write_bytes(snapshot_manifest) + for directory in sorted({(staging / name).parent + for name, _ in files} | {staging}): + fsync_directory(directory) + os.rename(staging, target) + fsync_directory(target.parent) + except BaseException: + discard_destination(str(staging)) + raise + + +def _check_readable(index: Path, manifest: dict) -> None: + from diskgraph import SamplerEngine + + try: + SamplerEngine(str(index), num_threads=1) + except Exception as error: + built = (manifest.get('environment') or {}).get('kumo_engine') + raise RuntimeError( + f'This DiskGraph cannot open the snapshot exported by kumo-engine ' + f'{built} ({error}). Rebuild it with this kumo-engine or load it ' + 'with the exporting one.') from error diff --git a/packages/kumo-engine/kumo/engine/_tracking.py b/packages/kumo-engine/kumo/engine/_tracking.py new file mode 100644 index 000000000..6464c8b40 --- /dev/null +++ b/packages/kumo-engine/kumo/engine/_tracking.py @@ -0,0 +1,445 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +r"""Platform-neutral run records and the `Tracking` protocol.""" + +import dataclasses +import json +import logging +import os +import queue +import threading +import time +import warnings +from dataclasses import dataclass, field +from pathlib import Path +from typing import ( + TYPE_CHECKING, + Any, + Callable, + Dict, + Iterator, + List, + Mapping, + Optional, + Protocol, + Sequence, + Tuple, +) + +from kumo.engine._events import Event, _finite_floats + +if TYPE_CHECKING: + from kumo.engine.graph import GraphSnapshot + from kumo.engine.trainer.trainer import TrainingResult + +logger = logging.getLogger(__name__) + +DROPPED_PARAMS = 'kumo.dropped_params' +DROPPED_METRICS = 'kumo.dropped_metrics' +# Tracker calls run on one background thread so a slow tracking server +# never stalls training; these bound how long the fit waits for it. +QUEUE_SIZE = 1000 +FINISH_SECONDS = 30.0 +MODEL_SECONDS = 3600.0 + + +@dataclass(frozen=True) +class FitContext: + r"""What a fit trains, known once its model plan is resolved.""" + + run_id: str + query: str + model_plan: Mapping[str, Any] + inputs: Mapping[str, Mapping[str, Optional[str]]] = field( + default_factory=dict) + + +@dataclass(frozen=True) +class RunRecord: + r"""One tracked run; ``steps`` are per-epoch ``(epoch, value)`` + histories, ``artifacts`` map file names to text and ``inputs`` map + source table names to their ``connector``, ``location``, ``version`` and + ``version_kind``.""" + + name: Optional[str] = None + params: Mapping[str, str] = field(default_factory=dict) + metrics: Mapping[str, float] = field(default_factory=dict) + steps: Mapping[str, Tuple[Tuple[int, float], + ...]] = field(default_factory=dict) + tags: Mapping[str, str] = field(default_factory=dict) + artifacts: Mapping[str, str] = field(default_factory=dict) + inputs: Mapping[str, Mapping[str, Optional[str]]] = field( + default_factory=dict) + + +INPUT_FIELDS = ('connector', 'location', 'version', 'version_kind') + + +def fit_inputs(snapshot_dir: Path, + table_dir: Path) -> Dict[str, Dict[str, Optional[str]]]: + r"""Source tables recorded in the snapshot's manifest, plus a custom + training table's source, if those manifests exist.""" + from kumo.engine._manifest import ( + SnapshotManifest, + TrainingTableManifest, + read_manifest, + ) + from kumo.engine.graph import SNAPSHOT_MANIFEST_FILE + from kumo.engine.pquery.training_table import STAGING_MANIFEST_FILE + + try: + snapshot = read_manifest(snapshot_dir / SNAPSHOT_MANIFEST_FILE, + SnapshotManifest) + table = read_manifest(table_dir / STAGING_MANIFEST_FILE, + TrainingTableManifest) + tables = {} if snapshot is None else dict( + snapshot.parents.source_tables) + if table is not None and table.parents.source_table is not None: + tables['training_table'] = table.parents.source_table + return { + str(name): {key: _text(getattr(entry, key)) + for key in INPUT_FIELDS} + for name, entry in tables.items() + } + except Exception: + logger.warning('Could not read the fit\'s source tables', + exc_info=True) + return {} + + +def _text(value: Any) -> Optional[str]: + return None if value is None else str(value) + + +def run_record(context: FitContext, result: Optional['TrainingResult'] = None, + fit_ended: Optional[Event] = None, *, + tags: Optional[Mapping[str, str]] = None) -> RunRecord: + r"""The fit run's record; ``tags`` are the user's extra tags.""" + metrics: Dict[str, float] = {} + pointers = _scheduler_tags() + artifacts = { + 'query.pql': context.query, + 'model_plan.json': json.dumps(dict(context.model_plan), indent=2, + default=str), + } + if result is not None: + metrics.update(result.metrics) + if fit_ended is not None and 'seconds' in fit_ended.metrics: + metrics['seconds'] = fit_ended.metrics['seconds'] + if getattr(result, 'bundle_id', None): + pointers['kumo.bundle_id'] = str(result.bundle_id) + provenance = Path(result.bundle_dir) / 'provenance.json' + if provenance.is_file(): + artifacts['provenance.json'] = provenance.read_text( + encoding='utf-8') + params = { + key: _param_value(value) + for key, value in _flatten(context.model_plan) + } + return RunRecord( + params=params, + metrics=_finite_floats(metrics), + tags={ + **_user_tags(tags, params), + 'kumo.kind': 'fit', + 'kumo.run_id': context.run_id, + **pointers, + **_error_tag(fit_ended), + }, + artifacts=artifacts, + inputs=dict(context.inputs), + ) + + +def _user_tags(tags: Optional[Mapping[str, Any]], + params: Mapping[str, str]) -> Dict[str, str]: + dropped = sorted(key for key in tags or {} + if key.startswith('kumo.') or key in params) + if dropped: + warnings.warn( + f'Tracking tags {dropped} are reserved and not recorded.', + UserWarning, stacklevel=3) + return { + key: str(value) + for key, value in (tags or {}).items() if key not in dropped + } + + +def _scheduler_tags() -> Dict[str, str]: + from kumo.engine._provenance import scheduler_context + + scheduler = scheduler_context() + if scheduler is None: + return {} + tags = { + 'kumo.scheduler': scheduler.platform, + 'kumo.job_id': scheduler.job_id, + 'kumo.job_run_id': scheduler.job_run_id, + } + return {key: str(value) for key, value in tags.items() if value} + + +def _error_tag(ended: Optional[Event]) -> Dict[str, str]: + if ended is None or 'error' not in ended.params: + return {} + return {'kumo.error': str(ended.params['error'])} + + +def trial_record(context: FitContext, events: Sequence[Event]) -> RunRecord: + r"""One trial attempt's record from its events so far.""" + started = next(event for event in events if event.name == 'trial_started') + trial, attempt = started.trial, int(started.params['attempt']) + params: Dict[str, str] = { + key: _param_value(value) + for key, value in started.params.items() + } + metrics: Dict[str, float] = {} + steps: Dict[str, List[Tuple[int, float]]] = {} + ended: Dict[str, str] = {} + for event in events: + if event.name == 'epoch_ended' and event.epoch is not None: + for name, value in event.metrics.items(): + steps.setdefault(name, []).append((event.epoch, value)) + elif event.name == 'trial_ended': + ended = { + 'kumo.status': str(event.params.get('status')), + **_error_tag(event), + } + metrics.update({ + name if name == 'seconds' else f'best_{name}': value + for name, value in event.metrics.items() + }) + return RunRecord( + name=f'trial-{trial}-attempt-{attempt}', + params=params, + metrics=metrics, + steps={ + name: tuple(history) + for name, history in steps.items() + }, + tags={ + 'kumo.kind': 'trial', + 'kumo.run_id': f'{context.run_id}-trial-{trial}-attempt-{attempt}', + 'kumo.parent_run': context.run_id, + **ended, + }, + ) + + +class Tracking(Protocol): + r"""Records one `Trainer.fit` in an experiment tracker. + + Calls come from one background thread. `start` gets the fit run's + record without metrics, `on_trial` a trial attempt's record so far with + the event that updated it, and `finish` the fit run's final record + (`None` if tracking stopped early, which should end the run as failed). + Optional members: ``prepare()``, called first on the fit's own thread; + ``log_model(result, snapshot)``, called after a fit that succeeded and + before `finish`, which the fit waits for while the tracking thread is + alive, up to an hour (an error, a stopped thread or the time limit ends + the run as failed); and ``tags`` (extra fit-run tags). + """ + def start(self, record: RunRecord) -> None: + ... + + def on_trial(self, record: RunRecord, event: Event) -> None: + ... + + def finish(self, record: Optional[RunRecord], + error: Optional[BaseException]) -> None: + ... + + +class TrackingSession: + def __init__(self, tracking: Optional[Tracking]) -> None: + self._tracker = tracking if _rank() == 0 else None + self._context: Optional[FitContext] = None + self._trials: Dict[int, List[Event]] = {} + self._fit_ended: Optional[Event] = None + self._jobs: 'queue.Queue[Optional[Callable[[], Any]]]' = queue.Queue( + QUEUE_SIZE) + self._worker: Optional[threading.Thread] = None + self._behind = False + + @property + def enabled(self) -> bool: + return self._tracker is not None + + def start(self, context: Callable[[], FitContext]) -> None: + if self._tracker is None: + return + try: + value = context() + except Exception: + logger.warning('Experiment tracking failed; the fit continues', + exc_info=True) + self._tracker = None + return + self._context = value + self._call(lambda tracker: getattr(tracker, 'prepare', _none)()) + self._worker = threading.Thread(target=self._work, + name='kumo-tracking', daemon=True) + self._worker.start() + self._submit(lambda: self._call(lambda tracker: tracker.start( + run_record(value, tags=getattr(tracker, 'tags', None))))) + + def on_event(self, event: Event) -> None: + if self._context is None: + return + if event.name == 'fit_ended': + self._fit_ended = event + if event.trial is None: + return + if event.name == 'trial_started': + self._trials[event.trial] = [] + events = self._trials.setdefault(event.trial, []) + if 'attempt' not in event.params and events: + event = dataclasses.replace(event, params={ + **event.params, 'attempt': events[0].params['attempt'] + }) + events.append(event) + record = trial_record(self._context, events) + self._submit(lambda: self._call( + lambda tracker: tracker.on_trial(record, event))) + + def log_model(self, result: 'TrainingResult', + snapshot: 'GraphSnapshot') -> Optional[BaseException]: + r"""Runs the tracker's ``log_model`` on the tracking thread and waits + for it; returns the error it raised.""" + if (self._worker is None + or not hasattr(self._tracker, 'log_model')): + return None + done = threading.Event() + failed: List[BaseException] = [] + + def log() -> None: + try: + tracker = self._tracker + if tracker is not None: + getattr(tracker, 'log_model')(result, snapshot) + except BaseException as error: + failed.append(error) + finally: + done.set() + + worker = self._worker + deadline = time.monotonic() + MODEL_SECONDS + try: + self._jobs.put(log, timeout=MODEL_SECONDS) + except queue.Full: + pass + while not done.wait(1.0): + if worker.is_alive() and time.monotonic() < deadline: + continue + failed.append( + TimeoutError(f'it did not finish within {MODEL_SECONDS:.0f}s') + if worker.is_alive() else + RuntimeError('the tracking thread stopped')) + tracker, self._tracker = self._tracker, None + if tracker is not None: + _abandon(tracker) + break + if failed: + warnings.warn(f'Could not log the model: {failed[0]}', + UserWarning, stacklevel=3) + return failed[0] + return None + + def finish(self, result: Optional['TrainingResult'], + error: Optional[BaseException]) -> None: + worker = self._worker + if self._context is None or worker is None: + return + try: + record: Optional[RunRecord] = run_record(self._context, result, + self._fit_ended) + except Exception: + logger.warning('Could not build the tracking record', + exc_info=True) + record = None + + def finish_tracker() -> None: + tracker, self._tracker = self._tracker, None + if tracker is not None: + tracker.finish(record, error) + + deadline = time.monotonic() + FINISH_SECONDS + try: + for job in (finish_tracker, None): + self._jobs.put(job, timeout=max(0.0, + deadline - time.monotonic())) + except queue.Full: + pass + worker.join(max(0.0, deadline - time.monotonic())) + if worker.is_alive(): + logger.warning( + 'Experiment tracking did not finish within %ss; it goes on ' + 'in the background and the fit returns', FINISH_SECONDS) + + def _submit(self, job: Callable[[], Any]) -> bool: + if self._behind: + return False + try: + self._jobs.put_nowait(job) + return True + except queue.Full: + self._behind = True + logger.warning('Experiment tracking is falling behind; the rest ' + 'of this fit is not tracked') + return False + + def _work(self) -> None: + while True: + job = self._jobs.get() + if job is None: + return + try: + job() + except Exception: + logger.warning('Experiment tracking failed', exc_info=True) + + def _call(self, call: Callable[[Tracking], Any]) -> None: + tracker = self._tracker + if tracker is None: + return + try: + call(tracker) + except Exception: + logger.warning( + 'Experiment tracking failed; it is off for the rest of this ' + 'run', exc_info=True) + self._tracker = None + _abandon(tracker) + + +def _none(*args: Any) -> None: + return None + + +def _abandon(tracker: Tracking) -> None: + try: + tracker.finish(None, None) + except Exception: + logger.debug('Could not end the run of %r', tracker, exc_info=True) + + +def _rank() -> int: + for name in ('RANK', 'LOCAL_RANK'): + if os.environ.get(name, '').strip() not in ('', '0'): + return 1 + return 0 + + +def _flatten(value: Any, prefix: str = '') -> Iterator[Tuple[str, Any]]: + if isinstance(value, Mapping): + for key, item in value.items(): + yield from _flatten(item, f'{prefix}{key}.') + else: + yield prefix[:-1], value + + +def _param_value(value: Any) -> str: + if isinstance(value, str): + return value + return json.dumps(value, default=str) diff --git a/packages/kumo-engine/kumo/engine/connector/__init__.py b/packages/kumo-engine/kumo/engine/connector/__init__.py index 8f7373203..48408d085 100644 --- a/packages/kumo-engine/kumo/engine/connector/__init__.py +++ b/packages/kumo-engine/kumo/engine/connector/__init__.py @@ -36,6 +36,25 @@ from kumo.engine.connector.snowflake import SnowflakeConnector from kumo.engine.connector.source_table import SourceTable, SQLSourceTable + +def connector_for_root(uri: str) -> Connector: + r"""A connector rooted at ``uri``: ``s3://``, ``gs://``, ``file://`` or a + local path.""" + from urllib.parse import urlparse + from urllib.request import url2pathname + + parsed = urlparse(uri) + if parsed.scheme == 's3': + return S3Connector(root_dir=uri) + if parsed.scheme == 'gs': + return GCSConnector(root_dir=uri) + if parsed.scheme not in ('', 'file'): + raise ValueError(f'{uri!r} is not s3://, gs://, file:// or a local ' + f'path.') + return FileConnector( + root_dir=url2pathname(parsed.path) if parsed.scheme else uri) + + __all__ = [ "Connector", "FileConnector", @@ -52,4 +71,5 @@ "GCSIdentity", "BigQueryIdentity", "DatabricksIdentity", + "connector_for_root", ] diff --git a/packages/kumo-engine/kumo/engine/connector/file.py b/packages/kumo-engine/kumo/engine/connector/file.py index 802da205a..df6d233dc 100644 --- a/packages/kumo-engine/kumo/engine/connector/file.py +++ b/packages/kumo-engine/kumo/engine/connector/file.py @@ -121,6 +121,11 @@ def get_table_config(self, connector: Union[Connector, SourceConnector], LocalFileSourceTable(path=str(Path(path).resolve()), file_type=FileType(kind.upper()))) + def table_location(self, connector: Union[Connector, SourceConnector], + table_name): + reader, name = self._reader(connector, table_name) + return str(Path(reader.resolve(name)[0]).resolve()) + def prepare_ingest(self, connector, table_name, dest_dir, *args, **kwargs): _disjoint(Path(connector.config.root_dir), Path(dest_dir)) return super().prepare_ingest(connector, table_name, dest_dir, *args, diff --git a/packages/kumo-engine/kumo/engine/connector/snowflake.py b/packages/kumo-engine/kumo/engine/connector/snowflake.py index 0e75a9947..bdfc97387 100644 --- a/packages/kumo-engine/kumo/engine/connector/snowflake.py +++ b/packages/kumo-engine/kumo/engine/connector/snowflake.py @@ -11,7 +11,7 @@ import os from functools import cached_property from pathlib import Path -from typing import Dict, Optional +from typing import Any, Dict, Optional from kumoapi.data_source import ( DataSourceType, @@ -100,6 +100,12 @@ def sink_backend(self) -> ArtifactSink: "snowflake").SnowflakeSinkBackend return SnowflakeSinkBackend(identity=self.identity) + def session(self) -> Any: + r"""The Snowpark session this connector reads through: its identity, + pointed at its warehouse, database and schema, kept alive.""" + with self.source_backend._session(self) as session: + return session + def source_table_for(self, table_name: str) -> SnowflakeSourceTable: return SnowflakeSourceTable( snowflake_connector_id=self.id, diff --git a/packages/kumo-engine/kumo/engine/graph.py b/packages/kumo-engine/kumo/engine/graph.py index 53ebcd199..966bde62f 100644 --- a/packages/kumo-engine/kumo/engine/graph.py +++ b/packages/kumo-engine/kumo/engine/graph.py @@ -61,24 +61,39 @@ """ import json +import threading import warnings -from dataclasses import dataclass, field -from functools import cached_property +from dataclasses import dataclass, field, replace +from functools import cached_property, partial from pathlib import Path -from typing import Any, Dict, Mapping, Sequence +from typing import Any, Dict, Mapping, Optional, Sequence from kumo_core.pquery import NativeGraph, normalize_graph from kumoapi.graph import Edge +from kumo.engine._manifest import ( + RECONSTRUCTED, + SnapshotManifest, + SnapshotParents, + SourceTableEntry, + read_manifest, + stamp, + write_manifest, +) +from kumo.engine._provenance import source_table from kumo.engine._scratch import ( fingerprint, get_scratch_dir, snapshot_scratch_dir, ) from kumo.engine._storage import network_storage +from kumo.engine.connector.base import Connector from kumo.engine.table import Table +SNAPSHOT_MANIFEST_FILE = 'snapshot-manifest.json' + _warned_scratch: set[Path] = set() +_warned_backfill: set[Path] = set() def _warn_if_network_scratch(root: Path) -> None: @@ -123,6 +138,14 @@ def _warn_low_fk_match_rates(stats: Dict[str, Any], ) +def _ingest_stats(path: Any) -> Dict[str, Any]: + try: + stats = json.loads(Path(path).read_text())['stats'] + except (OSError, ValueError, KeyError, TypeError): + return {} + return stats if isinstance(stats, dict) else {} + + def _warn_ingest_data_loss(result: Any, graph_config: Dict[str, Any]) -> None: from kumo_core.diskgraph.validation import parse_failure_warning @@ -146,10 +169,8 @@ def _warn_ingest_data_loss(result: Any, graph_config: Dict[str, Any]) -> None: warnings.warn(parse_failure_warning(name, column, count), UserWarning, stacklevel=3) - try: - report = json.loads(Path(result.ingest_stats_path).read_text()) - stats = report['stats'] - except (AttributeError, OSError, ValueError, KeyError, TypeError): + stats = _ingest_stats(getattr(result, 'ingest_stats_path', None)) + if not stats: return tables = {table['table_name']: table for table in graph_config['tables']} duplicates: Dict[str, int] = {} @@ -184,6 +205,78 @@ def _warn_ingest_data_loss(result: Any, graph_config: Dict[str, Any]) -> None: ) +def _store() -> Any: + from kumo_core.diskgraph.snapshot import ( + DiskgraphLocalConfig, + LocalSnapshotStore, + ) + + return LocalSnapshotStore( + DiskgraphLocalConfig(data_root=str(get_scratch_dir() / 'snapshots'))) + + +def _lock(graph_key: str) -> Any: + from filelock import FileLock + + persist_dir = snapshot_scratch_dir(graph_key) / 'source_tables' + persist_dir.parent.mkdir(parents=True, exist_ok=True) + return FileLock(f'{persist_dir}.lock') + + +def _loaded(store: Any, handle: Any) -> Optional[Any]: + from diskgraph import DiskgraphCorruptionError, read_success_marker + from kumo_core.diskgraph.snapshot import rebased + + path = store.path(handle) + if read_success_marker(path) is None: + return None + try: + store.verify(handle) + return rebased(store, handle) + except (DiskgraphCorruptionError, ValueError) as error: + warnings.warn(f"Reloading the damaged graph index at {path}: {error}", + UserWarning, stacklevel=3) + return None + + +def _snapshot_id(graph_key: str, data_version: str) -> str: + return f'{graph_key}@{data_version[:16]}' + + +_COMPARABLE = ('exact', 'timestamp', 'listing') + + +def _source_change(table: Table, recorded: Optional[SourceTableEntry], + record_sources: bool) -> Dict[str, Any]: + from kumoapi.source_table import CustomSQLSourceTable + + before = None if recorded is None else recorded.version + current, same_source = None, False + comparable = (recorded is not None + and recorded.version_kind in _COMPARABLE) + try: + source = table.definition.source_table + if not isinstance(source, CustomSQLSourceTable): + connector = table.connector + stamp = connector.source_backend.get_data_version( + connector, table.table_name) + now = source_table(connector, table.table_name, None, stamp, + record_sources=record_sources) + current = now.version + comparable = comparable and now.version_kind in _COMPARABLE + same_source = recorded is not None and ( + recorded.connector, recorded.location) == (now.connector, + now.location) + except Exception: + current = None + if not (same_source and comparable and before is not None + and current is not None): + status = 'unknown' + else: + status = 'unchanged' if before == current else 'changed' + return {'status': status, 'recorded': before, 'current': current} + + @dataclass(frozen=True) class GraphSnapshot: r"""The materialized artifact produced by `Graph.snapshot()`: an on-disk @@ -210,11 +303,160 @@ class GraphSnapshot: `TrainingWorkflowRequest.from_spec`/`materialize(...)`). Computed once here, at snapshot time, rather than being rebuilt on every `Trainer.fit`/`predict` call. + record_sources: the `record_sources` choice this snapshot was taken + with; `None` when unknown (built elsewhere or loaded). """ id: str index_path: Path graph_config: Dict[str, Any] + record_sources: bool | None = field(default=None, compare=False) + + @property + def graph_key(self) -> str: + """The `Graph.id` this snapshot was built from.""" + return self.id.partition('@')[0] + + @property + def data_version(self) -> str: + """The full fingerprint of the source versions the index was built + from; `id` carries its first 16 characters.""" + return Path(self.index_path).name + + @cached_property + def native_graph(self) -> NativeGraph: + """`graph_config` in `kumo_core`'s normalized form, as + `Graph.native_graph`.""" + return normalize_graph(self.graph_config) + + def export(self, connector: Connector, path: str) -> str: + r"""Copy the index to `//` and + return that location. Any earlier export there is replaced; it is not + loadable until the new one completes. + + Args: + connector: an S3, GCS or file connector (on Databricks or + Snowflake, a `FileConnector` on a mounted volume). + """ + from kumo.engine._snapshot_io import export + + return export(self, connector, path) + + @classmethod + def load_from_connector(cls, connector: Connector, *, path: str, + graph: Optional['Graph'] = None) -> 'GraphSnapshot': + r"""Copy the snapshot exported at `path` into the local store. + + Args: + path: the location `export` returned. + graph: if given, warn when its source tables changed since the + snapshot was built (see `source_changes`). + """ + from kumo_core.diskgraph.snapshot import DiskgraphSnapshotHandle + + from kumo.engine._snapshot_io import load_into, read + + root, manifest = read(connector, path) + graph_key = manifest['graph_key'] + if graph is not None and graph.id != graph_key: + warnings.warn(f"Snapshot at {path} was taken of another graph.", + UserWarning, stacklevel=2) + store = _store() + handle = DiskgraphSnapshotHandle( + graph_key=graph_key, graph_data_version=manifest['data_version']) + with _lock(graph_key): + result = _loaded(store, handle) + if result is None: + load_into(connector, path, root, manifest, + Path(store.path(handle))) + result = _loaded(store, handle) + snapshot = cls(id=manifest['snapshot_id'], + index_path=Path(result.index_path), + graph_config=manifest['graph_config']) + if graph is not None: + changed = [ + alias + for alias, change in snapshot.source_changes(graph).items() + if change['status'] == 'changed' + ] + if changed: + warnings.warn( + f"Snapshot at {path} predates source table(s) " + f"{', '.join(changed)}; see `snapshot.source_changes`.", + UserWarning, stacklevel=2) + return snapshot + + def source_changes(self, graph: 'Graph') -> Dict[str, Dict[str, Any]]: + r"""Each of `graph`'s tables' source status since this snapshot: + `unchanged`, `changed` or `unknown`. Never raises.""" + try: + manifest = self.manifest + recorded = ({} if manifest is None or manifest.reconstructed else + dict(manifest.parents.source_tables)) + record_sources = (manifest is None + or manifest.record_sources is not False) + except Exception: + recorded, record_sources = {}, True + from kumo_core.sources.parallel import run_in_parallel + + tables = graph._named_tables + changes = run_in_parallel( + [ + partial(_source_change, table, recorded.get(alias), + record_sources) + for alias, table in tables.items() + ], workers=max(1, min(len(tables), 16)), stop=threading.Event()) + return dict(zip(tables, changes)) + + @property + def manifest(self) -> SnapshotManifest | None: + """The snapshot manifest stored next to the index, if any.""" + return read_manifest(Path(self.index_path) / SNAPSHOT_MANIFEST_FILE, + SnapshotManifest) + + @property + def counts(self) -> Dict[str, Dict[str, int]]: + """Source rows per table and edges per edge type in this index.""" + from kumo_core.diskgraph.snapshot import DiskgraphIngestResult + + stats = _ingest_stats( + DiskgraphIngestResult.for_index(self.index_path).ingest_stats_path) + return { + 'rows': dict(stats.get('source_rows_by_table', {})), + 'edges': dict(stats.get('edge_counts', {})), + } + + +def _write_snapshot_manifest(result: Any, snapshot_id: str, + source_tables: Dict[str, SourceTableEntry], + record_sources: bool, *, fresh: bool) -> None: + path = Path(result.index_path) / SNAPSHOT_MANIFEST_FILE + try: + if not fresh and path.exists(): + existing = read_manifest(path, SnapshotManifest) + if existing is None or (existing.record_sources + is not False) == record_sources: + return + manifest = replace( + existing, record_sources=record_sources, + parents=replace(existing.parents, + source_tables=source_tables)) + else: + manifest = SnapshotManifest( + id=snapshot_id, + parents=SnapshotParents(source_tables=source_tables), + record_sources=record_sources, + **(stamp() if fresh else RECONSTRUCTED)) + write_manifest(path, manifest) + except Exception as error: + if fresh or path not in _warned_backfill: + if not fresh: + _warned_backfill.add(path) + warnings.warn( + f"Could not write the snapshot manifest at {path}: {error}", + UserWarning, + stacklevel=4, + ) @dataclass(frozen=True) @@ -346,7 +588,7 @@ def validate(self) -> None: """ self.native_graph - def snapshot(self) -> GraphSnapshot: + def snapshot(self, *, record_sources: bool = True) -> GraphSnapshot: r"""Materialize this graph: read each table's connector, run the result through `kumo-core` to build column stats and the on-disk index. The index is stored under a directory keyed by both `self.id` @@ -355,9 +597,12 @@ def snapshot(self) -> GraphSnapshot: and `GraphSnapshot`'s docstring for the resulting `id` shape. If a valid build already exists for that exact key, ingest is skipped and it's reused as is. - """ - from filelock import FileLock + Args: + record_sources: if False, the snapshot manifest records each + source as connector and table name only (no bucket, path or + account) and Custom SQL as a hash only (no SQL text). + """ self.validate() for table in self._named_tables.values(): table.connector.validate_scratch_dir(get_scratch_dir()) @@ -365,7 +610,6 @@ def snapshot(self) -> GraphSnapshot: graph_key = self.id persist_dir = snapshot_scratch_dir(graph_key) / 'source_tables' - persist_dir.parent.mkdir(parents=True, exist_ok=True) # Guards two things, both shared across every `snapshot()` call for # this `graph_key` -- different processes, or different threads in @@ -383,11 +627,12 @@ def snapshot(self) -> GraphSnapshot: # mid-build -- a plain in-memory lock wouldn't coordinate across # processes, and a lock-file-existence mutex would leave a stale # lock behind after a crash. - with FileLock(f'{persist_dir}.lock'): - return self._snapshot_no_lock(graph_key, persist_dir) + with _lock(graph_key): + return self._snapshot_no_lock(graph_key, persist_dir, + record_sources) - def _snapshot_no_lock(self, graph_key: str, persist_dir: Path) \ - -> GraphSnapshot: + def _snapshot_no_lock(self, graph_key: str, persist_dir: Path, + record_sources: bool = True) -> GraphSnapshot: r"""`snapshot()`'s body, run under its `FileLock`. Not meant to be called on its own -- `persist_dir`, unlike everything else here, isn't a pure function of `self` alone (it's `snapshot()`'s lock key, @@ -401,9 +646,7 @@ def _snapshot_no_lock(self, graph_key: str, persist_dir: Path) \ from kumo_core.diskgraph.runner import IngestExecutionConfig from kumo_core.diskgraph.snapshot import ( DiskgraphIngestResult, - DiskgraphLocalConfig, DiskgraphSnapshotHandle, - LocalSnapshotStore, check_diskgraph_output, ) from kumo_core.sources.layout import staging_dir @@ -413,12 +656,11 @@ def _snapshot_no_lock(self, graph_key: str, persist_dir: Path) \ from kumo.engine.connector.source_table import checked_sql graph_config = self._graph_config - store = LocalSnapshotStore( - DiskgraphLocalConfig(data_root=str(get_scratch_dir() / - 'snapshots'))) + store = _store() table_data_versions: Dict[str, str] = {} source_paths: Dict[str, str] = {} + source_tables: Dict[str, SourceTableEntry] = {} try: for alias, table in self._named_tables.items(): @@ -452,6 +694,9 @@ def _snapshot_no_lock(self, graph_key: str, persist_dir: Path) \ "limitation in kumoai/backend/local/graph.py's " "_stage_sources, which raises the same way.") table_data_versions[alias] = prepared.version + source_tables[alias] = source_table( + connector, name, sql, prepared.version, + record_sources=record_sources) paths = prepared.source.get('paths') or [] if paths: @@ -500,11 +745,16 @@ def _snapshot_no_lock(self, graph_key: str, persist_dir: Path) \ ), ) + snapshot_id = _snapshot_id(graph_key, data_version) + _write_snapshot_manifest(result, snapshot_id, source_tables, + record_sources, fresh=existing is None) + _warn_ingest_data_loss(result, graph_config) return GraphSnapshot( - id=f'{graph_key}@{data_version[:16]}', + id=snapshot_id, index_path=Path(result.index_path), graph_config=graph_config, + record_sources=record_sources, ) finally: discard_destination(str(persist_dir)) diff --git a/packages/kumo-engine/kumo/engine/integrations/__init__.py b/packages/kumo-engine/kumo/engine/integrations/__init__.py new file mode 100644 index 000000000..f2aa5b939 --- /dev/null +++ b/packages/kumo-engine/kumo/engine/integrations/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/packages/kumo-engine/kumo/engine/integrations/_mlflow_model.py b/packages/kumo-engine/kumo/engine/integrations/_mlflow_model.py new file mode 100644 index 000000000..314964044 --- /dev/null +++ b/packages/kumo-engine/kumo/engine/integrations/_mlflow_model.py @@ -0,0 +1,140 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +import tempfile +from pathlib import Path +from typing import Any, Dict, Mapping, Optional, Tuple + +import pandas as pd + +from kumo.engine.connector._extras import import_provider +from kumo.engine.connector.base import Connector +from kumo.engine.graph import GraphSnapshot + +_pyfunc = import_provider('mlflow.pyfunc', 'mlflow') +_schema = import_provider('mlflow.types.schema', 'mlflow') +_signature = import_provider('mlflow.models.signature', 'mlflow') + +_DTYPES = { + 'PREDICTED': 'boolean', + 'SCORE': 'double', + 'TIMESTAMP': 'datetime', + 'False_PROB': 'double', + 'True_PROB': 'double', +} + + +class KumoModel(_pyfunc.PythonModel): + def load_context(self, context: Any) -> None: + from kumo.engine.trainer.trainer import Trainer + + self._trainer = Trainer.load(context.artifacts['bundle']) + self._columns = list(context.model_config['columns']) + self._snapshots: Dict[str, GraphSnapshot] = {} + self._default = context.model_config.get('snapshot') + if self._default: + self._snapshot(self._default) + + def predict(self, context: Any, model_input: pd.DataFrame, + params: Optional[Dict[str, Any]] = None) -> pd.DataFrame: + params = params or {} + if len(model_input) == 0 and not params.get('population'): + return pd.DataFrame(columns=self._columns) + uri = params.get('snapshot') or self._default + if not uri: + raise ValueError("Pass params={'snapshot': }.") + return predict(self._trainer, self._snapshot(uri), model_input, + params) + + def _snapshot(self, uri: str) -> GraphSnapshot: + if uri not in self._snapshots: + connector, path = snapshot_location(uri) + self._snapshots[uri] = GraphSnapshot.load_from_connector( + connector, path=path) + return self._snapshots[uri] + + +def predict(trainer: Any, snapshot: GraphSnapshot, rows: pd.DataFrame, + params: Mapping[str, Any]) -> pd.DataFrame: + from kumo.engine._scratch import scoped_scratch_dir + from kumo.engine.connector.file import FileConnector + + time = params.get('prediction_time') + time = None if time is None else pd.Timestamp(time) + options = dict( + binary_classification_threshold=params.get( + 'binary_classification_threshold'), + num_classes_to_return=params.get('num_classes_to_return'), + ) + trainer.scratch_dir.mkdir(parents=True, exist_ok=True) + with tempfile.TemporaryDirectory(dir=trainer.scratch_dir) as tmp, \ + scoped_scratch_dir(Path(tmp) / 'scratch'): + if len(rows) == 0: + return trainer.predict_snapshot( + snapshot, prediction_time=None + if time is None else time.to_pydatetime(), **options).frame + rows = rows.rename(columns={'ENTITY': 'ENTITY1'}) + if time is not None and 'TIMESTAMP' not in rows: + rows = rows.assign(TIMESTAMP=time) + (Path(tmp) / 'rows').mkdir() + rows.to_parquet(Path(tmp) / 'rows' / 'rows.parquet', index=False) + return trainer.predict_snapshot( + snapshot, + entities=FileConnector(root_dir=str(Path(tmp) / 'rows'))[ + 'rows.parquet'], **options).frame + + +def snapshot_location(uri: str) -> Tuple[Connector, str]: + from kumo.engine.connector import connector_for_root + + parent, slash, name = uri.rstrip('/').rpartition('/') + return connector_for_root(parent or slash or '.'), name + + +def signature(snapshot: GraphSnapshot, query: str) -> Tuple[Any, pd.DataFrame]: + from kumo_core.train.formatting import prediction_output_columns + from kumo_core.types import TaskType + + from kumo.engine.pquery import PredictiveQuery + + bound = PredictiveQuery.from_snapshot(snapshot, query) + parsed = bound._parsed + entity = snapshot.native_graph.table(parsed.entity_table) + entity_type = _entity_type(str(entity.column(entity.primary_key).dtype)) + task_type = TaskType(parsed.task_type) + dtypes = { + **_DTYPES, 'TARGET_PRED': + 'boolean' if task_type == TaskType.BINARY_CLASSIFICATION else 'double' + } + inputs = _schema.Schema([ + _schema.ColSpec(entity_type, 'ENTITY'), + _schema.ColSpec('datetime', 'TIMESTAMP', required=False), + ]) + outputs = _schema.Schema([ + _schema.ColSpec(entity_type, 'ENTITY'), + *(_schema.ColSpec(dtypes.get(name) or _schema.AnyType(), name) + for name in prediction_output_columns(task_type)), + ]) + params = _schema.ParamSchema([ + _schema.ParamSpec('snapshot', 'string', None), + _schema.ParamSpec('prediction_time', 'string', None), + _schema.ParamSpec('binary_classification_threshold', 'double', None), + _schema.ParamSpec('num_classes_to_return', 'long', None), + _schema.ParamSpec('population', 'boolean', False), + ]) + example = pd.DataFrame({ + 'ENTITY': { + 'long': [0, 1], + 'double': [0.0, 1.0], + }.get(entity_type, ['entity-0', 'entity-1']) + }) + return _signature.ModelSignature(inputs, outputs, params), example + + +def _entity_type(dtype: str) -> str: + if dtype.startswith('int') or dtype == 'byte': + return 'long' + if dtype.startswith('float'): + return 'double' + return 'string' diff --git a/packages/kumo-engine/kumo/engine/integrations/_mlflow_model_code.py b/packages/kumo-engine/kumo/engine/integrations/_mlflow_model_code.py new file mode 100644 index 000000000..03cedae08 --- /dev/null +++ b/packages/kumo-engine/kumo/engine/integrations/_mlflow_model_code.py @@ -0,0 +1,7 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +from kumo.engine.connector._extras import import_provider +from kumo.engine.integrations._mlflow_model import KumoModel + +import_provider('mlflow.models', 'mlflow').set_model(KumoModel()) diff --git a/packages/kumo-engine/kumo/engine/integrations/mlflow.py b/packages/kumo-engine/kumo/engine/integrations/mlflow.py new file mode 100644 index 000000000..c32dde826 --- /dev/null +++ b/packages/kumo-engine/kumo/engine/integrations/mlflow.py @@ -0,0 +1,506 @@ +# 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]'``).""" + +import hashlib +import json +import logging +import os +import posixpath +import re +import time +import warnings +from contextlib import ExitStack, contextmanager +from pathlib import Path +from typing import ( + TYPE_CHECKING, + Any, + Callable, + Dict, + FrozenSet, + Iterator, + Mapping, + Optional, + Sequence, + Set, + Tuple, +) + +from kumoapi.plugin import PublishedModel, check_model_alias + +from kumo.engine._events import Event +from kumo.engine._manifest import BundleProvenance +from kumo.engine._tracking import DROPPED_PARAMS, INPUT_FIELDS, RunRecord +from kumo.engine._version import __version__ +from kumo.engine.connector._extras import import_provider + +if TYPE_CHECKING: + from kumo_core.embedding_models import EmbeddingModelStore + + from kumo.engine.graph import GraphSnapshot + from kumo.engine.trainer.trainer import Trainer, TrainingResult + +logger = logging.getLogger(__name__) + +_MAX_PARAM_KEY = 250 +_MAX_PARAM_VALUE = 6000 +# MLflow's portable name charset (its Windows rule, which omits ':'). +_PARAM_KEY = re.compile(r'[/\w.\- ]+') +# MLflow's defaults (7 retries, 120 s timeout) can hold a call for minutes. +_HTTP_LIMITS = { + 'MLFLOW_HTTP_REQUEST_MAX_RETRIES': '2', + 'MLFLOW_HTTP_REQUEST_TIMEOUT': '30', +} +# Model upload only, and only where the limits above apply. +_UPLOAD_HTTP = {'MLFLOW_HTTP_REQUEST_TIMEOUT': '600'} +_PATH_CONNECTORS = ('LOCAL_FILE', 'S3', 'GCS', 'ADLS') + + +class MlflowTracking: + r"""Records `Trainer.fit` in MLflow: one run per fit (nested under the + active run, if any) with a child run per trial attempt. Source tables + are logged as run inputs: Databricks tables as Delta 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. + + Args: + experiment_name: experiment for the fit run, created if missing. + `None` uses the current experiment. + run_name: name of the fit run. `None` lets MLflow pick one. + log_system_metrics: sample CPU, memory, disk, network and (with + ``nvidia-ml-py``) GPU usage on the fit run. + tags: extra tags for the fit run. + log_model: log the fitted model on the fit run (see `log_model`) + before the fit returns. + registered_model_name: also register it; implies ``log_model``. + aliases: registered-model aliases to point at the new version. + serving_snapshot: exported snapshot URI the served model loads at + startup (see `log_model`). + + ``run_id`` is the MLflow id of the latest fit run, once it has started. + After a fit, ``registered`` is the registered model version, else None. + + Raises: + ValueError: on aliases without a registered model name, or invalid. + """ + + def __init__( + self, + experiment_name: Optional[str] = None, + run_name: Optional[str] = None, + log_system_metrics: bool = True, + tags: Optional[Mapping[str, str]] = None, + log_model: bool = False, + registered_model_name: Optional[str] = None, + aliases: Sequence[str] = (), + serving_snapshot: Optional[str] = None, + ) -> None: + self._mlflow = import_provider('mlflow', 'mlflow') + _check_registration(registered_model_name, aliases) + self.experiment_name = experiment_name + self.run_name = run_name + self.log_system_metrics = log_system_metrics + self.tags = dict(tags or {}) + self.registered_model_name = registered_model_name + self.aliases = tuple(aliases) + self.serving_snapshot = serving_snapshot + self.logs_model = log_model or registered_model_name is not None + self.registered: Optional[PublishedModel] = None + self._outer: Optional[str] = None + self.run_id: Optional[str] = None + self._reset() + + def _reset(self) -> None: + self._client: Any = None + self._parent: Optional[str] = None + self._experiment_id: Optional[str] = None + self._children: Dict[Tuple[int, int], str] = {} + self._artifacts: Set[str] = set() + self._model_id: Optional[str] = None + self._http_owned: FrozenSet[str] = frozenset() + self._guard = ExitStack() + + def prepare(self) -> None: + active = self._mlflow.active_run() + self._outer = active.info.run_id if active is not None else None + + def start(self, record: RunRecord) -> None: + self._reset() + self.run_id = None + self.registered = None + self._http_owned = frozenset(name for name in _HTTP_LIMITS + if name not in os.environ) + self._guard.enter_context(_limited_http(self._http_owned)) + try: + autologging = import_provider('mlflow.utils.autologging_utils', + 'mlflow') + # Not reentrant: leaving it clears the flag. + if not getattr(autologging, '_AUTOLOGGING_GLOBALLY_DISABLED', + False): + self._guard.enter_context(autologging.disable_autologging()) + except Exception: + logger.warning('Could not suspend MLflow autologging', + exc_info=True) + mlflow = self._mlflow + self._client = mlflow.MlflowClient() + experiment_id = None + if self.experiment_name is not None: + experiment = self._client.get_experiment_by_name( + self.experiment_name) + experiment_id = (experiment.experiment_id + if experiment else self._client.create_experiment( + self.experiment_name)) + run = mlflow.start_run( + experiment_id=experiment_id, + run_name=self.run_name, + parent_run_id=self._outer, + tags=record.tags, + log_system_metrics=self.log_system_metrics, + ) + self._parent = self.run_id = run.info.run_id + self._experiment_id = run.info.experiment_id + self._log_params(self._parent, record.params) + self._log_artifacts(record.artifacts) + _log_inputs(mlflow, self._client, self._parent, record.inputs) + + def on_trial(self, record: RunRecord, event: Event) -> None: + key = (event.trial or 0, int(event.params['attempt'])) + if event.name == 'trial_started': + run = self._client.create_run( + self._experiment_id, run_name=record.name, + tags={**record.tags, 'mlflow.parentRunId': self._parent}) + self._children[key] = run.info.run_id + self._log_params(run.info.run_id, record.params) + elif event.name == 'epoch_ended': + self._log_metrics(self._children.get(key), event.metrics, + event.epoch or 0) + elif event.name == 'trial_ended': + run_id = self._children.pop(key) + self._log_metrics(run_id, record.metrics) + self._log_tags(run_id, record.tags) + succeeded = event.params.get('status') == 'succeeded' + self._client.set_terminated(run_id, + 'FINISHED' if succeeded else 'FAILED') + + def log_model(self, result: 'TrainingResult', + snapshot: 'GraphSnapshot') -> None: + run_id = self._parent + if not self.logs_model or run_id is None: + return + upload = { + name: value + for name, value in _UPLOAD_HTTP.items() if name in self._http_owned + } + os.environ.update(upload) + try: + info = _log_model(self._mlflow, result, snapshot, name='model', + run_id=run_id, + serving_snapshot=self.serving_snapshot) + finally: + if self._parent == run_id: + os.environ.update({name: _HTTP_LIMITS[name] for name in upload}) + if self._parent != run_id: + logger.warning('The fit run %s ended before its model was ' + 'logged; the model is not registered', run_id) + return + self._model_id = info.model_id + if self.registered_model_name is not None: + self.registered = _register(self._mlflow, info, + self.registered_model_name, + self.aliases) + + def finish(self, record: Optional[RunRecord], + error: Optional[BaseException]) -> None: + status = 'FAILED' + try: + if record is not None and self._parent is not None: + self._log_metrics(self._parent, record.metrics, + model_id=self._model_id) + self._log_tags(self._parent, record.tags) + self._log_artifacts(record.artifacts) + if error is None: + status = 'FINISHED' + finally: + _quietly(self._mlflow.flush_async_logging) + for run_id in self._children.values(): + _quietly(self._client.set_terminated, run_id, 'FAILED') + if self._parent is not None: + active = self._mlflow.active_run() + if active is not None and active.info.run_id == self._parent: + _quietly(self._mlflow.end_run, status) + else: + _quietly(self._client.set_terminated, self._parent, status) + self._guard.close() + self._reset() + + def _log_artifacts(self, artifacts: Mapping[str, str]) -> None: + for name, text in artifacts.items(): + if name not in self._artifacts: + self._client.log_text(self._parent, text, name) + self._artifacts.add(name) + + def _log_tags(self, run_id: str, tags: Mapping[str, str]) -> None: + entities = self._mlflow.entities + self._client.log_batch( + run_id, tags=[ + entities.RunTag(key, value) + for key, value in tags.items() + ]) + + def _log_params(self, run_id: str, params: Mapping[str, str]) -> None: + entities = self._mlflow.entities + kept = _mlflow_params(params) + dropped = len(params) - len(kept) + tags = [entities.RunTag(DROPPED_PARAMS, str(dropped))] if dropped else [] + self._client.log_batch( + run_id, params=[ + entities.Param(key, value) + for key, value in kept.items() + ], tags=tags) + + def _log_metrics(self, run_id: Optional[str], + metrics: Mapping[str, float], step: int = 0, + model_id: Optional[str] = None) -> None: + if run_id is None or not metrics: + return + entities = self._mlflow.entities + timestamp = int(time.time() * 1000) + self._client.log_batch( + run_id, metrics=[ + entities.Metric(key, float(value), timestamp, step, + model_id=model_id) + for key, value in metrics.items() + ]) + + +def log_model( + trainer: 'Trainer', + *, + snapshot: 'GraphSnapshot', + name: str = 'model', + registered_model_name: Optional[str] = None, + aliases: Sequence[str] = (), + run_id: Optional[str] = None, + serving_snapshot: Optional[str] = None, +) -> Any: + r"""Log a trainer's bundle as an MLflow pyfunc model in ``run_id`` (or + the active run, else a new one) with its source tables as run inputs; + returns its `ModelInfo`. Scores with + ``predict(entities, params={'snapshot': uri})``; with + ``serving_snapshot`` (an exported snapshot URI), each serving replica + downloads that snapshot once at startup and scores against it when no + ``snapshot`` param is given.""" + mlflow = import_provider('mlflow', 'mlflow') + _check_registration(registered_model_name, aliases) + result = trainer._training_result + if result is None: + raise ValueError('log_model requires a fitted or loaded Trainer') + with _in_run(mlflow, run_id) as active: + _log_inputs(mlflow, mlflow.MlflowClient(), active, + _provenance_inputs(result.provenance)) + info = _log_model(mlflow, result, snapshot, name=name, run_id=active, + serving_snapshot=serving_snapshot) + if result.metrics: + mlflow.log_metrics(dict(result.metrics), run_id=info.run_id, + model_id=info.model_id) + if registered_model_name is not None: + _register(mlflow, info, registered_model_name, aliases) + return info + + +def load_model( + model_uri: str, + embedding_models: Optional['EmbeddingModelStore'] = None, +) -> 'Trainer': + r"""Restore the `Trainer` of a model `log_model` logged.""" + from kumo.engine._scratch import new_run_id, trainer_scratch_dir + from kumo.engine.trainer.trainer import Trainer, warn_if_other_bundle + + mlflow = import_provider('mlflow', 'mlflow') + destination = trainer_scratch_dir(new_run_id()) / 'mlflow-model' + destination.mkdir(parents=True) + local = Path( + mlflow.artifacts.download_artifacts(artifact_uri=model_uri, + dst_path=str(destination))) + model = mlflow.models.Model.load(str(local)) + bundle = model.flavors['python_function']['artifacts']['bundle']['path'] + trainer = Trainer.load(local / bundle, embedding_models) + warn_if_other_bundle(model.metadata, trainer) + return trainer + + +def _log_model(mlflow: Any, result: 'TrainingResult', + snapshot: Optional['GraphSnapshot'], *, name: str, run_id: str, + serving_snapshot: Optional[str]) -> Any: + from kumo.engine.integrations import _mlflow_model, _mlflow_model_code + from kumo.engine.trainer._spec import read_training_spec + from kumo.engine.trainer.trainer import registry_pointers + + if snapshot is None: + raise ValueError('log_model needs the snapshot the model trained on') + spec = read_training_spec(result.bundle_dir) + signature, example = _mlflow_model.signature(snapshot, + spec.predictive_query) + provenance = result.bundle_dir / 'provenance.json' + model_logger = logging.getLogger('mlflow.models.model') + model_logger.addFilter(_no_serving_check) + try: + with _in_run(mlflow, run_id): + return mlflow.pyfunc.log_model( + name=name, + python_model=_mlflow_model_code.__file__, + model_config={ + 'columns': signature.outputs.input_names(), + 'snapshot': serving_snapshot, + }, + artifacts={ + 'bundle': str(result.bundle_dir), + **({'provenance': str(provenance)} + if provenance.is_file() else {}), + }, + signature=signature, + input_example=example, + metadata=registry_pointers(result), + pip_requirements=[f'kumo-engine[s3,gcs]=={__version__}'], + ) + finally: + model_logger.removeFilter(_no_serving_check) + + +def _provenance_inputs( + provenance: Optional[BundleProvenance] +) -> Dict[str, Mapping[str, Any]]: + if provenance is None: + return {} + tables = dict(provenance.source_tables or {}) + custom = (None if provenance.training_table is None else + provenance.training_table.source_table) + if custom is not None and custom.location: + tables['training_table'] = custom + return { + name: {key: getattr(entry, key) + for key in INPUT_FIELDS} + for name, entry in tables.items() + } + + +def _log_inputs(mlflow: Any, client: Any, run_id: str, + inputs: Mapping[str, Mapping[str, Optional[str]]]) -> None: + try: + meta = import_provider('mlflow.data.meta_dataset', 'mlflow') + entities = mlflow.entities + datasets = [] + for name, entry in inputs.items(): + source = _dataset_source(entry) + if source is None: + continue + digest = hashlib.sha256( + json.dumps(dict(entry), sort_keys=True).encode()).hexdigest() + dataset = meta.MetaDataset(source, name=name, digest=digest[:8]) + datasets.append( + entities.DatasetInput( + dataset._to_mlflow_entity(), + tags=[entities.InputTag('mlflow.data.context', 'train')])) + if datasets: + client.log_inputs(run_id, datasets=datasets) + except Exception: + logger.warning('Could not log the fit\'s source tables to MLflow', + exc_info=True) + + +def _register(mlflow: Any, info: Any, name: str, + aliases: Sequence[str]) -> PublishedModel: + version = mlflow.register_model(info.model_uri, name) + client = mlflow.MlflowClient() + for alias in aliases: + client.set_registered_model_alias(name, alias, version.version) + return PublishedModel(name=name, version=str(version.version), + uri=f'models:/{name}/{version.version}') + + +def _no_serving_check(record: logging.LogRecord) -> bool: + # MLflow's input-example check predicts without a snapshot; no switch. + return not record.getMessage().startswith( + 'Failed to validate serving input example') + + +@contextmanager +def _in_run(mlflow: Any, run_id: Optional[str]) -> Iterator[str]: + active = mlflow.active_run() + if active is not None and run_id in (None, active.info.run_id): + yield active.info.run_id + return + with mlflow.start_run(run_id=run_id, nested=active is not None) as run: + yield run.info.run_id + + +def _check_registration(registered_model_name: Optional[str], + aliases: Sequence[str]) -> None: + if isinstance(aliases, str): + raise TypeError('aliases must be a list of names, not a string') + if aliases and registered_model_name is None: + raise ValueError('aliases require registered_model_name') + for alias in aliases: + check_model_alias(alias) + + +def _dataset_source(entry: Mapping[str, Optional[str]]) -> Any: + connector = (entry.get('connector') or '').upper() + location = entry.get('location') + if not location or location.startswith('sql:'): + return None + if connector == 'DATABRICKS': + delta = import_provider('mlflow.data.delta_dataset_source', 'mlflow') + return delta.DeltaDatasetSource( + delta_table_name=_unquoted(location), + delta_table_version=_delta_version(entry)) + if connector in _PATH_CONNECTORS: + registry = import_provider('mlflow.data.dataset_source_registry', + 'mlflow') + with warnings.catch_warnings(): + warnings.simplefilter('ignore') + return registry.resolve_dataset_source(location) + return None + + +def _unquoted(name: str) -> str: + parts = re.findall(r'`((?:[^`]|``)*)`', name) + return '.'.join(part.replace('``', '`') for part in parts) or name + + +def _delta_version(entry: Mapping[str, Optional[str]]) -> Optional[int]: + version = entry.get('version') or '' + head = version.split('_', 1)[0] + if entry.get('version_kind') == 'exact' and head.isdigit(): + return int(head) + return None + + +def _mlflow_params(params: Mapping[str, str]) -> Dict[str, str]: + return { + key: value + for key, value in params.items() + if len(key) <= _MAX_PARAM_KEY and _PARAM_KEY.fullmatch(key) is not None + and posixpath.normpath(key) == key and not key.startswith('/') + and len(value) <= _MAX_PARAM_VALUE + } + + +def _quietly(function: Callable[..., Any], *args: Any) -> None: + try: + function(*args) + except Exception: + logger.debug('MLflow cleanup call failed', exc_info=True) + + +@contextmanager +def _limited_http(names: FrozenSet[str]) -> Iterator[None]: + os.environ.update({name: _HTTP_LIMITS[name] for name in names}) + try: + yield + finally: + for name in names: + os.environ.pop(name, None) diff --git a/packages/kumo-engine/kumo/engine/integrations/snowflake.py b/packages/kumo-engine/kumo/engine/integrations/snowflake.py new file mode 100644 index 000000000..4d78d5b56 --- /dev/null +++ b/packages/kumo-engine/kumo/engine/integrations/snowflake.py @@ -0,0 +1,251 @@ +# 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]'``).""" + +import logging +import re +import tempfile +from contextlib import ExitStack +from pathlib import Path +from typing import Any, Dict, List, Mapping, Optional, Set, Tuple + +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.connector._extras import import_provider + +logger = logging.getLogger(__name__) + +_MAX_METRICS = 200 +_MAX_PARAMS = 1000 +# Metric name budgets: per-epoch metrics give way first, then trial summaries. +_EPOCH_METRICS = _MAX_METRICS - 60 +_SUMMARY_METRICS = _MAX_METRICS - 30 +_FINISH_PARAMS = 20 + + +class SnowflakeTracking: + r"""Records `Trainer.fit` in Snowflake Experiment Tracking as one run; + 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. + + Args: + session: Snowpark session; needs CREATE EXPERIMENT on the schema, or + MODIFY on an existing experiment. + experiment_name: experiment for the fit run, created if missing. + database: `None` uses the session's current database. + schema: `None` uses the session's current schema, or ``PUBLIC`` when + `database` is given or the session has none. + run_name: name of the fit run. `None` lets Snowflake pick one; a + taken name gets a suffix from the kumo run id. + tags: extra tags for the fit run. + native_model_lineage: keep the run open as a context for the whole + fit, so models registered meanwhile are linked to it. + + ``run_id`` is the latest fit run's name (Snowflake runs are identified + by name), once it has started. + """ + + def __init__( + self, + session: Any, + *, + experiment_name: str, + database: Optional[str] = None, + schema: Optional[str] = None, + run_name: Optional[str] = None, + tags: Optional[Mapping[str, str]] = None, + native_model_lineage: bool = False, + ) -> None: + self._module = import_provider('snowflake.ml.experiment', 'snowflake') + self.session = session + self.experiment_name = experiment_name + self.database = database + self.schema = schema + self.run_name = run_name + self.tags = dict(tags or {}) + self.native_model_lineage = native_model_lineage + self.run_id: Optional[str] = None + self._reset() + + def _reset(self) -> None: + self._tracking: Any = None + self._run: Any = None + self._run_context = ExitStack() + self._trial_tags: Dict[Tuple[int, int], Set[str]] = {} + self._metrics: Set[str] = set() + self._params = 0 + self._start_tags: Set[str] = set() + self._artifacts: Set[str] = set() + self._dropped_metrics: Set[str] = set() + self._dropped_params = 0 + + def start(self, record: RunRecord) -> None: + self._reset() + self.run_id = None + cls = self._module.ExperimentTracking + tracking = getattr(cls, '_instance', None) + if tracking is None: + tracking = cls(self.session) + 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') + return + self._tracking = tracking + database = self.database or self.session.get_current_database() + schema = self.schema or (None if self.database else + self.session.get_current_schema()) + experiment = tracking.set_experiment(self.experiment_name, + database_name=database, + schema_name=schema or 'PUBLIC') + self._run = tracking.start_run( + self._new_run_name(tracking, experiment.name, + record.tags['kumo.run_id'])) + self.run_id = str(self._run.name) + if self.native_model_lineage: + self._run_context.enter_context(self._run) + self._start_tags = set(record.tags) + self._log_params({**record.params, **record.tags}) + self._log_artifacts(record.artifacts) + + def on_trial(self, record: RunRecord, event: Event) -> None: + if self._run is None: + return + prefix = f'{record.name}/' + key = (event.trial or 0, int(event.params['attempt'])) + if event.name == 'trial_started': + self._trial_tags[key] = set(record.tags) + self._log_params({prefix + k: v + for k, v in {**record.params, + **record.tags}.items()}) + elif event.name == 'epoch_ended': + self._log_metrics({prefix + k: v + for k, v in event.metrics.items()}, + _EPOCH_METRICS, step=event.epoch or 0) + elif event.name == 'trial_ended': + self._log_metrics({prefix + k: v + for k, v in record.metrics.items()}, + _SUMMARY_METRICS) + logged = self._trial_tags.pop(key, set()) + self._log_params({prefix + k: v + for k, v in record.tags.items() + if k not in logged}) + + def finish(self, record: Optional[RunRecord], + error: Optional[BaseException]) -> None: + status = 'FAILED' + try: + if self._run is not None and record is not None: + self._log_metrics(record.metrics, _MAX_METRICS) + self._log_params({k: v + for k, v in record.tags.items() + if k not in self._start_tags}, _MAX_PARAMS) + self._log_artifacts(record.artifacts) + self._log_dropped() + if error is None: + status = 'FINISHED' + finally: + try: + if self._run is not None: + self._tracking.end_run(status=status) + except Exception: + logger.warning('Could not end the Snowflake run %s', status, + exc_info=True) + finally: + self._run_context.close() + self._reset() + + def _new_run_name(self, tracking: Any, experiment: Any, + run_id: str) -> Optional[str]: + if self.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 + # that is neither. + existing = [ + str(row['name']) + for row in tracking._sql_client.show_runs_in_experiment( + experiment_name=experiment) + ] + short = run_id.rsplit('-', 1)[-1] + for name in (self.run_name, _suffixed(self.run_name, short[:8]), + _suffixed(self.run_name, short)): + if not _taken(name, existing): + break + else: + name = None + if name != self.run_name: + logger.warning( + 'Snowflake run %s already exists in experiment %s; recording ' + 'this fit as %s', self.run_name, self.experiment_name, + name or 'a generated run name') + return name + + def _log_artifacts(self, artifacts: Mapping[str, str]) -> None: + with tempfile.TemporaryDirectory() as directory: + for name, text in artifacts.items(): + if name in self._artifacts: + continue + path = Path(directory) / name + path.write_text(text, encoding='utf-8') + self._tracking.log_artifact(str(path)) + self._artifacts.add(name) + + def _log_dropped(self) -> None: + dropped = { + DROPPED_METRICS: len(self._dropped_metrics), + DROPPED_PARAMS: self._dropped_params, + } + dropped = {name: count for name, count in dropped.items() if count} + if dropped: + self._tracking.log_params(dropped) + + def _log_params(self, params: Mapping[str, str], + limit: int = _MAX_PARAMS - _FINISH_PARAMS) -> None: + kept = {} + for name, value in params.items(): + if '$$' in name or '$$' in value or self._params >= limit: + self._dropped_params += 1 + else: + kept[name] = value + self._params += 1 + if kept: + self._tracking.log_params(kept) + + def _log_metrics(self, metrics: Mapping[str, float], limit: int, + step: int = 0) -> None: + kept = {} + for name, value in metrics.items(): + if '$$' in name: + self._dropped_metrics.add(name) + continue + if name not in self._metrics: + if len(self._metrics) >= limit: + self._dropped_metrics.add(name) + continue + self._metrics.add(name) + kept[name] = value + if kept: + self._tracking.log_metrics(kept, step=step) + + +def _suffixed(name: str, suffix: str) -> str: + if len(name) > 1 and name[0] == name[-1] == '"': + return f'{name[:-1]}_{suffix}"' + return f'{name}_{suffix}' + + +def _taken(name: str, existing: List[str]) -> bool: + resolved = resolve_snowflake_identifier(name) + regex = ''.join('.*' if c == '%' else '.' if c == '_' else re.escape(c) + for c in name) + return any( + other == resolved + or re.fullmatch(regex, other, re.IGNORECASE | re.DOTALL) is not None + for other in existing) diff --git a/packages/kumo-engine/kumo/engine/pquery/_custom_source.py b/packages/kumo-engine/kumo/engine/pquery/_custom_source.py index a98c74975..74ea2e9d9 100644 --- a/packages/kumo-engine/kumo/engine/pquery/_custom_source.py +++ b/packages/kumo-engine/kumo/engine/pquery/_custom_source.py @@ -8,6 +8,7 @@ import os import shutil from pathlib import Path +from typing import Optional, Tuple import pyarrow.dataset as ds import pyarrow.parquet as pq @@ -15,8 +16,9 @@ from kumo.engine.connector.source_table import SourceTable -def stage_source_table(source_table: SourceTable, destination: Path) -> Path: - """Read ``source_table`` through its connector and retain its parquet parts.""" +def stage_source_table(source_table: SourceTable, + destination: Path) -> Tuple[Path, Optional[str]]: + """Stage ``source_table``'s parquet parts; returns them and their version.""" if not isinstance(source_table, SourceTable): raise TypeError("source_table must be a kumo.engine SourceTable") destination.mkdir(parents=True, exist_ok=False) @@ -62,7 +64,7 @@ def stage_source_table(source_table: SourceTable, destination: Path) -> Path: errno.EPERM, errno.EACCES, errno.EOPNOTSUPP, errno.ENOSYS): raise shutil.copyfile(part, target) - return output + return output, prepared.version def _validate_source(path, query, graph, *, weight_col, split_col, diff --git a/packages/kumo-engine/kumo/engine/pquery/prediction_table.py b/packages/kumo-engine/kumo/engine/pquery/prediction_table.py index d0f96bc01..f64912109 100644 --- a/packages/kumo-engine/kumo/engine/pquery/prediction_table.py +++ b/packages/kumo-engine/kumo/engine/pquery/prediction_table.py @@ -104,7 +104,8 @@ def from_source_table( pq._require_matching_snapshot(snapshot) run_id = new_run_id() run_dir = prediction_table_scratch_dir(run_id) - rows_path = stage_source_table(source_table, run_dir / "source-table") + rows_path, _ = stage_source_table(source_table, + run_dir / "source-table") request, _ = _build_prediction_table_request( snapshot, pq._parsed, diff --git a/packages/kumo-engine/kumo/engine/pquery/predictive_query.py b/packages/kumo-engine/kumo/engine/pquery/predictive_query.py index 20fb51a9e..1e738fc64 100644 --- a/packages/kumo-engine/kumo/engine/pquery/predictive_query.py +++ b/packages/kumo-engine/kumo/engine/pquery/predictive_query.py @@ -6,15 +6,17 @@ `PredictiveQuery(graph, query)` binds a `Graph` at construction, matching the cloud SDK's `kumo.PredictiveQuery(graph=graph, query=QUERY)` -- not a -`GraphSnapshot`. Parsing (`kumo_core.pquery.parse_and_validate_query` against -`self.graph.native_graph` -- `Graph`'s own `cached_property`, needing no -materialized index) happens lazily, on first access to the `_parsed` -`cached_property`, and only once: `validate()`/`get_task_type()` both read -that same property rather than re-parsing on every call. Confirmed against -the existing local backend's own `_LocalPQueryBackend.validate` -(`kumoai/backend/local/pquery.py`), which parses against -`to_portable(definition)` -- the schema-level conversion, not an ingested -snapshot -- that this doesn't need a `GraphSnapshot` either. +`GraphSnapshot`. `PredictiveQuery.from_snapshot(snapshot, query)` binds a +snapshot instead and parses against `snapshot.graph_config`. + +Parsing (`kumo_core.pquery.parse_and_validate_query` against +`self.graph.native_graph`, needing no materialized index) happens lazily, on first +access to the `_parsed` `cached_property`, and only once: +`validate()`/`get_task_type()` both read that same property rather than +re-parsing on every call. Confirmed against the existing local backend's own +`_LocalPQueryBackend.validate` (`kumoai/backend/local/pquery.py`), which parses +against `to_portable(definition)` -- the schema-level conversion, not an +ingested snapshot -- that this doesn't need a `GraphSnapshot` either. Construction itself stays I/O-free, consistent with `Table`/`Graph`: nothing parses until `_parsed` is actually touched (by `validate()`, `get_task_type()`, @@ -25,8 +27,8 @@ `suggest_model_plan` uses resolved graph/query metadata. Methods that generate training or prediction tables also read the materialized index. These methods -accept an explicit `GraphSnapshot` produced by `self.graph.snapshot()` and -never trigger graph ingest themselves. +accept an explicit `GraphSnapshot` of this query's graph and never trigger +graph ingest themselves. Supported, matching `kumo.engine.pquery.predictive_query.PredictiveQuery` in the cloud SDK: `validate`, `get_task_type`, `suggest_model_plan`, @@ -90,22 +92,34 @@ PredictionTable, _compile_prediction_table, ) -from kumo.engine.pquery.training_table import TrainingTable +from kumo.engine.pquery.training_table import TrainingTable, _record_manifest logger = logging.getLogger(__name__) @dataclass(frozen=True) class PredictiveQuery: - r"""An immutable PQL query bound to a `Graph`.""" + r"""An immutable PQL query bound to a `Graph` or a `GraphSnapshot`.""" - graph: Graph + graph: Graph | GraphSnapshot query: str + @classmethod + def from_snapshot(cls, snapshot: GraphSnapshot, + query: str) -> 'PredictiveQuery': + r"""Bind `query` to `snapshot`'s graph without its `Graph` object. + + Snapshots passed to this query's methods must come from the same + graph (same `graph_key`). That only checks the snapshot belongs to + this query's graph; it says nothing about compatibility with a + trained model. + """ + return cls(graph=snapshot, query=query) + @cached_property def _parsed(self) -> ValidatedNativeQuery: - r"""Parses and validates `self.query` against `self.graph. - native_graph`, via `kumo_core.pquery.parse_and_validate_query` -- + r"""Parses and validates `self.query` against `self.graph.native_graph`, + via `kumo_core.pquery.parse_and_validate_query` -- the same parser `Trainer._predict` already calls directly, and what `kumo_core.training.workflow.TrainingWorkflowRequest.from_spec` uses internally for `fit`. No materialized `GraphSnapshot` needed @@ -125,23 +139,21 @@ def _parsed(self) -> ValidatedNativeQuery: `task_type`/`task_family`/`query_type` on the returned `ValidatedNativeQuery` are all plain `str`; `get_task_type` below wraps `task_type` into the shared, public `kumoapi.task.TaskType`. - - Reuses `self.graph.native_graph` (itself a `cached_property`) - rather than building its own `graph_config`/`NativeGraph` -- one - canonical conversion per `Graph`, shared by `Graph.validate()` and - every `PredictiveQuery` bound to it. """ return parse_and_validate_query(self.query, self.graph.native_graph) def _require_matching_snapshot(self, snapshot: GraphSnapshot) -> None: - # snapshot.id is 'graph_key@data_version_prefix' (see - # Graph.snapshot()); self.graph.id is the graph_key alone, so - # membership is a prefix check, not equality. - if not snapshot.id.startswith(f'{self.graph.id}@'): + if isinstance(self.graph, GraphSnapshot): + if snapshot.id == self.graph.id: + return + graph_key = self.graph.graph_key + else: + graph_key = self.graph.id + if snapshot.graph_key != graph_key: raise ValueError( f"snapshot {snapshot.id!r} was not materialized from " - f"self.graph (id {self.graph.id!r}) -- call " - "self.graph.snapshot() and pass its result.") + f"this query's graph (id {graph_key!r}) -- pass a " + "snapshot of that graph.") def validate(self, verbose: bool = True) -> ValidationResponse: r"""Return validation warnings and information for a valid query. @@ -195,7 +207,7 @@ def suggest_model_plan( train_table_spec: Optional[TrainingTableSpec] = None, ) -> ModelPlan: r"""Suggests a modeling plan given this query and a materialized - `snapshot` of `self.graph`. Mirrors the cloud SDK's + `snapshot` of this query's graph. Mirrors the cloud SDK's `PredictiveQuery.suggest_model_plan`. Uses resolved graph/query metadata without generating a training table @@ -221,7 +233,7 @@ def generate_training_table( plan: Optional[TrainingTableGenerationPlan] = None, ) -> TrainingTable: r"""Compiles this query into a `TrainingTable` against a materialized - `snapshot` of `self.graph`. Mirrors the cloud SDK's + `snapshot` of this query's graph. Mirrors the cloud SDK's `PredictiveQuery.generate_training_table`. `Trainer.fit` consumes the returned `TrainingTable` directly. Its @@ -229,7 +241,7 @@ def generate_training_table( settings are retained without keeping a graph-bound task plan. Args: - snapshot: a snapshot of `self.graph` to compile against. + snapshot: a snapshot of this query's graph to compile against. plan: if not given, uses `kumo_core`'s generation defaults, matching the local backend's existing `suggest_training_table_plan` default. @@ -308,7 +320,7 @@ def generate_training_table( or {})), ) - return TrainingTable( + table = TrainingTable( query=self.query, train_table_path=Path(generated.train_table_path), manifest_path=Path(generated.manifest_path), @@ -316,6 +328,8 @@ def generate_training_table( required_stages=task_plan.task_dataset.required_stages, generation_plan=generation_plan, ) + _record_manifest(table, snapshot, task_plan.train_table) + return table def generate_prediction_table( self, @@ -325,7 +339,7 @@ def generate_prediction_table( plan: Optional[TrainingTableGenerationPlan] = None, ) -> PredictionTable: r"""Compiles this query into a `PredictionTable` against a - materialized `snapshot` of `self.graph`. Mirrors the cloud SDK's + materialized `snapshot` of this query's graph. Mirrors the cloud SDK's `PredictiveQuery.generate_prediction_table`, and parallels `generate_training_table` above. @@ -337,7 +351,7 @@ def generate_prediction_table( docstring). Args: - snapshot: a snapshot of `self.graph` to compile against. + snapshot: a snapshot of this query's graph to compile against. prediction_time: the anchor timestamp for a temporal query; when omitted, uses the latest source-time maximum across the query's target aggregations in the snapshot. diff --git a/packages/kumo-engine/kumo/engine/pquery/training_table.py b/packages/kumo-engine/kumo/engine/pquery/training_table.py index ae58aac63..990b85dcb 100644 --- a/packages/kumo-engine/kumo/engine/pquery/training_table.py +++ b/packages/kumo-engine/kumo/engine/pquery/training_table.py @@ -2,21 +2,36 @@ # All rights reserved. # SPDX-License-Identifier: Apache-2.0 +import json +import warnings from dataclasses import dataclass, replace from pathlib import Path -from typing import TYPE_CHECKING, Optional +from typing import TYPE_CHECKING, Any, Dict, Optional from kumo_core.model_plan import TrainingTableGenerationPlan from kumoapi.model_plan import \ TrainingTableGenerationPlan as PublicGenerationPlan from kumoapi.train import TrainingTableSpec +from kumo.engine import _provenance as provenance +from kumo.engine._manifest import ( + SourceTableEntry, + TrainingTableManifest, + TrainingTableParents, + read_manifest, + stamp, + write_manifest, +) from kumo.engine.connector.source_table import SourceTable from kumo.engine.graph import GraphSnapshot if TYPE_CHECKING: from kumo.engine.trainer.output_config import TrainingTableOutputConfig +STAGING_MANIFEST_FILE = 'training-table-manifest.json' +EXPORT_MANIFEST_FILE = '_kumo_manifest.json' +_RANDOM_SPLITS = ('RandomSplit', 'RandomLinkSplit') + @dataclass(frozen=True) class TrainingTable: @@ -38,6 +53,27 @@ class TrainingTable: #: Effective generation settings carried forward into ``Trainer.fit``. generation_plan: TrainingTableGenerationPlan + @property + def tt_id(self) -> str: + """The id DiskGraph assigned to this table at generation.""" + from kumo_core.training.export import read_spec + + return read_spec(str(self.staging_dir))['tt_id'] + + @property + def manifest(self) -> Optional[TrainingTableManifest]: + """Lineage recorded at generation; None for older tables.""" + return read_manifest(self.staging_dir / STAGING_MANIFEST_FILE, + TrainingTableManifest) + + @property + def source_mode(self) -> Optional[str]: + """``pquery`` or ``custom_tt``, as DiskGraph recorded at generation.""" + try: + return json.loads(self.manifest_path.read_text())['source_mode'] + except (OSError, ValueError, KeyError, TypeError): + return None + def export(self, output_config: "TrainingTableOutputConfig") -> str: r"""Write the staged rows with a ``SPLIT`` column (``0`` train, ``1`` validation, ``2`` test) for :meth:`from_source_table`. Returns the @@ -63,7 +99,27 @@ def export(self, output_config: "TrainingTableOutputConfig") -> str: with tempfile.TemporaryDirectory(dir=root) as output: files = export_train_table(str(self.staging_dir), output) return write_files(connector, output_config.table_name, files, - output_config.write_mode, export_id=export_id) + output_config.write_mode, export_id=export_id, + **self._export_sidecars(connector, + output_config)) + + def _export_sidecars(self, connector, output_config) -> Dict[str, Any]: + from kumo_core.object_store_sink import ObjectStoreSinkBackend + from kumo_core.sources import layout + + manifest = self.staging_dir / STAGING_MANIFEST_FILE + name = output_config.table_name.rstrip('/') + try: + if (not isinstance(connector.sink_backend, ObjectStoreSinkBackend) + or layout.file_type(name) == layout.CSV + or not manifest.is_file()): + return {} + return {'sidecars': {EXPORT_MANIFEST_FILE: manifest.read_bytes()}} + except Exception as error: + warnings.warn( + f"Exporting without {EXPORT_MANIFEST_FILE}: {error}", + UserWarning, stacklevel=3) + return {} @classmethod def from_source_table( @@ -126,7 +182,8 @@ def from_source_table( ) run_id = new_run_id() run_dir = train_table_scratch_dir(run_id) - rows_path = stage_source_table(source_table, run_dir / "source-table") + rows_path, version = stage_source_table(source_table, + run_dir / "source-table") generation_plan = from_json_dict( to_json_dict(public_plan), TrainingTableGenerationPlan, @@ -183,7 +240,7 @@ def from_source_table( year_over_year=("yoy_offsets_secs" in (task_plan.train_table.compiler_options or {})), ) - return cls( + table = cls( query=pq.query, train_table_path=Path(generated.train_table_path), manifest_path=Path(generated.manifest_path), @@ -191,3 +248,45 @@ def from_source_table( required_stages=task_plan.task_dataset.required_stages, generation_plan=generation_plan, ) + _record_manifest( + table, snapshot, task_plan.train_table, + source_table=provenance.source_table( + source_table.connector, source_table.table_name, + source_table.sql, version, + record_sources=_record_sources(snapshot))) + return table + + +def _record_sources(snapshot: GraphSnapshot) -> bool: + if snapshot.record_sources is not None: + return snapshot.record_sources + manifest = snapshot.manifest + return manifest is not None and manifest.record_sources is True + + +def _record_manifest(table: TrainingTable, snapshot: GraphSnapshot, + request: Any, *, + source_table: Optional[SourceTableEntry] = None) -> None: + path = table.staging_dir / STAGING_MANIFEST_FILE + try: + from kumo_core.io.json import to_json_dict + + seed = None + if request.split_family in _RANDOM_SPLITS: + seed = (request.split_config or {}).get('rnd_seed') or 0 + write_manifest( + path, + TrainingTableManifest( + id=table.tt_id, + parents=TrainingTableParents(snapshot=snapshot.id, + source_table=source_table), + query=table.query, + generation_plan=to_json_dict(table.generation_plan), + task_type=request.task_type, + split_seed=seed, + **stamp(), + )) + except Exception as error: + warnings.warn( + f"Could not write the training-table manifest at {path}: " + f"{error}", UserWarning, stacklevel=3) diff --git a/packages/kumo-engine/kumo/engine/tracking.py b/packages/kumo-engine/kumo/engine/tracking.py new file mode 100644 index 000000000..a466674b0 --- /dev/null +++ b/packages/kumo-engine/kumo/engine/tracking.py @@ -0,0 +1,11 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +r"""API for writing an experiment-tracking adapter for ``fit(tracking=...)``.""" + +from kumo.engine._tracking import RunRecord, Tracking + +__all__ = [ + 'Tracking', + 'RunRecord', +] diff --git a/packages/kumo-engine/kumo/engine/trainer/__init__.py b/packages/kumo-engine/kumo/engine/trainer/__init__.py index 3301763b5..b768c814c 100644 --- a/packages/kumo-engine/kumo/engine/trainer/__init__.py +++ b/packages/kumo-engine/kumo/engine/trainer/__init__.py @@ -2,18 +2,26 @@ # All rights reserved. # SPDX-License-Identifier: Apache-2.0 +from kumo.engine.trainer._scoring import ScoringDefaults from kumo.engine.trainer.output_config import ( ModelOutputConfig, PredictionOutputConfig, TrainingTableOutputConfig, ) -from kumo.engine.trainer.trainer import PredictionResult, Trainer, TrainingResult +from kumo.engine.trainer.trainer import ( + PredictionResult, + Trainer, + TrainingResult, + set_model_alias, +) __all__ = [ "Trainer", "TrainingResult", "PredictionResult", + "ScoringDefaults", "ModelOutputConfig", "PredictionOutputConfig", "TrainingTableOutputConfig", + "set_model_alias", ] diff --git a/packages/kumo-engine/kumo/engine/trainer/_provenance.py b/packages/kumo-engine/kumo/engine/trainer/_provenance.py new file mode 100644 index 000000000..f7dec8807 --- /dev/null +++ b/packages/kumo-engine/kumo/engine/trainer/_provenance.py @@ -0,0 +1,105 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Fit record in ``provenance.json``; excluded from the bundle id.""" + +from dataclasses import replace +from datetime import datetime, timezone +from pathlib import Path +from typing import Dict, Optional + +from kumo.engine._manifest import ( + RECONSTRUCTED, + BundleParents, + BundleProvenance, + BundleTrainingTable, + FitTiming, + SourceTableEntry, + TrackingRef, + read_manifest, + stamp, + write_manifest, +) +from kumo.engine.graph import GraphSnapshot +from kumo.engine.pquery.training_table import TrainingTable +from kumo.engine.trainer._spec import PROVENANCE_ROLE, listed_file + + +def write_provenance( + bundle_dir: Path, + *, + bundle_id: Optional[str], + train_table: TrainingTable, + warm_start_bundle_id: Optional[str], + started_at: datetime, + fit_mode: Optional[str], + snapshot: Optional[GraphSnapshot] = None, + tracked_run_id: Optional[str] = None, +) -> None: + from kumo_core.artifacts import PROVENANCE_FILE, add_bundle_file + + finished_at = datetime.now(timezone.utc) + write_manifest( + bundle_dir / PROVENANCE_FILE, + BundleProvenance( + id=bundle_id, + parents=BundleParents(training_table=train_table.tt_id, + warm_start=warm_start_bundle_id), + fit=FitTiming( + started_at=started_at.isoformat(timespec='seconds'), + finished_at=finished_at.isoformat(timespec='seconds'), + seconds=round((finished_at - started_at).total_seconds(), 3), + ), + training_table=_training_table(train_table), + source_tables=_source_tables(snapshot), + fit_mode=fit_mode, + tracking=None if tracked_run_id is None else TrackingRef( + run_id=tracked_run_id), + **stamp(), + )) + add_bundle_file(bundle_dir, PROVENANCE_ROLE, PROVENANCE_FILE) + + +def _source_tables( + snapshot: Optional[GraphSnapshot] +) -> Optional[Dict[str, SourceTableEntry]]: + manifest = None if snapshot is None else snapshot.manifest + if manifest is None or manifest.reconstructed: + return None + return manifest.parents.source_tables + + +def _training_table(train_table: TrainingTable) -> BundleTrainingTable: + manifest = train_table.manifest + return BundleTrainingTable( + source=train_table.source_mode, + snapshot=None if manifest is None else manifest.parents.snapshot, + source_table=None + if manifest is None else manifest.parents.source_table) + + +def record_derivation(bundle_dir: Path, parent_bundle_id: Optional[str]) -> None: + from kumo_core.artifacts import ( + PROVENANCE_FILE, + add_bundle_file, + bundle_id_from_manifest, + ) + + provenance = read_provenance(bundle_dir) + if provenance is None: + if listed_file(bundle_dir, PROVENANCE_ROLE) is not None: + raise ValueError('its provenance.json is unreadable or from a ' + 'newer kumo-engine, so it is left unchanged') + provenance = BundleProvenance(id=None, parents=BundleParents(), + **RECONSTRUCTED) + write_manifest( + bundle_dir / PROVENANCE_FILE, + replace(provenance, id=bundle_id_from_manifest(bundle_dir), + parents=replace(provenance.parents, + derived_from=parent_bundle_id))) + add_bundle_file(bundle_dir, PROVENANCE_ROLE, PROVENANCE_FILE) + + +def read_provenance(bundle_dir: Path) -> Optional[BundleProvenance]: + path = listed_file(bundle_dir, PROVENANCE_ROLE) + return None if path is None else read_manifest(path, BundleProvenance) diff --git a/packages/kumo-engine/kumo/engine/trainer/_runtime.py b/packages/kumo-engine/kumo/engine/trainer/_runtime.py index 15d372e10..8ce71b3d8 100644 --- a/packages/kumo-engine/kumo/engine/trainer/_runtime.py +++ b/packages/kumo-engine/kumo/engine/trainer/_runtime.py @@ -10,8 +10,10 @@ def local_training_bindings(training_input: Any, run_name: Optional[str] = None, - embedding_models: Any = None) -> Any: + embedding_models: Any = None, + instrumentation: Any = None) -> Any: from kumo_core.training.runtime import local_training_bindings as bind return bind(training_input, run_name=run_name, - embedding_models=embedding_models) + embedding_models=embedding_models, + instrumentation=instrumentation) diff --git a/packages/kumo-engine/kumo/engine/trainer/_scoring.py b/packages/kumo-engine/kumo/engine/trainer/_scoring.py new file mode 100644 index 000000000..daa0d3e6f --- /dev/null +++ b/packages/kumo-engine/kumo/engine/trainer/_scoring.py @@ -0,0 +1,78 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +r"""Scoring defaults stored as ``scoring-defaults.json`` in a model bundle.""" + +import json +import warnings +from dataclasses import asdict, dataclass +from numbers import Integral, Real +from pathlib import Path +from typing import Any, Mapping, Optional + +from kumo.engine._manifest import write_manifest +from kumo.engine.trainer._spec import SCORING_DEFAULTS_ROLE, listed_file + +SCORING_DEFAULTS_FILE = 'scoring-defaults.json' + + +@dataclass(frozen=True) +class ScoringDefaults: + """Prediction arguments a bundle applies when ``predict`` omits them.""" + + binary_classification_threshold: Optional[float] = None + num_classes_to_return: Optional[int] = None + + +def normalized_scoring_defaults(defaults: ScoringDefaults, + task_type: Any) -> ScoringDefaults: + from kumo_core.prediction_options import ( + resolve_classification_threshold, + resolve_num_classes_to_return, + ) + + threshold = defaults.binary_classification_threshold + count = defaults.num_classes_to_return + if threshold is not None: + if isinstance(threshold, bool) or not isinstance(threshold, Real): + raise ValueError('binary_classification_threshold must be a ' + f'number, but got {threshold!r}') + threshold = float(threshold) + if resolve_classification_threshold(task_type, threshold) is None: + raise ValueError( + 'binary_classification_threshold applies only to binary and ' + f'multilabel classification, not {task_type.value}') + if isinstance(count, Integral) and not isinstance(count, bool): + count = int(count) + if count is not None: + resolve_num_classes_to_return(task_type, count) + return ScoringDefaults(threshold, count) + + +def write_scoring_defaults(bundle_dir: Path, defaults: ScoringDefaults) -> None: + from kumo_core.artifacts import add_bundle_file, remove_bundle_file + + if defaults == ScoringDefaults(): + remove_bundle_file(bundle_dir, SCORING_DEFAULTS_ROLE) + return + write_manifest(bundle_dir / SCORING_DEFAULTS_FILE, asdict(defaults)) + add_bundle_file(bundle_dir, SCORING_DEFAULTS_ROLE, SCORING_DEFAULTS_FILE) + + +def read_scoring_defaults(bundle_dir: Path, + artifact_paths: Mapping[str, str]) -> ScoringDefaults: + from kumo_core.artifacts import bundle_task_type + + path = listed_file(bundle_dir, SCORING_DEFAULTS_ROLE) + if path is None: + return ScoringDefaults() + try: + data = json.loads(path.read_text(encoding='utf-8')) + return normalized_scoring_defaults( + ScoringDefaults(data.get('binary_classification_threshold'), + data.get('num_classes_to_return')), + bundle_task_type(artifact_paths)) + except Exception as error: + warnings.warn(f'Ignoring the scoring defaults in {bundle_dir}: {error}', + UserWarning, stacklevel=3) + return ScoringDefaults() diff --git a/packages/kumo-engine/kumo/engine/trainer/_spec.py b/packages/kumo-engine/kumo/engine/trainer/_spec.py index 0115532b7..757390902 100644 --- a/packages/kumo-engine/kumo/engine/trainer/_spec.py +++ b/packages/kumo-engine/kumo/engine/trainer/_spec.py @@ -9,14 +9,12 @@ separate and must be supplied through a compatible prediction snapshot. """ -import hashlib import json -import os -import tempfile from dataclasses import asdict, dataclass from pathlib import Path -from typing import Any, Dict +from typing import Any, Dict, Optional, Tuple +from kumo_core.io.filesystem import atomic_write_bytes from kumoapi.model_plan import ModelPlan, TrainingTableGenerationPlan #: Bundle metadata filename and schema version. @@ -24,10 +22,11 @@ SPEC_VERSION = 1 #: What `kumo_core`'s `VitalArtifactPublicationPolicy` names its own -#: manifest. `write_training_spec` amends this file's `artifacts` list so -#: `training-spec.json` is covered the same way every other bundle file is. +#: manifest; `training-spec.json` is listed in it like every other file. MANIFEST_FILE = 'publication-manifest.json' SPEC_ROLE = 'training_spec' +PROVENANCE_ROLE = 'provenance' +SCORING_DEFAULTS_ROLE = 'scoring_defaults' @dataclass(frozen=True) @@ -46,36 +45,6 @@ class TrainingSpec: training_table_generation: Dict[str, Any] -def _sha256(path: Path) -> str: - digest = hashlib.sha256() - digest.update(path.read_bytes()) - return digest.hexdigest() - - -def _atomic_write(path: Path, payload: bytes) -> None: - """Replace one metadata file without exposing partial JSON.""" - temporary = None - try: - with tempfile.NamedTemporaryFile(dir=path.parent, - prefix=f'.{path.name}.', - delete=False) as out: - temporary = Path(out.name) - out.write(payload) - out.flush() - os.fsync(out.fileno()) - os.replace(temporary, path) - temporary = None - if hasattr(os, 'O_DIRECTORY'): - descriptor = os.open(path.parent, os.O_RDONLY | os.O_DIRECTORY) - try: - os.fsync(descriptor) - finally: - os.close(descriptor) - finally: - if temporary is not None: - temporary.unlink(missing_ok=True) - - def write_training_spec( bundle_dir: Path, *, @@ -85,6 +54,7 @@ def write_training_spec( training_table_generation: TrainingTableGenerationPlan, ) -> None: r"""Write an integrity-covered training spec into a completed bundle.""" + from kumo_core.artifacts import add_bundle_file from kumo_core.io.json import to_json_dict spec = TrainingSpec( @@ -95,28 +65,12 @@ def write_training_spec( training_table_generation=to_json_dict(training_table_generation), ) spec_path = bundle_dir / SPEC_FILE - _atomic_write( + atomic_write_bytes( spec_path, (json.dumps(asdict(spec), indent=2, sort_keys=True) + '\n').encode('utf-8'), ) - - manifest_path = bundle_dir / MANIFEST_FILE - with open(manifest_path, encoding='utf-8') as handle: - manifest = json.load(handle) - manifest['artifacts'] = [ - entry for entry in manifest.get('artifacts') or [] - if entry.get('relative_path') != SPEC_FILE - ] + [{ - 'role': SPEC_ROLE, - 'relative_path': SPEC_FILE, - 'size_bytes': spec_path.stat().st_size, - 'sha256': _sha256(spec_path), - }] - _atomic_write( - manifest_path, - (json.dumps(manifest, indent=2, sort_keys=True) + '\n').encode('utf-8'), - ) + add_bundle_file(bundle_dir, SPEC_ROLE, SPEC_FILE) def read_training_spec(bundle_dir: Path) -> TrainingSpec: @@ -150,11 +104,23 @@ def read_training_spec(bundle_dir: Path) -> TrainingSpec: ) -def read_bundle_artifact_paths(bundle_dir: Path) -> Dict[str, str]: - r"""Rebuild the verified ``role -> path`` artifact mapping.""" +def listed_file(bundle_dir: Path, role: str) -> Optional[Path]: + try: + manifest = json.loads( + (bundle_dir / MANIFEST_FILE).read_text(encoding='utf-8')) + entries = manifest.get('artifacts') or [] + except (OSError, ValueError, AttributeError): + return None + return next((bundle_dir / entry['relative_path'] + for entry in entries if entry.get('role') == role), None) + + +def read_bundle_artifact_paths( + bundle_dir: Path) -> Tuple[Dict[str, str], str]: + r"""Rebuild the verified ``role -> path`` mapping and the bundle id.""" from kumo_core.artifacts import verify_model_bundle - verify_model_bundle(bundle_dir) + bundle_id = verify_model_bundle(bundle_dir) manifest_path = bundle_dir / MANIFEST_FILE if not manifest_path.is_file(): raise ValueError( @@ -175,5 +141,6 @@ def read_bundle_artifact_paths(bundle_dir: Path) -> Dict[str, str]: return { entry['role']: str(bundle_dir / entry['relative_path']) for entry in artifacts - if entry['role'] != SPEC_ROLE - } + if entry['role'] not in (SPEC_ROLE, PROVENANCE_ROLE, + SCORING_DEFAULTS_ROLE) + }, bundle_id diff --git a/packages/kumo-engine/kumo/engine/trainer/output_config.py b/packages/kumo-engine/kumo/engine/trainer/output_config.py index 4a325126d..b17641c6c 100644 --- a/packages/kumo-engine/kumo/engine/trainer/output_config.py +++ b/packages/kumo-engine/kumo/engine/trainer/output_config.py @@ -24,6 +24,9 @@ class PredictionOutputConfig: ``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. + ``lineage_columns`` writes constant ``KUMO_RUN_ID``, ``KUMO_BUNDLE_ID`` + and ``KUMO_SNAPSHOT_ID`` columns into every output row, only when passed + to ``Trainer.predict``; enable it only on a new table or directory. """ connector: Connector @@ -33,6 +36,7 @@ class PredictionOutputConfig: extra_fields: List[MetadataField] = field(default_factory=list) output_paths: Optional[Mapping[str, str]] = None + lineage_columns: bool = False def __post_init__(self) -> None: if self.output_paths is not None: diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index e44504897..8eb17ff3d 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -7,14 +7,18 @@ import hashlib import json import logging +import os +import shutil import time -from dataclasses import dataclass, replace +import warnings +from dataclasses import dataclass, field, replace from datetime import datetime, timezone from pathlib import Path -from typing import Dict, Mapping, Optional, cast +from typing import Any, Callable, Dict, Mapping, Optional, Sequence, cast import pandas as pd from kumo_core.embedding_models import EmbeddingModelStore +from kumo_core.instrumentation import format_error from kumo_core.model_plan import ModelPlan as CoreModelPlan from kumo_core.model_plan import inferred_or_missing_type_to_none from kumo_core.model_plan_conversion import public_model_plan_from_internal @@ -28,10 +32,13 @@ from kumo_core.types import ExplainMethod, PredictionArtifactType from kumoapi.model_plan import ModelPlan, TrainingTableGenerationPlan from kumoapi.plugin import ( + BUNDLE_ID_KEY, + ModelAliasingSink, ModelFetchingSink, ModelRegistryLocation, PublishedModel, SourceConnector, + check_model_alias, ) from kumoapi.train import TrainingStage from typing_extensions import Self @@ -41,11 +48,21 @@ require_embedding_models, require_trained_embedding_models, ) +from kumo.engine._events import Event, EventEmitter, query_sha256 +from kumo.engine._manifest import BundleProvenance from kumo.engine._scratch import new_run_id, trainer_scratch_dir +from kumo.engine._tracking import ( + FitContext, + Tracking, + TrackingSession, + fit_inputs, +) from kumo.engine.connector.base import Connector +from kumo.engine.connector.source_table import SourceTable from kumo.engine.graph import GraphSnapshot from kumo.engine.pquery.prediction_table import PredictionTable from kumo.engine.pquery.training_table import TrainingTable +from kumo.engine.trainer._scoring import ScoringDefaults, read_scoring_defaults from kumo.engine.trainer._spec import ( read_bundle_artifact_paths, read_training_spec, @@ -66,11 +83,12 @@ class TrainingResult: `published_model` records the destination's assigned name, version and URI when `fit()` exports the bundle. It is `None` without an export. - `holdout_path` is the holdout Parquet file written by `fit()`, or `None` - when the task produced none (also for results of `Trainer.load`). It is - not part of the published model bundle. The file is capped at 1 GiB in - memory; `holdout_truncated` is true when rows were dropped to fit, and - `None` when the size record could not be read. + `holdout_path` is the published model's holdout Parquet file written by + `fit()`, or `None` when it has none (no holdout task, a model refit on the + full data, or results of `Trainer.load`). It is not part of the published + model bundle. The file is capped at 1 GiB in memory; `holdout_truncated` is + true when rows were dropped to fit, and `None` when the size record could + not be read. """ metrics: Dict[str, float] @@ -79,6 +97,14 @@ class TrainingResult: published_model: Optional[PublishedModel] = None holdout_path: Optional[Path] = None holdout_truncated: Optional[bool] = False + bundle_id: Optional[str] = None + + @property + def provenance(self) -> Optional[BundleProvenance]: + """The bundle's recorded fit provenance; None for older bundles.""" + from kumo.engine.trainer._provenance import read_provenance + + return read_provenance(self.bundle_dir) def holdout_df(self) -> pd.DataFrame: r"""Read the holdout predictions as a pandas DataFrame. @@ -105,10 +131,13 @@ 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. + ``bundle_id`` and ``snapshot_id`` identify the model and graph scored. """ selection: PredictionSelectionRef artifact_type: PredictionArtifactType + bundle_id: Optional[str] = field(default=None, compare=False) + snapshot_id: Optional[str] = field(default=None, compare=False) @property def num_rows(self) -> int: @@ -141,8 +170,19 @@ 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=...)') + if output_config.lineage_columns and not self._has_lineage_columns(): + warnings.warn( + 'lineage_columns are written at prediction time; this result ' + 'has none, so the export has none either.', UserWarning, + stacklevel=2) self._export(output_config) + def _has_lineage_columns(self) -> bool: + from pyarrow import parquet + files = self.files + return bool(files) and LINEAGE_RUN_ID in parquet.read_schema( + files[0]).names + def _export(self, output_config: PredictionOutputConfig) -> None: write_files = getattr(output_config.connector.sink_backend, 'write_files', None) @@ -160,6 +200,9 @@ def _export(self, output_config: PredictionOutputConfig) -> None: export_id=f'{self.selection.job_id}-{kind.value}') +LINEAGE_RUN_ID = 'KUMO_RUN_ID' + + def _find_holdout(log_dir: str) -> tuple[Optional[Path], Optional[bool]]: from kumo_core.train.output import EvalOutput path = Path(EvalOutput.get_df_path(log_dir, TrainingStage.TEST)) @@ -176,6 +219,17 @@ def _find_holdout(log_dir: str) -> tuple[Optional[Path], Optional[bool]]: return path, truncated +def _published_holdout(trainer: Any) -> tuple[Optional[Path], Optional[bool]]: + from kumo_core.training.config import FitMode + + published = trainer.published_fit_mode() + if published is FitMode.refit_full: + logger.info('The published model is refit on the full data, including ' + 'the test split, so it has no holdout predictions.') + return None, False + return _find_holdout(trainer.get_log_dir(fit_mode=published)) + + 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. @@ -193,6 +247,56 @@ def _to_core_model_plan(model_plan: ModelPlan) -> CoreModelPlan: return internal_model_plan_from_public(model_plan) +def _link_or_copy(source: str, destination: str) -> None: + try: + os.link(source, destination) + except OSError: + shutil.copy2(source, destination) + + +def _require_publisher(export_config: ModelOutputConfig) -> None: + if not callable( + getattr(export_config.connector.sink_backend, 'publish_model', + None)): + raise NotImplementedError( + f'{type(export_config.connector).__name__} cannot publish ' + f'models') + if export_config.location and export_config.location.alias: + check_model_alias(export_config.location.alias) + + +def _publish(training_result: TrainingResult, + export_config: ModelOutputConfig) -> PublishedModel: + spec = read_training_spec(training_result.bundle_dir) + metadata = { + name: str(getattr(spec, name)) + for name in ('spec_version', 'predictive_query', 'snapshot_id') + } + metadata.update(registry_pointers(training_result)) + location = export_config.location + alias = location.alias if location else None + if location and alias: + location = replace(location, alias=None) + published = export_config.connector.sink_backend.publish_model( + export_config.connector, + str(training_result.bundle_dir), + path=export_config.path, + location=location, + metadata=metadata, + metrics=training_result.metrics, + ) + if location and alias: + try: + set_model_alias(export_config.connector, + replace(location, version=published.version), + alias) + except Exception as error: + warnings.warn( + f"Published {published.name} version {published.version} " + f"without alias {alias}: {error}", UserWarning, stacklevel=4) + return published + + def _stage_model_bundle(artifacts: Mapping[str, str], dest: Path) -> Path: from kumo_core.artifacts import stage_model_bundle return stage_model_bundle(artifacts, dest) @@ -223,6 +327,7 @@ def __init__( self._training_query: Optional[str] = None self._training_table_generation_plan: Optional[ TrainingTableGenerationPlan] = None + self._scoring_defaults = ScoringDefaults() @property def model_plan(self) -> ModelPlan: @@ -239,10 +344,69 @@ def training_table_generation_plan( """Effective table-generation settings saved by the last fit/load.""" return self._training_table_generation_plan + @property + def training_query(self) -> Optional[str]: + """Predictive query of the last fit/load.""" + return self._training_query + @property def scratch_dir(self) -> Path: return trainer_scratch_dir(self._id) + @property + def bundle_id(self) -> Optional[str]: + """Content id of the fitted or loaded model bundle.""" + return getattr(self._training_result, 'bundle_id', None) + + @property + def provenance(self) -> Optional[BundleProvenance]: + """Fit provenance of the fitted or loaded bundle, if recorded.""" + result = self._training_result + return None if result is None else result.provenance + + @property + def scoring_defaults(self) -> ScoringDefaults: + """Prediction defaults carried by the fitted or loaded bundle.""" + return self._scoring_defaults + + def _fitted(self, method: str) -> TrainingResult: + if self._training_result is None: + raise RuntimeError( + f"Trainer.{method}() requires a prior Trainer.fit() call on " + "this same Trainer instance.") + return self._training_result + + def with_scoring_defaults(self, defaults: ScoringDefaults) -> Self: + """Return a trainer on a copy of this bundle that carries ``defaults``.""" + from kumo_core.artifacts import bundle_task_type + + from kumo.engine.trainer._provenance import record_derivation + from kumo.engine.trainer._scoring import ( + normalized_scoring_defaults, + write_scoring_defaults, + ) + + training_result = self._fitted('with_scoring_defaults') + defaults = normalized_scoring_defaults( + defaults, bundle_task_type(training_result.artifact_paths)) + bundle_dir = trainer_scratch_dir(new_run_id()) / 'derived-model' + shutil.copytree(training_result.bundle_dir, bundle_dir, + copy_function=_link_or_copy) + write_scoring_defaults(bundle_dir, defaults) + try: + record_derivation(bundle_dir, training_result.bundle_id) + except Exception as error: + warnings.warn( + f"Could not record the parent bundle in provenance.json: " + f"{error}", UserWarning, stacklevel=2) + return self._load_bundle(bundle_dir, self._embedding_models) + + def export(self, config: ModelOutputConfig) -> PublishedModel: + """Publish the fitted or loaded model bundle.""" + training_result = self._fitted('export') + _require_publisher(config) + return _publish(training_result, config) + def fit( self, train_table: TrainingTable, @@ -250,6 +414,8 @@ def fit( snapshot: GraphSnapshot, export_config: Optional[ModelOutputConfig] = None, warm_start: Optional[Self] = None, + callbacks: Sequence[Callable[[Event], None]] = (), + tracking: Optional[Tracking] = None, ) -> TrainingResult: r"""Train a model against `train_table`. @@ -275,12 +441,63 @@ def fit( state, fitted encoders, architecture and neighbor sampling; this trainer's training controls still apply. Task types must match. Forecasting warm starts are unsupported. + callbacks: called with each progress `Event`; a callback that + raises is logged and not called again. + tracking: a `kumo.engine.tracking.Tracking`, e.g. + `kumo.engine.integrations.mlflow.MlflowTracking` or + `kumo.engine.integrations.snowflake.SnowflakeTracking`; + tracking errors are logged, never raised. """ + if export_config is not None: + _require_publisher(export_config) + run_id = f'{self._id}-fit-{new_run_id()}' + session = TrackingSession(tracking) + events = EventEmitter( + run_id, + [*callbacks, session.on_event] if session.enabled else callbacks) + fit_start = time.monotonic() + try: + started = { + 'snapshot_id': snapshot.id, + 'query_sha256': query_sha256(train_table.query), + } + except Exception: + started = {} + events.emit('fit_started', params=started) + try: + result = self._fit(train_table, snapshot=snapshot, + export_config=export_config, + warm_start=warm_start, run_id=run_id, + events=events, tracking=session) + except BaseException as error: + events.emit('fit_ended', params={ + 'status': 'failed', + 'error': format_error(error), + }) + session.finish(None, error) + raise + events.emit('fit_ended', metrics={ + **result.metrics, 'seconds': time.monotonic() - fit_start + }, params={'status': 'succeeded'}) + session.finish(result, session.log_model(result, snapshot)) + return result + + def _fit( + self, + train_table: TrainingTable, + *, + snapshot: GraphSnapshot, + export_config: Optional[ModelOutputConfig], + warm_start: Optional[Self], + run_id: str, + events: EventEmitter, + tracking: TrackingSession, + ) -> TrainingResult: from kumo_core.automl.runtime import ( inherit_warm_start_model_plan, validate_warm_start_task_type, ) - from kumo_core.io.json import load_json + from kumo_core.io.json import load_json, to_json_dict from kumo_core.training.orchestration import ( CallableTrainingOptions, ResolvedTrainingInput, @@ -297,24 +514,20 @@ def fit( from kumo.engine.trainer._runtime import local_training_bindings - if export_config is not None and not callable( - getattr(export_config.connector.sink_backend, 'publish_model', - None)): - raise NotImplementedError( - f'{type(export_config.connector).__name__} cannot publish ' - f'models') - run_id = f'{self._id}-fit-{new_run_id()}' + started_at = datetime.now(timezone.utc) run_dir = self.scratch_dir / 'runs' / run_id model_plan = _to_core_model_plan(self._model_plan) warm_start_artifacts = None + warm_start_bundle_id = None effective_plan = self._model_plan if warm_start is not None: if warm_start._training_result is None: raise ValueError( "warm_start requires a fitted or loaded Trainer") - warm_start_artifacts = read_bundle_artifact_paths( - warm_start._training_result.bundle_dir) + warm_start_artifacts, warm_start_bundle_id = ( + read_bundle_artifact_paths( + warm_start._training_result.bundle_dir)) inherit_warm_start_model_plan( model_plan, load_json(warm_start_artifacts['resolved_model_plan'], @@ -336,11 +549,11 @@ def fit( search=AutoMLSearchOptions(), ) if warm_start_artifacts is not None: - source_schema = json.loads( - Path(warm_start_artifacts['prediction_schema']).read_text()) + from kumo_core.artifacts import bundle_task_type + validate_warm_start_task_type( TaskType(workflow.query.task_type), - TaskType(source_schema['task_type']), + bundle_task_type(warm_start_artifacts), ) from kumo_core.pquery import ( query_is_temporal, @@ -408,6 +621,13 @@ def fit( tt_id=ingest_tt_id, ) del task_dataset + tracking.start(lambda: FitContext( + run_id=run_id, + query=train_table.query, + model_plan=to_json_dict(materialized.model_plan), + inputs=fit_inputs(Path(snapshot.index_path), + Path(train_table.staging_dir)), + )) training_input = ResolvedTrainingInput( task=materialized.task, model_plan=materialized.model_plan, @@ -419,7 +639,8 @@ def fit( training_input, bindings=local_training_bindings( training_input, run_name=run_id, - embedding_models=self._embedding_models), + embedding_models=self._embedding_models, + instrumentation=events), options=CallableTrainingOptions(), ) @@ -445,8 +666,10 @@ def fit( 'failed for run %s. The trained model remains available.', run_id) - holdout_path, holdout_truncated = _find_holdout( - result.best_trainer.log_dir) + holdout_path, holdout_truncated = _published_holdout( + result.best_trainer) + fit_mode = result.best_trainer.published_fit_mode() + fit_mode = None if fit_mode is None else fit_mode.value training_result = TrainingResult( metrics=dict(result.metrics), artifact_paths=dict(result.artifact_paths), @@ -463,6 +686,11 @@ def fit( training_table_generation=train_table.generation_plan, export_config=export_config, model_plan=effective_plan, + provenance=dict(train_table=train_table, + warm_start_bundle_id=warm_start_bundle_id, + started_at=started_at, fit_mode=fit_mode, + snapshot=snapshot, + tracked_run_id=run_id if tracking.enabled else None), ) def _finish_fit( @@ -474,8 +702,13 @@ def _finish_fit( training_table_generation: TrainingTableGenerationPlan, export_config: Optional[ModelOutputConfig], model_plan: Optional[ModelPlan] = None, + provenance: Optional[Mapping[str, Any]] = None, ) -> TrainingResult: """Commit durable fit outputs before making this trainer usable.""" + from kumo_core.artifacts import bundle_id_from_manifest + + from kumo.engine.trainer._provenance import write_provenance + effective_plan = model_plan if model_plan is not None else self._model_plan write_training_spec( training_result.bundle_dir, @@ -484,30 +717,83 @@ def _finish_fit( snapshot_id=snapshot_id, training_table_generation=training_table_generation, ) + try: + training_result = replace( + training_result, + bundle_id=bundle_id_from_manifest(training_result.bundle_dir)) + except Exception as error: + warnings.warn(f"Could not compute the bundle id: {error}", + UserWarning, stacklevel=3) + if provenance is not None: + try: + write_provenance(training_result.bundle_dir, + bundle_id=training_result.bundle_id, + **provenance) + except Exception as error: + warnings.warn( + f"Publishing the model bundle without provenance.json: " + f"{error}", UserWarning, stacklevel=3) if export_config is not None: - spec = read_training_spec(training_result.bundle_dir) - published = export_config.connector.sink_backend.publish_model( - export_config.connector, - str(training_result.bundle_dir), - path=export_config.path, - location=export_config.location, - metadata={ - name: str(getattr(spec, name)) - for name in ('spec_version', 'predictive_query', - 'snapshot_id') - }, - metrics=training_result.metrics, - ) - training_result = replace(training_result, - published_model=published) + training_result = replace( + training_result, + published_model=_publish(training_result, export_config)) self._training_result = training_result + self._scoring_defaults = ScoringDefaults() self._snapshot_id = snapshot_id self._training_query = predictive_query self._training_table_generation_plan = training_table_generation return training_result + def predict_snapshot( + self, + snapshot: GraphSnapshot, + *, + entities: Optional[SourceTable] = None, + prediction_time: Optional[datetime] = None, + output_config: Optional[PredictionOutputConfig] = None, + num_classes_to_return: Optional[int] = None, + num_workers: int = 1, + binary_classification_threshold: Optional[float] = None, + ) -> PredictionResult: + r"""Score ``snapshot`` with the query and table-generation plan saved + by the last fit/load. + + Without ``entities``, scores the query's population at + ``prediction_time`` (as :meth:`PredictiveQuery.generate_prediction_table`). + ``entities`` is a source table of ``ENTITY1`` (and ``TIMESTAMP`` where + the task needs one) rows to score instead. Other arguments are as in + :meth:`predict`. Warns when ``snapshot`` was taken of a different + graph than the one the model was trained on. + """ + from kumo.engine.pquery.predictive_query import PredictiveQuery + + if self._training_query is None: + raise RuntimeError( + "Trainer.predict_snapshot() requires a prior Trainer.fit() or " + "Trainer.load() call on this same Trainer instance.") + if entities is not None and prediction_time is not None: + raise ValueError("Pass prediction_time or entities, not both; " + "put anchor times in the entities' TIMESTAMP.") + trained = (self._snapshot_id or '').partition('@')[0] + if trained and snapshot.graph_key != trained: + warnings.warn( + f"Snapshot {snapshot.id} was taken of graph " + f"{snapshot.graph_key}, but this model was trained on a " + f"snapshot of graph {trained}.", UserWarning, stacklevel=2) + query = PredictiveQuery.from_snapshot(snapshot, self._training_query) + plan = self._training_table_generation_plan + table = (query.generate_prediction_table( + snapshot, prediction_time=prediction_time, plan=plan) + if entities is None else PredictionTable.from_source_table( + entities, pq=query, snapshot=snapshot, plan=plan)) + 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) + def predict( self, prediction_table: PredictionTable, @@ -537,6 +823,8 @@ def predict( Binary classification requires an explicit threshold. Multilabel classification defaults to 0.5. Both accept thresholds from 0 to 1; thresholds change predicted labels without changing class scores. + + Omitted arguments fall back to the bundle's :attr:`scoring_defaults`. """ from kumo_core.prediction_options import ( validate_num_classes_to_return, @@ -545,10 +833,10 @@ def predict( validate_prediction_workers(num_workers) validate_num_classes_to_return(num_classes_to_return) - if self._training_result is None: - raise RuntimeError( - "Trainer.predict() requires a prior Trainer.fit() call on " - "this same Trainer instance.") + training_result = self._fitted('predict') + if binary_classification_threshold is None: + binary_classification_threshold = ( + self._scoring_defaults.binary_classification_threshold) from kumo_core.pquery import normalize_graph, parse_and_validate_query from kumo_core.prediction_options import ( prediction_query_contract, @@ -582,6 +870,8 @@ def predict( # population and anchors still come from the immutable PredictionTable. native_graph = normalize_graph(graph_config) 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 = self._scoring_defaults.num_classes_to_return num_classes_to_return = resolve_num_classes_to_return( task_type, num_classes_to_return, @@ -608,7 +898,7 @@ def predict( from kumo_core.io.json import load_json # Loaded once: the embedding check and the in-process scorer share it. - encoders_path = self._training_result.artifact_paths.get('encoders') + encoders_path = training_result.artifact_paths.get('encoders') encoders = (load_json(encoders_path, EncodersConfig) if encoders_path and Path(encoders_path).is_file() else None) require_trained_embedding_models(self._embedding_models, encoders) @@ -672,17 +962,26 @@ def predict( model_dir = str( _stage_model_bundle( - self._training_result.artifact_paths, + training_result.artifact_paths, run_dir / 'model_bundle', )) extra_fields = ([f.value for f in output_config.extra_fields] if output_config is not None else []) + lineage = {} + if output_config is not None and output_config.lineage_columns: + lineage = { + name: value + for name, value in ((LINEAGE_RUN_ID, run_id), + ('KUMO_BUNDLE_ID', self.bundle_id), + ('KUMO_SNAPSHOT_ID', snapshot.id)) + if value is not None + } source_identity = [ snapshot.id, prediction_table.query, timestamps, sorted(kind.value for kind in artifacts), num_classes_to_return, threshold, extra_fields - ] + ] + ([lineage] if lineage else []) selection = score_prediction( materialized=materialized, model_dir=model_dir, @@ -704,11 +1003,14 @@ def predict( forecast_options=forecast_options, embedding_models=self._embedding_models, encoders=encoders, + constant_columns=lineage, ) del encoders _collect_garbage('predict') - result = PredictionResult(selection, artifact_type) + result = PredictionResult(selection, artifact_type, + bundle_id=self.bundle_id, + snapshot_id=snapshot.id) if output_config is not None: result._export(output_config) return result @@ -753,9 +1055,17 @@ def load_from_connector( if not isinstance(sink, ModelFetchingSink): raise TypeError( 'model_registry.sink_backend must implement ModelFetchingSink') - sink.fetch_model(cast(SourceConnector, model_registry), str(bundle_dir), - path=address.path, location=address.location) - return cls._load_bundle(bundle_dir, embedding_models) + alias = None if location is None else location.alias + if alias is not None and not isinstance(sink, ModelAliasingSink): + raise NotImplementedError( + f'{type(model_registry).__name__} cannot fetch models by ' + f'alias') + published = sink.fetch_model(cast(SourceConnector, model_registry), + str(bundle_dir), path=address.path, + location=address.location) + trainer = cls._load_bundle(bundle_dir, embedding_models) + warn_if_other_bundle(getattr(published, 'metadata', None), trainer) + return trainer @classmethod def _load_bundle( @@ -764,7 +1074,7 @@ def _load_bundle( embedding_models: Optional[EmbeddingModelStore] = None, ) -> Self: bundle_dir = bundle_dir.resolve() - paths = read_bundle_artifact_paths(bundle_dir) + paths, bundle_id = read_bundle_artifact_paths(bundle_dir) spec = read_training_spec(bundle_dir) from kumo_core.io.json import from_json_dict plan = from_json_dict(spec.model_plan, ModelPlan) @@ -776,5 +1086,31 @@ def _load_bundle( trainer._training_query = spec.predictive_query trainer._training_table_generation_plan = from_json_dict( spec.training_table_generation, TrainingTableGenerationPlan) - trainer._training_result = TrainingResult(metrics, paths, bundle_dir) + trainer._training_result = TrainingResult(metrics, paths, bundle_dir, + bundle_id=bundle_id) + trainer._scoring_defaults = read_scoring_defaults(bundle_dir, paths) return trainer + + +def set_model_alias(connector: Connector, location: ModelRegistryLocation, + alias: str) -> None: + """Point a registry alias at ``location``'s exact version.""" + sink = connector.sink_backend + if not isinstance(sink, ModelAliasingSink): + raise NotImplementedError( + f'{type(connector).__name__} cannot set model aliases') + sink.set_model_alias(cast(SourceConnector, connector), location, alias) + + +def warn_if_other_bundle(metadata: Optional[Mapping[str, Any]], + trainer: Trainer) -> None: + recorded = (metadata or {}).get(BUNDLE_ID_KEY) + if recorded and trainer.bundle_id and recorded != trainer.bundle_id: + warnings.warn( + f"The registry records bundle {recorded} for this version, but " + f"the fetched bundle is {trainer.bundle_id}.", UserWarning, + stacklevel=3) + + +def registry_pointers(result: TrainingResult) -> Dict[str, str]: + return {BUNDLE_ID_KEY: result.bundle_id} if result.bundle_id else {} diff --git a/packages/kumo-engine/pyproject.toml b/packages/kumo-engine/pyproject.toml index 03dd919fe..eb6effd9c 100644 --- a/packages/kumo-engine/pyproject.toml +++ b/packages/kumo-engine/pyproject.toml @@ -62,6 +62,11 @@ text=[ gcs=[ "kumo-gcs>=0.1.0.dev1", ] +# nvidia-ml-py adds GPU system metrics. +mlflow=[ + "mlflow>=3.13,<4", + "nvidia-ml-py", +] # Reports this wheel's release version as the Kumo version. The provider is a # top-level module (see `[tool.setuptools] py-modules`) because importing diff --git a/packages/kumo-engine/setup.py b/packages/kumo-engine/setup.py index 125b39632..2bbf67013 100644 --- a/packages/kumo-engine/setup.py +++ b/packages/kumo-engine/setup.py @@ -2,10 +2,13 @@ # All rights reserved. # SPDX-License-Identifier: Apache-2.0 +import os import re +import subprocess from pathlib import Path from setuptools import setup +from setuptools.command.build_py import build_py kumo_api_release_ver = '3.0.0' api_pyproject = Path(__file__).parent.parent / 'kumo-api' / 'pyproject.toml' @@ -22,7 +25,34 @@ kumo_api = f'kumo-api=={api_ver}' + +def _build_commit(): + commit = os.environ.get('KUMO_ENGINE_BUILD_COMMIT') + if commit: + return commit + command = ['git', '-C', str(Path(__file__).parent), 'rev-parse', + '--short=12', 'HEAD'] + try: + return subprocess.run(command, capture_output=True, text=True, + check=True, timeout=10).stdout.strip() or None + except (OSError, subprocess.SubprocessError): + return None + + +class _StampBuildCommit(build_py): + def run(self): + super().run() + target = Path(self.build_lib) / 'kumo' / 'engine' / '_version.py' + if target.is_file(): + text = target.read_text(encoding='utf-8') + target.unlink() + target.write_text( + f'{text}BUILD_COMMIT = {_build_commit()!r}\n', + encoding='utf-8') + + setup( + cmdclass={'build_py': _StampBuildCommit}, install_requires=[ "pandas", "pyarrow>=8.0.0", diff --git a/packages/kumo-engine/test/_snowflake_experiment_fake.py b/packages/kumo-engine/test/_snowflake_experiment_fake.py new file mode 100644 index 000000000..27b1cc778 --- /dev/null +++ b/packages/kumo-engine/test/_snowflake_experiment_fake.py @@ -0,0 +1,192 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# 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).""" + +import os +import re +import sys +import types +from typing import Any, Dict, List, Optional, Tuple + +from kumo_core.sources.sql import resolve_snowflake_identifier as _identifier + +STATUSES = ('FAILED', 'FINISHED') + + +def _like(pattern: str, name: str) -> bool: + regex = ''.join('.*' if c == '%' else '.' if c == '_' else re.escape(c) + for c in pattern) + return re.fullmatch(regex, name, re.IGNORECASE | re.DOTALL) is not None + + +class FakeRun: + def __init__(self, tracking: 'FakeExperimentTracking', experiment: str, + name: str) -> None: + self._tracking = tracking + self.experiment_name = experiment + self.name = name + + def __enter__(self) -> 'FakeRun': + self._tracking.calls.append(('run.__enter__', self.name)) + self._tracking.lineage_stack.append((self.experiment_name, self.name)) + return self + + def __exit__(self, exc_type: Any, exc: Any, tb: Any) -> None: + self._tracking.calls.append(('run.__exit__', self.name)) + self._tracking.lineage_stack.pop() + if self._tracking._run is self: + self._tracking.end_run( + status='FAILED' if exc_type is not None else None) + + +class FakeExperimentTracking: + _instance: Optional['FakeExperimentTracking'] = None + + def __new__(cls, *args: Any, **kwargs: Any) -> 'FakeExperimentTracking': + if cls._instance is None: + cls._instance = super().__new__(cls) + return cls._instance + + def __init__(self, session: Any, *, database_name: Optional[str] = None, + schema_name: Optional[str] = None, + capture_source_info: bool = True, model_registry: Any = None, + warnings_as_errors: bool = False) -> None: + self.capture_source_info = capture_source_info + self.warnings_as_errors = warnings_as_errors + if hasattr(self, '_initialized'): + self.calls.append(('__init__', database_name, schema_name)) + return + self._session = session + self.database = _identifier(database_name + or session.get_current_database()) + self.schema = _identifier(schema_name or session.get_current_schema() + or 'PUBLIC') + self.calls: List[Tuple[Any, ...]] = [('__init__', database_name, + schema_name)] + self.fail_on: set = set() + self.experiment: Optional[str] = None + self._run: Optional[FakeRun] = None + self.runs: Dict[str, Dict[str, Any]] = {} + self.lineage_stack: List[Tuple[str, str]] = [] + self._sql_client = types.SimpleNamespace( + show_runs_in_experiment=self._show_runs) + self._initialized = True + + def _show_runs(self, *, experiment_name: str, + like: Optional[str] = None) -> List[Dict[str, Any]]: + self._record('show_runs', experiment_name, like) + return self._matching(experiment_name, like) + + def _matching(self, experiment: str, + like: Optional[str]) -> List[Dict[str, Any]]: + return [{'name': name} for name, run in self.runs.items() + if run['experiment'] == experiment + and (like is None or _like(like, name))] + + def _record(self, *call: Any) -> None: + self.calls.append(call) + if call[0] in self.fail_on: + raise RuntimeError(f'fake Snowflake error in {call[0]}') + + def set_experiment(self, experiment_name: str, + database_name: Optional[str] = None, + schema_name: Optional[str] = None) -> Any: + self._record('set_experiment', experiment_name, database_name, + schema_name) + if database_name is not None and schema_name is None: + schema_name = 'PUBLIC' + location = (_identifier(database_name or self.database), + _identifier(schema_name or self.schema)) + name = _identifier(experiment_name) + if (name, *location) != (self.experiment, self.database, self.schema): + self._run = None + self.experiment = name + self.database, self.schema = location + return types.SimpleNamespace(name=name) + + def start_run(self, run_name: Optional[str] = None) -> FakeRun: + self._record('start_run', run_name) + if self._run is not None: + raise RuntimeError('A run is already active. Please end the ' + 'current run before starting a new one.') + if self.experiment is None: + self.experiment = 'DEFAULT' + matches = self._matching(self.experiment, + run_name) if run_name else [] + if matches: + name = matches[0]['name'] + if self.runs[name]['status'] != 'RUNNING': + raise RuntimeError(f'Run {run_name} exists but cannot be ' + f'resumed as it is no longer running.') + else: + name = _identifier(run_name or f'RUN_{len(self.runs) + 1}') + self.runs[name] = {'experiment': self.experiment, + 'status': 'RUNNING', 'params': {}, + 'metrics': [], 'artifacts': {}} + self._run = FakeRun(self, self.experiment, name) + return self._run + + def end_run(self, run_name: Optional[str] = None, + status: Optional[str] = None) -> None: + self._record('end_run', run_name, status) + if self.experiment is None: + raise RuntimeError('No experiment set.') + if self._run is None: + raise RuntimeError('No run is active.') + if status is not None and status not in STATUSES: + raise ValueError(f"Unsupported status value: '{status}'.") + self.runs[self._run.name]['status'] = status or 'FINISHED' + self._run = None + + def _active(self) -> Dict[str, Any]: + if self._run is None: + self.start_run() + return self.runs[self._run.name] + + def log_params(self, params: Dict[str, Any]) -> None: + self._record('log_params', dict(params)) + self._active()['params'].update( + {key: str(value) for key, value in params.items()}) + + def log_metrics(self, metrics: Dict[str, float], step: int = 0) -> None: + self._record('log_metrics', dict(metrics), step) + self._active()['metrics'].extend( + (key, float(value), step) for key, value in metrics.items()) + + def log_artifact(self, local_path: str, + artifact_path: Optional[str] = None) -> None: + self._record('log_artifact', local_path, artifact_path) + with open(local_path, encoding='utf-8') as file: + self._active()['artifacts'][os.path.basename(local_path)] = ( + file.read()) + + +class FakeSession: + def __init__(self, database: Optional[str] = 'DB', + schema: Optional[str] = 'PUBLIC') -> None: + self.database, self.schema = database, schema + self.statements: List[str] = [] + + def sql(self, query: str) -> Any: + self.statements.append(query) + return types.SimpleNamespace(collect=lambda: []) + + def get_current_database(self) -> Optional[str]: + return self.database + + def get_current_schema(self) -> Optional[str]: + return self.schema + + +def install(monkeypatch: Any) -> type: + """Makes `import snowflake.ml.experiment` return a fresh fake.""" + FakeExperimentTracking._instance = None + module = types.ModuleType('snowflake.ml.experiment') + module.ExperimentTracking = FakeExperimentTracking + for name in ('snowflake', 'snowflake.ml'): + if name not in sys.modules: + monkeypatch.setitem(sys.modules, name, types.ModuleType(name)) + monkeypatch.setitem(sys.modules, 'snowflake.ml.experiment', module) + return FakeExperimentTracking diff --git a/packages/kumo-engine/test/conftest.py b/packages/kumo-engine/test/conftest.py index e7601b972..4b60addea 100644 --- a/packages/kumo-engine/test/conftest.py +++ b/packages/kumo-engine/test/conftest.py @@ -59,6 +59,34 @@ def identity(filesystem, path): return probes +@pytest.fixture +def graph(tmp_path): + rng = np.random.default_rng(0) + data = tmp_path / 'data' + data.mkdir() + pd.DataFrame({'ID': np.arange(50)}).to_parquet(data / 'users.parquet') + pd.DataFrame({ + 'EID': + np.arange(500), + 'USER_ID': + rng.integers(0, 50, 500), + 'TS': + pd.Timestamp('2024-01-01') + + pd.to_timedelta(rng.integers(0, 120, 500), unit='D'), + }).to_parquet(data / 'events.parquet') + connector = kumo.FileConnector(root_dir=str(data)) + return kumo.Graph( + tables={ + 'users': + kumo.Table.from_source_table(connector['users.parquet'], + primary_key='ID'), + 'events': + kumo.Table.from_source_table( + connector['events.parquet'], primary_key='EID', + time_column='TS').with_column('USER_ID', stype=Stype.ID), + }, edges=[Edge(src_table='events', fkey='USER_ID', dst_table='users')]) + + @pytest.fixture def snapshot(tmp_path): config = { diff --git a/packages/kumo-engine/test/connector/test_custom_source.py b/packages/kumo-engine/test/connector/test_custom_source.py index aaa2de4e1..87080104a 100644 --- a/packages/kumo-engine/test/connector/test_custom_source.py +++ b/packages/kumo-engine/test/connector/test_custom_source.py @@ -40,11 +40,12 @@ def prepare_ingest(connector, table_name, destination, sql): "paths": [str(first), str(second)], }, credentials=None, + version=None, ) connector = SimpleNamespace(source_backend=Backend()) source = SourceTable("curated", connector) - output = stage_source_table(source, tmp_path / "staged") + output, _ = stage_source_table(source, tmp_path / "staged") assert output == tmp_path / "staged" / "parts" assert pq.read_table(output)["ENTITY1"].to_pylist() == [1, 2] @@ -62,7 +63,7 @@ def test_stage_source_table_accepts_partitioned_file_connector_directory( pq.write_table(pa.table({"ENTITY1": [2]}), parts / "two.parquet") connector = FileConnector(root_dir=str(root)) - output = stage_source_table(connector["parts"], tmp_path / "staged") + output, _ = stage_source_table(connector["parts"], tmp_path / "staged") assert output == tmp_path / "staged" / "parts" assert sorted(pq.read_table(output)["ENTITY1"].to_pylist()) == [1, 2] @@ -96,10 +97,10 @@ def refuse_link(*args, **kwargs): source={ "type": "parquet", "paths": [str(p) for p in paths] - }, credentials=None) + }, credentials=None, version=None) source = SourceTable("curated", Mock(source_backend=backend)) - output = stage_source_table(source, tmp_path / "staged") + output, _ = stage_source_table(source, tmp_path / "staged") retained = sorted(output.glob("*.parquet")) assert len(retained) == 2 @@ -124,7 +125,7 @@ def test_stage_source_table_casts_to_what_the_table_reader_reads(tmp_path): "TARGET1": pa.array(["a", "b"]).dictionary_encode(), }), root / "curated.parquet") - output = stage_source_table( + output, _ = stage_source_table( FileConnector(root_dir=str(root))["curated"], tmp_path / "staged") assert pq.read_table(output).schema == pa.schema([("ENTITY1", pa.int64()), @@ -137,7 +138,7 @@ def test_stage_source_table_retains_single_file_in_parts_directory(tmp_path): pq.write_table(pa.table({"ENTITY1": [1, 2]}), root / "table.parquet") connector = FileConnector(root_dir=str(root)) - output = stage_source_table(connector["table.parquet"], tmp_path / "staged") + output, _ = stage_source_table(connector["table.parquet"], tmp_path / "staged") assert output == tmp_path / "staged" / "parts" parts = list(output.glob("*.parquet")) @@ -171,8 +172,8 @@ def test_stage_source_table_retains_reencoded_float64_weight_parts( pq.write_table(table, root / "weighted.parquet") connector = FileConnector(root_dir=str(root)) - output = stage_source_table(connector["weighted.parquet"], - tmp_path / "staged") + output, _ = stage_source_table(connector["weighted.parquet"], + tmp_path / "staged") assert output == tmp_path / "staged" / "parts" assert len(list((tmp_path / "staged" / "source").rglob("*.parquet"))) > 1 diff --git a/packages/kumo-engine/test/test_bundle_provenance.py b/packages/kumo-engine/test/test_bundle_provenance.py new file mode 100644 index 000000000..8fd0ba3b4 --- /dev/null +++ b/packages/kumo-engine/test/test_bundle_provenance.py @@ -0,0 +1,144 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import json +import re +import shutil +from datetime import datetime, timezone +from types import SimpleNamespace + +import pandas as pd +import pytest +from kumo_core.artifacts import add_bundle_file +from kumoapi.model_plan import ModelPlan, TrainingTableGenerationPlan + +import kumo.engine as kumo +from kumo.engine._manifest import ( + BundleProvenance, + BundleTrainingTable, + SourceTableEntry, + TrainingTableParents, + read_manifest, + write_manifest, +) +from kumo.engine.trainer._spec import MANIFEST_FILE +from kumo.engine.trainer.trainer import TrainingResult + + +def test_fit_records_provenance(fitted): + trainer, _, snapshot = fitted + result = trainer._training_result + provenance = result.provenance + parents = provenance.parents + assert parents.training_table and parents.warm_start is None + assert parents.extra == {} + assert provenance.training_table == BundleTrainingTable( + source='pquery', snapshot=snapshot.id) + assert provenance.source_tables == snapshot.manifest.parents.source_tables + assert provenance.fit_mode == 'refit_trainval' + assert provenance.tracking is None + assert re.fullmatch('[0-9a-f]{64}', trainer.bundle_id) + assert trainer.bundle_id == result.bundle_id == provenance.id + assert kumo.Trainer.load(result.bundle_dir).bundle_id == provenance.id + + +def test_provenance_round_trips_unchanged(fitted, tmp_path): + trainer, _, _ = fitted + newer = json.loads( + (trainer._training_result.bundle_dir / 'provenance.json').read_text()) + newer['future'] = True + newer['fit']['gpu_seconds'] = 1.5 + newer['training_table'].update(connector='S3', location='s3://b/tt', + version='1', version_kind='exact') + source = tmp_path / 'newer.json' + source.write_text(json.dumps(newer, indent=2, sort_keys=True) + '\n') + + provenance = read_manifest(source, BundleProvenance) + write_manifest(tmp_path / 'out.json', provenance) + + assert (tmp_path / 'out.json').read_bytes() == source.read_bytes() + assert provenance.extra == {'future': True} + assert provenance.fit.extra == {'gpu_seconds': 1.5} + assert provenance.training_table.source_table.location == 's3://b/tt' + + +def test_bundle_without_provenance_loads_and_predicts_identically( + fitted, tmp_path): + trainer, prediction_table, snapshot = fitted + result = trainer._training_result + loaded = kumo.Trainer.load(result.bundle_dir) + assert (loaded.bundle_id, loaded.provenance) == (result.bundle_id, + result.provenance) + + old = tmp_path / 'old' + shutil.copytree(result.bundle_dir, old) + manifest = json.loads((old / MANIFEST_FILE).read_text()) + manifest['artifacts'] = [ + entry for entry in manifest['artifacts'] + if entry['role'] != 'provenance' + ] + (old / MANIFEST_FILE).write_text(json.dumps(manifest)) + (old / 'provenance.json').unlink() + restored = kumo.Trainer.load(old) + assert restored.bundle_id == result.bundle_id + assert restored.provenance is None + + def predict(model): + return model.predict( + prediction_table, snapshot=snapshot, + binary_classification_threshold=.5).frame.sort_values( + 'ENTITY').reset_index(drop=True) + + pd.testing.assert_frame_equal(predict(restored), predict(trainer)) + + +def _finish_stub_fit(bundle, trainer, manifest=None, source_mode=None): + (bundle / 'model').mkdir(parents=True) + (bundle / 'model' / 'config.json').write_text('{}') + (bundle / MANIFEST_FILE).write_text(json.dumps({'artifacts': []})) + add_bundle_file(bundle, 'trainer_config', 'model/config.json') + table = SimpleNamespace(tt_id='tt-1', manifest=manifest, + source_mode=source_mode) + return trainer._finish_fit( + TrainingResult({}, {}, bundle), predictive_query='PREDICT x', + snapshot_id='snap', export_config=None, + training_table_generation=TrainingTableGenerationPlan(), + provenance=dict(train_table=table, warm_start_bundle_id=None, + started_at=datetime.now(timezone.utc), fit_mode=None)) + + +def test_provenance_records_custom_or_missing_table_source(tmp_path): + source_table = SourceTableEntry(connector='S3', + location='s3://bucket/exports/tt', + version='17', version_kind='timestamp') + manifest = SimpleNamespace(parents=TrainingTableParents( + snapshot='snap', source_table=source_table)) + custom = _finish_stub_fit(tmp_path / 'custom', kumo.Trainer(ModelPlan()), + manifest, 'custom_tt') + assert custom.provenance.training_table == BundleTrainingTable( + source='custom_tt', snapshot='snap', source_table=source_table) + assert json.loads((tmp_path / 'custom' / + 'provenance.json').read_text())['training_table'] == { + 'source': 'custom_tt', + 'snapshot': 'snap', + **source_table.to_dict() + } + older = _finish_stub_fit(tmp_path / 'older', kumo.Trainer(ModelPlan())) + assert older.provenance.training_table == BundleTrainingTable( + source=None, snapshot=None) + + +def test_provenance_failure_warns_and_still_publishes(tmp_path, monkeypatch): + from kumo.engine.trainer import _provenance + + def fail(*args): + raise RuntimeError('unreadable manifest') + + monkeypatch.setattr(_provenance, '_training_table', fail) + trainer = kumo.Trainer(ModelPlan()) + with pytest.warns(UserWarning, match='unreadable manifest'): + result = _finish_stub_fit(tmp_path / 'bundle', trainer) + assert result.provenance is None + assert re.fullmatch('[0-9a-f]{64}', result.bundle_id) + assert trainer._training_result is result diff --git a/packages/kumo-engine/test/test_connector_for_root.py b/packages/kumo-engine/test/test_connector_for_root.py new file mode 100644 index 000000000..93073db79 --- /dev/null +++ b/packages/kumo-engine/test/test_connector_for_root.py @@ -0,0 +1,25 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import pytest + +from kumo.engine.connector import ( + FileConnector, + GCSConnector, + S3Connector, + connector_for_root, +) + + +def test_connector_for_root(tmp_path): + s3 = connector_for_root('s3://my-bucket/runs/1') + assert isinstance(s3, S3Connector) + assert (s3.bucket, s3.prefix) == ('my-bucket', 'runs/1') + assert isinstance(connector_for_root('gs://my-bucket/runs'), GCSConnector) + for uri in (str(tmp_path), tmp_path.as_uri()): + local = connector_for_root(uri) + assert isinstance(local, FileConnector) + assert local.config.root_dir == str(tmp_path) + with pytest.raises(ValueError, match='not s3://'): + connector_for_root('abfss://c@a.dfs.core.windows.net/x') diff --git a/packages/kumo-engine/test/test_custom_sql.py b/packages/kumo-engine/test/test_custom_sql.py index b2f80c231..498b045ce 100644 --- a/packages/kumo-engine/test/test_custom_sql.py +++ b/packages/kumo-engine/test/test_custom_sql.py @@ -131,7 +131,7 @@ def spy(graph_config, spec): def test_custom_training_rows_stage_the_query_result(warehouse, tmp_path): - staged = stage_source_table(warehouse.sql(SQL), tmp_path / 'staged') + staged, _ = stage_source_table(warehouse.sql(SQL), tmp_path / 'staged') assert pq.read_table(staged).to_pydict() == { 'user_id': [1, 3], 'age': [30, 50] diff --git a/packages/kumo-engine/test/test_fit_events.py b/packages/kumo-engine/test/test_fit_events.py new file mode 100644 index 000000000..849a53760 --- /dev/null +++ b/packages/kumo-engine/test/test_fit_events.py @@ -0,0 +1,152 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import json + +import numpy as np +import pandas as pd +import pytest +from kumoapi.model_plan import ModelPlan + +import kumo.engine as kumo +from kumo.engine._scratch import get_scratch_dir + + +@pytest.fixture +def table_and_snapshot(tmp_path): + previous_scratch = get_scratch_dir() + kumo.configure(scratch_dir=tmp_path / 'scratch') + try: + rng = np.random.default_rng(0) + (tmp_path / 'data').mkdir() + pd.DataFrame({ + 'ID': np.arange(400), + 'Y': rng.random(400), + 'F1': rng.normal(size=400), + }).to_parquet(tmp_path / 'data' / 'ents.parquet') + connector = kumo.FileConnector(root_dir=str(tmp_path / 'data')) + ents = kumo.Table.from_source_table(connector['ents.parquet'], + primary_key='ID').infer_metadata() + graph = kumo.Graph(tables={'ents': ents}, edges=[]) + snapshot = graph.snapshot() + query = kumo.PredictiveQuery(graph=graph, + query='PREDICT ents.Y FOR EACH ents.ID') + yield query.generate_training_table(snapshot), snapshot + finally: + kumo.configure(scratch_dir=previous_scratch) + + +def test_fit_reports_progress_events(table_and_snapshot, monkeypatch, + caplog): + from kumo_core.automl.optuna_searcher import TPESearcher + + update = TPESearcher.update + + def fail_second_trial(self, *args, **kwargs): + if self._num_updates: + raise RuntimeError('searcher down') + self._num_updates = 1 + return update(self, *args, **kwargs) + + monkeypatch.setattr(TPESearcher, '_num_updates', 0, raising=False) + monkeypatch.setattr(TPESearcher, 'update', fail_second_trial) + table, snapshot = table_and_snapshot + plan = ModelPlan() + plan.training_job.num_experiments = 2 + plan.training_job.disable_explain = True + plan.optimization.max_epochs = 2 + plan.optimization.max_steps_per_epoch = 2 + events = [] + broken_calls = [] + + def broken(event): + broken_calls.append(event) + raise RuntimeError('callback bug') + + result = kumo.Trainer(plan).fit(table, snapshot=snapshot, + callbacks=[broken, events.append]) + + assert len(broken_calls) == 1 + assert 'callback bug' in caplog.text + names = [event.name for event in events] + assert names == [ + 'fit_started', + 'trial_started', 'epoch_ended', 'epoch_ended', 'trial_ended', + 'trial_started', 'epoch_ended', 'epoch_ended', 'trial_ended', + 'fit_ended', + ] + assert len({event.run_id for event in events}) == 1 + assert set(events[0].params) == {'snapshot_id', 'query_sha256'} + assert [(e.trial, e.epoch) for e in events if e.name == 'epoch_ended' + ] == [(1, 1), (1, 2), (2, 1), (2, 2)] + assert any(name.startswith('val_') for name in events[2].metrics) + assert events[1].params['attempt'] == 1 + assert events[4].params == {'status': 'succeeded', 'attempt': 1} + assert {'val_mae', 'seconds'} <= set(events[4].metrics) + assert events[8].params == { + 'status': 'failed', + 'attempt': 1, + 'error': 'RuntimeError: searcher down' + } + assert events[-1].metrics == {**result.metrics, + 'seconds': events[-1].metrics['seconds']} + assert events[-1].params == {'status': 'succeeded'} + assert [json.loads(json.dumps(event.to_dict()))['name'] + for event in events] == names + + +def test_failed_fit_reports_fit_ended(table_and_snapshot, monkeypatch): + import kumo.engine.trainer.trainer as trainer_module + + def fail(*args, **kwargs): + raise ValueError('stop') + + monkeypatch.setattr(trainer_module, 'build_generated_training_task_dataset', + fail) + table, snapshot = table_and_snapshot + events = [] + + with pytest.raises(ValueError, match='stop'): + kumo.Trainer(ModelPlan()).fit(table, snapshot=snapshot, + callbacks=[events.append]) + + assert [event.name for event in events] == ['fit_started', 'fit_ended'] + assert events[-1].params == { + 'status': 'failed', + 'error': 'ValueError: stop' + } + + +def test_event_metrics_keep_finite_numbers(): + from kumo.engine._events import EventEmitter + + received = [] + EventEmitter('run', [received.append]).emit( + 'epoch_ended', metrics={'val_loss': 'n/a', 'val_auc': float('nan'), + 'val_mae': 0.5}) + + assert received[0].metrics == {'val_mae': 0.5} + + +def test_epoch_event_without_recorded_metrics_does_not_raise(tmp_path): + from datetime import datetime, timezone + from types import SimpleNamespace + + from kumo_core.callbacks.train_progress_logger import ( + AutoTrainingProgressLogger, + ) + + received = [] + logger = AutoTrainingProgressLogger( + str(tmp_path), 1, on_event=lambda *a, **k: received.append(a)) + progress = SimpleNamespace(num_epochs_completed=1, metrics=[], + estimated_train_time=0.0) + logger.set_current_trainer( + SimpleNamespace(training_progress_logger=SimpleNamespace( + progress=progress)), trial=1) + logger.progress.start_time = datetime.now(timezone.utc) + + logger.on_train_epoch_end(None, None) + + assert received == [] diff --git a/packages/kumo-engine/test/test_holdout_df.py b/packages/kumo-engine/test/test_holdout_df.py index 8c03af648..70bfe82c6 100644 --- a/packages/kumo-engine/test/test_holdout_df.py +++ b/packages/kumo-engine/test/test_holdout_df.py @@ -2,6 +2,8 @@ # All rights reserved. # SPDX-License-Identifier: Apache-2.0 +import re + import numpy as np import pandas as pd import pytest @@ -65,13 +67,16 @@ def graph(tmp_path_factory): QUERY = 'PREDICT COUNT(events.*, 0, 7, days) > 0 FOR EACH ents.ID' -def _fit(graph, query): +def _fit(graph, query, refit_trainval=True, refit_full=False, explain=False): graph, snapshot = graph pquery = kumo.PredictiveQuery(graph=graph, query=query) plan = ModelPlan() - plan.training_job.disable_explain = True + plan.training_job.disable_explain = not explain + plan.training_job.refit_trainval = refit_trainval + plan.training_job.refit_full = refit_full plan.optimization.max_epochs = 1 - plan.optimization.max_steps_per_epoch = 2 + # refit_full splits each epoch's batches across 3 stages. + plan.optimization.max_steps_per_epoch = 3 if refit_full else 2 trainer = kumo.Trainer(model_plan=plan) return trainer.fit(pquery.generate_training_table(snapshot), snapshot=snapshot) @@ -90,6 +95,23 @@ def test_holdout_df_after_fit(graph): assert not (result.bundle_dir / 'test_df.parquet').exists() +@pytest.mark.parametrize('refit_trainval,refit_full,explain,suffix', [ + (False, False, False, ''), + (True, False, True, '_refit_trainval'), + (True, True, False, None), +]) +def test_holdout_is_the_published_models(graph, refit_trainval, refit_full, + explain, suffix): + result = _fit(graph, QUERY, refit_trainval, refit_full, explain) + + if suffix is None: + assert result.holdout_path is None + else: + assert re.fullmatch(rf'version_\d+{suffix}', + result.holdout_path.parent.name) + assert len(result.holdout_df()) > 0 + + def test_no_holdout(graph): result = _fit( graph, 'PREDICT LIST_DISTINCT(events.ITEM_ID, 0, 7, days) ' diff --git a/packages/kumo-engine/test/test_job.py b/packages/kumo-engine/test/test_job.py new file mode 100644 index 000000000..671f3baef --- /dev/null +++ b/packages/kumo-engine/test/test_job.py @@ -0,0 +1,132 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import json +import os +import signal +import warnings + +import pytest +from kumo_core.sources.errors import SourceNotFoundError, TransientSourceError + +import kumo.engine as kumo +from kumo.engine._manifest import SchedulerContext, stamp + + +@pytest.fixture(autouse=True) +def no_scheduler(monkeypatch): + for name in ('KUMO_SCHEDULER', 'KUMO_JOB_ID', 'KUMO_JOB_RUN_ID'): + monkeypatch.delenv(name, raising=False) + + +def _summary(capsys): + out = capsys.readouterr().out.splitlines() + assert len(out) == 1 + return json.loads(out[0]) + + +def test_summary_records_ids_warnings_and_scheduler(capsys): + with kumo.job(job_id='job-7', job_run_id='run-7') as run: + scheduler = stamp()['scheduler'] + run.record(bundle_id='b-1', snapshot_id='s-1') + run.record(run_id='r-1', outputs={'predictions': 's3://out/p'}) + warnings.warn('threshold not set', UserWarning) + warnings.warn('use the new API', FutureWarning) + for i in range(30): + warnings.warn(f'{i}' * 400, RuntimeWarning) + warnings.warn(f'{i}' * 400, RuntimeWarning) + + assert scheduler == SchedulerContext(platform='other', job_id='job-7', + job_run_id='run-7') + summary = _summary(capsys) + assert summary == { + 'event': 'kumo.engine.job', + 'status': 'succeeded', + 'exit_code': 0, + 'scheduler': scheduler.to_dict(), + 'bundle_id': 'b-1', + 'snapshot_id': 's-1', + 'run_id': 'r-1', + 'outputs': { + 'predictions': 's3://out/p' + }, + 'warnings': summary['warnings'], + } + assert summary['warnings'][0] == 'threshold not set' + assert 'use the new API' not in summary['warnings'] + assert len(summary['warnings']) == 20 + assert len(set(summary['warnings'])) == 20 + assert max(map(len, summary['warnings'])) == 300 + assert 'KUMO_JOB_ID' not in os.environ + + +def test_warning_survives_a_patcher_that_wraps_kumos_showwarning(capsys): + original = {} + + def patched(*args, **kwargs): + original['show'](*args, **kwargs) + + with pytest.warns(UserWarning, match='threshold not set'): + original['show'] = warnings.showwarning + warnings.showwarning = patched + with kumo.job(): + original['show'] = warnings.showwarning + warnings.showwarning = patched + warnings.warn('threshold not set', UserWarning) + + assert _summary(capsys)['warnings'] == ['threshold not set'] + + +def _wrapped(cause): + try: + raise RuntimeError('wrapped') from cause + except RuntimeError as error: + return error + + +@pytest.mark.parametrize('error, kind', [ + (_wrapped(TransientSourceError('busy')), 'transient'), + (ConnectionResetError(), 'transient'), + (SourceNotFoundError('missing'), 'permanent'), + (ValueError('x' * 2000), 'unknown'), +]) +def test_failure_exits_non_zero_with_a_retry_hint(capsys, error, kind): + with pytest.raises(SystemExit) as exited: + with kumo.job(job_id='job-7') as run: + run.record(snapshot_id='s-1') + raise error + assert exited.value.code == 1 + summary = _summary(capsys) + assert (summary['status'], summary['exit_code']) == ('failed', 1) + assert summary['snapshot_id'] == 's-1' + assert summary['scheduler']['job_id'] == 'job-7' + assert summary['error']['type'] == type(error).__name__ + assert summary['error']['kind'] == kind + assert len(summary['error']['message']) <= 500 + assert 'KUMO_JOB_ID' not in os.environ + + +@pytest.mark.parametrize('error, code', [ + (KeyboardInterrupt(), 130), + (SystemExit(3), 3), + (SystemExit(0), 0), +]) +def test_interrupt_and_exit_still_print_the_summary(capsys, error, code): + with pytest.raises(type(error)): + with kumo.job(job_id='job-7'): + raise error + summary = _summary(capsys) + assert summary['exit_code'] == code + assert summary['status'] == ('succeeded' if code == 0 else 'failed') + assert 'KUMO_JOB_ID' not in os.environ + + +def test_sigterm_prints_the_summary_and_exits_143(capsys): + previous = signal.getsignal(signal.SIGTERM) + with pytest.raises(SystemExit) as exited: + with kumo.job(): + os.kill(os.getpid(), signal.SIGTERM) + assert exited.value.code == 143 + assert _summary(capsys)['status'] == 'failed' + assert signal.getsignal(signal.SIGTERM) is previous 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..2ee6d6860 --- /dev/null +++ b/packages/kumo-engine/test/test_job_templates.py @@ -0,0 +1,151 @@ +# 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 + +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 test_databricks_tasks_run_the_job_scripts(): + yaml = pytest.importorskip('yaml') + bundle = yaml.safe_load( + (_examples() / 'databricks' / 'databricks.yml').read_text()) + root = '${workspace.file_path}' + tasks = [ + task['spark_python_task'] for job in bundle['resources']['jobs'].values() + for task in job['tasks'] + ] + assert tasks == [{ + 'python_file': f'{root}/{script}', + 'parameters': [ + '--job-id', '{{job.id}}', '--job-run-id', '{{job.run_id}}' + ] + } for script in ('fit_job.py', 'score_job.py')] + + +def test_sagemaker_and_snowflake_name_the_job_scripts(): + pipeline = (_examples() / 'sagemaker' / 'pipeline.py').read_text() + register = (_examples() / 'snowflake' / 'register_jobs.py').read_text() + for script in ('fit_job.py', 'score_job.py'): + assert (_examples() / script).is_file() + assert f"/{script}'" in pipeline + assert f"'entrypoint': '{script}'" in register + + +def _data(root: Path) -> None: + root.mkdir() + rng = np.random.default_rng(0) + pd.DataFrame({'USER_ID': np.arange(50)}).to_parquet(root / 'users.parquet') + pd.DataFrame({'ITEM_ID': np.arange(20)}).to_parquet(root / '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 / 'orders.parquet') + + +def _summary(capsys): + lines = capsys.readouterr().out.splitlines() + return json.loads(lines[-1]) + + +def test_fit_then_score_locally(tmp_path, monkeypatch, capsys): + _data(tmp_path / 'data') + 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) + for name, value in (('KUMO_SOURCE', 'file'), + ('KUMO_DATA_ROOT', str(tmp_path / 'data')), + ('KUMO_ARTIFACT_ROOT', str(tmp_path / 'artifacts')), + ('KUMO_SCRATCH_DIR', str(tmp_path / 'scratch')), + ('KUMO_QUERY', 'PREDICT COUNT(orders.*, 0, 7, days) > 0 ' + 'FOR EACH users.USER_ID')): + monkeypatch.setenv(name, value) + for name in ('KUMO_REGISTRY', 'KUMO_PREDICTIONS_TABLE', 'KUMO_SNAPSHOT', + 'KUMO_MODEL_PATH', + 'KUMO_SCHEDULER', 'KUMO_JOB_ID', 'KUMO_JOB_RUN_ID'): + monkeypatch.delenv(name, raising=False) + monkeypatch.syspath_prepend(str(_examples())) + monkeypatch.delitem(sys.modules, 'jobs', raising=False) + + _module('fit_job').main(['--job-id', 'local', '--job-run-id', 'run-1']) + fit = _summary(capsys) + assert fit['status'] == 'succeeded', fit + assert fit['scheduler']['job_run_id'] == 'run-1' + assert fit['outputs']['model']['uri'] + assert fit['outputs']['model']['version'] is None + + monkeypatch.setenv('KUMO_MODEL_PATH', 'runs/run-1/model') + _module('score_job').main(['--job-id', 'local', '--job-run-id', 'run-2']) + score = _summary(capsys) + assert score['status'] == 'succeeded', score + assert (score['bundle_id'], score['snapshot_id']) != (None, None) + assert score['bundle_id'] == fit['bundle_id'] + assert score['outputs']['predictions'] == str( + tmp_path / 'artifacts' / 'runs' / 'run-2' / 'predictions') + '/' + predictions = pd.read_parquet(score['outputs']['predictions']) + assert len(predictions) == 50 + assert 'ANCHOR_TIMESTAMP' in predictions + + +def test_sagemaker_pipeline_needs_an_instance_type(capsys): + pipeline = _module('sagemaker/pipeline') + args = [ + '--role', 'r', '--image', 'i', '--code', 'c', '--data-root', 'd', + '--artifact-root', 'a', '--tracking-arn', 't' + ] + with pytest.raises(SystemExit): + pipeline.main(args) + assert '--instance-type' in capsys.readouterr().err + + +@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']) diff --git a/packages/kumo-engine/test/test_mlflow_model.py b/packages/kumo-engine/test/test_mlflow_model.py new file mode 100644 index 000000000..bedad396a --- /dev/null +++ b/packages/kumo-engine/test/test_mlflow_model.py @@ -0,0 +1,301 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import json +import warnings +from types import SimpleNamespace + +import pandas as pd +import pytest +from kumoapi.model_plan import ModelPlan +from kumoapi.plugin import PublishedModel +from test_mlflow_tracking import ( # noqa: F401 + _runs, + table_and_snapshot, + tracking_uri, +) + +import kumo.engine as kumo +from kumo.engine._scratch import get_scratch_dir + +mlflow = pytest.importorskip('mlflow') + + +def _plan(): + plan = ModelPlan() + plan.training_job.num_experiments = 1 + plan.training_job.disable_explain = True + plan.optimization.max_epochs = 2 + plan.optimization.max_steps_per_epoch = 2 + return plan + + +def _sorted(frame): + return frame.sort_values('ENTITY').reset_index(drop=True) + + +def test_fit_logs_registers_and_serves_the_model(table_and_snapshot, + tmp_path, caplog, + monkeypatch): + from kumo.engine.integrations.mlflow import MlflowTracking, load_model + + monkeypatch.setenv('KUMO_SCHEDULER', 'airflow') + monkeypatch.setenv('KUMO_JOB_RUN_ID', 'run-7') + table, snapshot = table_and_snapshot + plan = _plan() + tracker = MlflowTracking(log_system_metrics=False, + registered_model_name='m', aliases=('champion', )) + result = kumo.Trainer(plan).fit(table, snapshot=snapshot, tracking=tracker) + + assert 'MLflow tracking failed' not in caplog.text + parent, = [ + run for run in _runs(mlflow.MlflowClient()) + if 'mlflow.parentRunId' not in run.data.tags + ] + model, = mlflow.search_logged_models(output_format='list') + assert model.source_run_id == parent.info.run_id + metrics = {metric.key: metric.value for metric in model.metrics} + assert metrics == {**result.metrics, 'seconds': metrics['seconds']} + version = mlflow.MlflowClient().get_model_version_by_alias('m', 'champion') + assert version.source == model.model_uri + assert {k: v for k, v in parent.data.tags.items() + if k.startswith('kumo.')} == { + 'kumo.kind': 'fit', + 'kumo.run_id': parent.data.tags['kumo.run_id'], + 'kumo.bundle_id': result.bundle_id, + 'kumo.scheduler': 'airflow', + 'kumo.job_run_id': 'run-7', + } + assert tracker.registered == PublishedModel( + name='m', version=str(version.version), + uri=f'models:/m/{version.version}') + + info = mlflow.models.get_model_info('models:/m@champion') + assert info.metadata == {'kumo_bundle_id': result.bundle_id} + provenance = json.loads(mlflow.artifacts.load_text( + f'{model.artifact_location}/artifacts/provenance.json')) + assert provenance == result.provenance.to_dict() + assert provenance['tracking'] == { + 'run_id': parent.data.tags['kumo.run_id'] + } + assert [(c['name'], c['type']) for c in info.signature.inputs.to_dict() + ] == [('ENTITY', 'long'), ('TIMESTAMP', 'datetime')] + assert [(c['name'], c['type']) for c in info.signature.outputs.to_dict() + ] == [('ENTITY', 'long'), ('TARGET_PRED', 'double')] + + location = snapshot.export( + kumo.FileConnector(root_dir=str(tmp_path / 'exports')), 'snapshots') + served = mlflow.pyfunc.load_model('models:/m@champion').predict( + pd.DataFrame({'ENTITY': range(0, 400, 7)}), + params={'snapshot': str(tmp_path / 'exports' / location)}) + assert list(served) == ['ENTITY', 'TARGET_PRED'] + assert len(served) == len(range(0, 400, 7)) + assert load_model('models:/m@champion').bundle_id == result.bundle_id + + original = mlflow.models.Model.load + + def other_bundle(path): + model = original(path) + model.metadata['kumo_bundle_id'] = 'other' + return model + + with pytest.MonkeyPatch.context() as patch, pytest.warns( + UserWarning, match='records bundle other'): + patch.setattr(mlflow.models.Model, 'load', other_bundle) + load_model('models:/m@champion') + + +def test_logged_model_scores_the_query_population_and_timed_rows( + fitted, tracking_uri, tmp_path, monkeypatch): + from kumo.engine.integrations.mlflow import log_model + + trainer, table, snapshot = fitted + location = snapshot.export( + kumo.FileConnector(root_dir=str(tmp_path / 'exports')), 'snapshots') + + info = log_model(trainer, snapshot=snapshot, + serving_snapshot=str(tmp_path / 'exports' / location)) + + assert mlflow.get_run(info.run_id).info.status == 'FINISHED' + requirements = mlflow.pyfunc.get_model_dependencies(info.model_uri) + assert f'kumo-engine[gcs,s3]=={kumo.__version__}' in open( + requirements).read().splitlines() + assert {metric.key + for metric in mlflow.get_logged_model(info.model_id).metrics + } == set(trainer._training_result.metrics) + assert [c['name'] for c in info.signature.outputs.to_dict() + ] == ['ENTITY', 'TARGET_PRED', 'False_PROB', 'True_PROB'] + pyfunc = mlflow.pyfunc.load_model(info.model_uri) + params = { + 'snapshot': str(tmp_path / 'exports' / location), + 'prediction_time': str(pd.Timestamp(table.timestamps[0])), + 'binary_classification_threshold': 0.5, + } + expected = _sorted( + trainer.predict(table, snapshot=snapshot, + binary_classification_threshold=0.5).frame) + empty = pd.DataFrame({'ENTITY': pd.Series([], dtype='int64')}) + tables = get_scratch_dir() / 'prediction_tables' + before = set(tables.glob('*')) + predict = kumo.Trainer.predict + + def predict_beside_another_call(self, *args, **kwargs): + (tables / 'other-call').mkdir(parents=True, exist_ok=True) + return predict(self, *args, **kwargs) + + monkeypatch.setattr(kumo.Trainer, 'predict', predict_beside_another_call) + + assert pyfunc.predict(empty, params=params).empty + served = pyfunc.unwrap_python_model()._snapshots + assert list(served) == [params['snapshot']] + pd.testing.assert_frame_equal( + _sorted( + pyfunc.predict(empty, params={ + **params, 'snapshot': None, 'population': True + })), expected) + for rows, population in ((empty, True), + (pd.DataFrame({'ENTITY': range(50)}), False)): + pd.testing.assert_frame_equal( + _sorted( + pyfunc.predict(rows, + params={**params, 'population': population})), + expected) + scratch = pyfunc.unwrap_python_model()._trainer.scratch_dir + assert list(scratch.iterdir()) == [] + assert set(tables.glob('*')) == before | {tables / 'other-call'} + + +@pytest.mark.parametrize('options', [ + dict(registered_model_name='m', aliases=('SageMakerLifecycle-prod', )), + dict(log_model=True, aliases=('champion', )), +]) +def test_invalid_registration_is_refused(tracking_uri, options): + from kumo.engine.integrations.mlflow import MlflowTracking + + with pytest.raises(ValueError): + MlflowTracking(**options) + + +def test_fit_waits_for_model_logging_and_fails_the_run_if_it_fails( + monkeypatch): + import time + + from kumo.engine import _tracking + from kumo.engine._tracking import FitContext, TrackingSession + + monkeypatch.setattr(_tracking, 'FINISH_SECONDS', 0.01) + + class Tracker: + def __init__(self, fail): + self.fail, self.calls = fail, [] + + def start(self, record): + self.calls.append('start') + + def on_trial(self, record, event): + pass + + def log_model(self, result, snapshot): + time.sleep(0.2) + self.calls.append('log_model') + if self.fail: + raise RuntimeError('upload failed') + + def finish(self, record, error): + self.calls.append(('finish', error)) + + result = SimpleNamespace(metrics={}, bundle_dir='missing', bundle_id='b') + for fail in (False, True): + tracker = Tracker(fail) + session = TrackingSession(tracker) + session.start(lambda: FitContext('r', 'q', {})) + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter('always') + error = session.log_model(result, None) + session.finish(result, error) + session._worker.join(5) + assert tracker.calls[:2] == ['start', 'log_model'] + assert tracker.calls[2] == ('finish', error) + assert (error is not None) is fail + assert any('upload failed' in str(w.message) for w in caught) is fail + + +@pytest.mark.parametrize('stall', ['slow upload', 'dead thread']) +def test_model_logging_wait_is_bounded(monkeypatch, stall): + import threading + + from kumo.engine import _tracking + from kumo.engine._tracking import FitContext, TrackingSession + + monkeypatch.setattr(_tracking, 'MODEL_SECONDS', 0.5) + release = threading.Event() + + class Tracker: + calls = [] + + def start(self, record): + pass + + def on_trial(self, record, event): + pass + + def log_model(self, result, snapshot): + release.wait(30) + + def finish(self, record, error): + self.calls.append((record, error)) + + tracker = Tracker() + session = TrackingSession(tracker) + session.start(lambda: FitContext('r', 'q', {})) + if stall == 'dead thread': + session._submit(lambda: (_ for _ in ()).throw(SystemExit())) + session._worker.join(5) + with pytest.warns(UserWarning, match='Could not log the model'): + error = session.log_model(SimpleNamespace(), None) + release.set() + session.finish(None, error) + + assert type(error) is (TimeoutError + if stall == 'slow upload' else RuntimeError) + assert tracker.calls == [(None, None)] + + +def test_model_upload_gets_a_longer_http_timeout_and_a_late_one_is_dropped( + tracking_uri, monkeypatch, caplog): + import os + + from test_tracking_contract import CONTEXT + + from kumo.engine._tracking import TrackingSession + from kumo.engine.integrations import mlflow as integration + + for name in integration._HTTP_LIMITS: + monkeypatch.delenv(name, raising=False) + tracker = integration.MlflowTracking(log_system_metrics=False, + registered_model_name='m') + seen = [] + + def log_model(mlflow, result, snapshot, **kwargs): + seen.append(os.environ.get('MLFLOW_HTTP_REQUEST_TIMEOUT')) + if len(seen) == 2: + tracker.finish(None, None) + return SimpleNamespace(model_id=None, model_uri='runs:/x/model') + + monkeypatch.setattr(integration, '_log_model', log_model) + monkeypatch.setattr(integration, '_register', lambda *args: 'registered') + for _ in range(2): + session = TrackingSession(tracker) + session.start(lambda: CONTEXT) + assert session.log_model(SimpleNamespace(), None) is None + during = os.environ.get('MLFLOW_HTTP_REQUEST_TIMEOUT') + registered, tracker.registered = tracker.registered, None + session.finish(None, None) + session._worker.join(5) + assert registered == ('registered' if len(seen) == 1 else None) + assert during == ('30' if len(seen) == 1 else None) + + assert seen == ['600', '600'] + assert 'MLFLOW_HTTP_REQUEST_TIMEOUT' not in os.environ + assert 'the model is not registered' in caplog.text diff --git a/packages/kumo-engine/test/test_mlflow_tracking.py b/packages/kumo-engine/test/test_mlflow_tracking.py new file mode 100644 index 000000000..0babf6919 --- /dev/null +++ b/packages/kumo-engine/test/test_mlflow_tracking.py @@ -0,0 +1,310 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import json +import os +import socket +import threading +import time + +import numpy as np +import pandas as pd +import pytest +from kumoapi.encoder import Numerical +from kumoapi.model_plan import ModelPlan +from test_tracking_contract import CONTEXT, join_tracking_threads, replay + +import kumo.engine as kumo +from kumo.engine._scratch import get_scratch_dir +from kumo.engine._tracking import run_record + +mlflow = pytest.importorskip('mlflow') +autologging_utils = pytest.importorskip('mlflow.utils.autologging_utils') + +QUERY = 'PREDICT ents.Y FOR EACH ents.ID' + + +@pytest.fixture +def tracking_uri(tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) # the default artifact root is ./mlruns + previous = mlflow.get_tracking_uri() + uri = f'sqlite:///{tmp_path / "mlflow.db"}' + mlflow.set_tracking_uri(uri) + try: + yield uri + finally: + while mlflow.active_run() is not None: + mlflow.end_run() + mlflow.set_tracking_uri(previous) + + +@pytest.fixture +def table_and_snapshot(tmp_path, tracking_uri): + previous_scratch = get_scratch_dir() + kumo.configure(scratch_dir=tmp_path / 'scratch') + try: + rng = np.random.default_rng(0) + (tmp_path / 'data').mkdir() + pd.DataFrame({ + 'ID': np.arange(400), + 'Y': rng.random(400), + 'F1': rng.normal(size=400), + 'amount ($)': rng.normal(size=400), + }).to_parquet(tmp_path / 'data' / 'ents.parquet') + connector = kumo.FileConnector(root_dir=str(tmp_path / 'data')) + ents = kumo.Table.from_source_table(connector['ents.parquet'], + primary_key='ID').infer_metadata() + graph = kumo.Graph(tables={'ents': ents}, edges=[]) + snapshot = graph.snapshot() + query = kumo.PredictiveQuery(graph=graph, query=QUERY) + yield query.generate_training_table(snapshot), snapshot + finally: + kumo.configure(scratch_dir=previous_scratch) + + +def _runs(client): + return client.search_runs( + [e.experiment_id for e in client.search_experiments()]) + + +def test_fit_logs_a_run_per_fit_and_per_trial(table_and_snapshot, tmp_path, + caplog): + from kumo.engine.integrations.mlflow import MlflowTracking + + table, snapshot = table_and_snapshot + client = mlflow.MlflowClient() + plan = ModelPlan() + plan.training_job.num_experiments = 2 + plan.training_job.disable_explain = True + plan.optimization.max_epochs = 2 + plan.optimization.max_steps_per_epoch = 2 + plan.optimization.base_lr = [0.01, 0.001] + plan.column_processing.encoder_overrides = { + 'ents.amount ($)': Numerical('standard') + } + mlflow.pytorch.autolog() + try: + with mlflow.start_run() as outer: + tracking = MlflowTracking(experiment_name='ents') + result = kumo.Trainer(plan).fit(table, snapshot=snapshot, + tracking=tracking) + assert mlflow.active_run().info.run_id == outer.info.run_id + assert autologging_utils._AUTOLOGGING_GLOBALLY_DISABLED is False + finally: + mlflow.pytorch.autolog(disable=True) + + runs = _runs(client) + parent, = [ + run for run in runs + if run.data.tags.get('mlflow.parentRunId') == outer.info.run_id + ] + assert tracking.run_id == parent.info.run_id + children = sorted( + (run for run in runs + if run.data.tags.get('mlflow.parentRunId') == parent.info.run_id), + key=lambda run: run.info.run_name) + assert len(runs) == 4 + assert client.get_experiment(parent.info.experiment_id).name == 'ents' + assert parent.info.status == 'FINISHED' + assert {k for k in parent.data.tags if k.startswith('kumo.')} == { + 'kumo.kind', 'kumo.run_id', 'kumo.bundle_id', 'kumo.dropped_params'} + assert parent.data.tags['kumo.bundle_id'] == result.bundle_id + assert int(parent.data.tags['kumo.dropped_params']) >= 1 + params = parent.data.params + assert json.loads(params['optimization.base_lr']) == [0.01, 0.001] + assert not any('$' in name for name in params) + metrics = { + name: value + for name, value in parent.data.metrics.items() + if not name.startswith('system/') + } + assert metrics == {**result.metrics, 'seconds': metrics['seconds']} + assert 'off for the rest of this run' not in caplog.text + ents, = parent.inputs.dataset_inputs + assert (ents.dataset.name, ents.dataset.source_type) == ('ents', 'local') + assert json.loads(ents.dataset.source) == { + 'uri': str((tmp_path / 'data' / 'ents.parquet').resolve()) + } + + assert [run.info.run_name + for run in children] == ['trial-1-attempt-1', 'trial-2-attempt-1'] + for child in children: + assert child.info.status == 'FINISHED' + assert child.data.tags['kumo.status'] == 'succeeded' + assert {'best_val_mae', 'seconds', + 'val_mae'} <= set(child.data.metrics) + history = client.get_metric_history(child.info.run_id, 'val_mae') + assert sorted(metric.step for metric in history) == [1, 2] + + +def test_failed_fit_marks_the_run_failed(table_and_snapshot, monkeypatch): + import kumo_core.training.orchestration as orchestration + + from kumo.engine.integrations.mlflow import MlflowTracking + + def fail(*args, **kwargs): + raise ValueError('stop') + + monkeypatch.setattr(orchestration, 'run_callable_training', fail) + table, snapshot = table_and_snapshot + + with pytest.raises(ValueError, match='stop'): + kumo.Trainer(ModelPlan()).fit(table, snapshot=snapshot, + tracking=MlflowTracking()) + + run, = _runs(mlflow.MlflowClient()) + assert run.info.status == 'FAILED' + assert run.data.tags['kumo.error'] == 'ValueError: stop' + assert mlflow.active_run() is None + + +def test_mlflow_error_ends_runs_failed_and_is_not_raised( + tracking_uri, monkeypatch, caplog): + from mlflow.tracking import MlflowClient + + from kumo.engine.integrations.mlflow import MlflowTracking + + original = MlflowClient.log_batch + + def down(self, *args, metrics=(), **kwargs): + if metrics: + raise RuntimeError('tracking server down') + return original(self, *args, **kwargs) + + monkeypatch.setattr(MlflowClient, 'log_batch', down) + replay(MlflowTracking(log_system_metrics=False)) + + assert caplog.text.count('off for the rest of this run') == 1 + assert mlflow.active_run() is None + statuses = {run.info.status for run in _runs(mlflow.MlflowClient())} + assert statuses == {'FAILED'} + + +def test_user_disabled_autologging_stays_disabled(tracking_uri): + from kumo.engine.integrations.mlflow import MlflowTracking + + tracker = MlflowTracking(log_system_metrics=False) + with autologging_utils.disable_autologging(): + tracker.start(run_record(CONTEXT)) + tracker.finish(None, None) + assert autologging_utils._AUTOLOGGING_GLOBALLY_DISABLED is True + + +@pytest.fixture +def hung_server(): + server = socket.create_server(('127.0.0.1', 0)) + held = [] + + def accept(): + while True: + try: + held.append(server.accept()[0]) + except OSError: + return + + threading.Thread(target=accept, daemon=True).start() + try: + yield f'http://127.0.0.1:{server.getsockname()[1]}' + finally: + server.close() + for connection in held: + connection.close() + join_tracking_threads() + + +def test_a_hung_server_does_not_stall_the_fit(hung_server, monkeypatch): + import kumo.engine._tracking as tracking + from kumo.engine.integrations.mlflow import MlflowTracking + + monkeypatch.setattr(tracking, 'FINISH_SECONDS', 0.5) + monkeypatch.setenv('MLFLOW_HTTP_REQUEST_MAX_RETRIES', '0') + monkeypatch.delenv('MLFLOW_HTTP_REQUEST_TIMEOUT', raising=False) + previous = mlflow.get_tracking_uri() + mlflow.set_tracking_uri(hung_server) + try: + began = time.monotonic() + replay(MlflowTracking(log_system_metrics=False)) + elapsed = time.monotonic() - began + limits = {name: os.environ.get(name) for name in ( + 'MLFLOW_HTTP_REQUEST_MAX_RETRIES', 'MLFLOW_HTTP_REQUEST_TIMEOUT')} + finally: + mlflow.set_tracking_uri(previous) + + assert elapsed < 1.5 + assert limits == {'MLFLOW_HTTP_REQUEST_MAX_RETRIES': '0', + 'MLFLOW_HTTP_REQUEST_TIMEOUT': '30'} + + +def test_source_tables_are_logged_as_native_inputs( + tracking_uri, monkeypatch): + from mlflow.data.delta_dataset_source import DeltaDatasetSource + + from kumo.engine.integrations.mlflow import MlflowTracking + + monkeypatch.setattr(DeltaDatasetSource, '_lookup_table_id', + lambda self, name: 'table-id') + replay(MlflowTracking(log_system_metrics=False)) + + run, = [run for run in _runs(mlflow.MlflowClient()) + if run.data.tags['kumo.kind'] == 'fit'] + inputs = { + item.dataset.name: (item.dataset.source_type, + json.loads(item.dataset.source)) + for item in run.inputs.dataset_inputs + } + assert inputs == { + 'ents': ('s3', {'uri': 's3://bucket/ents.parquet'}), + 'orders': ('delta_table', { + 'delta_table_name': 'main.sales.orders', + 'delta_table_version': 7, + 'is_databricks_uc_table': True, + 'delta_table_id': 'table-id', + }), + } + orders = next(item for item in run.inputs.dataset_inputs + if item.dataset.name == 'orders') + assert mlflow.data.get_source(orders).delta_table_id == 'table-id' + + +def test_fit_inputs_come_from_the_manifests(tmp_path): + from kumo.engine._manifest import ( + RECONSTRUCTED, + SnapshotManifest, + SnapshotParents, + SourceTableEntry, + TrainingTableManifest, + TrainingTableParents, + write_manifest, + ) + from kumo.engine._tracking import fit_inputs + + snapshot, table = tmp_path / 'index', tmp_path / 'table' + snapshot.mkdir() + table.mkdir() + assert fit_inputs(snapshot, table) == {} + + entry = {'connector': 'LOCAL_FILE', 'location': '/data/ents.parquet', + 'version': '1790000000', 'version_kind': 'timestamp'} + write_manifest( + snapshot / 'snapshot-manifest.json', + SnapshotManifest(id='s', parents=SnapshotParents(source_tables={ + 'ents': SourceTableEntry(**entry, sql='SELECT 1')}), + **RECONSTRUCTED)) + write_manifest( + table / 'training-table-manifest.json', + TrainingTableManifest( + id='t', parents=TrainingTableParents( + snapshot='s', + source_table=SourceTableEntry(**{**entry, 'version': None})), + query='q', generation_plan={}, task_type='x', split_seed=None, + **RECONSTRUCTED)) + assert fit_inputs(snapshot, table) == { + 'ents': entry, + 'training_table': {**entry, 'version': None}, + } + + (snapshot / 'snapshot-manifest.json').write_text('not json') + with pytest.warns(UserWarning, match='unreadable manifest'): + assert fit_inputs(snapshot, table) == { + 'training_table': {**entry, 'version': None}} diff --git a/packages/kumo-engine/test/test_optional_imports.py b/packages/kumo-engine/test/test_optional_imports.py index c70368fed..c18d0310d 100644 --- a/packages/kumo-engine/test/test_optional_imports.py +++ b/packages/kumo-engine/test/test_optional_imports.py @@ -26,6 +26,7 @@ 's3fs', 'google', 'gcsfs', + 'mlflow', ) PROVIDER_PACKAGES = ( 'kumo_snowflake', @@ -114,6 +115,22 @@ def find_spec(self, name, path=None, target=None): # and only the stand-in `self` is missing. Fine: nothing # broke at import time. pass + + from kumo.engine.integrations.mlflow import MlflowTracking + try: + MlflowTracking() + except ImportError as e: + assert "kumo-engine[mlflow]" in str(e), e + else: + raise AssertionError("MlflowTracking() needs mlflow") + + from kumo.engine.integrations.snowflake import SnowflakeTracking + try: + SnowflakeTracking(None, experiment_name="ents") + except ImportError as e: + assert "kumo-engine[snowflake]" in str(e), e + else: + raise AssertionError("SnowflakeTracking() needs snowflake-ml-python") print("ok") ''') diff --git a/packages/kumo-engine/test/test_prediction_output_columns.py b/packages/kumo-engine/test/test_prediction_output_columns.py index 84eeb2631..7c6e66703 100644 --- a/packages/kumo-engine/test/test_prediction_output_columns.py +++ b/packages/kumo-engine/test/test_prediction_output_columns.py @@ -324,3 +324,69 @@ def test_multilabel_thresholds_preserve_scores(tmp_path, num_workers): binary_classification_threshold=threshold) finally: kumo.configure(scratch_dir=previous) + + +def test_lineage_columns_are_written_by_the_scorer(graph, classifier, + tmp_path): + import pyarrow as pa + import pyarrow.parquet as pq + + _, snapshot, _ = graph + trainer, _, table = classifier + connector = kumo.FileConnector(root_dir=str(tmp_path)) + kinds = ('predictions', 'embeddings') + + def config(name, lineage): + return PredictionOutputConfig( + connector=connector, + output_paths={kind: f'{name}/{kind}/' for kind in kinds}, + lineage_columns=lineage) + + def read(name, kind): + return pa.concat_tables( + pq.read_table(path) + for path in sorted((tmp_path / name / kind).glob('*.parquet'))) + + result = trainer.predict(table, snapshot=snapshot, + output_config=config('on', True), + binary_classification_threshold=0.5) + trainer.predict(table, snapshot=snapshot, + output_config=config('off', False), + binary_classification_threshold=0.5) + result.export(config('copy', False)) + lineage = { + 'KUMO_RUN_ID': result.selection.job_id, + 'KUMO_BUNDLE_ID': trainer.bundle_id, + 'KUMO_SNAPSHOT_ID': snapshot.id, + } + assert trainer.bundle_id is not None + assert list(result.frame.columns[-3:]) == list(lineage) + + for kind in kinds: + on, off, copy = (read(name, kind) for name in ('on', 'off', 'copy')) + assert on.num_rows == off.num_rows > 0 + assert on.schema.names == off.schema.names + list(lineage) + for field in off.schema: + assert on.schema.field(field.name).type == field.type + for name, value in lineage.items(): + assert on.column(name).type == pa.string() + assert set(on.column(name).to_pylist()) == {value} + assert copy.equals(on) + + +def test_lineage_export_of_a_result_without_them_warns( + graph, classifier, tmp_path): + from kumo.engine.trainer.trainer import PredictionResult + + _, snapshot, _ = graph + trainer, _, table = classifier + result = trainer.predict(table, snapshot=snapshot, + binary_classification_threshold=0.5) + legacy = PredictionResult(result.selection, result.artifact_type) + assert legacy == result and hash(legacy) == hash(result) + + with pytest.warns(UserWarning, match='lineage_columns'): + result.export(PredictionOutputConfig( + connector=kumo.FileConnector(root_dir=str(tmp_path)), + output_paths={'predictions': 'preds/'}, lineage_columns=True)) + assert not {'KUMO_RUN_ID', 'KUMO_BUNDLE_ID'} & set(result.frame) diff --git a/packages/kumo-engine/test/test_predictive_query.py b/packages/kumo-engine/test/test_predictive_query.py index 7f3794266..2f77622a3 100644 --- a/packages/kumo-engine/test/test_predictive_query.py +++ b/packages/kumo-engine/test/test_predictive_query.py @@ -465,7 +465,7 @@ def test_custom_training_table_materializes_weight_and_manual_split( ) monkeypatch.setattr('kumo.engine.pquery._custom_source.stage_source_table', - lambda source_table, destination: rows) + lambda source_table, destination: (rows, None)) monkeypatch.setattr( 'kumo.engine.pquery._custom_source.validate_training_source', lambda *args: None) @@ -512,7 +512,7 @@ def test_custom_training_table_uses_the_normal_inferred_split( staging = tmp_path / 'staging' captured = {} monkeypatch.setattr('kumo.engine.pquery._custom_source.stage_source_table', - lambda source_table, destination: rows) + lambda source_table, destination: (rows, None)) monkeypatch.setattr( 'kumo.engine.pquery._custom_source.validate_training_source', lambda *args: None) @@ -637,7 +637,7 @@ def test_custom_prediction_table_materializes_exact_source_rows( ) captured = {} monkeypatch.setattr('kumo.engine.pquery._custom_source.stage_source_table', - lambda source_table, destination: rows) + lambda source_table, destination: (rows, None)) monkeypatch.setattr( 'kumo.engine.pquery._custom_source.validate_prediction_source', lambda *args: None, diff --git a/packages/kumo-engine/test/test_query_from_snapshot.py b/packages/kumo-engine/test/test_query_from_snapshot.py new file mode 100644 index 000000000..da611aa38 --- /dev/null +++ b/packages/kumo-engine/test/test_query_from_snapshot.py @@ -0,0 +1,108 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from dataclasses import replace + +import pandas as pd +import pytest +from kumoapi.task import TaskType + +import kumo.engine as kumo +from kumo.engine._scratch import get_scratch_dir + + +def _scores(trainer, table, snapshot): + frame = trainer.predict(table, snapshot=snapshot, + binary_classification_threshold=0.5).frame + return frame.sort_values('ENTITY').reset_index(drop=True) + + +def test_snapshot_bound_query_predicts_like_graph_bound(fitted, tmp_path): + trainer, graph_table, snapshot = fitted + assert trainer.training_query == graph_table.query + query = kumo.PredictiveQuery.from_snapshot(snapshot, trainer.training_query) + expected = _scores(trainer, graph_table, snapshot) + + assert query.get_task_type() == TaskType.BINARY_CLASSIFICATION + table = query.generate_prediction_table( + snapshot, plan=trainer.training_table_generation_plan) + assert table.timestamps == graph_table.timestamps + pd.testing.assert_frame_equal(_scores(trainer, table, snapshot), expected) + + anchor = pd.Timestamp(graph_table.timestamps[0]) + pd.DataFrame({ + 'ENTITY1': range(50), + 'TIMESTAMP': [anchor] * 50, + }).to_parquet(tmp_path / 'rows.parquet') + supplied = kumo.PredictionTable.from_source_table( + kumo.FileConnector(root_dir=str(tmp_path))['rows.parquet'], pq=query, + snapshot=snapshot) + pd.testing.assert_frame_equal(_scores(trainer, supplied, snapshot), + expected) + + +def test_predict_snapshot_scores_with_the_saved_query(fitted, tmp_path): + trainer, graph_table, snapshot = fitted + expected = _scores(trainer, graph_table, snapshot) + + def scores(**options): + return trainer.predict_snapshot( + snapshot, binary_classification_threshold=0.5, + **options).frame.sort_values('ENTITY').reset_index(drop=True) + + pd.testing.assert_frame_equal(scores(), expected) + pd.DataFrame({ + 'ENTITY1': range(50), + 'TIMESTAMP': [pd.Timestamp(graph_table.timestamps[0])] * 50, + }).to_parquet(tmp_path / 'rows.parquet') + rows = kumo.FileConnector(root_dir=str(tmp_path))['rows.parquet'] + pd.testing.assert_frame_equal(scores(entities=rows), expected) + with pytest.raises(ValueError, match='not both'): + scores(entities=rows, prediction_time=graph_table.timestamps[0]) + other = replace(snapshot, id='other@' + snapshot.id.split('@')[1]) + with pytest.warns(UserWarning, match='trained on a snapshot of graph'): + scored = trainer.predict_snapshot( + other, binary_classification_threshold=0.5).frame + pd.testing.assert_frame_equal( + scored.sort_values('ENTITY').reset_index(drop=True), expected) + + +def test_exported_snapshot_scores_in_a_fresh_scratch_without_a_graph( + fitted, tmp_path): + trainer, graph_table, snapshot = fitted + expected = _scores(trainer, graph_table, snapshot) + exports = kumo.FileConnector(root_dir=str(tmp_path / 'exports')) + location = snapshot.export(exports, 'snapshots') + bundle = trainer._training_result.bundle_dir + previous = get_scratch_dir() + kumo.configure(scratch_dir=tmp_path / 'fresh') + try: + loaded = kumo.GraphSnapshot.load_from_connector(exports, path=location) + assert loaded.id == snapshot.id + assert loaded.index_path.is_relative_to(tmp_path / 'fresh') + model = kumo.Trainer.load(bundle) + scored = model.predict_snapshot( + loaded, binary_classification_threshold=0.5).frame + pd.testing.assert_frame_equal( + scored.sort_values('ENTITY').reset_index(drop=True), expected) + query = kumo.PredictiveQuery.from_snapshot(loaded, model.training_query) + table = query.generate_prediction_table( + loaded, plan=model.training_table_generation_plan) + pd.testing.assert_frame_equal(_scores(model, table, loaded), expected) + finally: + kumo.configure(scratch_dir=previous) + + +def test_snapshot_bound_query_checks_snapshot_graph(fitted): + _, graph_table, snapshot = fitted + query = kumo.PredictiveQuery.from_snapshot(snapshot, graph_table.query) + newer = replace(snapshot, id=snapshot.id.split('@')[0] + '@newer') + other = replace(snapshot, id='other@' + snapshot.id.split('@')[1]) + plain = replace(snapshot, id='plain') + + query.suggest_model_plan(newer) + kumo.PredictiveQuery.from_snapshot(plain, graph_table.query) \ + .suggest_model_plan(plain) + with pytest.raises(ValueError, match='was not materialized'): + query.suggest_model_plan(other) diff --git a/packages/kumo-engine/test/test_scoring_defaults.py b/packages/kumo-engine/test/test_scoring_defaults.py new file mode 100644 index 000000000..dd764e2f3 --- /dev/null +++ b/packages/kumo-engine/test/test_scoring_defaults.py @@ -0,0 +1,156 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import hashlib +import json +import shutil +from dataclasses import replace +from types import SimpleNamespace + +import numpy as np +import pandas as pd +import pytest +from kumo_core.artifacts import add_bundle_file, remove_bundle_file +from kumo_core.types import TaskType +from kumoapi.model_plan import ModelPlan +from kumoapi.plugin import PublishedModel + +import kumo.engine as kumo +from kumo.engine.trainer._scoring import normalized_scoring_defaults + + +def _frame(trainer, table, snapshot, **kwargs): + result = trainer.predict(table, snapshot=snapshot, **kwargs) + return result.frame.sort_values('ENTITY').reset_index(drop=True) + + +def _digest(directory): + return { + str(path.relative_to(directory)): + hashlib.sha256(path.read_bytes()).hexdigest() + for path in sorted(directory.rglob('*')) if path.is_file() + } + + +def _copy(trainer, destination): + shutil.copytree(trainer._training_result.bundle_dir, destination) + return destination + + +def test_derived_bundle_thresholds_unless_explicit(fitted): + trainer, table, snapshot = fitted + parent_dir = trainer._training_result.bundle_dir + original = _digest(parent_dir) + probs = _frame(trainer, table, snapshot, + binary_classification_threshold=.5)['True_PROB'] + threshold = float(probs.median()) + child = trainer.with_scoring_defaults( + kumo.ScoringDefaults(binary_classification_threshold=threshold)) + + implicit = _frame(child, table, snapshot) + pd.testing.assert_frame_equal( + implicit, + _frame(trainer, table, snapshot, + binary_classification_threshold=threshold)) + assert implicit['TARGET_PRED'].any() + assert not _frame(child, table, snapshot, + binary_classification_threshold=1.0)['TARGET_PRED'].any() + + assert _digest(parent_dir) == original + assert trainer.scoring_defaults == kumo.ScoringDefaults() + assert child.bundle_id != trainer.bundle_id + assert child.provenance.id == child.bundle_id + assert trainer.provenance.id == trainer.bundle_id + assert child.provenance.parents == replace(trainer.provenance.parents, + derived_from=trainer.bundle_id) + cleared = child.with_scoring_defaults(kumo.ScoringDefaults()) + assert cleared.bundle_id == trainer.bundle_id + + +def test_equal_defaults_give_equal_bundle_ids(fitted): + trainer, *_ = fitted + children = [ + trainer.with_scoring_defaults( + kumo.ScoringDefaults(binary_classification_threshold=value)) + for value in (1, np.float32(1)) + ] + assert children[0].bundle_id == children[1].bundle_id + assert children[0].scoring_defaults.binary_classification_threshold == 1.0 + + +def test_derivation_from_old_or_newer_provenance(fitted, tmp_path): + trainer, *_ = fitted + defaults = kumo.ScoringDefaults(binary_classification_threshold=.5) + old = _copy(trainer, tmp_path / 'old') + remove_bundle_file(old, 'provenance') + child = kumo.Trainer.load(old).with_scoring_defaults(defaults) + assert child.provenance.reconstructed is True + assert child.provenance.id == child.bundle_id + recorded = json.loads( + (child._training_result.bundle_dir / 'provenance.json').read_text()) + assert recorded == { + 'kind': 'model_bundle', 'schema_version': 1, 'id': child.bundle_id, + 'parents': {'derived_from': trainer.bundle_id}, 'created_at': None, + 'environment': None, 'scheduler': None, 'reconstructed': True + } + + newer = _copy(trainer, tmp_path / 'newer') + record = json.dumps({'schema_version': 999, 'parents': {}}) + (newer / 'provenance.json').write_text(record) + add_bundle_file(newer, 'provenance', 'provenance.json') + with pytest.warns(UserWarning, match='left unchanged'): + child = kumo.Trainer.load(newer).with_scoring_defaults(defaults) + assert (child._training_result.bundle_dir / + 'provenance.json').read_text() == record + + +def test_defaults_survive_export_and_connector_load(fitted, tmp_path): + trainer, *_ = fitted + child = trainer.with_scoring_defaults( + kumo.ScoringDefaults(binary_classification_threshold=.4)) + registry = kumo.FileConnector(root_dir=str(tmp_path)) + child.export(kumo.ModelOutputConfig(registry, path='m1')) + loaded = kumo.Trainer.load_from_connector(registry, path='m1') + assert loaded.scoring_defaults == child.scoring_defaults + assert loaded.bundle_id == child.bundle_id + + +def test_derived_export_publishes_the_child_bundle_id(fitted): + trainer, *_ = fitted + child = trainer.with_scoring_defaults( + kumo.ScoringDefaults(binary_classification_threshold=.4)) + calls = [] + sink = SimpleNamespace(publish_model=lambda *a, **k: calls.append(k) or + PublishedModel(name='m', version='1', uri='u')) + child.export( + kumo.ModelOutputConfig(SimpleNamespace(sink_backend=sink), path='m')) + assert calls[0]['metadata']['kumo_bundle_id'] == child.bundle_id + + +def test_defaults_are_validated_against_the_task(fitted): + trainer, *_ = fitted + with pytest.raises(ValueError, match='num_classes_to_return'): + trainer.with_scoring_defaults( + kumo.ScoringDefaults(num_classes_to_return=2)) + with pytest.raises(RuntimeError, match='requires a prior'): + kumo.Trainer(ModelPlan()).with_scoring_defaults(kumo.ScoringDefaults()) + assert normalized_scoring_defaults( + kumo.ScoringDefaults(.3, np.int64(2)), + TaskType.MULTILABEL_CLASSIFICATION) == kumo.ScoringDefaults(.3, 2) + + +def test_unlisted_or_unreadable_defaults_are_ignored(fitted, tmp_path): + trainer, table, snapshot = fitted + bundle = _copy(trainer, tmp_path / 'old') + (bundle / 'scoring-defaults.json').write_text('not json') + old = kumo.Trainer.load(bundle) + assert old.scoring_defaults == kumo.ScoringDefaults() + assert old.bundle_id == trainer.bundle_id + + add_bundle_file(bundle, 'scoring_defaults', 'scoring-defaults.json') + with pytest.warns(UserWarning, match='Ignoring the scoring defaults'): + ignored = kumo.Trainer.load(bundle) + assert ignored.scoring_defaults == kumo.ScoringDefaults() + with pytest.raises(ValueError, match='must be specified'): + ignored.predict(table, snapshot=snapshot) diff --git a/packages/kumo-engine/test/test_semantic_carry_forward.py b/packages/kumo-engine/test/test_semantic_carry_forward.py index 15a4e2330..9481d14d1 100644 --- a/packages/kumo-engine/test/test_semantic_carry_forward.py +++ b/packages/kumo-engine/test/test_semantic_carry_forward.py @@ -44,6 +44,11 @@ def _stub_generated_task_dataset_builder(monkeypatch): ) +def _best_trainer(log_dir): + return SimpleNamespace(get_log_dir=lambda fit_mode: str(log_dir), + published_fit_mode=lambda: None) + + def _write_prediction_shard(options, frame): from pathlib import Path @@ -406,7 +411,7 @@ def fake_materialize(request, task_dataset, generated, **kwargs): metrics={}, artifact_paths={}, explain_disabled=True, - best_trainer=SimpleNamespace(log_dir=str(tmp_path)), + best_trainer=_best_trainer(tmp_path), )) monkeypatch.setattr('kumo.engine.trainer.trainer.write_training_spec', lambda *args, **kwargs: None) @@ -536,7 +541,7 @@ def fake_materialize(request, task_dataset, generated, **kwargs): metrics={}, artifact_paths={}, explain_disabled=True, - best_trainer=SimpleNamespace(log_dir=str(tmp_path)), + best_trainer=_best_trainer(tmp_path), )) monkeypatch.setattr('kumo.engine.trainer.trainer.write_training_spec', lambda *args, **kwargs: None) @@ -612,7 +617,7 @@ def fake_materialize(request, task_dataset, generated, **kwargs): metrics={}, artifact_paths={}, explain_disabled=True, - best_trainer=SimpleNamespace(log_dir=str(tmp_path)), + best_trainer=_best_trainer(tmp_path), )) monkeypatch.setattr('kumo.engine.trainer.trainer.write_training_spec', lambda *args, **kwargs: None) @@ -791,7 +796,7 @@ def fake_materialize(request, task_dataset, generated, **kwargs): metrics={}, artifact_paths={}, explain_disabled=True, - best_trainer=SimpleNamespace(log_dir=str(tmp_path)), + best_trainer=_best_trainer(tmp_path), )) monkeypatch.setattr('kumo.engine.trainer.trainer.write_training_spec', lambda *args, **kwargs: None) @@ -869,10 +874,9 @@ def fake_train(training_input, **kwargs): both_workers.wait(timeout=5) time.sleep(.05) run_name = Path(training_input.home_dir).name - return SimpleNamespace( - metrics={'run_name': - run_name}, artifact_paths={}, explain_disabled=True, - best_trainer=SimpleNamespace(log_dir=str(tmp_path))) + return SimpleNamespace(metrics={'run_name': run_name}, + artifact_paths={}, explain_disabled=True, + best_trainer=_best_trainer(tmp_path)) monkeypatch.setattr( 'kumo.engine.trainer.trainer.build_generated_training_task_dataset', diff --git a/packages/kumo-engine/test/test_snapshot_export.py b/packages/kumo-engine/test/test_snapshot_export.py new file mode 100644 index 000000000..81c367199 --- /dev/null +++ b/packages/kumo-engine/test/test_snapshot_export.py @@ -0,0 +1,483 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import hashlib +import json +import os +import shutil +import threading +import warnings +import zlib +from contextlib import contextmanager +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import Mock + +import pandas as pd +import pytest +from pyarrow.fs import LocalFileSystem, SubTreeFileSystem + +import kumo.engine as kumo +import kumo.engine._snapshot_io as snapshot_io +from kumo.engine import GraphSnapshot +from kumo.engine._scratch import get_scratch_dir + + +@pytest.fixture +def scratch(tmp_path): + previous = get_scratch_dir() + + def use(name): + kumo.configure(scratch_dir=tmp_path / name) + + use('cpu') + yield use + kumo.configure(scratch_dir=previous) + + +@pytest.fixture +def exports(tmp_path): + return kumo.FileConnector(root_dir=str(tmp_path / 'exports')) + + +def _index_files(path): + return { + str(file.relative_to(path)): file.read_bytes() + for file in Path(path).rglob('*') if file.is_file() + and file.name not in ('_SUCCESS', 'snapshot-manifest.json') + and not str(file.relative_to(path)).startswith('train_tables/') + } + + +def _export(snapshot, exports, tmp_path): + location = snapshot.export(exports, 'one') + return location, tmp_path / 'exports' / location + + +def _edit_manifest(exported, edit): + path = exported / 'kumo_snapshot.json' + manifest = json.loads(path.read_text()) + edit(manifest) + path.write_text(json.dumps(manifest)) + + +def _edit_index_manifest(exported, edit): + path = exported / 'index/_MANIFEST.json' + manifest = json.loads(path.read_text()) + edit(manifest) + path.write_text(json.dumps(manifest)) + _edit_manifest( + exported, lambda m: m.update( + index_manifest=hashlib.sha256(path.read_bytes()).hexdigest())) + + +def _load(exports, location, **kwargs): + return GraphSnapshot.load_from_connector(exports, path=location, **kwargs) + + +def test_export_load_round_trip_relocates_and_is_reused( + tmp_path, scratch, graph, exports, monkeypatch): + monkeypatch.setattr(snapshot_io, '_CHUNK', 1024) + monkeypatch.setattr(snapshot_io, '_RANGE_BYTES', 4096) + snapshot = graph.snapshot() + location = snapshot.export(exports, 'one/') + graph_id, data_version = snapshot.id.split('@') + assert location == f'one/{graph_id}/{data_version}' + assert snapshot.export(exports, 'one') == location + exported = tmp_path / 'exports' / location + recorded = (snapshot.index_path / 'snapshot-manifest.json').read_bytes() + assert (exported / 'snapshot-manifest.json').read_bytes() == recorded + + shutil.rmtree(tmp_path / 'cpu') + scratch('gpu') + loaded = _load(exports, location, graph=graph) + assert (loaded.id, loaded.graph_config) == (snapshot.id, + snapshot.graph_config) + assert loaded.index_path.is_relative_to(tmp_path / 'gpu') + assert _index_files(loaded.index_path) == _index_files(exported / 'index') + assert (loaded.index_path / + 'snapshot-manifest.json').read_bytes() == recorded + reexported = tmp_path / 'exports' / loaded.export(exports, 'two') + assert (reexported / 'snapshot-manifest.json').read_bytes() == recorded + + monkeypatch.setattr( + 'kumo_core.diskgraph.ingest.execute_diskgraph_ingest', + Mock(side_effect=AssertionError('ingest must not run'))) + monkeypatch.setattr(snapshot_io, 'load_into', + Mock(side_effect=AssertionError('already loaded'))) + assert _load(exports, location) == loaded + assert graph.snapshot() == loaded + + +def test_damaged_index_is_reused_by_snapshot_but_reloaded_by_load( + tmp_path, scratch, graph, exports, monkeypatch): + snapshot = graph.snapshot() + location, _ = _export(snapshot, exports, tmp_path) + expected = _index_files(snapshot.index_path) + stats = snapshot.index_path / 'col_stats.json' + stats.write_bytes(stats.read_bytes()[:-1]) + with warnings.catch_warnings(): + warnings.simplefilter('error') + assert graph.snapshot() == snapshot + + with pytest.warns(UserWarning, match='Reloading the damaged'): + loaded = _load(exports, location) + assert _index_files(loaded.index_path) == expected + + +def test_interrupted_reexport_is_not_loadable_and_reruns_from_scratch( + tmp_path, scratch, graph, exports, monkeypatch): + @contextmanager + def four(found, worker_ceiling=None): + yield 4 + + monkeypatch.setattr('kumo_core.sources.object_store.copies', four) + snapshot = graph.snapshot() + location, exported = _export(snapshot, exports, tmp_path) + (exported / 'index/stray.bin').write_bytes(b'x') + open_output, finished = snapshot_io._open_output, threading.Event() + + def interrupted(client, path): + if path.endswith('diskgraph_config.json'): + finished.wait(10) + raise OSError('connection reset') + if path.endswith('col_stats.json'): + finished.set() + return open_output(client, path) + + with monkeypatch.context() as patch: + patch.setattr(snapshot_io, '_open_output', interrupted) + patch.setattr(snapshot_io, '_ATTEMPTS', 1) + with pytest.raises(OSError, match='connection reset'): + snapshot.export(exports, 'one') + with pytest.raises(FileNotFoundError, match='No complete snapshot'): + _load(exports, location) + assert snapshot.export(exports, 'one') == location + assert not (exported / 'index/stray.bin').exists() + scratch('gpu') + assert _index_files(_load(exports, location).index_path) == _index_files( + snapshot.index_path) + + +def test_concurrent_exports_of_one_snapshot_both_succeed( + tmp_path, scratch, graph, exports, monkeypatch): + snapshot = graph.snapshot() + transfer, raced = snapshot_io._transfer, [] + + def racing(*args): + if not raced: + raced.append(None) + raced[0] = snapshot.export(exports, 'one') + return transfer(*args) + + monkeypatch.setattr(snapshot_io, '_transfer', racing) + assert raced == [snapshot.export(exports, 'one')] + assert _load(exports, raced[0]).id == snapshot.id + + +def test_export_repairs_files_a_dying_concurrent_export_wiped( + tmp_path, scratch, graph, exports, monkeypatch): + snapshot = graph.snapshot() + write, wiped = snapshot_io._write, [] + + def wipe_before_commit(filesystem, path, payload): + if path.endswith('/index/_MANIFEST.json') and not wiped: + wiped.append(path) + shutil.rmtree(Path(path).parent) + return write(filesystem, path, payload) + + monkeypatch.setattr(snapshot_io, '_write', wipe_before_commit) + location, _ = _export(snapshot, exports, tmp_path) + assert wiped + scratch('gpu') + assert _index_files(_load(exports, location).index_path) == _index_files( + snapshot.index_path) + + +def test_export_gives_up_when_another_export_keeps_wiping_it( + scratch, graph, exports, monkeypatch): + write = snapshot_io._write + + def wipe_before_commit(filesystem, path, payload): + if path.endswith('/index/_MANIFEST.json'): + shutil.rmtree(Path(path).parent) + return write(filesystem, path, payload) + + monkeypatch.setattr(snapshot_io, '_write', wipe_before_commit) + with pytest.raises(snapshot_io.ConcurrentExportError, + match='keeps replacing it'): + graph.snapshot().export(exports, 'one') + + +def test_failed_upload_is_retried_with_a_fresh_client( + scratch, graph, exports, monkeypatch): + open_output, clients = snapshot_io._open_output, [] + + def expiring(client, path): + if path.endswith('col_stats.json'): + clients.append(client) + if len(clients) == 1: + raise OSError('ExpiredToken') + return open_output(client, path) + + monkeypatch.setattr(snapshot_io, '_open_output', expiring) + graph.snapshot().export(exports, 'one') + assert len(clients) == 2 and clients[0] is not clients[1] + + +def test_reexport_replaces_whatever_the_location_holds( + tmp_path, scratch, graph, exports): + snapshot = graph.snapshot() + graph_id, data_version = snapshot.id.split('@') + occupied = tmp_path / 'exports/one' / graph_id / data_version + occupied.mkdir(parents=True) + (occupied / 'other.parquet').write_bytes(b'x') + location, exported = _export(snapshot, exports, tmp_path) + assert not (exported / 'other.parquet').exists() + stats = exported / 'index/col_stats.json' + original = stats.read_bytes() + stats.write_bytes(original[:3]) + (exported / 'index/index_metadata.json').unlink() + _edit_manifest(exported, lambda m: m.update(snapshot_id='other')) + assert snapshot.export(exports, 'one') == location + assert stats.read_bytes() == original + assert _load(exports, location).id == snapshot.id + + +def test_failed_load_leaves_nothing_and_restarts_cleanly( + tmp_path, scratch, graph, exports): + snapshot = graph.snapshot() + location, exported = _export(snapshot, exports, tmp_path) + stats = exported / 'index/col_stats.json' + original = stats.read_bytes() + stats.write_bytes(original[:-1] + bytes([original[-1] ^ 1])) + + scratch('gpu') + with pytest.raises(ValueError, match='corrupt.*checksum'): + _load(exports, location) + graph_dir = tmp_path / 'gpu/snapshots' / snapshot.index_path.parent.name + assert not [path for path in graph_dir.iterdir() if path.suffix != '.lock'] + + stats.write_bytes(original) + (graph_dir / f'{snapshot.index_path.name}.loading').mkdir() + loaded = _load(exports, location) + assert _index_files(loaded.index_path) == _index_files(snapshot.index_path) + assert [ + path.name for path in graph_dir.iterdir() if path.suffix != '.lock' + ] == [snapshot.index_path.name] + + +def test_load_refuses_before_downloading_when_disk_is_too_small( + tmp_path, scratch, graph, exports, monkeypatch): + location, _ = _export(graph.snapshot(), exports, tmp_path) + scratch('gpu') + monkeypatch.setattr(snapshot_io.os, 'statvfs', + lambda path: SimpleNamespace(f_frsize=1, f_bavail=0)) + monkeypatch.setattr(snapshot_io, '_download', + Mock(side_effect=AssertionError('downloaded'))) + with pytest.raises(OSError, match='free bytes'): + _load(exports, location) + + +@pytest.mark.parametrize('index, edit', [ + (False, lambda manifest: manifest.update(data_version='../../victim')), + (False, lambda manifest: manifest.pop('crc32')), + (True, lambda manifest: manifest['files'][0].update( + rel_path='../../victim/x')), + (True, lambda manifest: manifest['files'].append(manifest['files'][0])), +]) +def test_load_rejects_a_hostile_or_malformed_manifest(tmp_path, scratch, graph, + exports, index, edit): + location, exported = _export(graph.snapshot(), exports, tmp_path) + victim = tmp_path / 'victim' + victim.mkdir() + (victim / 'keep').write_text('x') + (_edit_index_manifest if index else _edit_manifest)(exported, edit) + scratch('gpu/a/b') + with pytest.raises(ValueError): + _load(exports, location) + assert (victim / 'keep').read_text() == 'x' + + +def test_load_rejects_unknown_formats_and_accepts_unknown_fields( + tmp_path, scratch, graph, exports): + snapshot = graph.snapshot() + location, exported = _export(snapshot, exports, tmp_path) + scratch('gpu') + + _edit_manifest(exported, lambda m: m.update(format=2)) + with pytest.raises(ValueError, match='format 2'): + _load(exports, location) + + _edit_manifest(exported, lambda m: m.update(format=1, added_later={'x': 1})) + other = kumo.Graph(tables={'users': graph.tables['users']}) + with pytest.warns(UserWarning, match='another graph'): + loaded = _load(exports, location, graph=other) + assert loaded.id == snapshot.id + + +@pytest.mark.parametrize('recorded', ['newer', 'missing']) +def test_snapshot_manifest_travels_unchanged_or_is_reconstructed( + tmp_path, scratch, graph, exports, recorded): + snapshot = graph.snapshot() + path = snapshot.index_path / 'snapshot-manifest.json' + if recorded == 'newer': + path.write_bytes(b'{"schema_version": 99, "future": [1]}') + else: + path.unlink() + location, exported = _export(snapshot, exports, tmp_path) + scratch('gpu') + with warnings.catch_warnings(): + warnings.simplefilter('ignore') + loaded = _load(exports, location) + shipped = (exported / 'snapshot-manifest.json').read_bytes() + assert (loaded.index_path / 'snapshot-manifest.json').read_bytes() == shipped + if recorded == 'newer': + assert shipped == b'{"schema_version": 99, "future": [1]}' + else: + manifest = loaded.manifest + assert (manifest.id, manifest.reconstructed) == (snapshot.id, True) + assert manifest.parents.source_tables == {} + + +def test_warehouse_connectors_are_not_supported(scratch, graph): + connector = SimpleNamespace(sink_backend=object()) + with pytest.raises(NotImplementedError, match='FileConnector'): + graph.snapshot().export(connector, 'x') + with pytest.raises(NotImplementedError, match='FileConnector'): + _load(connector, 'x') + + +def test_load_explains_an_index_the_installed_diskgraph_cannot_open( + tmp_path, scratch, graph, exports): + location, exported = _export(graph.snapshot(), exports, tmp_path) + metadata = exported / 'index/index_metadata.json' + garbage = b'{' * metadata.stat().st_size + metadata.write_bytes(garbage) + + def unreadable(manifest): + manifest['crc32']['index_metadata.json'] = [zlib.crc32(garbage)] + manifest['environment']['kumo_engine'] = '0.0.1' + + _edit_manifest(exported, unreadable) + scratch('gpu') + with pytest.raises(RuntimeError, match='exported by kumo-engine 0.0.1'): + _load(exports, location) + + +@pytest.mark.parametrize('scheme', ['s3', 'gs']) +def test_object_store_round_trip(tmp_path, scratch, graph, monkeypatch, + scheme): + storage = tmp_path / 'storage' + storage.mkdir() + + class Bucket(SubTreeFileSystem): + def __init__(self, *args, **kwargs): + super().__init__(str(storage), LocalFileSystem()) + + monkeypatch.setattr('kumo_s3.filesystem', Bucket) + monkeypatch.setattr('kumo_gcs.backend.fs.GcsFileSystem', Bucket) + snapshot = graph.snapshot() + connector = kumo.S3Connector() if scheme == 's3' else kumo.GCSConnector() + location = snapshot.export(connector, f'{scheme}://bucket/snapshots/') + graph_id, data_version = snapshot.id.split('@') + assert location == f'{scheme}://bucket/snapshots/{graph_id}/{data_version}' + scratch('gpu') + loaded = _load(connector, location, graph=graph) + assert _index_files(loaded.index_path) == _index_files(snapshot.index_path) + + +def _set_mtime(path, seconds): + os.utime(path, (seconds, seconds)) + + +def _rewrite_events(data): + frame = pd.read_parquet(data / 'events.parquet') + frame.iloc[::-1].to_parquet(data / 'events.parquet') + _set_mtime(data / 'events.parquet', 1_700_009_123) + + +@pytest.fixture +def dated(tmp_path, graph): + data = tmp_path / 'data' + for name in ('users', 'events'): + _set_mtime(data / f'{name}.parquet', 1_700_000_123) + return data + + +@pytest.mark.parametrize('record_sources', [True, False]) +def test_source_changes_reads_stamps_without_staging( + scratch, graph, dated, monkeypatch, record_sources): + from kumo_core.sources.object_store import ObjectStoreSourceReader + + from kumo.engine.connector.file import _FileSourceBackend + + _set_mtime(dated / 'users.parquet', 1_700_002_800) + snapshot = graph.snapshot(record_sources=record_sources) + fail = Mock(side_effect=AssertionError('must not stage')) + monkeypatch.setattr('kumo_core.sources.prepare.read_table', fail) + monkeypatch.setattr(ObjectStoreSourceReader, 'prepare', fail) + monkeypatch.setattr(_FileSourceBackend, 'prepare_ingest', fail) + assert snapshot.source_changes(graph) == { + 'users': {'status': 'unchanged', 'recorded': '1700002800', + 'current': '1700002800'}, + 'events': {'status': 'unchanged', 'recorded': '1700000123', + 'current': '1700000123'}, + } + _rewrite_events(dated) + assert snapshot.source_changes(graph)['events'] == { + 'status': 'changed', 'recorded': '1700000123', 'current': '1700009123'} + + monkeypatch.setattr(_FileSourceBackend, 'get_data_version', + Mock(side_effect=OSError('listing failed'))) + assert snapshot.source_changes(graph)['events'] == { + 'status': 'unknown', 'recorded': '1700000123', 'current': None} + + +def test_source_changes_compares_directory_listings(tmp_path, scratch): + parts = tmp_path / 'dir' / 'users' + parts.mkdir(parents=True) + pd.DataFrame({'ID': range(5)}).to_parquet(parts / 'part-0.parquet') + connector = kumo.FileConnector(root_dir=str(tmp_path / 'dir')) + graph = kumo.Graph(tables={ + 'users': kumo.Table.from_source_table(connector['users'], + primary_key='ID') + }, edges=[]) + snapshot = graph.snapshot() + assert snapshot.source_changes(graph)['users']['status'] == 'unchanged' + pd.DataFrame({'ID': range(5, 9)}).to_parquet(parts / 'part-1.parquet') + assert snapshot.source_changes(graph)['users']['status'] == 'changed' + + +@pytest.mark.parametrize('edit', ['absent', 'reconstructed', 'hour_bucket']) +def test_source_changes_unknown_without_a_comparable_version( + scratch, graph, dated, edit): + snapshot = graph.snapshot() + path = snapshot.index_path / 'snapshot-manifest.json' + manifest = json.loads(path.read_text()) + if edit == 'absent': + path.unlink() + elif edit == 'reconstructed': + path.write_text(json.dumps({**manifest, 'reconstructed': True})) + else: + for entry in manifest['parents']['source_tables'].values(): + entry['version_kind'] = 'hour_bucket' + path.write_text(json.dumps(manifest)) + assert {change['status'] for change in snapshot.source_changes( + graph).values()} == {'unknown'} + + +def test_load_warns_only_with_graph_and_only_for_changed_tables( + tmp_path, scratch, graph, dated, exports): + location, _ = _export(graph.snapshot(), exports, tmp_path) + scratch('gpu') + with warnings.catch_warnings(): + warnings.simplefilter('error') + _load(exports, location, graph=graph) + _rewrite_events(dated) + _load(exports, location) + with pytest.warns(UserWarning) as caught: + _load(exports, location, graph=graph) + [message] = [str(warning.message) for warning in caught] + assert 'source table(s) events;' in message diff --git a/packages/kumo-engine/test/test_snapshot_export_live.py b/packages/kumo-engine/test/test_snapshot_export_live.py new file mode 100644 index 000000000..3c69b548a --- /dev/null +++ b/packages/kumo-engine/test/test_snapshot_export_live.py @@ -0,0 +1,44 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import contextlib +import os +import uuid + +import pytest + +import kumo.engine as kumo +from kumo.engine._scratch import get_scratch_dir +from kumo.engine._snapshot_io import _location + +ROOTS = { + 's3': os.environ.get('KUMO_SNAPSHOT_LIVE_S3_ROOT'), + 'gs': os.environ.get('KUMO_SNAPSHOT_LIVE_GCS_ROOT'), +} + + +@pytest.mark.parametrize('scheme', ['s3', 'gs']) +def test_snapshot_round_trip_through_object_store(tmp_path, graph, scheme): + root = ROOTS[scheme] + if not root: + pytest.skip('set KUMO_SNAPSHOT_LIVE_S3_ROOT / _GCS_ROOT to run') + connector = kumo.S3Connector() if scheme == 's3' else kumo.GCSConnector() + path = f"{root.rstrip('/')}/engine-snapshot-{uuid.uuid4().hex[:12]}" + previous = get_scratch_dir() + try: + kumo.configure(scratch_dir=tmp_path / 'cpu') + snapshot = graph.snapshot() + location = snapshot.export(connector, path) + kumo.configure(scratch_dir=tmp_path / 'gpu') + loaded = kumo.GraphSnapshot.load_from_connector( + connector, path=location, graph=graph) + finally: + kumo.configure(scratch_dir=previous) + filesystem, prefix = _location(connector, path) + filesystem.delete_dir_contents(prefix, missing_dir_ok=True) + with contextlib.suppress(OSError): + filesystem.delete_dir(prefix) + assert loaded.id == snapshot.id + assert (loaded.index_path / 'col_stats.json').read_bytes() == ( + snapshot.index_path / 'col_stats.json').read_bytes() diff --git a/packages/kumo-engine/test/test_snapshot_manifest.py b/packages/kumo-engine/test/test_snapshot_manifest.py new file mode 100644 index 000000000..10a6e163e --- /dev/null +++ b/packages/kumo-engine/test/test_snapshot_manifest.py @@ -0,0 +1,284 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import hashlib +import json +import os +from dataclasses import replace +from datetime import datetime +from types import SimpleNamespace + +import numpy as np +import pandas as pd +import pytest +from kumo_core.custom_sql import custom_sql_synthetic_name +from kumo_core.sources.identity import SourceVersion +from kumo_core.types import ConnectorType +from kumo_core.version import get_kumo_version +from kumoapi.graph import Edge +from kumoapi.typing import Stype + +import kumo.engine._provenance as provenance +from kumo.engine import ( + BigQueryConnector, + DatabricksConnector, + FileConnector, + Graph, + S3Connector, + SnowflakeConnector, + Table, +) +from kumo.engine._manifest import ( + SchedulerContext, + SnapshotManifest, + SourceTableEntry, + read_manifest, + write_manifest, +) +from kumo.engine.connector.identity import ( + DatabricksIdentity, + SnowflakeIdentity, +) +from kumo.engine.graph import SNAPSHOT_MANIFEST_FILE, GraphSnapshot + +_SCHEDULER_ENV = ('KUMO_SCHEDULER', 'KUMO_JOB_ID', 'KUMO_JOB_RUN_ID', + 'TRAINING_JOB_NAME', 'TRAINING_JOB_ARN', + 'SNOWFLAKE_SERVICE_NAME', 'DATABRICKS_RUNTIME_VERSION') + + +@pytest.fixture +def no_scheduler(tmp_path, monkeypatch): + for name in _SCHEDULER_ENV: + monkeypatch.delenv(name, raising=False) + monkeypatch.setattr(provenance, 'PROCESSING_JOB_CONFIG', + tmp_path / 'processingjobconfig.json') + return monkeypatch + + +@pytest.fixture +def make_snapshot(tmp_path, no_scheduler): + no_scheduler.setattr('kumo.engine._scratch._root', tmp_path / 'scratch') + data = tmp_path / 'data' + data.mkdir() + pd.DataFrame({'ID': np.arange(5)}).to_parquet(data / 'ents.parquet') + + def make(events: int | None = None, **options) -> GraphSnapshot: + if events is not None: + pd.DataFrame({ + 'EID': np.arange(events), + 'ENT_ID': np.arange(events) % 5, + }).to_parquet(data / 'events.parquet') + # Single-file versions are mtime seconds. + stamp = 1.7e9 + events + os.utime(data / 'events.parquet', (stamp, stamp)) + connector = FileConnector(root_dir=str(data)) + return Graph( + { + 'ents': + Table.from_source_table(connector['ents'], + primary_key='ID').infer_metadata(), + 'events': + Table.from_source_table( + connector['events.parquet'], + primary_key='EID').infer_metadata().with_column( + 'ENT_ID', stype=Stype.ID), + }, [Edge('events', 'ENT_ID', 'ents')]).snapshot(**options) + + return make + + +def test_snapshot_manifest_records_sources_and_diffs(make_snapshot, + tmp_path): + before = make_snapshot(30) + manifest = before.manifest + + assert manifest.kind == 'graph_snapshot' + assert manifest.id == before.id + data = (tmp_path / 'data').resolve() + assert {a: e.location + for a, e in manifest.parents.source_tables.items()} == { + 'ents': str(data / 'ents.parquet'), + 'events': str(data / 'events.parquet'), + } + assert {e.version_kind + for e in manifest.parents.source_tables.values() + } == {'timestamp'} + assert manifest.extra == {} + assert before.counts['rows'] == {'ents': 5, 'events': 30} + assert before.counts['edges']['events_ENT_ID_to_ents_ID'] == 30 + assert manifest.environment.kumo_engine == get_kumo_version() + assert manifest.environment.diskgraph + assert manifest.scheduler is None + assert datetime.fromisoformat(manifest.created_at).tzinfo is not None + assert str(tmp_path / 'scratch') not in json.dumps(manifest.to_dict()) + + after = make_snapshot(40) + assert after.counts['rows']['events'] == 40 + sources = manifest.parents.source_tables + new_sources = after.manifest.parents.source_tables + assert new_sources['ents'] == sources['ents'] + assert new_sources['events'].version != sources['events'].version + + +def test_index_without_manifest_is_reused_and_backfilled( + make_snapshot, monkeypatch): + import kumo_core.diskgraph.ingest as ingest + + snapshot = make_snapshot(30) + original = snapshot.manifest + (snapshot.index_path / SNAPSHOT_MANIFEST_FILE).unlink() + monkeypatch.setattr(ingest, 'execute_diskgraph_ingest', + lambda *args: pytest.fail('cached index re-ingested')) + + reused = make_snapshot() + + assert reused == snapshot + assert reused.manifest == replace(original, created_at=None, + environment=None, scheduler=None, + reconstructed=True) + + +def test_manifest_write_failure_leaves_snapshot_unchanged( + make_snapshot, monkeypatch): + def fail(path, manifest): + raise OSError('disk full') + + monkeypatch.setattr('kumo.engine.graph.write_manifest', fail) + with pytest.warns(UserWarning, match='disk full'): + snapshot = make_snapshot(30) + + assert snapshot.manifest is None + assert (snapshot.index_path / '_SUCCESS').is_file() + + +@pytest.mark.parametrize('env, processing, expected', [ + ({}, None, None), + ({'KUMO_SCHEDULER': 'databricks', 'KUMO_JOB_ID': '11', + 'KUMO_JOB_RUN_ID': '22', 'TRAINING_JOB_NAME': 'ignored'}, None, + ('databricks', '11', '22')), + ({'KUMO_JOB_RUN_ID': 'run-9', 'TRAINING_JOB_NAME': 'train-1'}, None, + ('sagemaker', 'train-1', 'run-9')), + ({}, {'ProcessingJobName': 'proc-1', 'ProcessingJobArn': 'arn:proc-1'}, + ('sagemaker', 'proc-1', 'arn:proc-1')), + ({'SNOWFLAKE_SERVICE_NAME': 'JOB_A1'}, None, ('snowflake', None, 'JOB_A1')), +]) +def test_scheduler_context(no_scheduler, env, processing, expected): + for name, value in env.items(): + no_scheduler.setenv(name, value) + if processing is not None: + provenance.PROCESSING_JOB_CONFIG.write_text(json.dumps(processing)) + + context = provenance.scheduler_context() + + assert context == (expected and SchedulerContext( + **dict(zip(('platform', 'job_id', 'job_run_id'), expected)))) + + +@pytest.mark.parametrize('content', [ + None, '{"schema_version": 2}', '{"schema_version": true}', '{"kind": ' +]) +def test_unsupported_manifest_is_ignored(tmp_path, content): + if content is not None: + (tmp_path / SNAPSHOT_MANIFEST_FILE).write_text(content) + with pytest.warns(UserWarning): + assert GraphSnapshot('s', tmp_path, {}).manifest is None + else: + assert GraphSnapshot('s', tmp_path, {}).manifest is None + + +@pytest.mark.parametrize('connector, name, sql, location', [ + (S3Connector(root_dir='s3://bucket-a/data/'), 'exports/tt/', None, + 's3://bucket-a/data/exports/tt'), + (S3Connector(root_dir='s3://bucket-a/data/'), 'q4#1 50%.parquet', None, + 's3://bucket-a/data/q4%231%2050%25.parquet'), + (S3Connector(), 's3://bucket-a/sales%23q4.parquet', None, + 's3://bucket-a/sales%23q4.parquet'), + (SnowflakeConnector(warehouse='wh', database='db', schema_name='sch', + identity=SnowflakeIdentity(account='org-acct')), + 'TT', None, 'org-acct/"DB"."SCH"."TT"'), + (SnowflakeConnector(warehouse='wh', database='db', schema_name='sch'), + 'TT', None, '"DB"."SCH"."TT"'), + (DatabricksConnector( + catalog='cat', schema_name='sch', warehouse_id='w', + identity=DatabricksIdentity(host='https://workspace.example.com', + client_id='c', client_secret='s')), + 'cat.sch.tt', None, '`cat`.`sch`.`tt`'), + (BigQueryConnector(project_id='data-project', dataset_id='analytics'), + 'tt', None, 'data-project.analytics.tt'), + (SimpleNamespace(source_type='OTHER', source_backend=None), 'TT', None, + 'TT'), + (None, 'synthetic', 'SELECT 1', 'sql:' + + hashlib.sha256(b'SELECT 1').hexdigest()), +]) +def test_source_location(connector, name, sql, location): + assert provenance.source_location(connector, name, sql) == location + + +def test_unrecordable_location_is_null(): + with pytest.warns(UserWarning, match='location'): + entry = provenance.source_table(None, 'users', None, 7) + assert entry == SourceTableEntry(connector=None, location=None, + version='7', version_kind='unknown') + + +@pytest.mark.parametrize('record_sources', [True, False]) +def test_custom_sql_records_text_unless_sources_are_off(record_sources): + version = SourceVersion('abc', 'hour_bucket', {'orders': '17'}) + entry = provenance.source_table( + SnowflakeConnector(warehouse='wh', database='db', schema_name='sch'), + 'q', 'SELECT * FROM orders', version, + record_sources=record_sources).to_dict() + expected = { + 'connector': 'SNOWFLAKE', + 'location': 'sql:' + hashlib.sha256(b'SELECT * FROM orders').hexdigest(), + 'version': 'abc', + 'version_kind': 'hour_bucket', + 'dependencies': {'orders': '17'}, + } + if record_sources: + expected.update(sql='SELECT * FROM orders', dialect='snowflake', + namespace=['db', 'sch']) + reformatted = provenance.source_table( + SnowflakeConnector(warehouse='wh', database='db', schema_name='sch'), + 'q', 'select *\n from ORDERS -- note', version) + assert entry.pop('normalized_sql_fingerprint') == ( + reformatted.normalized_sql_fingerprint) == custom_sql_synthetic_name( + 'SELECT * FROM orders', ConnectorType.SNOWFLAKE) + assert reformatted.location != entry['location'] + assert entry == expected + + +def test_record_sources_false_keeps_table_names_only(make_snapshot, + tmp_path): + names = {'ents': 'ents', 'events': 'events.parquet'} + redacted = make_snapshot(30, record_sources=False) + sources = redacted.manifest.parents.source_tables + assert {alias: e.location for alias, e in sources.items()} == names + assert str(tmp_path) not in json.dumps(redacted.manifest.to_dict()) + assert redacted.manifest.record_sources is False + assert redacted.record_sources is False + + full = make_snapshot(record_sources=True) + assert full == redacted + assert full.record_sources is True + assert full.manifest.parents.source_tables['ents'].location == str( + (tmp_path / 'data' / 'ents.parquet').resolve()) + assert full.manifest.created_at == redacted.manifest.created_at + + +def test_snapshot_manifest_round_trips_unchanged(make_snapshot, tmp_path): + written = make_snapshot(30).index_path / SNAPSHOT_MANIFEST_FILE + newer = json.loads(written.read_text()) + newer['future'] = {'kept': [1, 2]} + newer['environment']['cuda'] = '13.0' + newer['parents']['source_tables']['ents']['row_count'] = 5 + source = tmp_path / 'newer.json' + source.write_text(json.dumps(newer, indent=2, sort_keys=True) + '\n') + + manifest = read_manifest(source, SnapshotManifest) + write_manifest(tmp_path / 'out.json', manifest) + + assert (tmp_path / 'out.json').read_bytes() == source.read_bytes() + assert manifest.extra == {'future': {'kept': [1, 2]}} + assert manifest.environment.extra == {'cuda': '13.0'} diff --git a/packages/kumo-engine/test/test_snowflake_connector.py b/packages/kumo-engine/test/test_snowflake_connector.py index 1e4fcc687..b2b9254c8 100644 --- a/packages/kumo-engine/test/test_snowflake_connector.py +++ b/packages/kumo-engine/test/test_snowflake_connector.py @@ -108,3 +108,18 @@ def test_a_missing_target_field_with_no_override_is_reported( with pytest.raises(ValueError, match="names no warehouse"): SnowflakeConnector.from_spcs_env() + + +def test_session_is_the_readers_session(monkeypatch): + from kumo_snowflake import session as sessions + + connector = SnowflakeConnector(warehouse='WH', database='DB', + schema_name='S') + calls = [] + monkeypatch.setattr( + sessions._sessions, 'get', + lambda principal, target, fresh=False: calls.append(target) or 'live') + monkeypatch.setattr(type(connector.source_backend), '_principal', + lambda self, target: 'principal') + assert connector.session() == 'live' + assert calls == [connector] diff --git a/packages/kumo-engine/test/test_snowflake_tracking.py b/packages/kumo-engine/test/test_snowflake_tracking.py new file mode 100644 index 000000000..7eaf66e5e --- /dev/null +++ b/packages/kumo-engine/test/test_snowflake_tracking.py @@ -0,0 +1,216 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import inspect +import threading +import time +from unittest import mock + +import _snowflake_experiment_fake +import pytest +from _snowflake_experiment_fake import FakeExperimentTracking, FakeSession +from test_tracking_contract import ( + CONTEXT, + EVENTS, + RESULT, + join_tracking_threads, +) + +from kumo.engine._tracking import TrackingSession + +FAKED = ('__init__', 'set_experiment', 'start_run', 'end_run', 'log_params', + 'log_metrics', 'log_artifact') + + +@pytest.mark.parametrize('method', FAKED) +def test_fake_matches_installed_experiment_tracking(method): + experiment = pytest.importorskip('snowflake.ml.experiment') + real = inspect.signature(getattr(experiment.ExperimentTracking, method)) + fake = inspect.signature(getattr(FakeExperimentTracking, method)) + assert [(p.name, p.kind, p.default) for p in fake.parameters.values() + ] == [(p.name, p.kind, p.default) for p in real.parameters.values()] + + +def test_installed_singleton_exposes_what_the_adapter_reads(monkeypatch): + experiment = pytest.importorskip('snowflake.ml.experiment') + cls = experiment.ExperimentTracking + monkeypatch.setattr(cls, '_instance', None) + session = mock.MagicMock() + session.get_current_database.return_value = 'DB' + session.get_current_schema.return_value = 'PUBLIC' + + tracking = cls(session) + + assert cls(mock.MagicMock()) is tracking + assert tracking._run is None + assert tracking._session is session + assert 'experiment_name' in inspect.signature( + tracking._sql_client.show_runs_in_experiment).parameters + + +@pytest.fixture +def fake(monkeypatch): + return _snowflake_experiment_fake.install(monkeypatch) + + +def _tracking(session=None, **kwargs): + from kumo.engine.integrations.snowflake import SnowflakeTracking + return SnowflakeTracking(session or FakeSession(), experiment_name='ents', + **kwargs) + + +def _replay(tracking, result=RESULT, error=None, during=None): + session = TrackingSession(tracking) + session.start(lambda: CONTEXT) + for event in EVENTS: + session.on_event(event) + if during is not None: + caught_up = threading.Event() + session._submit(caught_up.set) + caught_up.wait(30) + during() + session.finish(result, error) + + +def test_native_model_lineage_holds_the_run_open_during_the_fit(fake): + lineage = [] + _replay( + _tracking(database='ml', schema='runs', run_name='fit', + native_model_lineage=True), + during=lambda: lineage.extend(fake._instance.lineage_stack)) + + assert lineage == [('ENTS', 'FIT')] + assert fake._instance.lineage_stack == [] + assert fake._instance.calls[1] == ('set_experiment', 'ents', 'ml', 'runs') + + +def test_snowflake_errors_warn_once_and_end_the_run_failed(fake, caplog): + session = FakeSession() + fake(session).fail_on = {'log_metrics'} + + _replay(_tracking(session)) + + tracking = fake._instance + assert caplog.text.count('off for the rest of this run') == 1 + run, = tracking.runs.values() + assert run['status'] == 'FAILED' + assert tracking._run is None + + +def test_failed_fit_ends_the_run_failed(fake): + _replay(_tracking(), None, ValueError('stop')) + + run, = fake._instance.runs.values() + assert run['status'] == 'FAILED' + + +@pytest.mark.parametrize('busy', ['active run', 'other session']) +def test_an_instance_in_use_is_left_alone(fake, busy, caplog): + session = FakeSession() + user = fake(session, warnings_as_errors=True) + user.set_experiment('mine') + if busy == 'active run': + user.start_run('mine_1') + else: + session = FakeSession() + calls = len(user.calls) + + _replay(_tracking(session)) + + assert user.calls[calls:] == [] + assert user.warnings_as_errors is True + assert caplog.text.count('not tracked in Snowflake') == 1 + + +def test_location_defaults_to_the_session_not_the_last_fit(fake): + session = FakeSession(database='DEV', schema='WORK') + _replay(_tracking(session, database='ml')) + _replay(_tracking(session)) + + assert [ + call[1:] for call in fake._instance.calls if call[0] == 'set_experiment' + ] == [('ents', 'ml', 'PUBLIC'), ('ents', 'DEV', 'WORK')] + assert not any(call[0].startswith('run.') for call in fake._instance.calls) + + +@pytest.mark.parametrize('taken', ['FIT', 'MYXRUN']) +def test_an_existing_run_is_never_resumed(fake, taken, caplog): + session = FakeSession() + user = fake(session) + user.set_experiment('ents') + user.start_run(taken) + if taken == 'FIT': + user.end_run() + else: + user._run = None # still running, in another process + run_name = 'fit' if taken == 'FIT' else 'my_run' + + tracking = _tracking(session, run_name=run_name) + _replay(tracking) + + new = f"{run_name}_{CONTEXT.run_id.rsplit('-', 1)[-1][:8]}".upper() + assert tracking.run_id == new + 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 + + +def test_metric_and_param_limits(fake, monkeypatch): + import kumo.engine.integrations.snowflake as snowflake + + monkeypatch.setattr(snowflake, '_EPOCH_METRICS', 2) + monkeypatch.setattr(snowflake, '_SUMMARY_METRICS', 3) + monkeypatch.setattr(snowflake, '_MAX_METRICS', 5) + result = type(RESULT)(metrics={**RESULT.metrics, 'r2': float('nan')}, + bundle_dir=RESULT.bundle_dir) + _replay(_tracking(tags={'note': 'a$$b', 'level': 2}), result) + + run, = fake._instance.runs.values() + assert sorted({name for name, _, _ in run['metrics']}) == [ + 'test_mae', 'trial-1-attempt-1/best_val_mae', + 'trial-1-attempt-1/train_loss', 'trial-1-attempt-1/val_mae', 'val_mae' + ] + assert run['params']['kumo.dropped_metrics'] == '6' + assert run['params']['kumo.dropped_params'] == '1' + assert 'note' not in run['params'] + assert run['params']['level'] == '2' + + +def test_quoted_run_names_are_taken_only_by_their_exact_name(): + from kumo.engine.integrations.snowflake import _taken + + assert _taken('"MyRun"', ['MyRun']) + assert not _taken('"MyRun"', ['MYRUN']) + assert _taken('my_run', ['MYXRUN']) + + +def test_the_users_session_is_not_altered(fake): + session = FakeSession() + + _replay(_tracking(session)) + + assert session.statements == [] + + +def test_a_hung_snowflake_call_does_not_stall_the_fit(fake, monkeypatch, + caplog): + import kumo.engine._tracking as tracking + + monkeypatch.setattr(tracking, 'FINISH_SECONDS', 0.5) + release = threading.Event() + monkeypatch.setattr(FakeExperimentTracking, 'log_metrics', + lambda *args, **kwargs: release.wait(30)) + began = time.monotonic() + try: + _replay(_tracking()) + elapsed = time.monotonic() - began + finally: + release.set() + join_tracking_threads() + + assert elapsed < 1.5 + assert 'did not finish within' in caplog.text + run, = fake._instance.runs.values() + assert run['status'] == 'FINISHED' diff --git a/packages/kumo-engine/test/test_tracking_contract.py b/packages/kumo-engine/test/test_tracking_contract.py new file mode 100644 index 000000000..5d04e8c2f --- /dev/null +++ b/packages/kumo-engine/test/test_tracking_contract.py @@ -0,0 +1,397 @@ +# 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.""" + +import json +import threading +import time +from dataclasses import dataclass, field +from pathlib import Path +from types import SimpleNamespace +from typing import Any, Callable, Dict, List, Mapping, Optional, Tuple + +import pytest + +from kumo.engine import Event +from kumo.engine._tracking import ( + DROPPED_METRICS, + DROPPED_PARAMS, + FitContext, + TrackingSession, + run_record, +) +from kumo.engine.tracking import RunRecord + +RUN = 'trainer-fit-1' +USER_TAGS = {'team': 'ents', 'kumo.run_id': 'overridden'} +QUERY = 'PREDICT ents.Y FOR EACH ents.ID' +INPUTS = { + 'ents': {'connector': 'S3', 'location': 's3://bucket/ents.parquet', + 'version': '1790000000', 'version_kind': 'timestamp'}, + 'orders': {'connector': 'DATABRICKS', + 'location': '`main`.`sales`.`orders`', + 'version': '7_1790000000', 'version_kind': 'exact'}, + 'users': {'connector': 'SNOWFLAKE', 'location': '"DB"."SCH"."USERS"', + 'version': '1790000000', 'version_kind': 'timestamp'}, +} +CONTEXT = FitContext( + run_id=RUN, query=QUERY, inputs=INPUTS, + model_plan={ + 'optimization': {'base_lr': [0.01, 0.001], 'max_epochs': 2}, + 'training_job': {'tune_metric': 'mae'}, + 'column_processing': { + 'encoder_overrides': {'ents.amount ($)': 'numerical'} + }, + }) + + +def _event(name, trial=None, epoch=None, metrics=None, **params): + return Event(name=name, run_id=RUN, time=0.0, trial=trial, epoch=epoch, + metrics=metrics or {}, params=params) + + +def _epochs(trial, *values): + return [ + _event('epoch_ended', trial, epoch, {'train_loss': loss, + 'val_mae': mae}) + for epoch, (loss, mae) in enumerate(values, 1) + ] + + +# Trial 2's first attempt fails; epoch events carry no attempt, as in core. +EVENTS = [ + _event('trial_started', 1, attempt=1, **{'optimization.base_lr': 0.01}), + *_epochs(1, (0.5, 0.4), (0.3, 0.35)), + _event('trial_ended', 1, metrics={'val_mae': 0.35, 'seconds': 3.0}, + status='succeeded', attempt=1), + _event('trial_started', 2, attempt=1, **{'optimization.base_lr': 0.001}), + _event('trial_ended', 2, status='failed', attempt=1, + error='OutOfMemoryError: CUDA out of memory'), + _event('trial_started', 2, attempt=2, **{'optimization.base_lr': 0.001, + 'optimization.batch_size': 256}), + *_epochs(2, (0.6, 0.5), (0.4, 0.3)), + _event('trial_ended', 2, metrics={'val_mae': 0.3, 'seconds': 4.0}, + status='succeeded', attempt=2), + _event('fit_ended', metrics={'val_mae': 0.3, 'test_mae': 0.32, + 'seconds': 9.0}, status='succeeded'), +] +RESULT = SimpleNamespace(metrics={'val_mae': 0.3, 'test_mae': 0.32}, + bundle_dir=Path('no-bundle')) + + +class _Recorder: + def __init__(self, tracker=None): + self.calls: List[Tuple[str, Any]] = [] + self.tracker = tracker + self.tags = getattr(tracker, 'tags', None) + + def prepare(self): + getattr(self.tracker, 'prepare', lambda: None)() + + def start(self, record): + self.calls.append(('start', record)) + if self.tracker is not None: + self.tracker.start(record) + + def on_trial(self, record, event): + self.calls.append(('on_trial', record)) + if self.tracker is not None: + self.tracker.on_trial(record, event) + + def finish(self, record, error): + self.calls.append(('finish', record)) + if self.tracker is not None: + self.tracker.finish(record, error) + + +def join_tracking_threads(): + for thread in threading.enumerate(): + if thread.name == 'kumo-tracking': + thread.join(30) + + +def replay(tracker=None, events=EVENTS, result=RESULT, + error=None) -> Dict[str, RunRecord]: + r"""Feeds the recorded fit to `tracker`; returns the records expected + of it.""" + recorder = _Recorder(tracker) + session = TrackingSession(recorder) + session.start(lambda: CONTEXT) + for event in events: + session.on_event(event) + session.finish(result, error) + records = { + record.tags['kumo.run_id']: record + for hook, record in recorder.calls if hook == 'on_trial' + } + records[RUN] = run_record(CONTEXT, result, events[-1], tags=USER_TAGS) + return records + + +@dataclass +class Observed: + name: Optional[str] + params: Mapping[str, str] + metrics: Mapping[str, float] + steps: Mapping[str, Tuple[Tuple[int, float], ...]] + tags: Mapping[str, str] + parent: Optional[str] = None + artifacts: Mapping[str, str] = field(default_factory=dict) + + +@dataclass(frozen=True) +class Adapter: + make: Callable[[], Any] + observe: Callable[[], Dict[str, Observed]] + kept_params: Callable[[Mapping[str, str]], Mapping[str, str]] = dict + nested_runs: bool = False + step_metrics: bool = False + artifacts: bool = False + + +def gaps(expected: Dict[str, RunRecord], observed: Dict[str, Observed], + adapter: Adapter) -> List[str]: + kept_params = adapter.kept_params + found = [] + if set(observed) != set(expected): + return [f'runs {sorted(observed)} != {sorted(expected)}'] + dropped = sum(len(record.params) - len(kept_params(record.params)) + for record in expected.values()) + counted = sum(int(seen.tags.get(DROPPED_PARAMS, 0)) + for seen in observed.values()) + if counted != dropped: + found.append(f'{DROPPED_PARAMS}: {counted} != {dropped}') + for run_id, record in expected.items(): + seen = observed[run_id] + + def differ(what, have, want): + if 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('params', dict(seen.params), dict(kept_params(record.params))) + differ('metrics', {k: seen.metrics.get(k) for k in record.metrics}, + dict(record.metrics)) + names = set(record.metrics) + if adapter.step_metrics: + names |= set(record.steps) + differ('steps', {k: seen.steps.get(k) for k in record.steps}, + dict(record.steps)) + differ('metric names', sorted(seen.metrics), sorted(names)) + differ('tags', {k: seen.tags.get(k) for k in record.tags}, + dict(record.tags)) + differ( + 'kumo tags', + sorted(k for k in seen.tags if k.startswith('kumo.') + and k not in (DROPPED_PARAMS, DROPPED_METRICS)), + sorted(k for k in record.tags if k.startswith('kumo.'))) + if adapter.nested_runs: + differ('parent', seen.parent, record.tags.get('kumo.parent_run')) + if adapter.artifacts: + differ('artifacts', dict(seen.artifacts), dict(record.artifacts)) + return found + + +@pytest.fixture +def mlflow_adapter(tmp_path, monkeypatch): + mlflow = pytest.importorskip('mlflow') + from kumo.engine.integrations.mlflow import MlflowTracking, _mlflow_params + + monkeypatch.chdir(tmp_path) + previous = mlflow.get_tracking_uri() + mlflow.set_tracking_uri(f'sqlite:///{tmp_path / "mlflow.db"}') + + def observe(): + client = mlflow.MlflowClient() + runs = client.search_runs( + [e.experiment_id for e in client.search_experiments()]) + kumo_ids = {r.info.run_id: r.data.tags['kumo.run_id'] for r in runs} + + def history(run_id, key): + return tuple((m.step, m.value) for m in sorted( + client.get_metric_history(run_id, key), + key=lambda m: (m.step, m.timestamp))) + + return { + kumo_ids[run.info.run_id]: Observed( + name=run.info.run_name, params=run.data.params, + metrics=run.data.metrics, tags=run.data.tags, + steps={k: history(run.info.run_id, k) + for k in run.data.metrics}, + parent=kumo_ids.get(run.data.tags.get('mlflow.parentRunId')), + artifacts={ + a.path: mlflow.artifacts.load_text( + f'{run.info.artifact_uri}/{a.path}') + for a in client.list_artifacts(run.info.run_id) + }) + for run in runs + } + + try: + yield Adapter( + make=lambda: MlflowTracking(log_system_metrics=False, + tags=USER_TAGS), observe=observe, + kept_params=_mlflow_params, nested_runs=True, step_metrics=True, + artifacts=True) + finally: + while mlflow.active_run() is not None: + mlflow.end_run() + mlflow.set_tracking_uri(previous) + + +@pytest.fixture +def snowflake_adapter(monkeypatch): + import _snowflake_experiment_fake + + fake = _snowflake_experiment_fake.install(monkeypatch) + from kumo.engine.integrations.snowflake import SnowflakeTracking + + def observe(): + observed = {} + for run in fake._instance.runs.values(): + names: Dict[Optional[str], Dict[str, Dict[str, Any]]] = {} + + def bucket(key): + prefix, _, name = key.rpartition('/') + return names.setdefault(prefix or None, { + 'params': {}, 'tags': {}, 'metrics': {}, 'steps': {} + }), name + + for key, value in run['params'].items(): + runs, name = bucket(key) + # Snowflake has no tags: they are params of the same name. + if name.startswith('kumo.') or name in USER_TAGS: + runs['tags'][name] = value + else: + runs['params'][name] = value + for key, value, step in run['metrics']: + runs, name = bucket(key) + runs['metrics'][name] = value + runs['steps'][name] = (*runs['steps'].get(name, ()), + (step, value)) + for prefix, runs in names.items(): + observed[runs['tags']['kumo.run_id']] = Observed( + name=prefix, + artifacts=run['artifacts'] if prefix is None else {}, + **runs) + return observed + + return Adapter( + make=lambda: SnowflakeTracking(_snowflake_experiment_fake.FakeSession(), + experiment_name='ents', tags=USER_TAGS), + observe=observe, step_metrics=True, artifacts=True) + + +@pytest.fixture(params=['mlflow', 'snowflake']) +def adapter(request) -> Adapter: + return request.getfixturevalue(f'{request.param}_adapter') + + +def test_records_of_the_recorded_fit(): + records = replay() + + fit = records[RUN] + assert fit.tags == {'team': 'ents', 'kumo.kind': 'fit', 'kumo.run_id': RUN} + assert fit.params == { + 'optimization.base_lr': '[0.01, 0.001]', + 'optimization.max_epochs': '2', + 'training_job.tune_metric': 'mae', + 'column_processing.encoder_overrides.ents.amount ($)': 'numerical', + } + with pytest.warns(UserWarning, + match="'kumo.bundle_id', 'training_job.tune_metric'"): + spoofed = run_record(CONTEXT, tags={'kumo.bundle_id': 'x', + 'training_job.tune_metric': 'y', + 'a': 1}) + assert spoofed.tags.keys() - fit.tags.keys() == {'a'} + assert fit.metrics == {'val_mae': 0.3, 'test_mae': 0.32, 'seconds': 9.0} + assert fit.artifacts['query.pql'] == QUERY + assert json.loads(fit.artifacts['model_plan.json']) == CONTEXT.model_plan + assert fit.inputs == INPUTS + + oom = records[f'{RUN}-trial-2-attempt-1'] + assert oom == RunRecord( + name='trial-2-attempt-1', + params={'optimization.base_lr': '0.001', 'attempt': '1'}, tags={ + 'kumo.kind': 'trial', + 'kumo.run_id': f'{RUN}-trial-2-attempt-1', + 'kumo.parent_run': RUN, + 'kumo.status': 'failed', + 'kumo.error': 'OutOfMemoryError: CUDA out of memory', + }) + retry = records[f'{RUN}-trial-2-attempt-2'] + assert retry.params['attempt'] == '2' + assert retry.metrics == {'best_val_mae': 0.3, 'seconds': 4.0} + assert retry.steps == {'train_loss': ((1, 0.6), (2, 0.4)), + 'val_mae': ((1, 0.5), (2, 0.3))} + + +def test_adapter_records_the_same_runs(adapter, tmp_path): + (tmp_path / 'provenance.json').write_text('{"bundle": "b"}') + result = SimpleNamespace(metrics=RESULT.metrics, bundle_dir=tmp_path) + expected = replay(adapter.make(), result=result) + + assert expected[RUN].artifacts['provenance.json'] == '{"bundle": "b"}' + assert gaps(expected, adapter.observe(), adapter) == [] + + +def test_a_raising_tracker_is_ended_and_not_called_again(caplog): + class RaisesOnTrial(_Recorder): + def on_trial(self, record, event): + super().on_trial(record, event) + raise RuntimeError('tracker down') + + broken = RaisesOnTrial() + replay(broken) + + assert [hook for hook, _ in broken.calls] == [ + 'start', 'on_trial', 'finish' + ] + assert broken.calls[-1] == ('finish', None) + assert caplog.text.count('off for the rest of this run') == 1 + + +def test_adapter_records_a_failed_fit(adapter): + failed = _event('fit_ended', status='failed', error='ValueError: stop') + expected = replay(adapter.make(), [*EVENTS[:-1], failed], None, + ValueError('stop')) + + assert expected[RUN].tags['kumo.error'] == 'ValueError: stop' + assert gaps(expected, adapter.observe(), adapter) == [] + + +def test_only_rank_zero_tracks(monkeypatch): + monkeypatch.setenv('LOCAL_RANK', '1') + recorder = _Recorder() + session = TrackingSession(recorder) + + session.start(lambda: CONTEXT) + session.finish(RESULT, None) + + assert recorder.calls == [] + + +def test_a_hung_tracker_does_not_stall_the_fit(monkeypatch, caplog): + import kumo.engine._tracking as tracking + + monkeypatch.setattr(tracking, 'FINISH_SECONDS', 0.5) + release = threading.Event() + + class Hung(_Recorder): + def start(self, record): + super().start(record) + release.wait(30) + + hung = Hung() + began = time.monotonic() + try: + replay(hung) + elapsed = time.monotonic() - began + finally: + release.set() + join_tracking_threads() + + assert elapsed < 1.5 + assert 'did not finish within' in caplog.text diff --git a/packages/kumo-engine/test/test_trainer_registry.py b/packages/kumo-engine/test/test_trainer_registry.py index f8a775afc..2c509eccd 100644 --- a/packages/kumo-engine/test/test_trainer_registry.py +++ b/packages/kumo-engine/test/test_trainer_registry.py @@ -141,11 +141,11 @@ def test_write_then_read_training_spec_round_trips(tmp_path: Path) -> None: def test_training_spec_manifest_update_is_atomic_and_retryable( tmp_path: Path, monkeypatch) -> None: - import kumo.engine.trainer._spec as spec_module + import kumo_core.io.filesystem as fs_module bundle = _fake_published_bundle(tmp_path) original_manifest = (bundle / MANIFEST_FILE).read_bytes() - real_replace = spec_module.os.replace + real_replace = fs_module.os.replace failed = False def fail_first_manifest_replace(source, destination): @@ -155,7 +155,7 @@ def fail_first_manifest_replace(source, destination): raise OSError('simulated interrupted manifest commit') return real_replace(source, destination) - monkeypatch.setattr(spec_module.os, 'replace', fail_first_manifest_replace) + monkeypatch.setattr(fs_module.os, 'replace', fail_first_manifest_replace) with pytest.raises(OSError, match='interrupted manifest commit'): write_training_spec(bundle, predictive_query='PREDICT COUNT(*)', model_plan=ModelPlan(), snapshot_id='snap-123', @@ -210,13 +210,13 @@ def test_finish_fit_persists_spec_without_export_before_exposing_state( trainer = Trainer(ModelPlan()) result = TrainingResult({'auc': .9}, {}, bundle) - assert trainer._finish_fit(result, predictive_query='PREDICT COUNT(*)', - snapshot_id='snap-123', - training_table_generation=_generation_plan(), - export_config=None) is result + finished = trainer._finish_fit( + result, predictive_query='PREDICT COUNT(*)', snapshot_id='snap-123', + training_table_generation=_generation_plan(), export_config=None) assert (bundle / SPEC_FILE).is_file() - assert trainer._training_result is result + assert finished.bundle_id == Trainer.load(bundle).bundle_id + assert trainer._training_result is finished assert trainer._snapshot_id == 'snap-123' assert trainer._training_query == 'PREDICT COUNT(*)' assert result.published_model is None @@ -238,7 +238,11 @@ def test_finish_fit_forwards_publication_details_and_returns_assigned_version( class Sink: def publish_model(self, connector, bundle_dir, **kwargs): spec = json.loads((Path(bundle_dir) / SPEC_FILE).read_text()) - assert kwargs['metadata'] == { + assert { + key: value + for key, value in kwargs['metadata'].items() + if not key.startswith('kumo_') + } == { name: str(spec[name]) for name in ('spec_version', 'predictive_query', 'snapshot_id') } @@ -258,6 +262,7 @@ def publish_model(self, connector, bundle_dir, **kwargs): assert calls[0][2]['path'] == path assert calls[0][2]['location'] == location assert calls[0][2]['metrics'] == {'val_auroc': .9} + assert calls[0][2]['metadata']['kumo_bundle_id'] == result.bundle_id assert result.published_model is published assert result.metrics == original.metrics assert result.bundle_dir == bundle @@ -567,3 +572,113 @@ def fake_predict(**kwargs) -> None: binary_classification_threshold=.5) assert result.frame['TARGET_PRED'].tolist() == [True] + + +def _export(sink, bundle, location=None, artifact_paths=None): + return Trainer(ModelPlan())._finish_fit( + TrainingResult({}, artifact_paths or {}, + bundle), predictive_query='PREDICT COUNT(*)', + snapshot_id='snap-123', training_table_generation=_generation_plan(), + export_config=SimpleNamespace( + connector=SimpleNamespace(sink_backend=sink), + path=None if location else 'm', location=location)) + + +def _recording_sink(calls): + return SimpleNamespace(publish_model=lambda *a, **k: calls.append(k) or + PublishedModel(name='m', version='1', uri='u')) + + +def test_publish_records_only_the_bundle_id_pointer(tmp_path: Path) -> None: + calls = [] + result = _export(_recording_sink(calls), _fake_published_bundle(tmp_path)) + assert {k: v + for k, v in calls[0]['metadata'].items() + if k.startswith('kumo_')} == {'kumo_bundle_id': result.bundle_id} + + +class _AliasingSink(_FakeSink): + def __init__(self, recorded: Optional[str] = None, + fail: bool = False) -> None: + super().__init__() + self.recorded = recorded + self.fail = fail + self.aliases: dict = {} + + def set_model_alias(self, connector, location, alias) -> None: + if self.fail: + raise RuntimeError('registry unavailable') + self.aliases[alias] = location.version + + def fetch_model(self, connector, local_dir, **kwargs) -> PublishedModel: + shutil.copytree(self.published['CHURN_GNN@latest'], local_dir, + dirs_exist_ok=True) + return PublishedModel(name='CHURN_GNN', version='V1', uri='fake://', + metadata={'kumo_bundle_id': self.recorded}) + + +@pytest.mark.parametrize('sink,warning', [ + (_AliasingSink(), None), + (_AliasingSink(fail=True), 'registry unavailable'), + (_FakeSink(), 'cannot set model aliases'), +]) +def test_export_with_alias_publishes_then_sets_it_best_effort( + tmp_path: Path, sink, warning) -> None: + import warnings + + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter('always') + result = _export(sink, _fake_published_bundle(tmp_path), + ModelRegistryLocation('CHURN_GNN', alias='prod')) + assert result.published_model.version == 'V1' + assert list(sink.published) == ['CHURN_GNN@latest'] + messages = [str(w.message) for w in caught if 'alias' in str(w.message)] + if warning is None: + assert sink.aliases == {'prod': 'V1'} and messages == [] + else: + assert len(messages) == 1 and warning in messages[0] + + +def test_load_by_alias_warns_on_recorded_bundle_mismatch( + tmp_path: Path) -> None: + bundle = _fake_published_bundle(tmp_path) + write_training_spec(bundle, predictive_query='PREDICT COUNT(*)', + model_plan=ModelPlan(), snapshot_id='snap-123', + training_table_generation=_generation_plan()) + sink = _AliasingSink('other') + sink.publish_model(None, str(bundle), + location=ModelRegistryLocation('CHURN_GNN')) + with pytest.warns(UserWarning, match='registry records bundle other'): + loaded = Trainer.load_from_connector( + _FakeConnector(sink), + location=ModelRegistryLocation('CHURN_GNN', alias='prod')) + assert loaded.bundle_id == Trainer.load(bundle).bundle_id + + +def test_alias_needs_an_aliasing_sink() -> None: + from kumo.engine import set_model_alias + + connector = _FakeConnector(_FakeSink()) + with pytest.raises(NotImplementedError, match='alias'): + Trainer.load_from_connector( + connector, location=ModelRegistryLocation('m', alias='prod')) + with pytest.raises(NotImplementedError, match='alias'): + set_model_alias(connector, ModelRegistryLocation('m', version='1'), + 'prod') + for kwargs in (dict(alias='prod', version='3'), dict(alias='')): + with pytest.raises(ValueError, match='alias'): + ModelRegistryLocation('m', **kwargs) + + +def test_invalid_export_alias_is_rejected_before_publishing( + tmp_path: Path) -> None: + sink = _AliasingSink() + trainer = Trainer(ModelPlan()) + trainer._training_result = TrainingResult({}, {}, + _fake_published_bundle(tmp_path)) + with pytest.raises(ValueError, match='model alias'): + trainer.export( + SimpleNamespace( + connector=SimpleNamespace(sink_backend=sink), path=None, + location=ModelRegistryLocation('CHURN_GNN', alias='LATEST'))) + assert sink.published == {} diff --git a/packages/kumo-engine/test/test_training_explanations.py b/packages/kumo-engine/test/test_training_explanations.py index 240c070cf..350c54444 100644 --- a/packages/kumo-engine/test/test_training_explanations.py +++ b/packages/kumo-engine/test/test_training_explanations.py @@ -49,7 +49,8 @@ def explain(**kwargs): best_trainer = SimpleNamespace( explain=Mock(side_effect=explain), subgraph_collector_dir=lambda: str(tmp_path / 'subgraph_collector'), - log_dir=str(tmp_path), + get_log_dir=lambda fit_mode: str(tmp_path), + published_fit_mode=lambda: None, ) def aggregate_features(*args): diff --git a/packages/kumo-engine/test/test_training_table_export.py b/packages/kumo-engine/test/test_training_table_export.py index cedee0e8c..7bd2e20f5 100644 --- a/packages/kumo-engine/test/test_training_table_export.py +++ b/packages/kumo-engine/test/test_training_table_export.py @@ -4,7 +4,9 @@ import json from collections import Counter +from dataclasses import replace from functools import partial +from types import SimpleNamespace import numpy as np import pandas as pd @@ -17,6 +19,13 @@ from kumoapi.typing import Stype import kumo.engine as kumo +from kumo.engine._manifest import ( + SourceTableEntry, + TrainingTableManifest, + TrainingTableParents, + read_manifest, + write_manifest, +) from kumo.engine._scratch import get_scratch_dir from kumo.engine.trainer.output_config import TrainingTableOutputConfig @@ -225,9 +234,9 @@ def test_custom_link_table_ingests_multiple_weighted_parts( stage_source_table = _custom_source.stage_source_table def stage_multipart_source(*args, **kwargs): - path = stage_source_table(*args, **kwargs) + path, version = stage_source_table(*args, **kwargs) assert len(list(path.glob('*.parquet'))) > 1 - return path + return path, version monkeypatch.setattr(_custom_source, 'stage_source_table', stage_multipart_source) @@ -292,3 +301,132 @@ def test_custom_training_table_declares_integer_timestamp_unit(graph, unit, for stage in seeds: columns = ['ENTITY1', 'TIMESTAMP', 'TARGET1'] assert _rows(seeds[stage], columns) == _rows(original_seeds[stage], columns) + + +def test_manifest_and_export_sidecar(graph, tmp_path): + from kumo_core.training.export import export_train_table + + _, snapshot, connector, out = graph + query = 'PREDICT ents.LABEL FOR EACH ents.ID' + pquery, table, _ = _export(graph, query, 'lineage') + manifest = table.manifest + assert manifest.id == table.tt_id == _spec(table)['tt_id'] + assert manifest.parents == TrainingTableParents(snapshot=snapshot.id) + assert manifest.query == query + assert manifest.extra == {} + assert table.source_mode == 'pquery' + assert manifest.split_seed == 0 + + directory = out / 'lineage' + baseline = tmp_path / 'baseline' + export_train_table(str(table.staging_dir), str(baseline)) + for part in baseline.glob('*.parquet'): + assert (directory / part.name).read_bytes() == part.read_bytes() + assert json.loads( + (directory / '_kumo_manifest.json').read_bytes()) == manifest.to_dict() + + reloaded = kumo.TrainingTable.from_source_table( + connector['lineage'], TrainingTableSpec(split_col='SPLIT'), pq=pquery, + snapshot=snapshot) + manifest = reloaded.manifest + assert reloaded.source_mode == 'custom_tt' + assert manifest.parents.source_table == SourceTableEntry( + connector='LOCAL_FILE', + location=str(directory.resolve()), + version=connector.source_backend.get_data_version( + connector, 'lineage'), + version_kind='listing', + ) + assert _seeds(reloaded)[1] == _seeds(table)[1] + + +def test_training_table_manifest_round_trips_unchanged(graph, tmp_path): + _, table, _ = _export(graph, 'PREDICT ents.LABEL FOR EACH ents.ID', + 'round_trip') + newer = json.loads( + (table.staging_dir / 'training-table-manifest.json').read_text()) + newer['future'] = 1 + newer['parents']['label_source'] = 'x' + source = tmp_path / 'newer.json' + source.write_text(json.dumps(newer, indent=2, sort_keys=True) + '\n') + + manifest = read_manifest(source, TrainingTableManifest) + write_manifest(tmp_path / 'out.json', manifest) + + assert (tmp_path / 'out.json').read_bytes() == source.read_bytes() + assert manifest.extra == {'future': 1} + assert manifest.parents.extra == {'label_source': 'x'} + + +def test_custom_table_follows_the_snapshots_record_sources( + graph, monkeypatch): + graph_, snapshot, connector, _ = graph + pquery, _, _ = _export(graph, 'PREDICT ents.LABEL FOR EACH ents.ID', + 'unlocated') + + def location(snapshot): + table = kumo.TrainingTable.from_source_table( + connector['unlocated'], TrainingTableSpec(split_col='SPLIT'), + pq=pquery, snapshot=snapshot) + return table.manifest.parents.source_table.location + + try: + assert location(graph_.snapshot(record_sources=False)) == 'unlocated' + finally: + graph_.snapshot() + assert location(snapshot) != 'unlocated' + monkeypatch.setattr(kumo.GraphSnapshot, 'manifest', property(lambda _: None)) + assert location(replace(snapshot, record_sources=None)) == 'unlocated' + + +def test_manifest_write_failure_keeps_generation_and_export( + graph, monkeypatch): + from kumo.engine.pquery import training_table + + _, _, connector, out = graph + + def fail(*args, **kwargs): + raise OSError('read-only scratch') + + monkeypatch.setattr(training_table, 'write_manifest', fail) + with pytest.warns(UserWarning, match='read-only scratch'): + _, table, exported = _export( + graph, 'PREDICT ents.LABEL FOR EACH ents.ID', 'no_manifest') + assert table.manifest is None + assert not (out / 'no_manifest' / '_kumo_manifest.json').exists() + assert exported.num_rows == sum(_seeds(table)[1].values()) + + +def test_unreadable_manifest_still_exports(graph): + graph_, snapshot, connector, out = graph + table = kumo.PredictiveQuery( + graph=graph_, query='PREDICT ents.LABEL FOR EACH ents.ID' + ).generate_training_table(snapshot) + manifest = table.staging_dir / 'training-table-manifest.json' + manifest.chmod(0) + try: + with pytest.warns(UserWarning, match='_kumo_manifest.json'): + table.export(TrainingTableOutputConfig(connector, 'unreadable')) + finally: + manifest.chmod(0o644) + assert list((out / 'unreadable').glob('*.parquet')) + assert not (out / 'unreadable' / '_kumo_manifest.json').exists() + + +def test_no_sidecar_for_warehouse_or_csv_exports(graph): + graph_, snapshot, connector, out = graph + table = kumo.PredictiveQuery( + graph=graph_, query='PREDICT ents.LABEL FOR EACH ents.ID' + ).generate_training_table(snapshot) + calls = [] + + def write_files(connector, table_name, files, write_mode, **kwargs): + calls.append(kwargs) + + warehouse = SimpleNamespace( + source_type='SNOWFLAKE', + sink_backend=SimpleNamespace(write_files=write_files)) + table.export(TrainingTableOutputConfig(warehouse, 'DB.SCHEMA.TT')) + assert calls == [{'export_id': f'{table.tt_id}:tt-export-v1'}] + table.export(TrainingTableOutputConfig(connector, 'rows.csv/')) + assert not (out / 'rows.csv' / '_kumo_manifest.json').exists() diff --git a/packages/kumo-engine/test/test_warm_start.py b/packages/kumo-engine/test/test_warm_start.py index e79130045..3370f632f 100644 --- a/packages/kumo-engine/test/test_warm_start.py +++ b/packages/kumo-engine/test/test_warm_start.py @@ -182,6 +182,7 @@ def configure(self): patch.setattr(BaseLightningModule, 'configure_optimizers', configure) second = destination.fit(table, snapshot=snapshot, warm_start=source) assert restored_states and configured_states + assert second.provenance.parents.warm_start == source.bundle_id assert destination.model_plan == configured_plan assert _hash_bundle(durable) == original_hashes assert _hash_bundle(source_bundle) == original_hashes diff --git a/packages/kumo-snowflake/pyproject.toml b/packages/kumo-snowflake/pyproject.toml index 960687071..47547c8b9 100644 --- a/packages/kumo-snowflake/pyproject.toml +++ b/packages/kumo-snowflake/pyproject.toml @@ -27,13 +27,15 @@ dependencies = [ [project.optional-dependencies] # What the SDK plugin half needs beyond the platform reader: Snowpark for the -# session it opens, the model registry for what a sink publishes, and pyyaml -# for the manifest that carries. The platform installs this distribution +# session it opens, the model registry for what a sink publishes (and, from +# 2.0, the experiment tracking kumo-engine records fits in; capped below 3 +# because that adapter reads ExperimentTracking internals), and pyyaml for the +# manifest that carries. The platform installs this distribution # without this extra and imports `SnowflakeSourceReader` alone, so none of it # is loaded there -- `__init__` exports the plugin lazily to keep that true. sdk = [ "snowflake-snowpark-python", - "snowflake-ml-python", + "snowflake-ml-python>=2.0,<3", "pyyaml", ] test = [ diff --git a/packages/kumo-snowflake/src/kumo_snowflake/identity.py b/packages/kumo-snowflake/src/kumo_snowflake/identity.py index 7b17c9a78..e62270e42 100644 --- a/packages/kumo-snowflake/src/kumo_snowflake/identity.py +++ b/packages/kumo-snowflake/src/kumo_snowflake/identity.py @@ -7,7 +7,11 @@ from typing import Any, Iterable, Mapping from kumo_core.sources.errors import PermanentSourceError, SourceNotFoundError -from kumo_core.sources.identity import epoch_second, hour_version +from kumo_core.sources.identity import ( + SourceVersion, + epoch_second, + hour_version, +) from kumo_core.sources.sql import resolve_snowflake_identifier from kumo_snowflake.sql import placeholder, quote_identifier @@ -95,7 +99,8 @@ def data_version( altered = {str(row[0]): row[1] for row in cursor.fetchall()} if table not in altered: return hour_version() - return epoch_second(pd.to_datetime(altered[table])) + return SourceVersion(epoch_second(pd.to_datetime(altered[table])), + "timestamp") class MissingColumnTypesError(PermanentSourceError): diff --git a/packages/kumo-snowflake/src/kumo_snowflake/registry.py b/packages/kumo-snowflake/src/kumo_snowflake/registry.py index fd220653e..16d14d983 100644 --- a/packages/kumo-snowflake/src/kumo_snowflake/registry.py +++ b/packages/kumo-snowflake/src/kumo_snowflake/registry.py @@ -33,11 +33,13 @@ """ import json +import logging import shutil from contextlib import contextmanager from pathlib import Path from typing import ( TYPE_CHECKING, + Any, ContextManager, Dict, Iterator, @@ -53,9 +55,13 @@ resolved, ) from kumoapi.plugin import ( + BUNDLE_ID_KEY, + ENGINE_MARKER, ModelRegistryLocation, PublishedModel, SourceConnector, + check_model_alias, + registry_lineage, ) if TYPE_CHECKING: @@ -74,8 +80,9 @@ #: merely happens to be named alike. _MANIFEST_FILE = 'publication-manifest.json' -#: Engine-bundle marker stored through the registry metadata API. -_ENGINE_MARKER = 'kumo_engine' +_SYSTEM_ALIASES = frozenset({'DEFAULT', 'FIRST', 'LAST'}) + +_LOG = logging.getLogger(__name__) #: Where a fetched version is unpacked before its bundle is lifted out. Named #: rather than a temporary directory so a half-finished import leaves something @@ -209,6 +216,40 @@ def _refuse_path(path: Optional[str], action: str) -> None: f'a path.') +def _version_record(metadata: Mapping[str, str]) -> Dict[str, str]: + record = {ENGINE_MARKER: '1'} + if metadata.get(BUNDLE_ID_KEY): + record[BUNDLE_ID_KEY] = str(metadata[BUNDLE_ID_KEY]) + return record + + +def _recorded_lineage(version: "ModelVersion") -> Dict[str, str]: + try: + return registry_lineage(dict(version.show_metrics() or {})) + except Exception as error: + _LOG.warning('Could not read registry metadata: %s', error) + return {} + + +def _alias_holders(model: Any) -> Dict[str, str]: + holders = {} + for row in model.show_versions().to_dict('records'): + for alias in json.loads(row.get('aliases') or '[]'): + holders[alias] = str(row['name']) + return holders + + +def _restore_alias(model: Any, version: str, alias: str) -> None: + try: + model.version(version).set_alias(alias) + except Exception as error: + _LOG.warning('Alias %s is now unset: restoring it on version %s ' + 'failed: %s', alias, version, error) + else: + _LOG.warning('Could not move alias %s; it stays on version %s.', + alias, version) + + class SnowflakeModelRegistryMixin: r"""The model half of :class:`SnowflakeSinkBackend`. @@ -267,8 +308,8 @@ def publish_model( the model, and optionally the version. A version of ``None`` lets the registry assign one -- the usual case, and the reason this reports back what it created. - metadata: Returned with the publication result. Training metadata - remains in the bundle and is not recorded in registry tags. + metadata: Returned with the publication result; its + ``kumo_bundle_id`` is also recorded in the version's metrics. metrics: Returned with the publication result. Training metrics remain in the bundle's metrics.json. project_id: Ignored; a local runtime has no projects. @@ -288,13 +329,14 @@ def publish_model( f"{bundle} does not look like a published model bundle: it " f"has no {_SPEC_FILE}, which an export writes into every one.") + record = _version_record(metadata) with self._registry(connector) as registry: version = registry.log_model( _bundle_model(str(bundle)), model_name=location.model_name, version_name=location.version, comment=location.comment, - metrics={_ENGINE_MARKER: '1'}, + metrics=record, signatures=_carrier_signature(), ) return _published(connector, version, metadata=metadata, @@ -308,10 +350,11 @@ def list_models( ) -> List[PublishedModel]: r"""Lists Kumo model versions visible in the connector's registry. - A Kumo export records only ``kumo_engine=1`` in registry metadata. + A Kumo export records ``kumo_engine=1`` in registry metadata. That marker makes discovery both cheap and precise: versions logged by another tool are omitted, and no bundle is downloaded merely to decide whether it belongs in this list. + ``metadata`` carries the recorded ``kumo_bundle_id``. """ del project_id found: List[PublishedModel] = [] @@ -319,9 +362,11 @@ def list_models( for model in registry.models(): for version in model.versions(): recorded = dict(version.show_metrics() or {}) - if recorded.get(_ENGINE_MARKER) != '1': + if recorded.get(ENGINE_MARKER) != '1': continue - found.append(_published(connector, version)) + found.append( + _published(connector, version, + metadata=registry_lineage(recorded))) return sorted(found, key=lambda item: (item.name, item.version)) def fetch_model( @@ -345,14 +390,15 @@ def fetch_model( though the caller always hands this a fresh, empty ``mkdtemp``. path: Refused, as in :meth:`publish_model`. - location: The model, and optionally the version. ``None`` takes the - registry's default version, which is what a caller who names - none means. + location: The model, and optionally the version or an alias. With + neither, takes the registry's default version, which is what a + caller who names none means. project_id: Ignored. Returns: The name, the version actually read -- resolved, when the caller - named none -- where it lives, and metrics read from metrics.json. + named none -- where it lives, metrics read from metrics.json, and + the version's recorded ``kumo_bundle_id`` as ``metadata``. """ del project_id _refuse_path(path, 'fetching') @@ -366,8 +412,18 @@ def fetch_model( # Resolved before anything on disk is touched: a name the registry # does not know must not cost the caller the bundle already there. model = registry.get_model(location.model_name) - version = (model.default if location.version is None else - model.version(location.version)) + alias = getattr(location, 'alias', None) + if alias is not None: + if resolved(alias) not in _alias_holders(model): + raise ValueError( + f'No version of {location.model_name} holds alias ' + f'{alias}.') + version = model.version(alias) + elif location.version is None: + version = model.default + else: + version = model.version(location.version) + recorded = _recorded_lineage(version) destination.parent.mkdir(parents=True, exist_ok=True) staged = destination.parent / f'{destination.name}{_EXPORT_DIR}' @@ -385,9 +441,52 @@ def fetch_model( isinstance(key, str) and isinstance(value, (int, float)) for key, value in metrics.items())): raise ValueError('Invalid training metrics in model bundle.') - published = _published(connector, version, metrics=metrics) + published = _published(connector, version, metadata=recorded, + metrics=metrics) if destination.exists(): shutil.rmtree(destination) shutil.move(str(bundle), str(destination)) shutil.rmtree(staged, ignore_errors=True) return published + + def set_model_alias( + self, + connector: SourceConnector, + location: ModelRegistryLocation, + alias: str, + ) -> None: + r"""Points :obj:`alias` at ``location.version``, moving it if held. + + Snowflake allows one alias per version. A failed move puts the alias + back on the version that held it. + """ + if location.version is None: + raise ValueError('Setting an alias needs a ModelRegistryLocation ' + 'naming an exact version.') + check_model_alias(alias) + wanted = resolved(alias) + with self._registry(connector) as registry: + model = registry.get_model(location.model_name) + target = model.version(location.version) + name = resolved(str(target.version_name)) + holders = { + held: version + for held, version in _alias_holders(model).items() + if held not in _SYSTEM_ALIASES + } + if holders.get(wanted) == name: + return + held = sorted(a for a, v in holders.items() if v == name) + if held: + raise ValueError( + f'Version {name} of {location.model_name} already holds ' + f'alias {held[0]}; unset it first.') + previous = holders.get(wanted) + if previous is not None: + model.version(previous).unset_alias(alias) + try: + target.set_alias(alias) + except Exception: + if previous is not None: + _restore_alias(model, previous, alias) + raise diff --git a/packages/kumo-snowflake/src/kumo_snowflake/session.py b/packages/kumo-snowflake/src/kumo_snowflake/session.py index 0e237c868..08b8c4ebf 100644 --- a/packages/kumo-snowflake/src/kumo_snowflake/session.py +++ b/packages/kumo-snowflake/src/kumo_snowflake/session.py @@ -461,6 +461,16 @@ def _principal(self, connector: SourceConnector) -> Principal: f'accepted, as {self._default_hint}.'), ) + def table_location(self, connector: SourceConnector, + table_name: str) -> str: + """``account/"DB"."SCHEMA"."TABLE"``; the account when it resolves.""" + try: + account = _snowflake_config(self._principal(connector)).account + except Exception: + account = None + qualified = self._qualified(connector, table_name) + return f'{account}/{qualified}' if account else qualified + @staticmethod def _qualified(connector: SourceConnector, table_name: str) -> str: return (f'{_quoted(resolved(_target_config(connector).database))}.' diff --git a/packages/kumo-snowflake/tests/test_identity.py b/packages/kumo-snowflake/tests/test_identity.py index f4ba9fd8d..eade78ae6 100644 --- a/packages/kumo-snowflake/tests/test_identity.py +++ b/packages/kumo-snowflake/tests/test_identity.py @@ -30,6 +30,7 @@ def stamp(recorder=None, table=TABLE, **overrides): def test_a_table_is_stamped_with_the_second_it_was_last_altered(): """A version carrying `.0` is a different path segment and a different key.""" assert stamp() == ALTERED_VERSION + assert stamp().kind == "timestamp" assert stamp(last_altered=ALTERED.replace( microsecond=900000)) == (ALTERED_VERSION) @@ -46,6 +47,7 @@ def test_a_differently_cased_table_is_not_substituted(): def test_a_view_of_either_kind_falls_back_to_the_hour_window(table_type): recorder = cursor(table_type=table_type) assert stamp(recorder) == hour_version() + assert stamp(cursor(table_type=table_type)).kind == "hour_bucket" assert recorder.matching("LAST_ALTERED") == [] diff --git a/packages/kumo-snowflake/tests/test_registry.py b/packages/kumo-snowflake/tests/test_registry.py index 9d07e6c28..9ed7f1c94 100644 --- a/packages/kumo-snowflake/tests/test_registry.py +++ b/packages/kumo-snowflake/tests/test_registry.py @@ -8,6 +8,7 @@ from pathlib import Path from types import SimpleNamespace +import pandas as pd import pytest from kumo_snowflake import registry as registry_module from kumo_snowflake.registry import SnowflakeModelRegistryMixin @@ -30,7 +31,11 @@ def export(directory): version = SimpleNamespace( model_name='CHURN', version_name='V1', fully_qualified_model_name='DB.SCHEMA.CHURN', export=export, - show_metrics=lambda: pytest.fail('read registry metrics')) + show_metrics=lambda: { + 'kumo_engine': '1', + 'kumo_bundle_id': 'b1', + 'auc': 0.5 + }) model = SimpleNamespace(default=version, version=lambda _: version) @contextmanager @@ -60,6 +65,7 @@ def test_fetch_returns_bundle_metrics(fetch, requested_version, metrics): payload = json.dumps(metrics) if metrics is not None else None published = run(payload, requested_version) assert published.metrics == (metrics or {}) + assert published.metadata == {'kumo_bundle_id': 'b1'} assert published.version == 'V1' assert (destination / 'weights.bin').read_bytes() == b'model weights' assert not (destination / 'previous').exists() @@ -75,9 +81,13 @@ def test_invalid_metrics_preserve_previous_destination(fetch, payload): assert (destination / 'previous').read_bytes() == b'previous bundle' -@pytest.mark.parametrize('metadata', [{}, {'predictive_query': 'x' * 300}]) -def test_publish_records_only_engine_marker_in_catalog(tmp_path, monkeypatch, - metadata): +@pytest.mark.parametrize('metadata', [{}, { + 'predictive_query': 'x' * 300, + 'kumo_bundle_id': 'b1', + 'kumo_engine': '0' +}]) +def test_publish_records_only_pointers_in_catalog( + tmp_path, monkeypatch, metadata): bundle = tmp_path / 'bundle' bundle.mkdir() (bundle / 'training-spec.json').write_text('{}') @@ -99,16 +109,20 @@ def registry(_): monkeypatch.setattr(registry_module, '_carrier_signature', lambda: {}) published = sink.publish_model(SimpleNamespace(), str(bundle), location=ModelRegistryLocation('CHURN'), - metadata=metadata, metrics={'auc': 0.92}) - assert calls[0]['metrics'] == {'kumo_engine': '1'} + metadata=metadata, + metrics={'auc': 0.92, 'loss': float('nan')}) + assert calls[0]['metrics'] == { + **({'kumo_bundle_id': 'b1'} if metadata else {}), 'kumo_engine': '1' + } assert published.metadata == metadata - assert published.metrics == {'auc': 0.92} + assert published.metrics['auc'] == 0.92 -def test_listing_filters_engine_marker_and_returns_empty_metrics(monkeypatch): +def test_listing_filters_engine_marker_and_returns_lineage(monkeypatch): versions = [] for name, recorded in [('V1', { 'kumo_engine': '1', + 'kumo_bundle_id': 'b1', 'auc': 0.92 }), ('V2', { 'other': '1', @@ -134,4 +148,112 @@ def registry(_): found = sink.list_models(SimpleNamespace()) assert [item.version for item in found] == ['V1', 'V4'] assert [item.metrics for item in found] == [{}, {}] - assert all(item.metadata == {} for item in found) + assert [item.metadata for item in found] == [{'kumo_bundle_id': 'b1'}, {}] + + +class _Model: + def __init__(self, aliases): + self.aliases = aliases + + def show_versions(self): + return pd.DataFrame([ + dict(name=name, aliases=json.dumps(held)) + for name, held in self.aliases.items() + ]) + + def version(self, version_or_alias): + name = next((name for name, held in self.aliases.items() + if version_or_alias.upper() in held), version_or_alias) + held = self.aliases[name] + + def export(directory): + bundle = Path(directory) / 'bundle' + bundle.mkdir() + (bundle / 'training-spec.json').write_text(name) + (bundle / 'publication-manifest.json').write_text('{}') + + return SimpleNamespace( + version_name=name, model_name='CHURN', + fully_qualified_model_name='DB.SCHEMA.CHURN', export=export, + set_alias=lambda alias: held.append(alias.upper()), + unset_alias=lambda alias: held.remove(alias.upper()), + show_metrics=lambda: { + 'kumo_engine': '1', + 'kumo_bundle_id': f'id-{name}' + }) + + +def _aliasing_sink(monkeypatch, model): + @contextmanager + def registry(_): + yield SimpleNamespace(get_model=lambda _: model) + + sink = SnowflakeModelRegistryMixin() + monkeypatch.setattr(sink, '_registry', registry) + return sink + + +def test_set_alias_moves_it_and_refuses_what_snowflake_cannot(monkeypatch): + from kumo_snowflake.sink import SnowflakeSinkBackend + from kumoapi.plugin import ModelAliasingSink + + assert isinstance(SnowflakeSinkBackend(), ModelAliasingSink) + model = _Model({'V1': ['FIRST'], 'V2': ['LAST', 'DEFAULT'], 'V3': []}) + sink = _aliasing_sink(monkeypatch, model) + + def set_alias(version, alias): + sink.set_model_alias(SimpleNamespace(), + ModelRegistryLocation('CHURN', version=version), + alias) + + set_alias('V2', 'prod') + set_alias('V3', 'prod') + set_alias('V3', 'prod') + assert model.aliases == { + 'V1': ['FIRST'], + 'V2': ['LAST', 'DEFAULT'], + 'V3': ['PROD'] + } + for version, alias, match in [('V3', 'staging', 'already holds'), + ('V1', 'default', 'reserved'), + (None, 'prod', 'exact version')]: + with pytest.raises(ValueError, match=match): + set_alias(version, alias) + assert model.aliases['V3'] == ['PROD'] + + +def test_fetch_by_alias_reads_the_aliased_version(monkeypatch, tmp_path): + sink = _aliasing_sink(monkeypatch, _Model({ + 'V1': ['DEFAULT'], + 'V2': ['PROD'] + })) + published = sink.fetch_model( + SimpleNamespace(), str(tmp_path / 'out'), + location=ModelRegistryLocation('CHURN', alias='prod')) + assert published.version == 'V2' + assert published.metadata == {'kumo_bundle_id': 'id-V2'} + assert (tmp_path / 'out' / 'training-spec.json').read_text() == 'V2' + with pytest.raises(ValueError, match='holds alias'): + sink.fetch_model(SimpleNamespace(), str(tmp_path / 'other'), + location=ModelRegistryLocation('CHURN', alias='V1')) + + +def test_failed_alias_move_puts_it_back(monkeypatch): + model = _Model({'V1': ['PROD'], 'V2': []}) + + def fail(alias): + raise RuntimeError('lost connection') + + def failing_version(name): + version = _Model.version(model, name) + if name == 'V2': + version.set_alias = fail + return version + + model.version = failing_version + sink = _aliasing_sink(monkeypatch, model) + with pytest.raises(RuntimeError, match='lost connection'): + sink.set_model_alias(SimpleNamespace(), + ModelRegistryLocation('CHURN', version='V2'), + 'prod') + assert model.aliases == {'V1': ['PROD'], 'V2': []} diff --git a/setup.py b/setup.py index 361976afc..64108b9a0 100644 --- a/setup.py +++ b/setup.py @@ -293,6 +293,8 @@ def parse_project_version(project_file_path: Path) -> str: "respx==0.22.0", "pytorch-frame==0.3.0", "pytorch-lightning==2.6.5", + # kumo-engine's `mlflow` extra; its wheel is tested with --no-deps. + "mlflow>=3.13,<4", ] dev_requires = test_requires + [ diff --git a/test/train/test_kumo_pred_ttng.py b/test/train/test_kumo_pred_ttng.py index d3baf2684..0380f8245 100644 --- a/test/train/test_kumo_pred_ttng.py +++ b/test/train/test_kumo_pred_ttng.py @@ -516,6 +516,7 @@ def _predict_lp(self, **kwargs): "ignore_time_col": False, "pred_time_col_name": None, "output_extra_fields": None, + "output_constant_columns": None, "max_output_file_size": 1024**3, "progress_log_dir": None, "progress_log_id": 0, diff --git a/uv.lock b/uv.lock index 9ba891169..318cf7d63 100644 --- a/uv.lock +++ b/uv.lock @@ -140,7 +140,7 @@ dependencies = [ { name = "frozenlist" }, { name = "multidict" }, { name = "propcache" }, - { name = "typing-extensions", marker = "python_full_version < '3.13' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "typing-extensions", marker = "python_full_version < '3.13'" }, { name = "yarl" }, ] sdist = { url = "https://files.pythonhosted.org/packages/6c/4c/bdccd81e9ee225b69c60e7766c9a5b05364f118f4d383713b89a682d772d/aiohttp-3.14.5.tar.gz", hash = "sha256:5558a7f5a05af9ecf744af91e5baefc436f93c9333e656c27ec253f9a6bbe178", size = 8078052, upload-time = "2026-10-11T01:05:12.408Z" } @@ -304,7 +304,7 @@ version = "1.4.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "frozenlist" }, - { name = "typing-extensions", marker = "python_full_version < '3.13' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "typing-extensions", marker = "python_full_version < '3.13'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/61/62/06741b579156360248d1ec624842ad0edf697050bbaf7c3e46394e106ad1/aiosignal-1.4.0.tar.gz", hash = "sha256:f47eecd9468083c2029cc99945502cb7708b082c232f9aca65da147157b251c7", size = 25007, upload-time = "2025-07-03T22:54:43.528Z" } wheels = [ @@ -376,7 +376,7 @@ version = "4.14.2" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "idna" }, - { name = "typing-extensions", marker = "python_full_version < '3.13' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "typing-extensions", marker = "python_full_version < '3.13'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/61/cc/a381afa6efea9f496eff839d4a6a1aed3bfafc7b3ab4b0d1b243a12573dd/anyio-4.14.2.tar.gz", hash = "sha256:cfa139f3ed1a23ee8f88a145ddb5ac7605b8bbfd8592baacd7ce3d8bb4313c7f", size = 260176, upload-time = "2026-07-12T20:29:07.082Z" } wheels = [ @@ -572,12 +572,21 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/88/c6/92fcd42f1ba33e1184263f25bfabf3d27c383410470f169e4b8163bf9c17/beautifulsoup4-4.15.0-py3-none-any.whl", hash = "sha256:d6f88de62e1d4e38ecb1077eb9724cd0eff29d2a08ca16a401e9b9e93f117cf9", size = 109924, upload-time = "2026-06-07T16:44:21.566Z" }, ] +[[package]] +name = "blinker" +version = "1.9.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/21/28/9b3f50ce0e048515135495f198351908d99540d69bfdc8c1d15b73dc55ce/blinker-1.9.0.tar.gz", hash = "sha256:b4ce2265a7abece45e7cc896e98dbebe6cead56bcf805a3d23136d145f5445bf", size = 22460, upload-time = "2024-11-08T17:25:47.436Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/10/cb/f2ad4230dc2eb1a74edf38f1a38b9b52277f75bef262d8908e60d957e13c/blinker-1.9.0-py3-none-any.whl", hash = "sha256:ba0efaa9080b619ff2f3459d1d500c57bddea4a6b424b60a91141db6fd2f08bc", size = 8458, upload-time = "2024-11-08T17:25:46.184Z" }, +] + [[package]] name = "blockbuster" version = "1.5.28" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "forbiddenfruit", marker = "implementation_name == 'cpython' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "forbiddenfruit", marker = "implementation_name == 'cpython'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/28/14/de0062a8e3758fe1880729c1828196e6e951b8cd3299dfdd1e68aeae6e63/blockbuster-1.5.28.tar.gz", hash = "sha256:5af2ed9de6c874339a76cc831d1d11db79800a974573580e52ad827a70803482", size = 98235, upload-time = "2026-09-18T10:51:11.263Z" } wheels = [ @@ -605,7 +614,7 @@ source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "botocore-stubs" }, { name = "types-s3transfer" }, - { name = "typing-extensions", marker = "python_full_version < '3.12' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "typing-extensions", marker = "python_full_version < '3.12'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/7b/94/9dd7426ce339c15ff6bae15b609977aad2f8f778244f9a06e99db821c974/boto3_stubs-1.42.90.tar.gz", hash = "sha256:d0ea4bc1352ba5ef8fae13cfb2fef511f372113ede0388b50cdd19811d39ce01", size = 102688, upload-time = "2026-04-16T21:25:36.455Z" } wheels = [ @@ -657,8 +666,8 @@ dependencies = [ { name = "numpy" }, { name = "pandas" }, { name = "scikit-learn" }, - { name = "scipy", version = "1.17.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, - { name = "scipy", version = "1.18.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "scipy", version = "1.17.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" }, + { name = "scipy", version = "1.18.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/6d/28/3325463e27238bacefce34902ce2e5883ea3bd4044ce93e744ef499a6d72/category_encoders-2.11.1.tar.gz", hash = "sha256:f21b99185800c19b562b2928a0fe67dbf5658d048cd70451bfe8811c4123fdb0", size = 76600, upload-time = "2026-09-08T15:29:09.823Z" } wheels = [ @@ -679,7 +688,7 @@ name = "cffi" version = "2.1.1" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "pycparser", marker = "implementation_name != 'PyPy' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "pycparser", marker = "implementation_name != 'PyPy'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/9e/ef/008a1939e372c06329a3fce4279c02f328488f3526744906eeec3da7ad5f/cffi-2.1.1.tar.gz", hash = "sha256:dd31f52ea1086513bb9df30f8fcee9b8918323ae067a3d5b78bc826a000712be", size = 530807, upload-time = "2026-08-03T21:21:18.939Z" } wheels = [ @@ -960,7 +969,7 @@ name = "colorlog" version = "6.12.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "colorama", marker = "sys_platform == 'win32' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "colorama", marker = "sys_platform == 'win32'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/8c/55/ba79756cb90c8d69d599d57785398ac87bba7b19c80e87f4e8a562197c93/colorlog-6.12.0.tar.gz", hash = "sha256:2a7924c1dadf18b22a0eb8b06d1c7b01d5341707ec1641eb6fcc4fde0c3e8e5f", size = 18151, upload-time = "2026-07-23T13:40:40.71Z" } wheels = [ @@ -981,7 +990,7 @@ resolution-markers = [ "python_full_version < '3.12' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu130' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests'", ] dependencies = [ - { name = "numpy", marker = "python_full_version < '3.12' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "numpy" }, ] sdist = { url = "https://files.pythonhosted.org/packages/58/01/1253e6698a07380cd31a736d248a3f2a50a7c88779a1813da27503cadc2a/contourpy-1.3.3.tar.gz", hash = "sha256:083e12155b210502d0bca491432bb04d56dc3432f95a979b429f2848c3dbe880", size = 13466174, upload-time = "2025-07-26T12:03:12.549Z" } wheels = [ @@ -1086,7 +1095,7 @@ resolution-markers = [ "python_full_version == '3.12.*' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu130' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests'", ] dependencies = [ - { name = "numpy", marker = "python_full_version >= '3.12' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "numpy" }, ] sdist = { url = "https://files.pythonhosted.org/packages/83/5a/a55177dd22553a277388e8a1b3220e92de91bacb28356cdc73caa240121d/contourpy-1.4.0.tar.gz", hash = "sha256:20156f5a1ac4f8ce02656e39a61e82164a3d359796dc8026f75b062783d500e1", size = 13323726, upload-time = "2026-09-11T19:05:05.808Z" } wheels = [ @@ -1277,7 +1286,7 @@ wheels = [ [package.optional-dependencies] toml = [ - { name = "tomli", marker = "python_full_version <= '3.11' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "tomli", marker = "python_full_version <= '3.11'" }, ] [[package]] @@ -1285,7 +1294,7 @@ name = "cryptography" version = "50.0.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "cffi", marker = "platform_python_implementation != 'PyPy' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "cffi", marker = "platform_python_implementation != 'PyPy'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/de/41/6cbdcf9142d00fe82836fbb51e503e58088575cf7a0fe1dbff6695bf0840/cryptography-50.0.0.tar.gz", hash = "sha256:eeac2acb5a20ed25e0ad6d1df9891a520b78b404266b6d11778f25d5d691a6c9", size = 880201, upload-time = "2026-07-31T14:25:10.11Z" } wheels = [ @@ -1341,7 +1350,7 @@ name = "cuda-bindings" version = "13.4.2" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "cuda-pathfinder", marker = "extra == 'extra-4-kumo-cu130' or extra == 'extra-4-kumo-cu132' or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra != 'extra-4-kumo-cpu' and extra != 'group-4-kumo-public-tests')" }, + { name = "cuda-pathfinder", marker = "extra == 'extra-4-kumo-cu130' or extra == 'extra-4-kumo-cu132' or (extra != 'extra-4-kumo-cpu' and extra != 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests')" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/19/6f/e00ffcebcad6405a2326a52612eff98240a67c637dfed516e61b881e2db3/cuda_bindings-13.4.2-cp311-cp311-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:fc0a18f65b26459cd3c49c0250a44b419cdf790cffcb096ff4d46a1973f6f4b7", size = 6484069, upload-time = "2026-09-17T22:05:52.594Z" }, @@ -1390,34 +1399,34 @@ wheels = [ [package.optional-dependencies] cudart = [ - { name = "nvidia-cuda-runtime", version = "13.0.96", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform == 'linux' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'linux' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-cuda-runtime", version = "13.0.96", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform == 'linux' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, ] cufft = [ - { name = "nvidia-cufft", version = "12.0.0.61", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform == 'linux' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'linux' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-cufft", version = "12.0.0.61", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform == 'linux' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, ] cufile = [ - { name = "nvidia-cufile", version = "1.15.1.6", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform == 'linux' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'linux' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-cufile", version = "1.15.1.6", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform == 'linux' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, ] cupti = [ - { name = "nvidia-cuda-cupti", version = "13.0.85", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform == 'linux' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'linux' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-cuda-cupti", version = "13.0.85", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform == 'linux' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, ] curand = [ - { name = "nvidia-curand", version = "10.4.0.35", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform == 'linux' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'linux' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-curand", version = "10.4.0.35", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform == 'linux' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, ] cusolver = [ - { name = "nvidia-cusolver", version = "12.0.4.66", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform == 'linux' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'linux' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-cusolver", version = "12.0.4.66", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform == 'linux' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, ] cusparse = [ - { name = "nvidia-cusparse", version = "12.6.3.3", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform == 'linux' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'linux' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-cusparse", version = "12.6.3.3", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform == 'linux' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, ] nvjitlink = [ - { name = "nvidia-nvjitlink", version = "13.0.88", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform == 'linux' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'linux' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-nvjitlink", version = "13.0.88", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform == 'linux' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, ] nvrtc = [ - { name = "nvidia-cuda-nvrtc", version = "13.0.88", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform == 'linux' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'linux' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-cuda-nvrtc", version = "13.0.88", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform == 'linux' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, ] nvtx = [ - { name = "nvidia-nvtx", version = "13.0.85", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform == 'linux' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'linux' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-nvtx", version = "13.0.85", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform == 'linux' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, ] [[package]] @@ -1436,43 +1445,43 @@ wheels = [ [package.optional-dependencies] cublas = [ - { name = "nvidia-cublas", version = "13.4.0.1", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'AMD64' and sys_platform == 'win32' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, - { name = "nvidia-cuda-nvrtc", version = "13.2.78", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'AMD64' and sys_platform == 'win32' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-cublas", version = "13.4.0.1", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'AMD64' and sys_platform == 'win32' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-cuda-nvrtc", version = "13.2.78", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'AMD64' and sys_platform == 'win32' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, ] cudart = [ - { name = "nvidia-cuda-runtime", version = "13.2.75", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'AMD64' and sys_platform == 'win32' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-cuda-runtime", version = "13.2.75", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'AMD64' and sys_platform == 'win32' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, ] cufft = [ - { name = "nvidia-cufft", version = "12.2.0.46", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'AMD64' and sys_platform == 'win32' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, - { name = "nvidia-nvjitlink", version = "13.4.92", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'AMD64' and sys_platform == 'win32' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-cufft", version = "12.2.0.46", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'AMD64' and sys_platform == 'win32' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-nvjitlink", version = "13.4.92", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'AMD64' and sys_platform == 'win32' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, ] cufile = [ - { name = "nvidia-cufile", version = "1.17.1.22", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'aarch64' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (platform_machine == 'aarch64' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'aarch64' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (platform_machine == 'aarch64' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'aarch64' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (platform_machine == 'x86_64' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (platform_machine == 'x86_64' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (platform_machine == 'x86_64' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform != 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'linux' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'linux' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-cufile", version = "1.17.1.22", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, ] cupti = [ - { name = "nvidia-cuda-cupti", version = "13.2.75", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'AMD64' and sys_platform == 'win32' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-cuda-cupti", version = "13.2.75", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'AMD64' and sys_platform == 'win32' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, ] curand = [ - { name = "nvidia-curand", version = "10.4.2.55", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'AMD64' and sys_platform == 'win32' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-curand", version = "10.4.2.55", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'AMD64' and sys_platform == 'win32' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, ] cusolver = [ - { name = "nvidia-cublas", version = "13.4.0.1", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'AMD64' and sys_platform == 'win32' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, - { name = "nvidia-cusolver", version = "12.2.0.1", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'AMD64' and sys_platform == 'win32' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, - { name = "nvidia-cusparse", version = "12.7.10.1", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'AMD64' and sys_platform == 'win32' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, - { name = "nvidia-nvjitlink", version = "13.4.92", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'AMD64' and sys_platform == 'win32' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-cublas", version = "13.4.0.1", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'AMD64' and sys_platform == 'win32' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-cusolver", version = "12.2.0.1", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'AMD64' and sys_platform == 'win32' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-cusparse", version = "12.7.10.1", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'AMD64' and sys_platform == 'win32' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-nvjitlink", version = "13.4.92", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'AMD64' and sys_platform == 'win32' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, ] cusparse = [ - { name = "nvidia-cusparse", version = "12.7.10.1", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'AMD64' and sys_platform == 'win32' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, - { name = "nvidia-nvjitlink", version = "13.4.92", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'AMD64' and sys_platform == 'win32' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-cusparse", version = "12.7.10.1", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'AMD64' and sys_platform == 'win32' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-nvjitlink", version = "13.4.92", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'AMD64' and sys_platform == 'win32' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, ] nvjitlink = [ - { name = "nvidia-nvjitlink", version = "13.4.92", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'AMD64' and sys_platform == 'win32' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-nvjitlink", version = "13.4.92", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'AMD64' and sys_platform == 'win32' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, ] nvrtc = [ - { name = "nvidia-cuda-nvrtc", version = "13.2.78", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'AMD64' and sys_platform == 'win32' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-cuda-nvrtc", version = "13.2.78", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'AMD64' and sys_platform == 'win32' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, ] nvtx = [ - { name = "nvidia-nvtx", version = "13.2.75", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'AMD64' and sys_platform == 'win32' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-nvtx", version = "13.2.75", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'AMD64' and sys_platform == 'win32' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, ] [[package]] @@ -1594,6 +1603,20 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/02/08/9c41fb51ab5b43eb21674aff13df270e8ba6c4b29c8624e328dc7a9482af/distlib-0.4.3-py2.py3-none-any.whl", hash = "sha256:4b0ce306c966eb73bc3a7b6abad017c556dadd92c44701562cd528ac7fde4d5b", size = 470628, upload-time = "2026-06-12T08:04:50.506Z" }, ] +[[package]] +name = "docker" +version = "7.2.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "pywin32", marker = "sys_platform == 'win32'" }, + { name = "requests" }, + { name = "urllib3" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/88/7f/731ff914b0255d3d065f45fd4e626d4b8c95dbcbaada049f337a6ac16410/docker-7.2.0.tar.gz", hash = "sha256:cebb93773d334f778e023a7ee352a8d6e13ab1bd3b863a4d4a59dec897df43ac", size = 118731, upload-time = "2026-07-09T14:53:46.39Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/75/23/529140fe1aab80fc6992f93a706deec709140a6397439139a054e1515c45/docker-7.2.0-py3-none-any.whl", hash = "sha256:a3f45fdeb9165e2d25d9a1d02ddf3bc70fb572cf5ebbf9b58558c22caf29b71f", size = 148775, upload-time = "2026-07-09T14:53:45.224Z" }, +] + [[package]] name = "docutils" version = "0.21.2" @@ -1715,6 +1738,36 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/35/f8/08d37b2cd89da306e3520bd27f8a85692122b42b56c0c2c3784ff09c022f/flake8-7.1.2-py2.py3-none-any.whl", hash = "sha256:1cbc62e65536f65e6d754dfe6f1bada7f5cf392d6f5db3c2b85892466c3e7c1a", size = 57745, upload-time = "2025-02-16T18:45:42.351Z" }, ] +[[package]] +name = "flask" +version = "3.1.3" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "blinker" }, + { name = "click" }, + { name = "itsdangerous" }, + { name = "jinja2" }, + { name = "markupsafe" }, + { name = "werkzeug" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/26/00/35d85dcce6c57fdc871f3867d465d780f302a175ea360f62533f12b27e2b/flask-3.1.3.tar.gz", hash = "sha256:0ef0e52b8a9cd932855379197dd8f94047b359ca0a78695144304cb45f87c9eb", size = 759004, upload-time = "2026-02-19T05:00:57.678Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/7f/9c/34f6962f9b9e9c71f6e5ed806e0d0ff03c9d1b0b2340088a0cf4bce09b18/flask-3.1.3-py3-none-any.whl", hash = "sha256:f4bcbefc124291925f1a26446da31a5178f9483862233b23c0c96a20701f670c", size = 103424, upload-time = "2026-02-19T05:00:56.027Z" }, +] + +[[package]] +name = "flask-cors" +version = "6.0.5" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "flask" }, + { name = "werkzeug" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/47/03/4e464a50860f9adf08b5c1d3479cb8ea1f12af2aa69535c7042c6e628135/flask_cors-6.0.5.tar.gz", hash = "sha256:30c5031552cd59f620ac0c8211dac45b345d3b2df310e7721879e4f46ef9c601", size = 101386, upload-time = "2026-06-08T20:20:17.765Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/49/55/5bb1a2d918e9f02f131e47a59032bae70e48050e986e941511fd737a935c/flask_cors-6.0.5-py3-none-any.whl", hash = "sha256:68fcf75693e961f3af26683b23c4b9a8fb6b64de17d20d0c37b95e8de7ab2ed8", size = 16692, upload-time = "2026-06-08T20:20:16.247Z" }, +] + [[package]] name = "fonttools" version = "4.65.0" @@ -1924,6 +1977,30 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/69/43/a0096d96d3271640bbda056b9d22249be141a54509859b99002ee6149a01/gcsfs-2026.6.0-py3-none-any.whl", hash = "sha256:4a81fe40fbdef450aeaacb534b42c04059f8f3b080e3e169137247a5bb023464", size = 89629, upload-time = "2026-06-18T03:16:07.407Z" }, ] +[[package]] +name = "gitdb" +version = "4.0.12" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "smmap" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/72/94/63b0fc47eb32792c7ba1fe1b694daec9a63620db1e313033d18140c2320a/gitdb-4.0.12.tar.gz", hash = "sha256:5ef71f855d191a3326fcfbc0d5da835f26b13fbcba60c32c21091c349ffdb571", size = 394684, upload-time = "2025-01-02T07:20:46.413Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a0/61/5c78b91c3143ed5c14207f463aecfc8f9dbb5092fb2869baf37c273b2705/gitdb-4.0.12-py3-none-any.whl", hash = "sha256:67073e15955400952c6565cc3e707c554a4eea2e428946f7a4c162fab9bd9bcf", size = 62794, upload-time = "2025-01-02T07:20:43.624Z" }, +] + +[[package]] +name = "gitpython" +version = "3.2.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "gitdb" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/6e/2d/6f6e649818da44d4499604802c89329b8d9799687a124e3a5e467a643336/gitpython-3.2.0.tar.gz", hash = "sha256:fb92310af6844d96adc95ca066ed2e617c00e1dbd146a326626c81e72e18cc2e", size = 238091, upload-time = "2026-09-30T09:10:32.65Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/fa/48/f36d233ce749d1bb880cb5e6f8bb171dc2effabc7435c31da604d831b3ec/gitpython-3.2.0-py3-none-any.whl", hash = "sha256:bd70c5ec05cd2b797423e7eb312147d2458d3cca92085888fba2213f85905537", size = 228537, upload-time = "2026-09-30T09:10:31.159Z" }, +] + [[package]] name = "google-api-core" version = "2.30.3" @@ -2149,6 +2226,42 @@ grpc = [ { name = "grpcio" }, ] +[[package]] +name = "graphene" +version = "3.4.3" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "graphql-core" }, + { name = "graphql-relay" }, + { name = "python-dateutil" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/cc/f6/bf62ff950c317ed03e77f3f6ddd7e34aaa98fe89d79ebd660c55343d8054/graphene-3.4.3.tar.gz", hash = "sha256:2a3786948ce75fe7e078443d37f609cbe5bb36ad8d6b828740ad3b95ed1a0aaa", size = 44739, upload-time = "2024-11-09T20:44:25.757Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/66/e0/61d8e98007182e6b2aca7cf65904721fb2e4bce0192272ab9cb6f69d8812/graphene-3.4.3-py2.py3-none-any.whl", hash = "sha256:820db6289754c181007a150db1f7fff544b94142b556d12e3ebc777a7bf36c71", size = 114894, upload-time = "2024-11-09T20:44:23.851Z" }, +] + +[[package]] +name = "graphql-core" +version = "3.2.13" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/4e/5e/aa0d4e701b50db0bab71b125dd19ddb0f98c638d008d6f7e3d8ce9cbc92e/graphql_core-3.2.13.tar.gz", hash = "sha256:bb81dd266d4ab7b591bd976f1b23639d97776cb9ac1a896b4a93c271e11ed618", size = 577471, upload-time = "2026-09-27T12:22:57.136Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/1f/b4/b52a324b0b297ddc49b8ea31a51d39fae63132f694a267b9381d94e2e0af/graphql_core-3.2.13-py3-none-any.whl", hash = "sha256:b0eb04f2c31556b2310a77c8fb53c74e8b56570f8fea594d89c6ef7827dbb497", size = 268191, upload-time = "2026-09-27T12:22:55.661Z" }, +] + +[[package]] +name = "graphql-relay" +version = "3.2.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "graphql-core" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/d1/13/98fbf8d67552f102488ffc16c6f559ce71ea15f6294728d33928ab5ff14d/graphql-relay-3.2.0.tar.gz", hash = "sha256:1ff1c51298356e481a0be009ccdff249832ce53f30559c1338f22a0e0d17250c", size = 50027, upload-time = "2022-04-16T11:03:45.447Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/74/16/a4cf06adbc711bd364a73ce043b0b08d8fa5aae3df11b6ee4248bcdad2e0/graphql_relay-3.2.0-py3-none-any.whl", hash = "sha256:c9b22bd28b170ba1fe674c74384a8ff30a76c8e26f88ac3aa1584dd3179953e5", size = 16940, upload-time = "2022-04-16T11:03:43.895Z" }, +] + [[package]] name = "graphviz" version = "0.20.1" @@ -2324,6 +2437,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/71/c4/3a77e4273e866b1b0c412afd80882e95941a37170032b5109d847c501124/grpcio_status-1.84.0-py3-none-any.whl", hash = "sha256:0c182ca0d6e60acbfd0e14499cf39a155e4827a1c3fd9f7638e49af15a74c30a", size = 14639, upload-time = "2026-09-14T07:10:15.175Z" }, ] +[[package]] +name = "gunicorn" +version = "26.2.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/d9/8a/e4ef6ee11701b6cd64702848415ffb69eeff85cb388a3c6c7fe86f22f3f8/gunicorn-26.2.0.tar.gz", hash = "sha256:62b864895d9ebff0b2f9867ba04fe811c93121596540830c9c916d0769668447", size = 787921, upload-time = "2026-08-24T15:05:59.3Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/fe/85/7522a52e5e2f42faf1a129113ab63e548c42e103e9af395b7bfe65e403e2/gunicorn-26.2.0-py3-none-any.whl", hash = "sha256:bd249d0b3f7972f7432f0a6b6ff3b3ee2d129f70cd1ff6c09a9dd9e29a2b88e3", size = 228389, upload-time = "2026-08-24T15:05:57.67Z" }, +] + [[package]] name = "h11" version = "0.16.0" @@ -2374,6 +2496,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/56/95/9377bcb415797e44274b51d46e3249eba641711cf3348050f76ee7b15ffc/httpx-0.27.2-py3-none-any.whl", hash = "sha256:7bb2708e112d8fdd7829cd4243970f0c223274051cb35ee80c03301ee29a3df0", size = 76395, upload-time = "2024-08-27T12:53:59.653Z" }, ] +[[package]] +name = "huey" +version = "3.4.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/57/25/1281d8f3febc30b5facbfc5e4163447241711591c094dac642d32c827eb0/huey-3.4.0.tar.gz", hash = "sha256:ac9b02e741f13da4092ea19bcaafc73669e97d3153914d2fed2868df144f6dfa", size = 630013, upload-time = "2026-09-04T02:04:34.521Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/57/31/cee585bd4d1e39947a426a6736f4e17d22b454952370fb9cf48eb3103085/huey-3.4.0-py3-none-any.whl", hash = "sha256:d0580762397744026f83baae5f2b97a62474a62b8565ac75bbafefe72f9c784f", size = 136493, upload-time = "2026-09-04T02:04:32.848Z" }, +] + [[package]] name = "humanize" version = "4.6.0" @@ -2471,6 +2602,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/0a/63/4036ae70eea279c63e2304b91ee0ac182f467f24f86394ecfe726092340b/isort-5.12.0-py3-none-any.whl", hash = "sha256:f84c2818376e66cf843d497486ea8fed8700b340f308f076c6fb1229dff318b6", size = 91198, upload-time = "2023-01-28T17:10:21.149Z" }, ] +[[package]] +name = "itsdangerous" +version = "2.2.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/9c/cb/8ac0172223afbccb63986cc25049b154ecfb5e85932587206f42317be31d/itsdangerous-2.2.0.tar.gz", hash = "sha256:e0050c0b7da1eea53ffaf149c0cfbb5c6e2e2b69c4bef22c81fa6eb73e5f6173", size = 54410, upload-time = "2024-04-16T21:28:15.614Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/04/96/92447566d16df59b2a776c0fb82dbc4d9e07cd95062562af01e408583fc4/itsdangerous-2.2.0-py3-none-any.whl", hash = "sha256:c6242fc49e35958c8b15141343aa660db5fc54d4f13a1db01a3f5891b98700ef", size = 16234, upload-time = "2024-04-16T21:28:14.499Z" }, +] + [[package]] name = "jinja2" version = "3.1.6" @@ -2749,8 +2889,8 @@ dependencies = [ { name = "requests" }, { name = "s3fs" }, { name = "scikit-learn" }, - { name = "scipy", version = "1.17.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, - { name = "scipy", version = "1.18.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "scipy", version = "1.17.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" }, + { name = "scipy", version = "1.18.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" }, { name = "setuptools" }, { name = "setuptools-scm" }, { name = "smart-open" }, @@ -2762,11 +2902,11 @@ dependencies = [ { name = "tenacity" }, { name = "tensorboard" }, { name = "thrift" }, - { name = "torch", version = "2.12.0", source = { registry = "https://download.pytorch.org/whl/cpu" }, marker = "(sys_platform == 'darwin' and extra == 'extra-4-kumo-cpu') or (sys_platform == 'darwin' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, + { name = "torch", version = "2.12.0", source = { registry = "https://download.pytorch.org/whl/cpu" }, marker = "(sys_platform == 'darwin' and extra == 'extra-4-kumo-cpu') or (sys_platform == 'darwin' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, { name = "torch", version = "2.12.0", source = { registry = "https://pypi.org/simple" }, marker = "(extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu130' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests')" }, - { name = "torch", version = "2.12.0+cpu", source = { registry = "https://download.pytorch.org/whl/cpu" }, marker = "(sys_platform != 'darwin' and extra == 'extra-4-kumo-cpu') or (sys_platform != 'darwin' and extra != 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "torch", version = "2.12.0+cpu", source = { registry = "https://download.pytorch.org/whl/cpu" }, marker = "(sys_platform != 'darwin' and extra == 'extra-4-kumo-cpu') or (sys_platform != 'darwin' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, { name = "torch", version = "2.12.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" }, marker = "extra == 'extra-4-kumo-cu130' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, - { name = "torch", version = "2.12.0+cu132", source = { registry = "https://download.pytorch.org/whl/cu132" }, marker = "extra == 'extra-4-kumo-cu132' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "torch", version = "2.12.0+cu132", source = { registry = "https://download.pytorch.org/whl/cu132" }, marker = "extra == 'extra-4-kumo-cu132' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, { name = "types-aiobotocore", extra = ["s3"] }, { name = "typing-extensions" }, { name = "urllib3" }, @@ -2786,7 +2926,7 @@ cu130 = [ { name = "torch", version = "2.12.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" } }, ] cu132 = [ - { name = "kumo-core", extra = ["faiss-gpu"], marker = "extra == 'extra-4-kumo-cu132' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "kumo-core", extra = ["faiss-gpu"], marker = "extra == 'extra-4-kumo-cu132' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, { name = "torch", version = "2.12.0+cu132", source = { registry = "https://download.pytorch.org/whl/cu132" } }, ] data-platform = [ @@ -2803,6 +2943,7 @@ dev = [ { name = "flake8" }, { name = "isort" }, { name = "kubernetes-stubs" }, + { name = "mlflow" }, { name = "moto" }, { name = "packaging" }, { name = "pick" }, @@ -2881,6 +3022,7 @@ rest = [ test = [ { name = "category-encoders" }, { name = "coverage" }, + { name = "mlflow" }, { name = "moto" }, { name = "pytest" }, { name = "pytest-asyncio" }, @@ -2904,7 +3046,7 @@ tracing = [ [package.dev-dependencies] public-tests = [ - { name = "torch", version = "2.12.0", source = { registry = "https://download.pytorch.org/whl/cpu" }, marker = "(sys_platform == 'darwin' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, + { name = "torch", version = "2.12.0", source = { registry = "https://download.pytorch.org/whl/cpu" }, marker = "(sys_platform == 'darwin' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, { name = "torch", version = "2.12.0+cpu", source = { registry = "https://download.pytorch.org/whl/cpu" }, marker = "(sys_platform != 'darwin' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, ] @@ -2986,6 +3128,8 @@ requires-dist = [ { name = "markdown", marker = "extra == 'rest'", specifier = "==3.10" }, { name = "matplotlib", marker = "extra == 'drawing'", specifier = ">=3.10.5" }, { name = "matplotlib", marker = "extra == 'full'", specifier = ">=3.10.5" }, + { name = "mlflow", marker = "extra == 'dev'", specifier = ">=3.13,<4" }, + { name = "mlflow", marker = "extra == 'test'", specifier = ">=3.13,<4" }, { name = "moto", marker = "extra == 'dev'", specifier = "==4.2.14" }, { name = "moto", marker = "extra == 'test'", specifier = "==4.2.14" }, { name = "mypy-boto3-s3", specifier = "==1.42.94" }, @@ -3210,11 +3354,11 @@ dependencies = [ { name = "sqlglot" }, { name = "tabulate" }, { name = "tensorboard" }, - { name = "torch", version = "2.12.0", source = { registry = "https://download.pytorch.org/whl/cpu" }, marker = "(sys_platform == 'darwin' and extra == 'extra-4-kumo-cpu') or (sys_platform == 'darwin' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, + { name = "torch", version = "2.12.0", source = { registry = "https://download.pytorch.org/whl/cpu" }, marker = "(sys_platform == 'darwin' and extra == 'extra-4-kumo-cpu') or (sys_platform == 'darwin' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, { name = "torch", version = "2.12.0", source = { registry = "https://pypi.org/simple" }, marker = "(extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu130' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests')" }, - { name = "torch", version = "2.12.0+cpu", source = { registry = "https://download.pytorch.org/whl/cpu" }, marker = "(sys_platform != 'darwin' and extra == 'extra-4-kumo-cpu') or (sys_platform != 'darwin' and extra != 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "torch", version = "2.12.0+cpu", source = { registry = "https://download.pytorch.org/whl/cpu" }, marker = "(sys_platform != 'darwin' and extra == 'extra-4-kumo-cpu') or (sys_platform != 'darwin' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, { name = "torch", version = "2.12.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" }, marker = "extra == 'extra-4-kumo-cu130' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, - { name = "torch", version = "2.12.0+cu132", source = { registry = "https://download.pytorch.org/whl/cu132" }, marker = "extra == 'extra-4-kumo-cu132' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "torch", version = "2.12.0+cu132", source = { registry = "https://download.pytorch.org/whl/cu132" }, marker = "extra == 'extra-4-kumo-cu132' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, { name = "torch-geometric" }, { name = "torchmetrics" }, { name = "typing-extensions" }, @@ -3339,6 +3483,10 @@ databricks = [ gcs = [ { name = "kumo-gcs" }, ] +mlflow = [ + { name = "mlflow" }, + { name = "nvidia-ml-py" }, +] s3 = [ { name = "kumo-s3" }, ] @@ -3352,11 +3500,11 @@ test = [ { name = "numpy" }, { name = "pytest" }, { name = "pytest-mock" }, - { name = "torch", version = "2.12.0", source = { registry = "https://download.pytorch.org/whl/cpu" }, marker = "(sys_platform == 'darwin' and extra == 'extra-4-kumo-cpu') or (sys_platform == 'darwin' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, + { name = "torch", version = "2.12.0", source = { registry = "https://download.pytorch.org/whl/cpu" }, marker = "(sys_platform == 'darwin' and extra == 'extra-4-kumo-cpu') or (sys_platform == 'darwin' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, { name = "torch", version = "2.12.0", source = { registry = "https://pypi.org/simple" }, marker = "(extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu130' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests')" }, - { name = "torch", version = "2.12.0+cpu", source = { registry = "https://download.pytorch.org/whl/cpu" }, marker = "(sys_platform != 'darwin' and extra == 'extra-4-kumo-cpu') or (sys_platform != 'darwin' and extra != 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "torch", version = "2.12.0+cpu", source = { registry = "https://download.pytorch.org/whl/cpu" }, marker = "(sys_platform != 'darwin' and extra == 'extra-4-kumo-cpu') or (sys_platform != 'darwin' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, { name = "torch", version = "2.12.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" }, marker = "extra == 'extra-4-kumo-cu130' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, - { name = "torch", version = "2.12.0+cu132", source = { registry = "https://download.pytorch.org/whl/cu132" }, marker = "extra == 'extra-4-kumo-cu132' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "torch", version = "2.12.0+cu132", source = { registry = "https://download.pytorch.org/whl/cu132" }, marker = "extra == 'extra-4-kumo-cu132' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, ] text = [ { name = "kumo-glove-6b-50d" }, @@ -3376,7 +3524,9 @@ requires-dist = [ { name = "kumo-glove-6b-50d", marker = "extra == 'text'", editable = "tools/glove-6b-50d" }, { name = "kumo-s3", marker = "extra == 's3'", editable = "packages/kumo-s3" }, { name = "kumo-snowflake", extras = ["sdk"], marker = "extra == 'snowflake'", editable = "packages/kumo-snowflake" }, + { name = "mlflow", marker = "extra == 'mlflow'", specifier = ">=3.13,<4" }, { name = "numpy", marker = "extra == 'test'", specifier = ">=1.26,<3" }, + { name = "nvidia-ml-py", marker = "extra == 'mlflow'" }, { name = "pandas" }, { name = "pyarrow", specifier = ">=8.0.0" }, { name = "pydantic", specifier = ">=1.10.21" }, @@ -3385,7 +3535,7 @@ requires-dist = [ { name = "torch", marker = "extra == 'test'", specifier = ">=2.9,<2.13" }, { name = "typing-extensions", specifier = ">=4.5.0" }, ] -provides-extras = ["test", "snowflake", "databricks", "bigquery", "s3", "text", "gcs"] +provides-extras = ["test", "snowflake", "databricks", "bigquery", "s3", "text", "gcs", "mlflow"] [[package]] name = "kumo-gcs" @@ -3531,7 +3681,7 @@ requires-dist = [ { name = "pyyaml", marker = "extra == 'sdk'" }, { name = "requests", specifier = ">=2.31" }, { name = "snowflake-connector-python", specifier = ">=4.3" }, - { name = "snowflake-ml-python", marker = "extra == 'sdk'" }, + { name = "snowflake-ml-python", marker = "extra == 'sdk'", specifier = ">=2.0,<3" }, { name = "snowflake-snowpark-python", marker = "extra == 'sdk'" }, ] provides-extras = ["sdk", "test"] @@ -3766,8 +3916,8 @@ name = "matplotlib" version = "3.11.2" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "contourpy", version = "1.3.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, - { name = "contourpy", version = "1.4.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "contourpy", version = "1.3.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" }, + { name = "contourpy", version = "1.4.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" }, { name = "cycler" }, { name = "fonttools" }, { name = "kiwisolver" }, @@ -3854,6 +4004,89 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/ec/9d/1832c79be2a87aece4ecbbddd3ae315b9cc90ae62c6755581afc06670867/mermaid_py-0.9.0-py3-none-any.whl", hash = "sha256:6abc53f9527118c147d59568852fe782dcdbe54d5512f23d0e3a1b5dec9a302b", size = 32952, upload-time = "2026-09-16T10:46:35.6Z" }, ] +[[package]] +name = "mlflow" +version = "3.17.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "aiohttp" }, + { name = "alembic" }, + { name = "cryptography" }, + { name = "docker" }, + { name = "flask" }, + { name = "flask-cors" }, + { name = "graphene" }, + { name = "gunicorn", marker = "sys_platform != 'win32'" }, + { name = "huey" }, + { name = "matplotlib" }, + { name = "mlflow-skinny" }, + { name = "mlflow-tracing" }, + { name = "numpy" }, + { name = "pandas" }, + { name = "pyarrow" }, + { name = "scikit-learn" }, + { name = "scipy", version = "1.17.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" }, + { name = "scipy", version = "1.18.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" }, + { name = "skops" }, + { name = "sqlalchemy" }, + { name = "waitress", marker = "sys_platform == 'win32'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/6e/30/953491ffbd83c90aba4a72c8a68c03ea16a8b016050557094d059268bf7d/mlflow-3.17.1.tar.gz", hash = "sha256:b41225845f200968bfeb015c48ad9e57a02f3bf4cee246fb04a75325728c21d8", size = 11158262, upload-time = "2026-10-13T09:00:18.681Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a4/64/ee5ca7500b3f4ca187b7c0c2f080a029eed1df7f22b334391587979b23dc/mlflow-3.17.1-py3-none-any.whl", hash = "sha256:726bf7feac25dbdb5b3098bf20b86574393bb75468eef721902ec025e724752f", size = 12004276, upload-time = "2026-10-13T09:00:14.749Z" }, +] + +[[package]] +name = "mlflow-skinny" +version = "3.17.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "anyio" }, + { name = "cachetools" }, + { name = "click" }, + { name = "cloudpickle" }, + { name = "databricks-sdk" }, + { name = "fastapi" }, + { name = "gitpython" }, + { name = "importlib-metadata" }, + { name = "opentelemetry-api" }, + { name = "opentelemetry-proto" }, + { name = "opentelemetry-sdk" }, + { name = "packaging" }, + { name = "protobuf" }, + { name = "pydantic" }, + { name = "python-dotenv" }, + { name = "pyyaml" }, + { name = "requests" }, + { name = "sqlparse" }, + { name = "starlette" }, + { name = "typing-extensions" }, + { name = "uvicorn" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/c7/b3/a72208dd21b7b8345d3d15eafa8d5d2897d36e21420e22d290fad5775aa4/mlflow_skinny-3.17.1.tar.gz", hash = "sha256:c4587083811f638d9c2e6e8ab7caabbccc51ccc09de2c87a503913b32bbfd7bf", size = 3483742, upload-time = "2026-10-13T08:58:49.651Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/0f/54/d0e5e2588013d7655f9cfb3d17fff4760c22652ce628b57ae4ac0fa913e6/mlflow_skinny-3.17.1-py3-none-any.whl", hash = "sha256:7374adf971e06f77f861dd8196c50e344557728b643192c665f0603c19a8d146", size = 4105632, upload-time = "2026-10-13T08:58:47.785Z" }, +] + +[[package]] +name = "mlflow-tracing" +version = "3.17.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "cachetools" }, + { name = "databricks-sdk" }, + { name = "opentelemetry-api" }, + { name = "opentelemetry-proto" }, + { name = "opentelemetry-sdk" }, + { name = "packaging" }, + { name = "protobuf" }, + { name = "pydantic" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/4c/da/e226d475789bbb8caf68aafa3cbf2b715e5f6eda00f7f659b7eda6fba1ef/mlflow_tracing-3.17.1.tar.gz", hash = "sha256:c154492045b5f3544ab07fd9ab9d7c07febfb53065d43f1a07f33a31c46b3505", size = 1688196, upload-time = "2026-10-13T08:59:23.98Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ce/19/2ba5a327efa9a45ea46c927103a7c57b24441375d8b3435f6b27d187ac91/mlflow_tracing-3.17.1-py3-none-any.whl", hash = "sha256:4f6405e48efdf9330101aa059383103e1e15559664ec792290b9b95556f9e851", size = 1981200, upload-time = "2026-10-13T08:59:22.314Z" }, +] + [[package]] name = "more-itertools" version = "11.1.0" @@ -4081,7 +4314,7 @@ name = "mypy-boto3-s3" version = "1.42.94" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "typing-extensions", marker = "python_full_version < '3.12' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "typing-extensions", marker = "python_full_version < '3.12'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/6d/55/2268b22037d1c6814c2229904d43f9cb3fcd9b9e9d01a3402b58a10270ce/mypy_boto3_s3-1.42.94.tar.gz", hash = "sha256:1d92d722cf00573b8111e98493ab386e0c1b59a1530b7fee4af77f2d9a1c477d", size = 77049, upload-time = "2026-04-22T21:31:00.857Z" } wheels = [ @@ -4279,7 +4512,7 @@ resolution-markers = [ "python_full_version < '3.12'", ] dependencies = [ - { name = "nvidia-cuda-nvrtc", version = "13.2.78", source = { registry = "https://pypi.org/simple" }, marker = "extra == 'extra-4-kumo-cu132' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-cuda-nvrtc", version = "13.2.78", source = { registry = "https://pypi.org/simple" }, marker = "extra == 'extra-4-kumo-cu132' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/d0/48/8571e55aaabbd112f06a39a71fbc3abf0b8a0d450921816a13a6d5e8fa0b/nvidia_cublas-13.4.0.1-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:705d7d214fbb20f134415ecadc488abf74f444c155a6555dec5687404afb18a9", size = 513051304, upload-time = "2026-04-13T09:49:32.758Z" }, @@ -4422,7 +4655,7 @@ version = "9.20.0.48" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "nvidia-cublas", version = "13.1.1.3", source = { registry = "https://pypi.org/simple" }, marker = "extra == 'extra-4-kumo-cu130' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests')" }, - { name = "nvidia-cublas", version = "13.4.0.1", source = { registry = "https://pypi.org/simple" }, marker = "extra == 'extra-4-kumo-cu132' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-cublas", version = "13.4.0.1", source = { registry = "https://pypi.org/simple" }, marker = "extra == 'extra-4-kumo-cu132' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/56/c5/83384d846b2fd17c44bd499b36c75a45ed4f095fbbb2252294e89cea5c5c/nvidia_cudnn_cu13-9.20.0.48-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:e31454ae00094b0c55319d9d15b6fa2fc50a9e1c0f5c8c80fb75258234e731e1", size = 444574296, upload-time = "2026-03-09T19:28:27.751Z" }, @@ -4460,7 +4693,7 @@ resolution-markers = [ "python_full_version < '3.12'", ] dependencies = [ - { name = "nvidia-nvjitlink", version = "13.4.92", source = { registry = "https://pypi.org/simple" }, marker = "extra == 'extra-4-kumo-cu132' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-nvjitlink", version = "13.4.92", source = { registry = "https://pypi.org/simple" }, marker = "extra == 'extra-4-kumo-cu132' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/be/b1/6e36381739b51e692d91646298ad5685c562929b899331c2bb6e9722a176/nvidia_cufft-12.2.0.46-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:2e8f3a22c2745f95327b7639ba2868ea2cbd5d3131ebe6313394e24164054f41", size = 218244743, upload-time = "2026-04-13T09:51:24.75Z" }, @@ -4572,9 +4805,9 @@ resolution-markers = [ "python_full_version < '3.12'", ] dependencies = [ - { name = "nvidia-cublas", version = "13.4.0.1", source = { registry = "https://pypi.org/simple" }, marker = "extra == 'extra-4-kumo-cu132' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, - { name = "nvidia-cusparse", version = "12.7.10.1", source = { registry = "https://pypi.org/simple" }, marker = "extra == 'extra-4-kumo-cu132' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, - { name = "nvidia-nvjitlink", version = "13.4.92", source = { registry = "https://pypi.org/simple" }, marker = "extra == 'extra-4-kumo-cu132' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-cublas", version = "13.4.0.1", source = { registry = "https://pypi.org/simple" }, marker = "extra == 'extra-4-kumo-cu132' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-cusparse", version = "12.7.10.1", source = { registry = "https://pypi.org/simple" }, marker = "extra == 'extra-4-kumo-cu132' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-nvjitlink", version = "13.4.92", source = { registry = "https://pypi.org/simple" }, marker = "extra == 'extra-4-kumo-cu132' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/94/21/21e5c491a8285a6ed039cbb8c60ec13ff3aff60c669e1caa4adaeb71f997/nvidia_cusolver-12.2.0.1-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:c9def09fcf300b9465cdf800142abf10149df28c4290bdf9e7b4a700480d31d7", size = 249744985, upload-time = "2026-04-13T09:54:33.757Z" }, @@ -4612,7 +4845,7 @@ resolution-markers = [ "python_full_version < '3.12'", ] dependencies = [ - { name = "nvidia-nvjitlink", version = "13.4.92", source = { registry = "https://pypi.org/simple" }, marker = "extra == 'extra-4-kumo-cu132' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-nvjitlink", version = "13.4.92", source = { registry = "https://pypi.org/simple" }, marker = "extra == 'extra-4-kumo-cu132' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/b5/09/9f1fa88dd305e838b71f2ca4e50064e7cbdf2e15cd6410d9e1b74066cc41/nvidia_cusparse-12.7.10.1-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:80cc98b38fdcf054a80bd9782fe2e63c66f90e202cac269f641357ccac5fe7e3", size = 168860233, upload-time = "2026-04-13T09:55:48.891Z" }, @@ -4630,6 +4863,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/31/83/f3647ce26916c94a6ca4ff1810623e2c405cff2dea6e78d29516b2514df9/nvidia_cusparselt_cu13-0.8.1-py3-none-win_amd64.whl", hash = "sha256:dccbd362f91a7b9024d1f55ee9f548ac065027ff15d8c8b0db889ab3a8f31215", size = 156885108, upload-time = "2025-09-05T18:51:35.958Z" }, ] +[[package]] +name = "nvidia-ml-py" +version = "13.615.71" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/fd/30/b25216758be3d3e2834825d8193609e2d71c770a8bd7984438c058c90268/nvidia_ml_py-13.615.71.tar.gz", hash = "sha256:bebe4e48f51b1dc75028c0815cb7bfa14a31a5bb80be70c9d980c6036953fc3d", size = 57485, upload-time = "2026-09-25T15:15:28.226Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/53/a1/1681dfa1c904d4e3e72e51b55a0ff012d50b766843ef832d584abe2113c6/nvidia_ml_py-13.615.71-py3-none-any.whl", hash = "sha256:959bf4adf6fe1308e4bd739e722236b0d1ec8392e2cefad33ff70c311380b9b6", size = 58132, upload-time = "2026-09-25T15:15:26.54Z" }, +] + [[package]] name = "nvidia-nccl-cu13" version = "2.29.7" @@ -5083,7 +5325,7 @@ name = "pick" version = "2.6.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "windows-curses", marker = "sys_platform == 'win32' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "windows-curses", marker = "sys_platform == 'win32'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/b6/f9/450cc44b21190d99b86df3e59f7573d85f3ef2c4f7670ce4f65cb2def3b4/pick-2.6.0.tar.gz", hash = "sha256:f24434a845ae9265ebf7657fc224b9a203ac75aa576b568e061484a72e27be82", size = 6266, upload-time = "2026-02-28T13:47:06.31Z" } wheels = [ @@ -5250,6 +5492,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/45/e2/bbb7129c9e7999a6b8ee9cca3b66486c25c423ab5a75f34071798b74ce94/pre_commit-4.6.2-py2.py3-none-any.whl", hash = "sha256:e2dde9a75d3bce11bd3831c26d134df00a2803c1d818be6a0383c3dcda25dc4e", size = 226202, upload-time = "2026-08-10T22:07:16.942Z" }, ] +[[package]] +name = "prettytable" +version = "3.18.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "wcwidth" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/81/74/ba08d81e668ccfe8658d7520a307e63c19862c08eb4ccb26f356c5239a7a/prettytable-3.18.0.tar.gz", hash = "sha256:439217116152244369caf3d9f1caf2f9fe29b03bd79e88d2928c8e718c95d680", size = 76373, upload-time = "2026-06-22T16:07:50.174Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/fe/be/2e6798ace5cc036f5d05d36b7b2fd85346f1a708c87060890b070d0ec607/prettytable-3.18.0-py3-none-any.whl", hash = "sha256:b3346e0e6f79180833aebaac088ae926340586cf6d7d991b9eb125b65f72313a", size = 37357, upload-time = "2026-06-22T16:07:48.595Z" }, +] + [[package]] name = "prometheus-client" version = "0.14.1" @@ -5763,7 +6017,7 @@ version = "26.4.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "cryptography" }, - { name = "typing-extensions", marker = "python_full_version < '3.13' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "typing-extensions", marker = "python_full_version < '3.13'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/3f/e8/7325d258199b159eb2c03fe32107533e2832e70e63f4fb88a6aa00023201/pyopenssl-26.4.0.tar.gz", hash = "sha256:28dfcce0162b9211413e26dfbfdf1d24317fbeba18fc93c12400a1856b2a0bc7", size = 182046, upload-time = "2026-08-01T19:50:50.512Z" } wheels = [ @@ -5796,7 +6050,7 @@ name = "pytest" version = "9.0.3" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "colorama", marker = "sys_platform == 'win32' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "colorama", marker = "sys_platform == 'win32'" }, { name = "iniconfig" }, { name = "packaging" }, { name = "pluggy" }, @@ -5813,7 +6067,7 @@ version = "1.3.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "pytest" }, - { name = "typing-extensions", marker = "python_full_version < '3.13' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "typing-extensions", marker = "python_full_version < '3.13'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/90/2c/8af215c0f776415f3590cac4f9086ccefd6fd463befeae41cd4d3f193e5a/pytest_asyncio-1.3.0.tar.gz", hash = "sha256:d7f52f36d231b80ee124cd216ffb19369aa168fc10095013c6b014a34d3ee9e5", size = 50087, upload-time = "2025-11-10T16:07:47.256Z" } wheels = [ @@ -5916,6 +6170,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/16/7d/e9ffbadfbf89c93848412d04594135c4ae8c1d37d9e053b9c3ed718fabc4/python_discovery-1.6.1-py3-none-any.whl", hash = "sha256:d43fcdef879fe795352bd13ccf8d185ba5a9f86f36cfcd00529f596e737442b3", size = 38664, upload-time = "2026-09-18T01:31:52.448Z" }, ] +[[package]] +name = "python-dotenv" +version = "1.2.4" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/74/26/2fbeedb218a787a5eea551c7532cac4e009f83d689dd2faa0d0353473f86/python_dotenv-1.2.4.tar.gz", hash = "sha256:f0d53e69935a851c0dcc78f3ab7aaccd8cabef0b92382b576b824212902873c0", size = 60824, upload-time = "2026-10-01T05:36:10Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/60/d1/38f3a3405989a89ac18390803e70c6ad7c7760da4f9b83cbeca0c44a0c72/python_dotenv-1.2.4-py3-none-any.whl", hash = "sha256:42269a8a5b3fd54ffa6f3d84b18abed50064717576b4ecf03dc4a55d8aa04fdc", size = 23266, upload-time = "2026-10-01T05:36:08.633Z" }, +] + [[package]] name = "python-logging-loki" version = "0.3.1" @@ -5947,11 +6210,11 @@ dependencies = [ { name = "pandas" }, { name = "pillow" }, { name = "pyarrow" }, - { name = "torch", version = "2.12.0", source = { registry = "https://download.pytorch.org/whl/cpu" }, marker = "(sys_platform == 'darwin' and extra == 'extra-4-kumo-cpu') or (sys_platform == 'darwin' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, + { name = "torch", version = "2.12.0", source = { registry = "https://download.pytorch.org/whl/cpu" }, marker = "(sys_platform == 'darwin' and extra == 'extra-4-kumo-cpu') or (sys_platform == 'darwin' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, { name = "torch", version = "2.12.0", source = { registry = "https://pypi.org/simple" }, marker = "(extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu130' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests')" }, - { name = "torch", version = "2.12.0+cpu", source = { registry = "https://download.pytorch.org/whl/cpu" }, marker = "(sys_platform != 'darwin' and extra == 'extra-4-kumo-cpu') or (sys_platform != 'darwin' and extra != 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "torch", version = "2.12.0+cpu", source = { registry = "https://download.pytorch.org/whl/cpu" }, marker = "(sys_platform != 'darwin' and extra == 'extra-4-kumo-cpu') or (sys_platform != 'darwin' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, { name = "torch", version = "2.12.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" }, marker = "extra == 'extra-4-kumo-cu130' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, - { name = "torch", version = "2.12.0+cu132", source = { registry = "https://download.pytorch.org/whl/cu132" }, marker = "extra == 'extra-4-kumo-cu132' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "torch", version = "2.12.0+cu132", source = { registry = "https://download.pytorch.org/whl/cu132" }, marker = "extra == 'extra-4-kumo-cu132' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, { name = "tqdm" }, ] sdist = { url = "https://files.pythonhosted.org/packages/22/64/4a4aaa2a90eab208bf9f929a3bcc3b17ceb99bcb9abecf1c2323818f2898/pytorch_frame-0.3.0.tar.gz", hash = "sha256:412b42002e1a21a09fc8b0c6e4f0d048bd4e5ebef2f7139bbbbae82d81040152", size = 104795, upload-time = "2025-11-02T23:22:07.936Z" } @@ -5968,11 +6231,11 @@ dependencies = [ { name = "lightning-utilities" }, { name = "packaging" }, { name = "pyyaml" }, - { name = "torch", version = "2.12.0", source = { registry = "https://download.pytorch.org/whl/cpu" }, marker = "(sys_platform == 'darwin' and extra == 'extra-4-kumo-cpu') or (sys_platform == 'darwin' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, + { name = "torch", version = "2.12.0", source = { registry = "https://download.pytorch.org/whl/cpu" }, marker = "(sys_platform == 'darwin' and extra == 'extra-4-kumo-cpu') or (sys_platform == 'darwin' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, { name = "torch", version = "2.12.0", source = { registry = "https://pypi.org/simple" }, marker = "(extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu130' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests')" }, - { name = "torch", version = "2.12.0+cpu", source = { registry = "https://download.pytorch.org/whl/cpu" }, marker = "(sys_platform != 'darwin' and extra == 'extra-4-kumo-cpu') or (sys_platform != 'darwin' and extra != 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "torch", version = "2.12.0+cpu", source = { registry = "https://download.pytorch.org/whl/cpu" }, marker = "(sys_platform != 'darwin' and extra == 'extra-4-kumo-cpu') or (sys_platform != 'darwin' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, { name = "torch", version = "2.12.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" }, marker = "extra == 'extra-4-kumo-cu130' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, - { name = "torch", version = "2.12.0+cu132", source = { registry = "https://download.pytorch.org/whl/cu132" }, marker = "extra == 'extra-4-kumo-cu132' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "torch", version = "2.12.0+cu132", source = { registry = "https://download.pytorch.org/whl/cu132" }, marker = "extra == 'extra-4-kumo-cu132' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, { name = "torchmetrics" }, { name = "tqdm" }, { name = "typing-extensions" }, @@ -5991,6 +6254,28 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/0f/7b/39c34ca613b0b198cb866466651b26b045e2009864c5183c979a3b83f383/pytz-2026.3.post1-py2.py3-none-any.whl", hash = "sha256:dd95840dd199baea12d9cc096a1d452caa6596a1c1e4b5f3dbd1541855d5e815", size = 508283, upload-time = "2026-07-25T15:12:05.782Z" }, ] +[[package]] +name = "pywin32" +version = "312" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/1f/f5/10a6e845a00fc5e7afd0a988b744f403d4d57162a28d160a093c4d9322f0/pywin32-312-cp311-cp311-win32.whl", hash = "sha256:17948aeadbdb091f0ced6ef0841620794e68327b94ee415571c1203594b7215c", size = 6362659, upload-time = "2026-06-04T07:49:21.349Z" }, + { url = "https://files.pythonhosted.org/packages/35/c4/dcd2d62b5944b6d5db53413a5899016ccd57ffcb7278f3f81655d25d2027/pywin32-312-cp311-cp311-win_amd64.whl", hash = "sha256:d11417d84412f859b722fad0841b3614459ed0047f7542d8362e77884f6b6e8a", size = 6928825, upload-time = "2026-06-04T07:49:23.934Z" }, + { url = "https://files.pythonhosted.org/packages/b7/56/3cbb433fe4501cdba2eb9040f56a4e1a8243faa4186b25295564d1a7a79d/pywin32-312-cp311-cp311-win_arm64.whl", hash = "sha256:b2200a054ca6d6625c4842fc56a4976a4b47f96b73dbe5538c3f813a80359f47", size = 6721875, upload-time = "2026-06-04T07:49:26.416Z" }, + { url = "https://files.pythonhosted.org/packages/83/ff/32aa7d2ed0ab12b323aaa64f9b75e6ad4f8fd09f9ccfc28c79414d46838d/pywin32-312-cp312-cp312-win32.whl", hash = "sha256:dab4f65ac9c4e48400a2a0530c46c3c579cd5905ecd11b80692373915269208b", size = 6371877, upload-time = "2026-06-04T07:49:28.836Z" }, + { url = "https://files.pythonhosted.org/packages/03/d9/77040d3b43df3f3be32ea289433d660d2727f5ba327bc73be835127d9d60/pywin32-312-cp312-cp312-win_amd64.whl", hash = "sha256:b457f6d628a47e8a7346ce22acb7e1a46a4a78b52e1d17e1af56871bd19a93bc", size = 6914841, upload-time = "2026-06-04T07:49:31.85Z" }, + { url = "https://files.pythonhosted.org/packages/e3/cc/7b1ec671775756020a0ee7f4feeaf3c568f0ab86bd3900088cf986937a92/pywin32-312-cp312-cp312-win_arm64.whl", hash = "sha256:6017c58e12f6809fbb0555b75df144c2922a9ffd18e4b9b5afa863b6c1a9d950", size = 6727901, upload-time = "2026-06-04T07:49:34.244Z" }, + { url = "https://files.pythonhosted.org/packages/2d/41/12fbfd7f36ed2146d8bc9de96c2741296bf0d490b98508496cff322e274c/pywin32-312-cp313-cp313-win32.whl", hash = "sha256:7a27df850933d16a8eabfbaeb73d52b273e2da667f80d70b01a89d1f6828d02c", size = 6370184, upload-time = "2026-06-04T07:49:36.253Z" }, + { url = "https://files.pythonhosted.org/packages/ba/db/36a78e3403099d31d9746d13fdcde5accc43c1155f375a34d15983a479a7/pywin32-312-cp313-cp313-win_amd64.whl", hash = "sha256:c53e878d15a1c44788082bfe712a905433473aa38f86375b7cf8b45e3acbaaf9", size = 6914298, upload-time = "2026-06-04T07:49:38.876Z" }, + { url = "https://files.pythonhosted.org/packages/84/37/c1697194092b76de9ed47ca124323f02c57ffc8a45c06f88a3d5acaf01eb/pywin32-312-cp313-cp313-win_arm64.whl", hash = "sha256:59aba5d5940842075343a5ddc6b11f1cdf0d1567fe745290359dfbcc7c2eb831", size = 6727640, upload-time = "2026-06-04T07:49:41.083Z" }, + { url = "https://files.pythonhosted.org/packages/fc/2b/1f3cded5822fd49c02f40544cbb5f58c7cfd6b1694869fd476cb6170ee97/pywin32-312-cp314-cp314-win32.whl", hash = "sha256:a77a90fbb6881238d2ca9c6fd797b25817f3768fe78d214a90137ff055a75f5b", size = 6468928, upload-time = "2026-06-04T07:49:43.188Z" }, + { url = "https://files.pythonhosted.org/packages/21/82/3bf86d2e2808902013132e1ce905a7da0da53790f3836c64bf44d55e24f3/pywin32-312-cp314-cp314-win_amd64.whl", hash = "sha256:a4dd3a848290ef724347b19f301045831d8e802fa4464f491b98b1e0a081432e", size = 7024157, upload-time = "2026-06-04T07:49:45.34Z" }, + { url = "https://files.pythonhosted.org/packages/a4/0e/73f6d6800b4f27655abd9e9f6aaeaefcddb2b946e4674efa2bab184a7f7b/pywin32-312-cp314-cp314-win_arm64.whl", hash = "sha256:9fce94568364e0155e6dfb781ac5d95903be8baf28670632beab1b523f300daa", size = 6839598, upload-time = "2026-06-04T07:49:47.613Z" }, + { url = "https://files.pythonhosted.org/packages/eb/61/caa39686032d2ebdd04ff0ab5cbe163126c0066d98e00c9018646e42393b/pywin32-312-cp315-cp315-win32.whl", hash = "sha256:5c1fbe4a937a73ae9297384a3da38518cbc694c68ad8a809b2e19acd350f03ed", size = 6471159, upload-time = "2026-06-04T07:49:50.035Z" }, + { url = "https://files.pythonhosted.org/packages/0f/cd/7e1de64a4a6f69c04214169657ccab0d93a670ea50e35eb8f489d7378249/pywin32-312-cp315-cp315-win_amd64.whl", hash = "sha256:c2f03a0f73f804a13c2735b99392b0cd426bb4f2c4d0178e5ac966a0f21618d5", size = 7025293, upload-time = "2026-06-04T07:49:54.857Z" }, + { url = "https://files.pythonhosted.org/packages/23/ed/4532e9388e65fa16b46776ef47ad631a64eda1631884488af707666350ed/pywin32-312-cp315-cp315-win_arm64.whl", hash = "sha256:a8597d28f267b39074aef51fa593530082b39cbe5a074226096857b1fed2dfb9", size = 6840337, upload-time = "2026-06-04T07:49:57.531Z" }, +] + [[package]] name = "pyyaml" version = "6.0.2" @@ -6207,8 +6492,8 @@ dependencies = [ { name = "joblib" }, { name = "narwhals" }, { name = "numpy" }, - { name = "scipy", version = "1.17.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, - { name = "scipy", version = "1.18.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "scipy", version = "1.17.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" }, + { name = "scipy", version = "1.18.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" }, { name = "threadpoolctl" }, ] sdist = { url = "https://files.pythonhosted.org/packages/d2/eb/eaf5e07fcc0da7149b0e084f24e54edd7441b9a89ce7e034032ae97fe3a0/scikit_learn-1.9.1.tar.gz", hash = "sha256:629cada3e33e2b9bf376cdc7614a47a4140b8aedc1d836579e359736fbd82977", size = 7786908, upload-time = "2026-09-10T18:34:04.679Z" } @@ -6271,7 +6556,7 @@ resolution-markers = [ "python_full_version < '3.12' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu130' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests'", ] dependencies = [ - { name = "numpy", marker = "python_full_version < '3.12' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "numpy" }, ] sdist = { url = "https://files.pythonhosted.org/packages/7a/97/5a3609c4f8d58b039179648e62dd220f89864f56f7357f5d4f45c29eb2cc/scipy-1.17.1.tar.gz", hash = "sha256:95d8e012d8cb8816c226aef832200b1d45109ed4464303e997c5b13122b297c0", size = 30573822, upload-time = "2026-02-23T00:26:24.851Z" } wheels = [ @@ -6365,7 +6650,7 @@ resolution-markers = [ "python_full_version == '3.12.*' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu130' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests'", ] dependencies = [ - { name = "numpy", marker = "python_full_version >= '3.12' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "numpy" }, ] sdist = { url = "https://files.pythonhosted.org/packages/7e/74/66de6258867beb2ef08f35f9f2ac017a52cacd5081714d239ff1a442d458/scipy-1.18.1.tar.gz", hash = "sha256:52c4b7422442aba924d03ad4019852b08a92e64ea187b933135687bfe2747307", size = 30781235, upload-time = "2026-08-21T23:28:50.599Z" } wheels = [ @@ -6476,6 +6761,23 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/b7/ce/149a00dd41f10bc29e5921b496af8b574d8413afcd5e30dfa0ed46c2cc5e/six-1.17.0-py2.py3-none-any.whl", hash = "sha256:4721f391ed90541fddacab5acf947aa0d3dc7d27b2e1e8eda2be8970586c3274", size = 11050, upload-time = "2024-12-04T17:35:26.475Z" }, ] +[[package]] +name = "skops" +version = "0.16.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "numpy" }, + { name = "packaging" }, + { name = "prettytable" }, + { name = "scikit-learn" }, + { name = "scipy", version = "1.17.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" }, + { name = "scipy", version = "1.18.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/52/4b/7fef0af1efcb06c0986762645fb414ad850314a891434a2bbaea12a962ea/skops-0.16.0.tar.gz", hash = "sha256:58ddb63fed9f2a8545baa1a91c45c26d44cf65c83e2fd217541c41e411cf1d74", size = 673544, upload-time = "2026-09-24T06:29:46.96Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/44/a3/08ae9ea7f7c16e20a447d9289dce40f06f52607ed9a39f554a8426d96ba3/skops-0.16.0-py3-none-any.whl", hash = "sha256:5b2eb1bc49392c74a38db04481e2d110f97d9a72d6ae38fbd698aa7a865a9070", size = 148650, upload-time = "2026-09-24T06:29:45.491Z" }, +] + [[package]] name = "smart-open" version = "5.2.1" @@ -6485,6 +6787,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/cd/11/05f68ea934c24ade38e95ac30a38407767787c4e3db1776eae4886ad8c95/smart_open-5.2.1-py3-none-any.whl", hash = "sha256:71d14489da58b60ce12fc3ecb823facc59a8b23cd1b58edb97175640350d3a62", size = 58638, upload-time = "2021-08-28T11:06:50.07Z" }, ] +[[package]] +name = "smmap" +version = "5.0.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/1f/ea/49c993d6dfdd7338c9b1000a0f36817ed7ec84577ae2e52f890d1a4ff909/smmap-5.0.3.tar.gz", hash = "sha256:4d9debb8b99007ae47165abc08670bd74cb74b5227dda7f643eccc4e9eb5642c", size = 22506, upload-time = "2026-03-09T03:43:26.1Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c1/d4/59e74daffcb57a07668852eeeb6035af9f32cbfd7a1d2511f17d2fe6a738/smmap-5.0.3-py3-none-any.whl", hash = "sha256:c106e05d5a61449cf6ba9a1e650227ecfb141590d2a98412103ff35d89fc7b2f", size = 24390, upload-time = "2026-03-09T03:43:24.361Z" }, +] + [[package]] name = "sniffio" version = "1.3.1" @@ -6582,8 +6893,8 @@ dependencies = [ { name = "pytimeparse" }, { name = "pyyaml" }, { name = "retrying" }, - { name = "scipy", version = "1.17.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, - { name = "scipy", version = "1.18.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "scipy", version = "1.17.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" }, + { name = "scipy", version = "1.18.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" }, { name = "snowflake-connector-python", extra = ["pandas"] }, { name = "snowflake-snowpark-python" }, { name = "sqlparse" }, @@ -6640,7 +6951,7 @@ source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "alabaster" }, { name = "babel" }, - { name = "colorama", marker = "sys_platform == 'win32' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "colorama", marker = "sys_platform == 'win32'" }, { name = "docutils" }, { name = "imagesize" }, { name = "jinja2" }, @@ -6794,7 +7105,7 @@ name = "sqlalchemy" version = "2.0.54" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "greenlet", marker = "platform_machine == 'AMD64' or platform_machine == 'WIN32' or platform_machine == 'aarch64' or platform_machine == 'amd64' or platform_machine == 'ppc64le' or platform_machine == 'win32' or platform_machine == 'x86_64' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "greenlet", marker = "platform_machine == 'AMD64' or platform_machine == 'WIN32' or platform_machine == 'aarch64' or platform_machine == 'amd64' or platform_machine == 'ppc64le' or platform_machine == 'win32' or platform_machine == 'x86_64'" }, { name = "typing-extensions" }, ] sdist = { url = "https://files.pythonhosted.org/packages/29/9c/271aa905cf2964f841371a97f3e63ab692bf51b4423d0491e67bc7f64037/sqlalchemy-2.0.54.tar.gz", hash = "sha256:baa8521e8ee9f24e75dfc7aaabc08020e551ef0d48d7c3e3536f5cddf277586b", size = 9969559, upload-time = "2026-09-15T21:06:57.337Z" } @@ -6861,7 +7172,7 @@ version = "1.3.1" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "anyio" }, - { name = "typing-extensions", marker = "python_full_version < '3.13' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "typing-extensions", marker = "python_full_version < '3.13'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/eb/e3/7c1dc7381d9f8ab7d854328ebfa884e62cb3f3d8549ddfd37c7814f42afa/starlette-1.3.1.tar.gz", hash = "sha256:05d0213193f2fbaae60e2ecb593b4add4262ad4e46536b54abe36f11a71724e0", size = 2703240, upload-time = "2026-06-12T09:23:11.602Z" } wheels = [ @@ -7067,14 +7378,14 @@ resolution-markers = [ "python_full_version < '3.12' and sys_platform == 'darwin'", ] dependencies = [ - { name = "filelock", marker = "(sys_platform == 'darwin' and extra == 'extra-4-kumo-cpu') or (sys_platform == 'darwin' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, - { name = "fsspec", extra = ["http"], marker = "(sys_platform == 'darwin' and extra == 'extra-4-kumo-cpu') or (sys_platform == 'darwin' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, - { name = "jinja2", marker = "(sys_platform == 'darwin' and extra == 'extra-4-kumo-cpu') or (sys_platform == 'darwin' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, - { name = "networkx", version = "3.6.1", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version < '3.12' and sys_platform == 'darwin' and extra == 'extra-4-kumo-cpu') or (python_full_version < '3.12' and sys_platform == 'darwin' and extra != 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (python_full_version >= '3.12' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (python_full_version >= '3.12' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (python_full_version >= '3.12' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (python_full_version >= '3.12' and extra != 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'darwin' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform != 'darwin' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'darwin' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'darwin' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'darwin' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'darwin' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (extra != 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra != 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, - { name = "networkx", version = "3.7", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version >= '3.12' and sys_platform == 'darwin' and extra == 'extra-4-kumo-cpu') or (python_full_version >= '3.12' and sys_platform == 'darwin' and extra != 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (python_full_version < '3.12' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (python_full_version < '3.12' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (python_full_version < '3.12' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (python_full_version < '3.12' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (python_full_version < '3.12' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (python_full_version < '3.12' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'darwin' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform != 'darwin' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'darwin' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'darwin' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'darwin' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'darwin' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (extra != 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra != 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, - { name = "setuptools", marker = "(sys_platform == 'darwin' and extra == 'extra-4-kumo-cpu') or (sys_platform == 'darwin' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, - { name = "sympy", marker = "(sys_platform == 'darwin' and extra == 'extra-4-kumo-cpu') or (sys_platform == 'darwin' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, - { name = "typing-extensions", marker = "(sys_platform == 'darwin' and extra == 'extra-4-kumo-cpu') or (sys_platform == 'darwin' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, + { name = "filelock", marker = "(sys_platform == 'darwin' and extra == 'extra-4-kumo-cpu') or (sys_platform == 'darwin' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "fsspec", extra = ["http"], marker = "(sys_platform == 'darwin' and extra == 'extra-4-kumo-cpu') or (sys_platform == 'darwin' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "jinja2", marker = "(sys_platform == 'darwin' and extra == 'extra-4-kumo-cpu') or (sys_platform == 'darwin' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "networkx", version = "3.6.1", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version < '3.12' and sys_platform == 'darwin' and extra == 'extra-4-kumo-cpu') or (python_full_version < '3.12' and sys_platform == 'darwin' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "networkx", version = "3.7", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version >= '3.12' and sys_platform == 'darwin' and extra == 'extra-4-kumo-cpu') or (python_full_version >= '3.12' and sys_platform == 'darwin' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "setuptools", marker = "(sys_platform == 'darwin' and extra == 'extra-4-kumo-cpu') or (sys_platform == 'darwin' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "sympy", marker = "(sys_platform == 'darwin' and extra == 'extra-4-kumo-cpu') or (sys_platform == 'darwin' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "typing-extensions", marker = "(sys_platform == 'darwin' and extra == 'extra-4-kumo-cpu') or (sys_platform == 'darwin' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, ] wheels = [ { url = "https://download-r2.pytorch.org/whl/cpu/torch-2.12.0-cp311-cp311-macosx_14_0_arm64.whl", hash = "sha256:10802fd383bbfed646212e765a72c37d2185205d4f26eb197a254e8ac7ddcb25", upload-time = "2026-05-12T16:20:07Z" }, @@ -7151,14 +7462,14 @@ resolution-markers = [ "python_full_version < '3.12' and sys_platform != 'darwin'", ] dependencies = [ - { name = "filelock", marker = "(sys_platform != 'darwin' and extra == 'extra-4-kumo-cpu') or (sys_platform != 'darwin' and extra != 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, - { name = "fsspec", extra = ["http"], marker = "(sys_platform != 'darwin' and extra == 'extra-4-kumo-cpu') or (sys_platform != 'darwin' and extra != 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, - { name = "jinja2", marker = "(sys_platform != 'darwin' and extra == 'extra-4-kumo-cpu') or (sys_platform != 'darwin' and extra != 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, - { name = "networkx", version = "3.6.1", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version < '3.12' and sys_platform != 'darwin' and extra == 'extra-4-kumo-cpu') or (python_full_version < '3.12' and sys_platform != 'darwin' and extra != 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (python_full_version >= '3.12' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (python_full_version >= '3.12' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (python_full_version >= '3.12' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (python_full_version >= '3.12' and extra != 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'darwin' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'darwin' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'darwin' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'darwin' and extra != 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (extra != 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra != 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, - { name = "networkx", version = "3.7", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version >= '3.12' and sys_platform != 'darwin' and extra == 'extra-4-kumo-cpu') or (python_full_version >= '3.12' and sys_platform != 'darwin' and extra != 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (python_full_version < '3.12' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (python_full_version < '3.12' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (python_full_version < '3.12' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (python_full_version < '3.12' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (python_full_version < '3.12' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (python_full_version < '3.12' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'darwin' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'darwin' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'darwin' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'darwin' and extra != 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (extra != 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra != 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, - { name = "setuptools", marker = "(sys_platform != 'darwin' and extra == 'extra-4-kumo-cpu') or (sys_platform != 'darwin' and extra != 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, - { name = "sympy", marker = "(sys_platform != 'darwin' and extra == 'extra-4-kumo-cpu') or (sys_platform != 'darwin' and extra != 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, - { name = "typing-extensions", marker = "(sys_platform != 'darwin' and extra == 'extra-4-kumo-cpu') or (sys_platform != 'darwin' and extra != 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "filelock", marker = "(sys_platform != 'darwin' and extra == 'extra-4-kumo-cpu') or (sys_platform != 'darwin' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "fsspec", extra = ["http"], marker = "(sys_platform != 'darwin' and extra == 'extra-4-kumo-cpu') or (sys_platform != 'darwin' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "jinja2", marker = "(sys_platform != 'darwin' and extra == 'extra-4-kumo-cpu') or (sys_platform != 'darwin' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "networkx", version = "3.6.1", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version < '3.12' and sys_platform != 'darwin' and extra == 'extra-4-kumo-cpu') or (python_full_version < '3.12' and sys_platform != 'darwin' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "networkx", version = "3.7", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version >= '3.12' and sys_platform != 'darwin' and extra == 'extra-4-kumo-cpu') or (python_full_version >= '3.12' and sys_platform != 'darwin' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "setuptools", marker = "(sys_platform != 'darwin' and extra == 'extra-4-kumo-cpu') or (sys_platform != 'darwin' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "sympy", marker = "(sys_platform != 'darwin' and extra == 'extra-4-kumo-cpu') or (sys_platform != 'darwin' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "typing-extensions", marker = "(sys_platform != 'darwin' and extra == 'extra-4-kumo-cpu') or (sys_platform != 'darwin' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, ] wheels = [ { url = "https://download-r2.pytorch.org/whl/cpu/torch-2.12.0%2Bcpu-cp311-cp311-linux_s390x.whl", hash = "sha256:aaa9c1f5e8c518d7be1e3c3e1b090ca7d63b6e353a1abd6cfdaf902405093467", upload-time = "2026-05-12T23:16:01Z" }, @@ -7201,21 +7512,21 @@ resolution-markers = [ "python_full_version < '3.12'", ] dependencies = [ - { name = "cuda-bindings", marker = "(sys_platform == 'linux' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, - { name = "cuda-toolkit", version = "13.0.2", source = { registry = "https://pypi.org/simple" }, extra = ["cudart", "cufft", "cufile", "cupti", "curand", "cusolver", "cusparse", "nvjitlink", "nvrtc", "nvtx"], marker = "(sys_platform == 'linux' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "cuda-bindings", marker = "(sys_platform == 'linux' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "cuda-toolkit", version = "13.0.2", source = { registry = "https://pypi.org/simple" }, extra = ["cudart", "cufft", "cufile", "cupti", "curand", "cusolver", "cusparse", "nvjitlink", "nvrtc", "nvtx"], marker = "(sys_platform == 'linux' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, { name = "filelock", marker = "extra == 'extra-4-kumo-cu130' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, { name = "fsspec", extra = ["http"], marker = "extra == 'extra-4-kumo-cu130' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, { name = "jinja2", marker = "extra == 'extra-4-kumo-cu130' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, { name = "networkx", version = "3.6.1", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version < '3.12' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, - { name = "networkx", version = "3.7", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version >= '3.12' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, - { name = "nvidia-cublas", version = "13.1.1.3", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform == 'linux' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, - { name = "nvidia-cudnn-cu13", marker = "(sys_platform == 'linux' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, - { name = "nvidia-cusparselt-cu13", marker = "(sys_platform == 'linux' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, - { name = "nvidia-nccl-cu13", marker = "(sys_platform == 'linux' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, - { name = "nvidia-nvshmem-cu13", marker = "(sys_platform == 'linux' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "networkx", version = "3.7", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version >= '3.12' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-cublas", version = "13.1.1.3", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform == 'linux' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-cudnn-cu13", marker = "(sys_platform == 'linux' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-cusparselt-cu13", marker = "(sys_platform == 'linux' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-nccl-cu13", marker = "(sys_platform == 'linux' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-nvshmem-cu13", marker = "(sys_platform == 'linux' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, { name = "setuptools", marker = "extra == 'extra-4-kumo-cu130' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, { name = "sympy", marker = "extra == 'extra-4-kumo-cu130' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, - { name = "triton", marker = "(sys_platform == 'linux' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "triton", marker = "(sys_platform == 'linux' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, { name = "typing-extensions", marker = "extra == 'extra-4-kumo-cu130' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, ] wheels = [ @@ -7250,21 +7561,21 @@ resolution-markers = [ "python_full_version < '3.12'", ] dependencies = [ - { name = "cuda-bindings", marker = "(sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, - { name = "cuda-toolkit", version = "13.2.1", source = { registry = "https://pypi.org/simple" }, extra = ["cublas", "cudart", "cufft", "cufile", "cupti", "curand", "cusolver", "cusparse", "nvjitlink", "nvrtc", "nvtx"], marker = "(sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, - { name = "filelock", marker = "extra == 'extra-4-kumo-cu132' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, - { name = "fsspec", extra = ["http"], marker = "extra == 'extra-4-kumo-cu132' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, - { name = "jinja2", marker = "extra == 'extra-4-kumo-cu132' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "cuda-bindings", marker = "(sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "cuda-toolkit", version = "13.2.1", source = { registry = "https://pypi.org/simple" }, extra = ["cublas", "cudart", "cufft", "cufile", "cupti", "curand", "cusolver", "cusparse", "nvjitlink", "nvrtc", "nvtx"], marker = "(sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "filelock", marker = "extra == 'extra-4-kumo-cu132' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, + { name = "fsspec", extra = ["http"], marker = "extra == 'extra-4-kumo-cu132' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, + { name = "jinja2", marker = "extra == 'extra-4-kumo-cu132' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, { name = "networkx", version = "3.6.1", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version < '3.12' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, - { name = "networkx", version = "3.7", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version >= '3.12' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, - { name = "nvidia-cudnn-cu13", marker = "(sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, - { name = "nvidia-cusparselt-cu13", marker = "(sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, - { name = "nvidia-nccl-cu13", marker = "(sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, - { name = "nvidia-nvshmem-cu13", marker = "(sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, - { name = "setuptools", marker = "extra == 'extra-4-kumo-cu132' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, - { name = "sympy", marker = "extra == 'extra-4-kumo-cu132' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, - { name = "triton", marker = "(sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, - { name = "typing-extensions", marker = "extra == 'extra-4-kumo-cu132' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "networkx", version = "3.7", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version >= '3.12' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-cudnn-cu13", marker = "(sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-cusparselt-cu13", marker = "(sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-nccl-cu13", marker = "(sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-nvshmem-cu13", marker = "(sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "setuptools", marker = "extra == 'extra-4-kumo-cu132' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, + { name = "sympy", marker = "extra == 'extra-4-kumo-cu132' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, + { name = "triton", marker = "(sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "typing-extensions", marker = "extra == 'extra-4-kumo-cu132' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, ] wheels = [ { url = "https://download-r2.pytorch.org/whl/cu132/torch-2.12.0%2Bcu132-cp311-cp311-manylinux_2_28_aarch64.whl", hash = "sha256:c8b75ec50e9a6bfc6db0f1e0bd297a1b94c9803c378c706c4e856c7cad331398", upload-time = "2026-05-13T00:04:15Z" }, @@ -7315,11 +7626,11 @@ dependencies = [ { name = "lightning-utilities" }, { name = "numpy" }, { name = "packaging" }, - { name = "torch", version = "2.12.0", source = { registry = "https://download.pytorch.org/whl/cpu" }, marker = "(sys_platform == 'darwin' and extra == 'extra-4-kumo-cpu') or (sys_platform == 'darwin' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, + { name = "torch", version = "2.12.0", source = { registry = "https://download.pytorch.org/whl/cpu" }, marker = "(sys_platform == 'darwin' and extra == 'extra-4-kumo-cpu') or (sys_platform == 'darwin' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, { name = "torch", version = "2.12.0", source = { registry = "https://pypi.org/simple" }, marker = "(extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu130' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests')" }, - { name = "torch", version = "2.12.0+cpu", source = { registry = "https://download.pytorch.org/whl/cpu" }, marker = "(sys_platform != 'darwin' and extra == 'extra-4-kumo-cpu') or (sys_platform != 'darwin' and extra != 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "torch", version = "2.12.0+cpu", source = { registry = "https://download.pytorch.org/whl/cpu" }, marker = "(sys_platform != 'darwin' and extra == 'extra-4-kumo-cpu') or (sys_platform != 'darwin' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, { name = "torch", version = "2.12.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" }, marker = "extra == 'extra-4-kumo-cu130' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, - { name = "torch", version = "2.12.0+cu132", source = { registry = "https://download.pytorch.org/whl/cu132" }, marker = "extra == 'extra-4-kumo-cu132' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "torch", version = "2.12.0+cu132", source = { registry = "https://download.pytorch.org/whl/cu132" }, marker = "extra == 'extra-4-kumo-cu132' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/78/1f/2cd9eb8f3390c3ec4693ac0871913d4b468964b3833638e4091a70817e0a/torchmetrics-1.8.1.tar.gz", hash = "sha256:04ca021105871637c5d34d0a286b3ab665a1e3d2b395e561f14188a96e862fdb", size = 580373, upload-time = "2025-08-07T20:44:44.631Z" } wheels = [ @@ -7331,7 +7642,7 @@ name = "tqdm" version = "4.70.1" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "colorama", marker = "sys_platform == 'win32' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "colorama", marker = "sys_platform == 'win32'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/0d/ea/b2a5bd54b28a324dae8211928b2d730b6547500342c7e6c6dea08bd0a485/tqdm-4.70.1.tar.gz", hash = "sha256:cefd0eca11b2a37a3aee776544d4f4ae913f02688135b5556b8788dfa474afc4", size = 171846, upload-time = "2026-09-11T07:25:16.601Z" } wheels = [ @@ -7399,7 +7710,7 @@ version = "3.7.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "botocore-stubs" }, - { name = "typing-extensions", marker = "python_full_version < '3.12' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "typing-extensions", marker = "python_full_version < '3.12'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/76/e8/ef1fcb876937dbdddc0f01b5df4ed53f33b166a6367d80a9014d5e5f091d/types_aiobotocore-3.7.0.tar.gz", hash = "sha256:fe35de52c12e5fdb89ca60b3989766e7fe827e3d2e95fcf4583e91581945205c", size = 87992, upload-time = "2026-05-10T03:19:32.353Z" } wheels = [ @@ -7416,7 +7727,7 @@ name = "types-aiobotocore-s3" version = "3.7.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "typing-extensions", marker = "python_full_version < '3.12' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "typing-extensions", marker = "python_full_version < '3.12'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/94/ce/5b7baa06cde79ab5f083f56f060ccd2b60f6249127d70bbb8a37aafbdcd4/types_aiobotocore_s3-3.7.0.tar.gz", hash = "sha256:6ec738853dbba9133707991b98ea4dab19f7c62e02b3cca016e6cd8e3d684576", size = 77662, upload-time = "2026-05-10T03:17:35.088Z" } wheels = [ @@ -7484,7 +7795,7 @@ name = "tzlocal" version = "5.4.4" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "tzdata", marker = "sys_platform == 'win32' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "tzdata", marker = "sys_platform == 'win32'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/81/5b/879b2f932adfa7a053c360d50bc896c977fa6426109185f7c12ebdd0cb9d/tzlocal-5.4.4.tar.gz", hash = "sha256:8dbb8660838688a7b6ba4fed31d18dedf842afb4d47ca050d6d891c2c15f3be4", size = 31170, upload-time = "2026-06-29T08:03:40.026Z" } wheels = [ @@ -7528,6 +7839,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/9d/e5/9446b7a2a3d73de5f9fcf8f825efe78b071b111729e89f4ac07649b9448b/virtualenv-21.9.0-py3-none-any.whl", hash = "sha256:8c7cde05f630fbedf61758258014a92e6d6d5b4b9c6e6be7279d233f89e09aac", size = 5477625, upload-time = "2026-09-20T00:29:18.065Z" }, ] +[[package]] +name = "waitress" +version = "3.0.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/bf/cb/04ddb054f45faa306a230769e868c28b8065ea196891f09004ebace5b184/waitress-3.0.2.tar.gz", hash = "sha256:682aaaf2af0c44ada4abfb70ded36393f0e307f4ab9456a215ce0020baefc31f", size = 179901, upload-time = "2024-11-16T20:02:35.195Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/8d/57/a27182528c90ef38d82b636a11f606b0cbb0e17588ed205435f8affe3368/waitress-3.0.2-py3-none-any.whl", hash = "sha256:c56d67fd6e87c2ee598b76abdd4e96cfad1f24cacdea5078d382b1f9d7b5ed2e", size = 56232, upload-time = "2024-11-16T20:02:33.858Z" }, +] + [[package]] name = "watchdog" version = "6.0.0" @@ -7572,6 +7892,42 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/23/41/36c7246c1c273fd316d9b8b81581074becc3137135f9ccd6307960a7879c/watchtower-3.0.1-py3-none-any.whl", hash = "sha256:566874a9dee34b90517b9a9b7c2d5e6b15fb9fd00dacb03bff8a7bf4c6295fe0", size = 17331, upload-time = "2023-01-29T15:33:52.494Z" }, ] +[[package]] +name = "wcwidth" +version = "0.9.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/f0/b4/7830542634bb2d3e62aa3b586a72d5b3b6c91c3168929e7000ef3fed041d/wcwidth-0.9.2.tar.gz", hash = "sha256:ae0ef90b90f6af38b54f1fe6d58662ec33b3cb4b8391958a62416d654231727b", size = 955039, upload-time = "2026-10-05T00:24:05.521Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/59/1e/4532a81fb9dfbf4114a816775e0a36c3a64ee1d1f4bba2094e2da50be5dc/wcwidth-0.9.2-cp310-abi3-macosx_10_9_x86_64.whl", hash = "sha256:7ef5a940bd5e30bac6e721f1a48fce0cd7bb3ece19e9c5d139e72c76c35cfd07", size = 605343, upload-time = "2026-10-05T00:23:22.649Z" }, + { url = "https://files.pythonhosted.org/packages/a0/07/cb6940e81134b7ed25fa312ee9ab536a63db0793b149f88a90e603ceace9/wcwidth-0.9.2-cp310-abi3-macosx_11_0_arm64.whl", hash = "sha256:ae0800c5339423cc53d33a266ad264b42ba8aaa16d4464f6e6b1bee607f50b17", size = 609960, upload-time = "2026-10-05T00:23:27.049Z" }, + { url = "https://files.pythonhosted.org/packages/a4/80/15ad05d40bfa99155639fb9e13b3d77083aa0fab893c816db2543d29005c/wcwidth-0.9.2-cp310-abi3-manylinux1_i686.manylinux2014_i686.manylinux_2_17_i686.manylinux_2_5_i686.whl", hash = "sha256:9e542f1f8475b78452a295495d7a5bc3ead565112e9446a64dc93462a41c2a79", size = 762566, upload-time = "2026-10-05T00:23:38.322Z" }, + { url = "https://files.pythonhosted.org/packages/bc/f0/b8ef7758003d66b60f093695831a86dcc726aac01ee6446ffcbda27b61e3/wcwidth-0.9.2-cp310-abi3-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:674b518af28d38ee645ff97b74f5760abee5fad4bac74413bfc4b881ef2ce724", size = 771632, upload-time = "2026-10-05T00:23:32.448Z" }, + { url = "https://files.pythonhosted.org/packages/db/6c/f940133c71427c208575910e981942bd78c98b1f7cd0d1425ca4b7457c04/wcwidth-0.9.2-cp310-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:751bef0ab404b6a1dc028b56b4b85d46486be1c55833f80da533e42dc691f389", size = 769763, upload-time = "2026-10-05T00:23:40.175Z" }, + { url = "https://files.pythonhosted.org/packages/92/8f/285f862826f721964ec7c42f81dc53d23afbd723a0f4cd989651f8218e25/wcwidth-0.9.2-cp310-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:c3d80f39ba4653a595edae9aa46a509d14883790a8fc23c5db221ceb207f64b7", size = 782134, upload-time = "2026-10-05T00:23:33.926Z" }, + { url = "https://files.pythonhosted.org/packages/c2/2d/64aa54882a5d556d3654c1f926d9118b797461033e23a158409941a37c8f/wcwidth-0.9.2-cp310-abi3-musllinux_1_2_i686.whl", hash = "sha256:0a47e03d8293590ecce66c45dc20ff7b4b885e3c78093722239585eca0d77ab2", size = 782851, upload-time = "2026-10-05T00:23:41.974Z" }, + { url = "https://files.pythonhosted.org/packages/59/39/52389f6de7fe2e9c14ceb8253dd99034bd86e1c87847ea3c100a97dded9a/wcwidth-0.9.2-cp310-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:67d901a4ad99249eb775b4ee4769ca97fa405d35a75f46e83166910a47003f04", size = 784621, upload-time = "2026-10-05T00:23:43.449Z" }, + { url = "https://files.pythonhosted.org/packages/b3/8b/20225500a076ace27bbcc8a6fd7c55125133c57a618816c7b7b8b73070b1/wcwidth-0.9.2-cp310-abi3-win32.whl", hash = "sha256:ee1fd0db9d9fd711a70f3e7765e0e04c05d26982fa05361456163062549d7da4", size = 593123, upload-time = "2026-10-05T00:23:55.953Z" }, + { url = "https://files.pythonhosted.org/packages/5a/d6/b0690f55ea0483530a18bac917fbadbf54f35122510446fc370f5f1c2453/wcwidth-0.9.2-cp310-abi3-win_amd64.whl", hash = "sha256:2a9746de704242bd4fdaabb31dd46b82f694a56a8d21081ad89b679a89da9fec", size = 597827, upload-time = "2026-10-05T00:23:57.489Z" }, + { url = "https://files.pythonhosted.org/packages/e5/11/6ecf4e9e268ab1a4ec617ffcccc2ee4a71301625f5490912dbaba462fa9c/wcwidth-0.9.2-cp310-abi3-win_arm64.whl", hash = "sha256:b9c6ab615e03723b7f8760ea2f27758d656e7e13b51515c9dca5c3e8b04612fa", size = 598415, upload-time = "2026-10-05T00:23:51.517Z" }, + { url = "https://files.pythonhosted.org/packages/4e/41/549eef1ab767032bdbdc1f0ab655d404b082b1e9a1dab1361dbba90f64ed/wcwidth-0.9.2-cp314-cp314t-macosx_10_15_x86_64.whl", hash = "sha256:eda88ffdc97c0fbf193d407114f2c7a54b379f67f6e52a7531ee3b9fe749eca7", size = 606555, upload-time = "2026-10-05T00:23:24.188Z" }, + { url = "https://files.pythonhosted.org/packages/9b/64/a875ed7ea71cacadc0ae11b5fd3fac3486efd58bb25e67a7344248dceadd/wcwidth-0.9.2-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:1bf361c8705576760623b4724ae564666d73b016f9a778bcfd1c7345378ef4ec", size = 610829, upload-time = "2026-10-05T00:23:28.563Z" }, + { url = "https://files.pythonhosted.org/packages/c6/98/513095e484fe79b6f2613d6a72f855f5d56b65e15c215c2a6746fbc638f5/wcwidth-0.9.2-cp314-cp314t-manylinux1_i686.manylinux2014_i686.manylinux_2_17_i686.manylinux_2_5_i686.whl", hash = "sha256:97b878d1e158da5ed9ac5aac53fa3a55e282103af6a09ec353865613d1a31a76", size = 775819, upload-time = "2026-10-05T00:23:45.116Z" }, + { url = "https://files.pythonhosted.org/packages/22/fc/c02f3eec57224731e78f84b68e272250f784b6205acc7e0dcef6a7c23a0e/wcwidth-0.9.2-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:59dab4049cbd982b478bca098528df2c79a9160636a3a163ffebffcbd7d1b892", size = 786915, upload-time = "2026-10-05T00:23:35.323Z" }, + { url = "https://files.pythonhosted.org/packages/3d/6f/b0529a79bac3fe8d94f32b4237a13dbc3f955508753f6a6f06c73d679dc2/wcwidth-0.9.2-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:bb08ceb501d6aaf94066c3ee122dd825b152df40ff0bd0df4dc27126233b948e", size = 784186, upload-time = "2026-10-05T00:23:46.366Z" }, + { url = "https://files.pythonhosted.org/packages/d5/bd/6357c84ca9a734bfc735b7c48dbe21336b3777fab8a4101d14976dfe49a7/wcwidth-0.9.2-cp314-cp314t-win32.whl", hash = "sha256:8b4e381590b9b7390e07e22b2c0c1bb96ce50e1d2243c866d9387600362d51ed", size = 600791, upload-time = "2026-10-05T00:23:59.398Z" }, + { url = "https://files.pythonhosted.org/packages/98/de/037591ca18d897cc2179559dde72e6efc6ce0c90e9cd1e6bca4e87c38b4b/wcwidth-0.9.2-cp314-cp314t-win_amd64.whl", hash = "sha256:f2f7b3bba5a5d5f31fc350fd36ce5b84b693c83b7eb95ee630b720da5a5ce06f", size = 605463, upload-time = "2026-10-05T00:24:01.049Z" }, + { url = "https://files.pythonhosted.org/packages/d0/07/c9d96e106d938d26f7ab639bc80b8199359a1645ba6e3498413313ab6f38/wcwidth-0.9.2-cp314-cp314t-win_arm64.whl", hash = "sha256:734aa9405b321d1042301aa19c943c4731ee9e3460e4f8feea3299c064c97a14", size = 605936, upload-time = "2026-10-05T00:23:52.765Z" }, + { url = "https://files.pythonhosted.org/packages/82/8a/a28d61d910005ac93dfe48be3a0ebaa49352d88cebd25323e69e6ff2f4a8/wcwidth-0.9.2-cp315-cp315t-macosx_10_15_x86_64.whl", hash = "sha256:42dbcb76ce8af39e2c9db410ac3f9bdf4e47eb41d6f44525952f172d3d98f724", size = 606575, upload-time = "2026-10-05T00:23:25.663Z" }, + { url = "https://files.pythonhosted.org/packages/01/c2/a3c66bd32766c8f4d6dc47d572532ba014fe5be30489f2576aff7cada363/wcwidth-0.9.2-cp315-cp315t-macosx_11_0_arm64.whl", hash = "sha256:138e1f8898e431b2f2d7881f8ca8d75591c1d3c21aa53f54e989bd6b39811da2", size = 611045, upload-time = "2026-10-05T00:23:30.421Z" }, + { url = "https://files.pythonhosted.org/packages/ec/8a/d39964f8f8c019d7d439b9b501d3e7bb42fee69f00354040ba0b27b5824c/wcwidth-0.9.2-cp315-cp315t-manylinux1_i686.manylinux2014_i686.manylinux_2_17_i686.manylinux_2_5_i686.whl", hash = "sha256:5175609bf8cc7398a5f48aa35207bd64ebf9f45e4c70df65f7fdc7a988041a3c", size = 778143, upload-time = "2026-10-05T00:23:47.7Z" }, + { url = "https://files.pythonhosted.org/packages/2f/53/525da13e8f9ff7b5b4e74ec6f8d68bdee63905796972e086c6b1b96670d2/wcwidth-0.9.2-cp315-cp315t-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:e5f669ae8c3d969c72032f9cdee019674b666e522d45e1e2099a2e9dda4a341d", size = 786630, upload-time = "2026-10-05T00:23:36.967Z" }, + { url = "https://files.pythonhosted.org/packages/ef/9f/d6a0c6df354b9d93466548a65cbf4ffcb48c719bbd307504cf3e76740837/wcwidth-0.9.2-cp315-cp315t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:196b47cf32f9df27ccda6dc513237f3c2429c4c659db428d60a5bc443d10f270", size = 784839, upload-time = "2026-10-05T00:23:49.88Z" }, + { url = "https://files.pythonhosted.org/packages/bf/d7/3021feed1ed7926021ec134943ad3b24a2f7ea742cc9976461171482ed77/wcwidth-0.9.2-cp315-cp315t-win32.whl", hash = "sha256:0cd4f7f2e53905dcb110d213a4c8529b6733fa3d232d8c717f946cc69a10349b", size = 600845, upload-time = "2026-10-05T00:24:02.497Z" }, + { url = "https://files.pythonhosted.org/packages/63/80/6a03356d8ee38261e3a78cf89ee03d8e7f12c572d969237be00869e2dc73/wcwidth-0.9.2-cp315-cp315t-win_amd64.whl", hash = "sha256:33df042f96c61ed3cd5fb3742fba427553a635bc578799857a48aa79f774a0b9", size = 605441, upload-time = "2026-10-05T00:24:04.052Z" }, + { url = "https://files.pythonhosted.org/packages/0c/48/1a308a86a833fd12ff7a08d0d2491ff4a72c8a92d12f5ead8317630f771e/wcwidth-0.9.2-cp315-cp315t-win_arm64.whl", hash = "sha256:48719a9bc76c2f84238693fe5013571fa5beffa3621cf228f1f3a9e30dae84b8", size = 605885, upload-time = "2026-10-05T00:23:54.274Z" }, + { url = "https://files.pythonhosted.org/packages/9c/b4/0bfa065af506540d9d558e3e5548cff00bc1f9b24e6e2a8512498e8628de/wcwidth-0.9.2-py3-none-any.whl", hash = "sha256:89ca642c5bf0101157a09366be69fad0379db1f700ae39a920e103234573670e", size = 301667, upload-time = "2026-10-05T00:23:21.097Z" }, +] + [[package]] name = "websocket-client" version = "1.9.2"