diff --git a/.github/workflows/testing.yml b/.github/workflows/testing.yml index fd347dc17..16d72e930 100644 --- a/.github/workflows/testing.yml +++ b/.github/workflows/testing.yml @@ -723,6 +723,41 @@ jobs: cd packages/kumo-sdk && python -m pytest -rs \ -W "ignore:You have an incompatible version of 'pyarrow' installed:UserWarning" + mlflow-floor: + # kumo-engine's `mlflow` extra allows the MLflow a Databricks Runtime + # ships; run the MLflow suites against the extra's floor. The tests left + # out read logs, which MLflow 3.5's SQLite migrations reconfigure. + needs: [diskgraph-wheel] + if: ${{ !cancelled() }} + runs-on: ubuntu-24.04 + permissions: + contents: read + steps: + - name: Checkout repository + uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6 + with: + ref: ${{ github.sha }} + + - name: Set up packages + uses: ./.github/actions/setup + with: + github_token: ${{ secrets.KUMO_GITHUB_BOT_TOKEN }} + diskgraph-wheel-artifact: diskgraph-wheel + + - name: Run the MLflow suites at the mlflow extra's floor + run: | + uv pip install --no-deps -e packages/kumo-engine + # mlflow 3.5.0 pins a pyarrow with no wheel for this Python; the + # extra installs mlflow-skinny, which holds the client and stores. + uv pip uninstall mlflow mlflow-skinny mlflow-tracing + uv pip install "mlflow-skinny==3.5.0" + python -c "import mlflow; assert mlflow.__version__ == '3.5.0'" + cd packages/kumo-engine && python -m pytest -rs \ + test/test_mlflow_target.py test/test_mlflow_model.py \ + test/test_mlflow_tracking.py -k "not (reuses_its_version or \ + late_one_is_dropped or names_its_cause_and_can_raise or \ + ends_runs_failed_and_is_not_raised)" + durations-aggregate: runs-on: ubuntu-24.04 needs: [pytest] diff --git a/examples/mlops/smoke.py b/examples/mlops/smoke.py new file mode 100644 index 000000000..3cd5c92ff --- /dev/null +++ b/examples/mlops/smoke.py @@ -0,0 +1,118 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Manual live check per platform; Databricks reads DATABRICKS_CLIENT_SECRET.""" + +import argparse +import os +import tempfile +import time +from pathlib import Path +from typing import Optional, Sequence + +import numpy as np +import pandas as pd + +import kumo.engine as kumo +from kumo.engine import ModelRegistryLocation +from kumo.engine.connector import ( + DatabricksConnector, + DatabricksIdentity, + S3Identity, + SnowflakeConnector, + SnowflakeIdentity, +) + +QUERY = 'PREDICT COUNT(orders.*, 0, 7, days) > 0 FOR EACH users.USER_ID' + + +def data(root: Path) -> kumo.Graph: + rng = np.random.default_rng(0) + root.mkdir(parents=True) + pd.DataFrame({'USER_ID': np.arange(200)}).to_parquet(root / 'users.parquet') + pd.DataFrame({ + 'ORDER_ID': np.arange(6000), + 'USER_ID': rng.integers(0, 200, 6000), + 'ORDER_TIME': pd.Timestamp('2024-01-01') + + pd.to_timedelta(rng.integers(0, 120, 6000), unit='D'), + }).to_parquet(root / 'orders.parquet') + files = kumo.FileConnector(root_dir=str(root)) + orders = kumo.Table.from_source_table( + files['orders.parquet'], primary_key='ORDER_ID', + time_column='ORDER_TIME').infer_metadata() + return kumo.Graph( + tables={ + 'users': kumo.Table.from_source_table( + files['users.parquet'], primary_key='USER_ID'), + 'orders': orders, + }, edges=[kumo.Edge(src_table='orders', fkey='USER_ID', + dst_table='users')]) + + +def tracking(args: argparse.Namespace, suffix: str): + from kumo.engine.integrations.mlflow import MlflowTracking, SageMakerMLflow + + if args.platform == 'databricks': + workspace = DatabricksConnector( + catalog=args.catalog, schema_name=args.schema, + warehouse_id=args.warehouse, identity=DatabricksIdentity( + host=args.host, client_id=args.client_id, + client_secret=os.environ['DATABRICKS_CLIENT_SECRET'])) + return (MlflowTracking(target=workspace, + experiment_name=args.experiment), + f'{args.catalog}.{args.schema}.kumo_smoke_{suffix}') + if args.platform == 'sagemaker': + identity = S3Identity(role_arn=args.role_arn) if args.role_arn else None + return (MlflowTracking( + target=SageMakerMLflow(args.arn, identity=identity), + experiment_name='kumo-smoke'), f'kumo-smoke-{suffix}') + from kumo.engine.integrations.snowflake import SnowflakeTracking + + connector = SnowflakeConnector( + warehouse=args.warehouse, database=args.database, + schema_name=args.schema, identity=SnowflakeIdentity( + account=args.account, username=args.user, + private_key=Path(args.private_key_file).read_text())) + return (SnowflakeTracking(connector, experiment_name='KUMO_SMOKE'), + f'KUMO_SMOKE_{suffix}') + + +def main(argv: Optional[Sequence[str]] = None) -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument('platform', + choices=('databricks', 'sagemaker', 'snowflake')) + for flag in ('--host', '--client-id', '--catalog', '--schema', + '--warehouse', '--experiment', '--arn', '--role-arn', + '--database', '--account', '--user', '--private-key-file'): + parser.add_argument(flag) + parser.add_argument('--artifact-root', + help='defaults to a local temporary directory') + args = parser.parse_args(argv) + + started = time.monotonic() + work = Path(tempfile.mkdtemp(prefix='kumo-smoke-')) + kumo.configure(scratch_dir=work / 'scratch') + graph = data(work / 'data') + tracker, model_name = tracking(args, time.strftime('%Y%m%d%H%M%S')) + location = ModelRegistryLocation(model_name, alias='candidate') + cfg = kumo.MLOps(args.artifact_root or str(work / 'root'), + tracking=tracker, registry=location, on_error='raise') + print(f'checked the target in {time.monotonic() - started:.0f} s') + with kumo.mlops(cfg): + snapshot = graph.snapshot() + query = kumo.PredictiveQuery(graph=graph, query=QUERY) + table = query.generate_training_table(snapshot) + plan = query.suggest_model_plan(snapshot) + plan.training_job.num_experiments = 1 + plan.optimization.max_epochs = 1 + result = kumo.Trainer(plan).fit(table, snapshot=snapshot) + print(f'fit run {result.tracking}, bundle {result.id}, ' + f'registered {result.registered_model}') + loaded = kumo.Trainer.load_from_connector(cfg, location=location) + assert loaded.result.id == result.id, (loaded.result.id, result.id) + print('loaded the registered model back') + print(f'done in {time.monotonic() - started:.0f} s') + + +if __name__ == '__main__': + main() diff --git a/packages/kumo-api/kumoapi/plugin.py b/packages/kumo-api/kumoapi/plugin.py index 2e855462e..709e4a59d 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, @@ -42,10 +44,13 @@ from kumoapi.typing import Dtype __all__ = [ + 'BUNDLE_ID_KEY', + 'ENGINE_MARKER', 'ArtifactLocation', 'ArtifactSink', 'FileStoreSink', 'IngestResult', + 'ModelAliasingSink', 'ModelFetchingSink', 'ModelListingSink', 'ModelPublishingSink', @@ -56,8 +61,16 @@ 'SourceBackend', 'SourceConnector', 'SourceCredentials', + 'check_model_alias', + 'registry_lineage', ] +_REGISTRY_KEY_PREFIX = 'kumo_' +ENGINE_MARKER = 'kumo_engine' +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: @@ -149,13 +162,15 @@ 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 on export; read on import. """ 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 @@ -166,6 +181,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) @@ -702,7 +721,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: @@ -712,6 +732,38 @@ def fetch_model( ... +@runtime_checkable +class ModelAliasingSink(ArtifactSink, Protocol): + r"""A registry sink that sets and fetches by 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 registry accepts ``alias``.""" + if (not isinstance(alias, str) or not _ALIAS.fullmatch(alias) + or _RESERVED_ALIAS.fullmatch(alias)): + raise ValueError( + f'{alias!r} cannot be a model alias: use letters, digits and ' + f'underscores, not starting with a digit, other than the reserved ' + f'DEFAULT, FIRST, LAST, LATEST and v.') + + +def registry_lineage(recorded: Mapping[str, Any]) -> Dict[str, str]: + r"""A version's ``kumo_*`` registry metadata, without the 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-core/src/kumo_core/artifact_export.py b/packages/kumo-core/src/kumo_core/artifact_export.py index cc5371652..cac3cfa28 100644 --- a/packages/kumo-core/src/kumo_core/artifact_export.py +++ b/packages/kumo-core/src/kumo_core/artifact_export.py @@ -304,10 +304,14 @@ def prepared_artifact_file(path, extra_columns=None, *, parquet_options=None, writer = None try: for batch in local_parquet_batches(path): - frame = batch.to_pandas() + # Pandas turns all-null bool columns into Arrow nulls. + table = pa.Table.from_batches([batch]) for name, value in (extra_columns or {}).items(): - frame[name] = value - table = pa.Table.from_pandas(frame, preserve_index=False) + column = pa.repeat(pa.scalar(value), table.num_rows) + table = (table.set_column( + table.column_names.index(name), name, column) + if name in table.column_names else + table.append_column(name, column)) if writer is None: writer = pq.ParquetWriter(staged, table.schema, **(parquet_options or {})) diff --git a/packages/kumo-core/src/kumo_core/artifacts.py b/packages/kumo-core/src/kumo_core/artifacts.py index 134463d8d..413a4eb93 100644 --- a/packages/kumo-core/src/kumo_core/artifacts.py +++ b/packages/kumo-core/src/kumo_core/artifacts.py @@ -14,6 +14,7 @@ from dataclasses import dataclass, field from pathlib import Path from typing import ( + TYPE_CHECKING, Any, ClassVar, Iterable, @@ -26,6 +27,9 @@ from .io.json import load_json, save_json from .manifest import RunManifest, RunPaths +if TYPE_CHECKING: + from .types import TaskType + PROVENANCE_FILE = '_provenance.json' @@ -378,6 +382,26 @@ def add_bundle_file(bundle_dir: Path, role: str, relative_path: str) -> None: _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' @@ -475,6 +499,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 12ebb575b..3dac36e15 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 @@ -460,6 +461,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. @@ -715,7 +717,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 @@ -849,7 +851,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() @@ -874,12 +876,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( @@ -887,7 +894,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 " @@ -904,9 +913,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) @@ -959,6 +980,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 109ead225..a758e1c44 100644 --- a/packages/kumo-core/src/kumo_core/automl/optuna_searcher.py +++ b/packages/kumo-core/src/kumo_core/automl/optuna_searcher.py @@ -28,11 +28,6 @@ def __init__( max_attempts: int = 50, **kwargs: Any, ) -> None: - warnings.filterwarnings( - "ignore", - message=".*plot_.*", - category=optuna.exceptions.ExperimentalWarning, - ) super().__init__(space, max_attempts) direction = "maximize" if self.tune_mode == "max" else "minimize" with warnings.catch_warnings(): @@ -66,6 +61,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 @@ -89,6 +85,12 @@ def sample_value(self, key: str, value: list[Any]) -> Any: return self.current_trial.suggest_categorical(key, value) def analyze(self, path: str) -> None: + with warnings.catch_warnings(): + warnings.simplefilter("ignore", + optuna.exceptions.ExperimentalWarning) + self._analyze(path) + + def _analyze(self, path: str) -> None: import matplotlib import matplotlib.pyplot as plt from optuna.visualization.matplotlib import ( 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..fa22a91d3 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,15 +88,22 @@ 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 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..8d17d417a 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,21 @@ 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, + 'train_epoch_seconds': metrics.time_taken, + }) + 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/diskgraph/lightning_data.py b/packages/kumo-core/src/kumo_core/diskgraph/lightning_data.py index 55ec1589d..4a730f3ce 100644 --- a/packages/kumo-core/src/kumo_core/diskgraph/lightning_data.py +++ b/packages/kumo-core/src/kumo_core/diskgraph/lightning_data.py @@ -5,6 +5,7 @@ from __future__ import annotations import ctypes +import errno import itertools import json import logging @@ -246,6 +247,11 @@ def open(cls, process: Any) -> _WorkerHandle: # The duplicated sentinel still gives an authoritative exit # observation for a worker that won the startup/exit race. pass + except OSError as error: + # SageMaker hosting refuses pidfd_open. + if error.errno not in (errno.ENOSYS, errno.EPERM): + os.close(sentinel_fd) + raise except BaseException: try: os.close(sentinel_fd) @@ -272,9 +278,7 @@ def send_signal(self, sig: int) -> None: try: if self.pidfd is not None: _pidfd_send_signal(self.pidfd, sig) - elif not sys.platform.startswith('linux'): - # Non-Linux development fallback. Production Linux/Python 3.14 - # uses pidfds, avoiding numeric-PID reuse entirely. + else: os.kill(self.pid, sig) except ProcessLookupError: pass diff --git a/packages/kumo-core/src/kumo_core/encoder_selection.py b/packages/kumo-core/src/kumo_core/encoder_selection.py index d0956b1cb..e0933c42e 100644 --- a/packages/kumo-core/src/kumo_core/encoder_selection.py +++ b/packages/kumo-core/src/kumo_core/encoder_selection.py @@ -440,12 +440,25 @@ def select_text_encoder( logger.warning( "Text column %r will not be encoded by default: its values are too " "short or have too few known words. To encode it anyway, set its " - "encoder to TextEmbedding().", + "encoder to TextEmbedding().%s", column_name, + _no_default_model(), ) return Null() +def _no_default_model() -> str: + from kumo_core.embedding_models import ( + DEFAULT_MODEL, + INSTALL_DEFAULT_MODEL, + _installed, + ) + if DEFAULT_MODEL in _installed(): + return "" + return (f" No word-vector model is installed for it either: " + f"{INSTALL_DEFAULT_MODEL}, or pass embedding_models= to Trainer.") + + def select_sequence_encoder( dtype: Dtype, stats: _EncoderStats, 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/prediction.py b/packages/kumo-core/src/kumo_core/prediction.py index ac89557c5..db85585ed 100644 --- a/packages/kumo-core/src/kumo_core/prediction.py +++ b/packages/kumo-core/src/kumo_core/prediction.py @@ -583,6 +583,7 @@ def _score_prediction_shard( runtime_num_workers=config.worker_count, runtime_worker_id=worker_id, task_manifest=manifest, + binary_labels=config.binary_threshold is not None, ) return _publish_scored_shard( collector, output_dir=config.output_dir, run_id=config.run_id, diff --git a/packages/kumo-core/src/kumo_core/summaries.py b/packages/kumo-core/src/kumo_core/summaries.py new file mode 100644 index 000000000..25c366db5 --- /dev/null +++ b/packages/kumo-core/src/kumo_core/summaries.py @@ -0,0 +1,54 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +r"""Typed, tracker-neutral results, named ``
/``.""" + +from dataclasses import dataclass +from typing import Any, Literal, Mapping, Optional, Tuple, Union + + +@dataclass(frozen=True) +class Scalar: + r"""A metric; ``role='model'`` also logs it on the model.""" + + name: str + value: float + role: Optional[Literal['model']] = None + + +@dataclass(frozen=True) +class Tag: + name: str + value: str + + +@dataclass(frozen=True) +class Table: + r"""A table; a ``brief`` one (at most 12 x 5) goes on the run card.""" + + name: str + columns: Tuple[str, ...] + rows: Tuple[Tuple[Any, ...], ...] + brief: bool = False + + +@dataclass(frozen=True) +class Note: + r"""A line of text for the run card, under ``section``.""" + + section: str + text: str + + +@dataclass(frozen=True) +class Dataset: + r"""A source table's location, schema, profile and context.""" + + name: str + source: Mapping[str, Optional[str]] + schema: Mapping[str, Mapping[str, str]] + profile: Mapping[str, Any] + context: str + + +Summary = Union[Scalar, Tag, Table, Note, Dataset] 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 ce07cbd03..4786678e8 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 @@ -1366,6 +1366,7 @@ def predict( artifact_stats_collector: PredictionArtifactStatsCollector | None = None, task_manifest: DiskgraphTaskManifest | None = None, + binary_labels: bool = True, ) -> Any: """TT-NG batch prediction over the wheel-backed predict stage. @@ -1542,6 +1543,10 @@ def _build_output(prediction: dict[str, Any]) -> PredOutput: out_timestamps=out_timestamps, array_targets=allow_array_targets(), ) + if (not binary_labels and self.task.task_type + == TaskType.BINARY_CLASSIFICATION): + df['TARGET_PRED'] = pd.Series(pd.NA, index=df.index, + dtype='boolean') embedding = None if return_embeddings: diff --git a/packages/kumo-core/src/kumo_core/train/formatting.py b/packages/kumo-core/src/kumo_core/train/formatting.py index 53ce12c3a..ecc129b7c 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,22 @@ 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.""" + 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 +309,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/training/runtime.py b/packages/kumo-core/src/kumo_core/training/runtime.py index ba53e5fe1..96d7a6768 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): @@ -206,6 +207,7 @@ def local_training_bindings( run_name: Optional[str] = None, embedding_models: Optional[EmbeddingModelStore] = None, accelerator: str = 'auto', + instrumentation: Optional[TrainingInstrumentation] = None, ) -> Any: r"""Binds one resolved training input to this process. @@ -224,7 +226,7 @@ def local_training_bindings( run_name=run_name, embedding_models=embedding_models, accelerator=accelerator), - 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/test_artifact_export.py b/packages/kumo-core/tests/test_artifact_export.py index b3701f00f..244a121ec 100644 --- a/packages/kumo-core/tests/test_artifact_export.py +++ b/packages/kumo-core/tests/test_artifact_export.py @@ -138,6 +138,46 @@ def test_file_metadata_uses_hosted_batch_bound_and_preserves_source(tmp_path): assert source.read_bytes() == before +def test_stamped_columns_keep_the_source_types(tmp_path): + import pyarrow as pa + import pyarrow.parquet as pq + from kumo_core.artifact_export import prepared_artifact_file + + source = tmp_path / 'scores.parquet' + schema = pa.schema([('ENTITY', pa.int64()), ('TARGET_PRED', pa.bool_()), + ('ANCHOR', pa.timestamp('ns', tz='UTC'))]) + pq.write_table(pa.table([[1, 2], [None, None], [None, 0]], schema=schema), + source) + with prepared_artifact_file(str(source), {'BUNDLE_ID': 'b'}) as staged: + assert pq.read_schema(staged) == schema.append( + pa.field('BUNDLE_ID', pa.string())) + + +def test_stamping_builds_arrow_columns_not_python_lists(tmp_path, monkeypatch): + import pyarrow as pa + import pyarrow.parquet as pq + from kumo_core.artifact_export import prepared_artifact_file + + source = tmp_path / 'scores.parquet' + pq.write_table(pa.table({'ENTITY': np.arange(200000, dtype=np.int64)}), + source) + real_array = pa.array + + def guarded(values, *args, **kwargs): + assert not isinstance(values, list) or len(values) < 1000 + return real_array(values, *args, **kwargs) + + monkeypatch.setattr(pa, 'array', guarded) + extra = {'BUNDLE_ID': 'b', 'RANK': 7, 'FLAG': True} + with prepared_artifact_file(str(source), extra) as staged: + got = pq.read_table(staged) + assert got.schema.types[1:] == [pa.string(), pa.int64(), pa.bool_()] + assert got.num_rows == 200000 + assert got['BUNDLE_ID'].unique().to_pylist() == ['b'] + assert got['RANK'].unique().to_pylist() == [7] + assert got['FLAG'].unique().to_pylist() == [True] + + @pytest.mark.parametrize('existing_empty', [False, True]) def test_multipart_failure_and_failed_recopy_remain_retryable( tmp_path, monkeypatch, existing_empty): diff --git a/packages/kumo-core/tests/test_artifacts.py b/packages/kumo-core/tests/test_artifacts.py index 17e5255b0..fb8d3c23f 100644 --- a/packages/kumo-core/tests/test_artifacts.py +++ b/packages/kumo-core/tests/test_artifacts.py @@ -17,6 +17,7 @@ VitalArtifactSources, add_bundle_file, bundle_id_from_manifest, + bundle_task_type, resolve_vital_artifact_sources, verify_model_bundle, ) @@ -263,3 +264,10 @@ def test_bundle_id_ignores_provenance_only(tmp_path: Path) -> None: (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_encoder_selection.py b/packages/kumo-core/tests/test_encoder_selection.py index 5edf9476b..a05c5b186 100644 --- a/packages/kumo-core/tests/test_encoder_selection.py +++ b/packages/kumo-core/tests/test_encoder_selection.py @@ -298,7 +298,14 @@ def test_target_selection_still_requires_numerical_stats(): select_target_encoder(Stype.numerical, {}) -def test_text_left_unencoded_is_reported_by_every_inference(caplog): +@pytest.mark.parametrize("installed", [True, False]) +def test_text_left_unencoded_is_reported_by_every_inference( + caplog, monkeypatch, installed): + from kumo_core import embedding_models + + monkeypatch.setattr( + embedding_models, "_installed", + lambda: {embedding_models.DEFAULT_MODEL: None} if installed else {}) stats = { stat_type: _EncoderStat(value=value, approx=True) for stat_type, value in ( @@ -314,7 +321,10 @@ def test_text_left_unencoded_is_reported_by_every_inference(caplog): assert [record.getMessage() for record in caplog.records] == [ "Text column 'short' will not be encoded by default: its values are " "too short or have too few known words. To encode it anyway, set its " - "encoder to TextEmbedding()." + "encoder to TextEmbedding()." + + ("" if installed else " No word-vector model is installed for it " + "either: pip install \"kumo-engine[text]\", or pass embedding_models= " + "to Trainer.") ] * 2 diff --git a/packages/kumo-core/tests/test_optuna_searcher.py b/packages/kumo-core/tests/test_optuna_searcher.py index 9d7f9f4f8..6dd82f8c2 100644 --- a/packages/kumo-core/tests/test_optuna_searcher.py +++ b/packages/kumo-core/tests/test_optuna_searcher.py @@ -40,3 +40,16 @@ def sampler(**kwargs: Any) -> optuna.samplers.BaseSampler: # deprecated and False is the default. assert ('warn_independent_sampling' in sampler_kwargs) == passed assert sampler_kwargs.get('warn_independent_sampling', False) is False + + +def test_analyze_emits_no_experimental_warnings(tmp_path) -> None: + searcher = TPESearcher(SPACE) + searcher.study.optimize(lambda trial: trial.suggest_float('lr', 0, 1), + n_trials=3) + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter('always') + searcher.analyze(str(tmp_path)) + assert [ + str(warning.message) for warning in caught + if issubclass(warning.category, optuna.exceptions.ExperimentalWarning) + ] == [] 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..af255f322 --- /dev/null +++ b/packages/kumo-core/tests/test_searched_params.py @@ -0,0 +1,36 @@ +# 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) diff --git a/packages/kumo-core/tests/test_worker_handles.py b/packages/kumo-core/tests/test_worker_handles.py new file mode 100644 index 000000000..03e3f454a --- /dev/null +++ b/packages/kumo-core/tests/test_worker_handles.py @@ -0,0 +1,46 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import errno +import multiprocessing as mp +import signal +import time + +import pytest +from kumo_core.diskgraph import lightning_data + + +@pytest.mark.parametrize('code', [errno.ENOSYS, errno.EPERM]) +def test_workers_are_stopped_where_pidfds_are_refused(monkeypatch, code): + def refused(pid): + raise OSError(code, 'pidfd_open refused') + + monkeypatch.setattr(lightning_data, '_pidfd_open', refused) + process = mp.get_context('spawn').Process(target=time.sleep, args=(60, )) + process.start() + handle = lightning_data._WorkerHandle.open(process) + try: + assert not handle.exited() + handle.send_signal(signal.SIGKILL) + process.join(10) + assert handle.exited() and process.exitcode == -signal.SIGKILL + finally: + handle.close() + process.kill() + process.join() + + +def test_other_pidfd_errors_still_fail(monkeypatch): + def broken(pid): + raise OSError(errno.EBADF, 'bad') + + monkeypatch.setattr(lightning_data, '_pidfd_open', broken) + process = mp.get_context('spawn').Process(target=time.sleep, args=(60, )) + process.start() + try: + with pytest.raises(OSError, match='bad'): + lightning_data._WorkerHandle.open(process) + finally: + process.kill() + process.join() diff --git a/packages/kumo-databricks/src/kumo_databricks/registry.py b/packages/kumo-databricks/src/kumo_databricks/registry.py index 1f7d2e352..a0d3f41ef 100644 --- a/packages/kumo-databricks/src/kumo_databricks/registry.py +++ b/packages/kumo-databricks/src/kumo_databricks/registry.py @@ -8,9 +8,9 @@ import json import logging import os +import re 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,13 +20,17 @@ 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' +_UC_FORBIDDEN = re.compile(r'[ /\x00-\x1f\x7f]') _LOG = logging.getLogger(__name__) @@ -42,6 +46,11 @@ def _address(connector: SourceConnector, path: Optional[str], schema, name = normalize_table_coordinate(location.model_name, catalog=config.catalog, schema=config.schema_name) + if len(name) > 255 or _UC_FORBIDDEN.search(name): + raise ValueError( + f'Unity Catalog cannot register a model named {name!r}: use up to ' + f'255 characters without spaces, "/" or ".", e.g. ' + f"'users_churn'.") return f'{config.catalog}.{schema}.{name}' @@ -57,6 +66,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], @@ -100,6 +120,20 @@ def _versions(client: WorkspaceClient, name: str) -> Iterator[dict[str, Any]]: continue +def _registered_version(client: WorkspaceClient, name: str, + bundle_id: str) -> Optional[dict[str, Any]]: + """The ready Kumo version of ``name`` that holds ``bundle_id``, if any.""" + try: + for version in _versions(client, name): + tags = _tags(version) + if (version.get('status') == 'READY' and ENGINE_MARKER in tags + and tags.get(BUNDLE_ID_KEY) == bundle_id): + return version + except NotFound: + pass + return None + + def _validate_catalog(catalog: object) -> None: if not isinstance(catalog, dict) or set(catalog) != {'metadata', 'metrics'}: raise ValueError('Invalid Kumo registry metadata or metrics.') @@ -134,7 +168,15 @@ 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]) + existing = _registered_version(client, name, tags[BUNDLE_ID_KEY]) + if existing is not None: + _LOG.info('%s version %s already holds bundle %s; reused it', name, + existing['version'], tags[BUNDLE_ID_KEY]) + return _published(name, existing, metadata=metadata, + metrics=metrics) try: client.registered_models.create(catalog_name=catalog_name, schema_name=schema, name=model) @@ -201,7 +243,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: @@ -220,13 +262,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.') @@ -234,7 +281,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) @@ -252,9 +299,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) @@ -278,13 +328,20 @@ 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))) class DatabricksModelRegistryMixin(_DatabricksSessions): """Publish, discover and reload bundles; prediction remains in Kumo.""" + def check_model_location(self, connector: SourceConnector, + location: ModelRegistryLocation) -> None: + """Raises ``ValueError`` unless this schema can hold ``location``.""" + _address(connector, None, location) + def publish_model( self, connector: SourceConnector, @@ -322,15 +379,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=location.alias) + + 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 579f648e0..d74fc3883 100644 --- a/packages/kumo-databricks/src/kumo_databricks/session.py +++ b/packages/kumo-databricks/src/kumo_databricks/session.py @@ -91,10 +91,12 @@ def _client_for(principal: Principal) -> WorkspaceClient: f"Identity {principal.name!r} holds " f"{type(credentials).__name__} credentials; Databricks " f"authenticates a service principal with an OAuth secret.") + # Explicit, so a DATABRICKS_TOKEN set for MLflow cannot leak in. return WorkspaceClient( host=config.host, client_id=config.client_id, client_secret=_secret(credentials.client_secret), + auth_type='oauth-m2m', ) diff --git a/packages/kumo-databricks/tests/test_registry.py b/packages/kumo-databricks/tests/test_registry.py index 050fcf18e..e42bc4019 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( @@ -97,10 +100,19 @@ def publish(resolved, name, bundle_dir, *, comment, metadata, metrics): @pytest.mark.parametrize( 'name', - ['other.analytics.churn', 'main.other.churn', 'other.churn', 'a.b.c.d', '']) + ['other.analytics.churn', 'main.other.churn', 'other.churn', 'a.b.c.d', '', + 'users churn', 'users/churn', 'x' * 256]) def test_cross_schema_or_invalid_names_are_refused(connector, name): with pytest.raises(ValueError): - registry._address(connector, None, ModelRegistryLocation(name)) + DatabricksSinkBackend().check_model_location( + connector, ModelRegistryLocation(name)) + + +@pytest.mark.parametrize('name', ['users_churn', 'users-churn', + 'main.analytics.Users-Churn']) +def test_unity_catalog_model_names_are_accepted(connector, name): + DatabricksSinkBackend().check_model_location(connector, + ModelRegistryLocation(name)) @pytest.mark.parametrize('kwargs', [ @@ -141,7 +153,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 +167,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 +186,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 +226,10 @@ 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 + children[child] = SimpleNamespace( + path=f"{path}/{child}{'/' if directory else ''}", + is_directory=directory) return children.values() def download(self, path): @@ -485,3 +506,67 @@ 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') + + +def test_publishing_a_registered_bundle_again_reuses_its_version( + client, bundle, caplog): + metadata = dict(kumo_bundle_id='ab' * 32) + first = registry.publish_bundle(client, 'main.analytics.churn', bundle, + metadata=metadata) + posts = sum(call[0] == 'POST' for call in client.calls) + with caplog.at_level('INFO', logger=registry.__name__): + again = registry.publish_bundle(client, 'main.analytics.churn', + bundle, metadata=metadata) + assert again == first and again.version == '7' + assert sum(call[0] == 'POST' for call in client.calls) == posts + assert 'already holds bundle' in caplog.text + other = registry.publish_bundle(client, 'main.analytics.churn', bundle, + metadata=dict(kumo_bundle_id='cd' * 32)) + assert sum(call[0] == 'POST' for call in client.calls) == posts + 1 + assert other.metadata == dict(kumo_bundle_id='cd' * 32) 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/docs/mlops.md b/packages/kumo-engine/docs/mlops.md index 1a0badcb9..83f4f7def 100644 --- a/packages/kumo-engine/docs/mlops.md +++ b/packages/kumo-engine/docs/mlops.md @@ -1,46 +1,128 @@ # MLOps with Kumo Engine +Turn MLOps on with one `kumo.MLOps` and the normal `kumo.engine` calls store their snapshots, training tables and models under an artifact root, record fits in your experiment tracker, and register models in its registry. + +Install the engine wheel by path, `pip install "/[mlflow]"` (add `s3`, `gcs`, `databricks`, `sagemaker` or `snowflake` as needed; Snowflake needs only `[snowflake]`): PyPI's `kumo-engine` 1.0.0 predates this API (see [Install](#install)). GPU, CUDA and Databricks Runtime notes are under [Install](#install). + ## Quickstart -With a `graph` and a `pquery` (see the README), turn MLOps on with an artifact root and every snapshot, training table and model is stored under it: +Each block assumes a `graph` and a `pquery` (see the README). It trains, stores and registers a model, then scores with it. + +### Laptop or VM + +```python +import kumo.engine as kumo +from kumo.engine.integrations.mlflow import MlflowTracking # where MLflow's own setting points; by default sqlite:///mlflow.db + +cfg = kumo.MLOps("artifacts", tracking=MlflowTracking("churn"), + registry=kumo.ModelRegistryLocation("users-churn", alias="candidate")) +with kumo.mlops(cfg): + snapshot = graph.snapshot() + result = kumo.Trainer(pquery.suggest_model_plan(snapshot)).fit(pquery.generate_training_table(snapshot), + snapshot=snapshot, scoring_defaults=kumo.ScoringDefaults(binary_classification_threshold=0.5)) + scores = result.predict(snapshot).frame + print(result.registered_model.version, result.tracking.run_id, len(scores)) +``` + +### Databricks + +In a notebook or job on Databricks Runtime 17.3 LTS ML or later, after `%pip install "/[databricks,mlflow,text]"` and `dbutils.library.restartPython()`: + +```python +import kumo.engine as kumo +from kumo.engine.integrations.mlflow import MlflowTracking + +cfg = kumo.MLOps("/Volumes/main/ml/kumo", tracking=MlflowTracking("/Users//churn"), + registry=kumo.ModelRegistryLocation("main.ml.users_churn", alias="candidate")) +with kumo.mlops(cfg): + snapshot = graph.snapshot() + result = kumo.Trainer(pquery.suggest_model_plan(snapshot)).fit(pquery.generate_training_table(snapshot), + snapshot=snapshot, scoring_defaults=kumo.ScoringDefaults(binary_classification_threshold=0.5)) + print(result.registered_model.version, result.tracking.url) +``` + +### SageMaker + +With the managed MLflow tracking server's ARN (setup: [SageMaker managed MLflow](#sagemaker-managed-mlflow)): + +```python +import kumo.engine as kumo +from kumo.engine.integrations.mlflow import MlflowTracking, SageMakerMLflow + +cfg = kumo.MLOps("s3:///kumo", tracking=MlflowTracking("users-churn", target=SageMakerMLflow("")), + registry=kumo.ModelRegistryLocation("users-churn", alias="candidate")) +with kumo.mlops(cfg): + snapshot = graph.snapshot() + result = kumo.Trainer(pquery.suggest_model_plan(snapshot)).fit(pquery.generate_training_table(snapshot), + snapshot=snapshot, scoring_defaults=kumo.ScoringDefaults(binary_classification_threshold=0.5)) + print(result.registered_model.version, result.tracking.url) +``` + +Tables on S3 are read through an `S3Connector`: + +```python +s3 = kumo.S3Connector(root_dir="s3:///data") +users = kumo.Table.from_source_table(s3["users.parquet"], primary_key="USER_ID").infer_metadata() +orders = kumo.Table.from_source_table(s3["orders.parquet"], primary_key="ORDER_ID", time_column="ORDER_TIME").infer_metadata() +graph = kumo.Graph(tables={"users": users, "orders": orders.with_column("USER_ID", stype="ID")}, + edges=[kumo.Edge(src_table="orders", fkey="USER_ID", dst_table="users")]) +``` + +### Snowflake + +Tracking and the Snowflake Model Registry through one connector (`kumo-engine[snowflake,s3]`): ```python import kumo.engine as kumo +from kumo.engine.integrations.snowflake import SnowflakeTracking -cfg = kumo.MLOps("artifacts") # or s3://..., gs://..., /Volumes/..., (connector, "@DB.SCHEMA.STAGE/kumo") +sf = kumo.SnowflakeConnector(warehouse="WH", database="ML", schema_name="MODELS") +cfg = kumo.MLOps("s3:///kumo", tracking=SnowflakeTracking(sf, experiment_name="CHURN"), + registry=kumo.ModelRegistryLocation("CHURN_GNN", alias="candidate")) with kumo.mlops(cfg): snapshot = graph.snapshot() result = kumo.Trainer(pquery.suggest_model_plan(snapshot)).fit(pquery.generate_training_table(snapshot), - snapshot=snapshot) + snapshot=snapshot, scoring_defaults=kumo.ScoringDefaults(binary_classification_threshold=0.5)) + print(result.registered_model.version, result.tracking.url) ``` +To serve only a better model, move its alias yourself: [Promote a model](#promote-a-model). + ## How it fits together - The **artifact root** holds every artifact a block produces, each in its own directory with a `_provenance.json` manifest that embeds its parents' manifests. +- The **tracker** records each fit as a run of its own; the **registry** is the tracker's: MLflow (Unity Catalog, SageMaker MLflow, any MLflow server) or the Snowflake Model Registry. | id | names | where to find it | |---|---|---| | `snapshot.id` | a graph snapshot | `/snapshots///`; `result.snapshot_id` | | `table.id` | a training table | `/training_tables//` | -| `result.id` | a model (its bundle id) | `/models//` | +| `result.id` | a model (its bundle id) | `/models//`; the `kumo.bundle_id` run tag; `kumo_bundle_id` on the registered version | +| `result.registered_model.version` | the registered version | the registry | +| `result.tracking` | the tracker's run of the fit: `platform`, `experiment`, `run_id`, `run_name`, `url` | the tracker's UI | + +Find runs with the platform's UI, or by tag: `kumo.kind` (`fit`, `trial`) and `kumo.bundle_id`. ## Turning MLOps on -`kumo.MLOps` says where artifacts go. Use it as a block, or pass it as `mlops=` on a single call (`fit`, `graph.snapshot`, `generate_training_table`): +`kumo.MLOps` says where artifacts go and what records them. Use it as a block, or pass it as `mlops=` on a single call (`fit`, `graph.snapshot`, `generate_training_table`): ```python -cfg = kumo.MLOps("s3://bucket/kumo", on_error="raise") +cfg = kumo.MLOps("s3:///kumo", tracking=..., registry=..., on_error="raise") with kumo.mlops(cfg): snapshot = graph.snapshot() trainer.fit(table, snapshot=snapshot, mlops=cfg) ``` -- **`artifact_root`** (required): a URI (`s3://`, `gs://`, `file://`, or a local path, including `/Volumes/...` on a Databricks cluster) or a `(connector, path)` pair for a store a URI cannot reach: a volume through a `DatabricksConnector` (`"/Volumes////kumo"`) or a named stage through a `SnowflakeConnector` (`"@DB.SCHEMA.STAGE/kumo"`) once the installed sink holds files there. +- **`artifact_root`** (required): a URI (`s3://`, `gs://`, `file://`, or a local path, including `/Volumes/...` on a Databricks cluster) or a `(connector, path)` pair for a store a URI cannot reach: a volume through a `DatabricksConnector` (`"/Volumes////kumo"`) or a named stage through a `SnowflakeConnector` (`"@DB.SCHEMA.STAGE/kumo"`) once the installed sink holds files there. The connector's `root_dir` must contain the path; `kumo.MLOps` raises otherwise, so a data connector rooted at `s3:///data` cannot hold `s3:///artifacts`. +- **`tracking`**: `MlflowTracking` or `SnowflakeTracking` (below), or your own [`Tracking`](#custom-trackers). +- **`registry`**: a `ModelRegistryLocation(name, alias=...)`; every fit registers its model there, in the tracker's registry. - **`store`**: the artifact kinds stored when produced: `snapshot`, `training_table` and `model` (default: all). -- **`on_error`**: `"warn"` (default) warns about a failed store and returns; `"raise"`, for scheduled jobs, raises once the call has finished, with the failure as the `__cause__`. After `fit`, `trainer.result` still holds the model. +- **`on_error`**: `"warn"` (default) warns about a failed store, tracking call or registration and returns; `"raise"`, for scheduled jobs, raises once the call has finished, with the failure as the `__cause__`. After `fit`, `trainer.result` still holds the model. A model whose bundle or assets could not be stored under the root is never registered. - **Precedence:** a call's `mlops=` replaces the enclosing block, and an inner block replaces an outer one. A block covers the calling thread or asyncio task. - **Off by default.** Without a block or `mlops=`, every call behaves as it always has; plain exports (`export(config)`, `export_config=`) never write to the artifact root. - **Per call, not per model.** MLOps applies only to a call given `mlops=` or run inside `with kumo.mlops(...)`. A later call outside the block, even on a model trained inside one, runs without MLOps, and `on_error="raise"` does not detect this: it covers only MLOps work that ran. There is no per-call off switch; to run one call without MLOps, call it outside the block. For scheduled jobs, wrap the whole job in one block. +- `kumo.MLOps(...)` checks the tracker and registry when it is built (one experiment lookup, one registry listing, the model name's rules), so wrong credentials or a bad name fail before any work. ### Reading artifacts back @@ -49,17 +131,130 @@ Pass the `kumo.MLOps` where a connector goes, with paths relative to the artifac ```python snapshot = kumo.GraphSnapshot.load_from_connector(cfg, id=result.snapshot_id) trainer = kumo.Trainer.load_from_connector(cfg, path=f"models/{result.id}") +champion = kumo.Trainer.load_from_connector(cfg, location=kumo.ModelRegistryLocation("users-churn", alias="champion")) table = kumo.TrainingTable.load_from_connector(cfg, path=f"training_tables/{table_id}", pquery=pquery, snapshot=snapshot) ``` -A loaded model is a `Trainer`; its `result` is the `TrainingResult` (`.id`, `.manifest`, `.metrics`, `.holdout_df()`). +`location=` loads the version an alias (or `version=`) names from the tracker's registry, one way on every platform. A loaded model is a `Trainer`; its `result` is the `TrainingResult` (`.id`, `.manifest`, `.metrics`, `.scoring_defaults`, `.holdout_df()`). + +## Train and register + +`fit` inside the block stores the bundle under `models//`, records the fit as one run (a child run per hyperparameter-search trial attempt on MLflow, `trial--attempt-/` names inside the run on Snowflake), and registers the model as a new version of the registry location, moving its alias to it. + +- `result.registered_model` is the version `MLOps(registry=)` registered; `result.published_model` the one `fit(export_config=...)` published. +- `result.tracking` is the tracker's run: `TrackingRef(platform="databricks", run_id="...", url="https://.../ml/experiments/.../runs/...")`. The model's `_provenance.json` records the same. +- `fit(scoring_defaults=kumo.ScoringDefaults(...))` stores prediction defaults in the bundle before it is stored and registered: `binary_classification_threshold` (binary and multilabel classification) and `num_classes_to_return` (multiclass, multilabel, link prediction). `predict` and the MLflow pyfunc apply them when a call omits them; explicit arguments and a query's `TOP k` win. A binary model always returns `True_PROB`; with no threshold passed or stored, its `TARGET_PRED` is null. `result.with_scoring_defaults(...)` retunes a fitted model as a new bundle whose manifest names the original in `parents.derived_from`. +- `result.export(cfg)` (or `trainer.export(cfg)` for a loaded model) registers a model at `cfg.registry` through its tracker. On MLflow and Unity Catalog, registering a bundle the registry already holds under that name returns its version instead of adding one (found by its bundle id tag, logged at info) and moves the alias to it; the Snowflake Model Registry adds a new version each time. +- If registration fails, `fit` warns with the full error and returns with `registered_model=None` (with `on_error="raise"`, it raises), and the fit run ends failed with `kumo.error`. + +**Model names** follow the registry's rules, checked when `kumo.MLOps` is built: + +| registry | rule | example | +|---|---|---| +| Unity Catalog | `catalog.schema.model`, each part ≤ 255 characters without space, `/` or control characters | `main.ml.users_churn` | +| SageMaker MLflow | letters, digits and hyphens, ≤ 57, starting and ending with a letter or digit | `users-churn` | +| Snowflake Model Registry | a Snowflake identifier | `USERS_CHURN` | + +An alias is a letter or underscore followed by letters, digits or underscores; `DEFAULT`, `FIRST`, `LAST`, `LATEST` and `v` are reserved. + +### Promote a model + +Register every fit under a candidate alias and move the serving alias yourself once the new model is better, comparing `result.metrics` (the holdout's `test_`, logged to MLflow as `test/`) with the champion's: + +```python +with cfg.tracking.client() as client: # the tracker's MLflow server, as its identity + first = "champion" not in client.get_registered_model(name).aliases + if first or result.metrics["test_auroc"] > kumo.Trainer.load_from_connector( + cfg, location=kumo.ModelRegistryLocation(name, alias="champion")).result.metrics["test_auroc"]: + client.set_registered_model_alias(name, "champion", result.registered_model.version) +``` + +In the Snowflake Model Registry a version holds at most one alias: take `champion` off the old version (there is none on the first promotion) and `candidate` off the new one, then set it: + +```sql +SHOW VERSIONS IN MODEL ML.MODELS.CHURN_GNN; -- or result.registered_model.version +ALTER MODEL ML.MODELS.CHURN_GNN VERSION champion UNSET ALIAS; -- skip on the first promotion +ALTER MODEL ML.MODELS.CHURN_GNN VERSION UNSET ALIAS; -- it holds candidate +ALTER MODEL ML.MODELS.CHURN_GNN VERSION SET ALIAS = champion; +``` + +Jobs that load `alias="champion"` pick up the new version on their next run. + +### Trackers + +`MlflowTracking(experiment_name=None, run_name=None, log_system_metrics=True, tags=None, log_model=False, *, target=None)`: +- With no `target`, it records where MLflow's own setting points (`MLFLOW_TRACKING_URI`, `mlflow.set_tracking_uri`; automatic on Databricks); MLflow's default is a local `sqlite:///mlflow.db`. `target` names another server with its credentials: a `DatabricksConnector`, a `SageMakerMLflow(arn)`, or a tracking URI. With a `target`, `experiment_name` is required to record runs. +- `log_model=True` logs the model on the fit run without registering it; `MLOps(registry=)` implies it. +- Building `kumo.MLOps` creates a missing experiment; if MLflow refuses (on Databricks, the experiment's folder must exist), it raises and says so. + +`SnowflakeTracking(source, *, experiment_name, database=None, schema_name=None, run_name=None, tags=None, native_model_lineage=False)`: +- `source` is a `SnowflakeConnector`, whose identity opens the session and whose Snowflake Model Registry `MLOps(registry=)` registers in, or a Snowpark session, which tracks only. +- `database` and `schema_name` default to the session's (`PUBLIC` when only `database` is given). The fit is in Snowsight under AI & ML » Experiments. +- The registered version carries the bundle; load it with `Trainer.load_from_connector`. ## Score `predict(snapshot)` scores the trained query's population at its prediction time; `entities=` scores given rows; `predict(prediction_table, snapshot=...)` scores a table generated beforehand. A scoring job needs no `Graph` code: `GraphSnapshot.load_from_connector(cfg, id=...)` and `Trainer.load_from_connector(cfg, ...)`, then `predict`. It warns when the snapshot was taken of a different graph than the model's, naming the schema changes. +### MLflow pyfunc + +A model registered through `MlflowTracking` is also an MLflow pyfunc. It is a platform adapter: a way to load and score a registered Kumo model through MLflow's model API (`mlflow.pyfunc.load_model(uri)`, `spark_udf`, or `mlflow models serve`). It does not replace `Trainer.predict`, the batch prediction API. It is not a low-latency online service: each call scores the rows it is given against a snapshot. + +```python +with cfg.tracking.client(): # MLflow's own calls reach the tracker's server, as its identity + model = mlflow.pyfunc.load_model("models:/main.ml.users_churn@champion") +scores = model.predict(pd.DataFrame({"ENTITY": [1, 2, 3]}), params={ + "snapshot": cfg.snapshot_uri(snapshot), # "/Volumes/main/ml/kumo/snapshots//" + "prediction_time": "2024-06-01"}) # temporal queries need it, or a TIMESTAMP column +``` + +- Input rows have an `ENTITY` column and optionally a `TIMESTAMP`; a temporal query needs one of `TIMESTAMP` or `prediction_time`. Scores come back in input order, one per scored row. Binary classification returns `True_PROB`; `TARGET_PRED` needs `params={"binary_classification_threshold": ...}` or a stored threshold, and is null without one. +- Each replica loads a snapshot on its first `predict` with that URI and keeps it. `log_model(trainer, snapshot=..., serving_snapshot=...)` logs a model whose replicas load that snapshot at startup and use it when no `snapshot` param is given; `serving_snapshot` takes a URI, or a `GraphSnapshot` or snapshot id stored under the block's artifact root. +- Snapshot URIs: `s3://`, `gs://`, `file://` or a local path, read with the ambient credentials. A `/Volumes/...` path off a Databricks cluster, or a stage path, is read through `model_config={"assets_from": connector}` (below), e.g. the connector of a `(connector, "/Volumes/...")` root. +- A model with stored assets (word vectors for text columns) on a volume, stage or bucket reads them, and snapshots there, with the identity it is given: `mlflow.pyfunc.load_model(uri, model_config={"assets_from": connector})`, a connector for that store with its credentials; an asset it cannot find raises, naming this option. +- `mlflow models serve` needs an image with kumo-engine installed: `mlflow models serve -m --env-manager local`. Its pip requirement is `kumo-engine[s3,gcs]` of the same version (plus `[text]` for text columns). + +## Credentials and running off-platform + +`MlflowTracking(target=...)` and `SnowflakeTracking(connector)` take the same identities as the connectors, so a fit on a laptop or any VM can record to Databricks, SageMaker or Snowflake. An identity with no credentials (or no identity) uses the platform's ambient chain (a notebook, job or instance role, `aws configure`, `databricks auth login`); one with credentials uses exactly those. + +```python +from kumo.engine import DatabricksConnector, DatabricksIdentity, S3Identity, SnowflakeConnector, SnowflakeIdentity + +workspace = DatabricksConnector(catalog="main", schema_name="ml", warehouse_id="", + identity=DatabricksIdentity(host="https://", client_id="", client_secret="")) +databricks = MlflowTracking("/Users//users_churn", target=workspace) +sagemaker = MlflowTracking("users-churn", target=SageMakerMLflow("", identity=S3Identity(role_arn=""))) +snowflake = SnowflakeTracking(SnowflakeConnector(warehouse="", database="ML", schema_name="RUNS", role="", + identity=SnowflakeIdentity(account="", username="", private_key="")), + experiment_name="USERS_CHURN") +``` + +`load_model(uri, target=...)` and `log_model(..., target=...)` take the same targets. A `DatabricksConnector` without an identity uses a notebook, a job or a `~/.databrickscfg` profile with no environment variables to set. MLflow has no per-client credentials, so with explicit credentials each MLflow call kumo-engine makes sets `DATABRICKS_HOST` and a short-lived `DATABRICKS_TOKEN` (never the secret), or the SageMaker tracking URI and AWS session, and restores them afterwards: don't call MLflow or Databricks yourself on other threads during a fit. For your own MLflow calls as that identity, use `with tracking.client() as client:`, an `MlflowClient` for the tracker's server with the same state set for the block (`mlflow.search_runs` inside it reaches that server too). A `SnowflakeConnector` without `role=` uses the user's default role. `examples/mlops/smoke.py databricks|sagemaker|snowflake` is a manual live check per platform. + +### SageMaker managed MLflow + +1. Create a tracking server once (15-25 minutes): `aws sagemaker create-mlflow-tracking-server --tracking-server-name --artifact-store-uri s3:///mlflow --role-arn `; the server role needs read and write on that prefix. A stopped server starts with `aws sagemaker start-mlflow-tracking-server --tracking-server-name ` (about 20 minutes); `kumo.MLOps` fails until `aws sagemaker describe-mlflow-tracking-server --tracking-server-name ` reports `Started` (`Created` for a new server). +2. Give the job's identity `sagemaker-mlflow:*` on the server ARN (an MLflow App instead needs `sagemaker:CallMlflowAppApi` on the app ARN; the sample policy on AWS's MLflow Apps setup page lists the Tracking Server actions, which don't apply) and S3 read and write on the artifact store and the kumo artifact root. Rights scoped to the root's prefix (`s3:GetObject`, `s3:PutObject` and `s3:DeleteObject` on `//*`, `s3:ListBucket` limited to that prefix) suffice: kumo-engine makes no bucket-level calls. +3. Install `kumo-engine[mlflow,sagemaker]` and pass `MlflowTracking(..., target=SageMakerMLflow(arn))`; the registry is the same server. +4. MLflow's run links do not open on SageMaker. A tracked fit warns with the command below, and `result.tracking.url` is `None`. To open the UI: + - **SageMaker AI Studio** (not Studio Classic): Applications » MLflow, then Open MLflow on the server's menu. Studio acts as the domain's execution role, which needs `sagemaker-mlflow:AccessUI` (grant it on `*` if a policy scoped to the server ARN is still denied) and `sagemaker:ListMlflowTrackingServers` to list it). + - **SageMaker Unified Studio**: in an IAM-mode domain (tested), the project's MLflow page only lists servers connected to the project: Connect Tracking Server with its ARN. AWS documents Create MLflow Tracking Server and Create MLflow App on that page for other domains, deleted with the project (untested here). + - **Anywhere**, with your own credentials: `aws sagemaker create-presigned-mlflow-tracking-server-url --tracking-server-name ` prints a one-time URL valid for 5 minutes; it needs `sagemaker:CreatePresignedMlflowTrackingServerUrl` and `sagemaker-mlflow:AccessUI`. + - **MLflow App** instead of a server: `aws sagemaker create-mlflow-app --name --artifact-store-uri s3:///mlflow --role-arn ` (ready at `Created`, about 5 minutes); pass its ARN as the target. The UI: `aws sagemaker create-presigned-mlflow-app-url --arn ` (needs `sagemaker:CreatePresignedMlflowAppUrl`), then pick the workspace `default`. +5. In a SageMaker notebook, use a Python 3.11+ kernel with the wheel installed by path (`pip install --find-links "/[mlflow,sagemaker,s3,text]" ipykernel ipywidgets`, then `python -m ipykernel install --user --name kumo-engine`); `SageMakerMLflow(arn)` uses the notebook's role. Notebook instances need the AL2023 platform (`notebook-al2023-v1`): the wheel needs glibc 2.34. + ## Reference +### Install + +- Install the engine wheel by path, `pip install "/[mlflow]"`. PyPI's `kumo-engine` is the 1.0.0 release, which predates the MLOps API, and a bare `kumo-engine[mlflow]` resolves to it over a dev wheel; with it installed, `kumo.MLOps` raises `AttributeError: module 'kumo.engine' has no attribute 'MLOps'` and `pip show kumo-engine` says 1.0.0. The wheel `scripts/build_wheel.sh` builds bundles kumo-core, kumo-api, kumo-pql, DiskGraph and the connector packages; with separately built wheels, add `--find-links `. +- **GPU:** on a plain GPU VM or SageMaker image, install a torch build that matches the NVIDIA driver first, e.g. `pip install --index-url https://download.pytorch.org/whl/cu126 "torch==2.12.*" torchvision`. PyPI's default torch is built for a recent CUDA (`torch.version.cuda`; CUDA 13 needs driver ≥ 580); cu126 runs on ≥ 525.60. With `accelerator="auto"`, a torch that cannot use the driver runs on CPU: check `torch.cuda.is_available()`. +- **Databricks:** Runtime 17.3 LTS ML or later keeps its torch, torchvision and pyarrow (kumo-engine accepts pyarrow 16 through 24; a custom CA bundle for S3 needs 21 or later). On 16.4 LTS ML pip replaces torch and breaks torchvision. The parent folder of an experiment path must exist. Registering to Unity Catalog from outside Databricks also needs `boto3`. +- **Disk:** `kumo.configure(scratch_dir=...)` holds snapshots, tables and bundles; Python's own temporary files (multiprocessing) follow `TMPDIR`, so point both at a large local disk. +- **MLflow:** the `mlflow` extra installs `mlflow-skinny` 3.5 or later (Databricks Runtime 17.3 ships 3.0.1, which `%pip` upgrades); a 3.x client works against SageMaker's MLflow 3.0 servers. +- **SageMaker serving:** MLflow's stock serving containers cannot install kumo-engine; on SageMaker endpoints serve with `-m /opt/ml/model -p 8080`. GPU endpoints need a driver for CUDA 12.6 (set `InferenceAmiVersion`, e.g. `al2023-ami-sagemaker-inference-gpu-4-1`); batch transform takes JSON (`dataframe_split` plus `params`). + ### Exporting and loading artifacts Every artifact has one way out and one way back. A `kumo.MLOps` stands in for the connector to read the artifact root. @@ -74,11 +269,28 @@ Loading needs a file destination; a table export is for consumers outside kumo-e ### Lineage and manifests -Every artifact carries `_provenance.json`, metadata only (the underscore keeps Spark and Hive readers off it). Parents are embedded; a parent without a manifest (a bundle from an older release) is an id stub. +Every artifact carries `_provenance.json`, metadata only (the underscore keeps Spark and Hive readers off it). Parents are embedded, minus their statistics and mutable fields; a parent without a manifest (a bundle from an older release) is an id stub. - **Snapshot:** each source table's connector, location and version, and the engine build. `graph.snapshot(record_sources=False)` keeps only connector and table names, a SHA-256 of each location, and Custom SQL as a hash. - **Training table:** the query, generation plan, task type and split seed, its snapshot, and for `TrainingTable.from_source_table` the source table. `table.id` is deterministic over snapshot, query, plan, seed and engine build; a source without a version gets a random id. -- **Model** (`ModelManifest`, `result.manifest`): the bundle id (a content hash excluding the manifest), fit times and fit mode, and the training table and warm-start model (an embedded model keeps only ids of its own model parents). The bundle carries the holdout predictions (`holdout.parquet`), so `holdout_df()` works on a loaded model. +- **Model** (`ModelManifest`, `result.manifest`): the bundle id (a content hash excluding the manifest), fit times and fit mode, the training table and warm-start model (an embedded model keeps only ids of its own model parents), `tracking` (the tracker's run) and `dependencies`. The bundle carries the holdout predictions (`holdout.parquet`), so `holdout_df()` works on a loaded model. + +### Artifacts, bundles and dependencies + +- A model's **bundle** holds only its own state: weights, encoders, the training spec and the holdout. Its manifest's **`parents`** are lineage; its **`dependencies`** are what `predict` needs besides the bundle and a snapshot: the `predictor` entry point, the `kumo-engine` `package` (with `[text]` when a text column is encoded) and content-addressed `asset`s. +- **Assets** are stored once under `/assets///` and shared by every model that uses them. `Trainer.load` and the pyfunc look for each at the location recorded, then under the enclosing block's root, then in the embedding-model store, check its sha256 and cache non-local copies under `~/.cache/kumo/assets/`. An asset on a stage, a volume or a bucket is read with the identity the loading call was given: `Trainer.load_from_connector(connector, ...)`, `load_model(uri, target=DatabricksConnector(...) or SageMakerMLflow(arn, identity=...))`, the block root's connector, or `mlflow.pyfunc.load_model(uri, model_config={"assets_from": connector})`. A stored asset that cannot be found raises `MissingDependencyError` naming the argument that passes such a connector. +- **Text encoders** are the first assets: word-vector models (`GloVeVectors`) are stored; other `TokenEmbeddingModel`s are recorded by fingerprint and must be registered wherever the model loads: + +```python +store = kumo.EmbeddingModelStore() +store.add("glove100", kumo.GloVeVectors("glove.6B.zip", dim=100)) # the official GloVe release .zip, unchanged +plan.column_processing.encoder_overrides = {"users.bio": kumo.TextEmbedding("glove100")} +with kumo.mlops(cfg): + result = kumo.Trainer(plan, embedding_models=store).fit(table, snapshot=snapshot) +trainer = kumo.Trainer.load_from_connector(cfg, path=f"models/{result.id}") # anywhere, no store +``` + +A text column `infer_metadata()` reads as categorical (short values from a small vocabulary) takes a text encoder once its stype is text: `users.with_column("bio", stype="text")`. ### Moving snapshots between machines @@ -87,3 +299,18 @@ Every artifact carries `_provenance.json`, metadata only (the underscore keeps S - **Reuse:** a load reuses the local copy only if it holds the export's bytes and passes DiskGraph's integrity check; an export skips the upload when the destination holds a complete export with the same bytes. - **Staleness:** `snapshot.source_changes(graph)` reports each table as `unchanged`, `changed` or `unknown`; `load_from_connector(..., graph=graph)` warns about another graph's schema or changed tables. - **Trust:** load only from paths whose writers you trust; checksums detect corruption, not tampering. A warehouse connector cannot hold a snapshot. + +### Custom trackers + +Other systems plug in through `kumo.engine.tracking.Tracking`, which receives ready-built `RunRecord`s so every tracker records the same names: `start(record)` with the fit run, `on_trial(record, event)` with a trial attempt after each of its events, and `finish(record, error)` with the final record (`None` when tracking stopped early). Optional members: `prepare()`, `tags`, `check(registry)` (when `kumo.MLOps` is built), `log_model(bundle_dir, manifest, snapshot, registry)` (logs and registers, returning the version) or `connector` (a connector whose model registry `registry` registers in), `load_model(location)` and `register(result, location)` (for loading and `result.export(cfg)`), `platform`/`run_id`/`run_url` (recorded as `result.tracking`). A tracker that raises is not called again for that fit and its run ends failed. The fit's data arrive in `finish` as `Dataset` summaries in `record.summaries`, each with its `name`, `source` (`connector`, `location`, `version`, `version_kind`), `schema`, `profile` and `context` (its role); a tracker's `tags` are added except `kumo.*` and param names. Skip summary types your tracker does not handle; new ones may be added. + +### Progress events + +`fit(callbacks=[...])` calls each callback with every `kumo.engine.tracking.Event`: `fit_started`, per trial `trial_started`, `epoch_ended` and `trial_ended`, then `fit_ended`. Ended events carry `params["status"]` and, on failure, `params["error"]`; a callback that raises is logged and not called again. For JSON lines: `callbacks=[lambda e: print(json.dumps(e.to_dict()))]`. `graph.snapshot(callbacks=[...])` reports `snapshot_started`, `table_staged` per table and `snapshot_ended` (with `snapshot_id` and `reused`), each with `params["seconds"]`, the time since the snapshot started, and logs each at INFO. `callbacks=[print]` prints each as its name and fields, e.g. `table_staged table=users seconds=12.0`. + +### Limits + +- Tracker calls run on a background thread; `fit` waits at most 30 s for tracking to finish, and up to an hour for model logging and registration (600 s MLflow HTTP timeout during the upload). During a fit MLflow HTTP calls use 2 retries and a 30 s timeout unless `MLFLOW_HTTP_REQUEST_MAX_RETRIES` / `MLFLOW_HTTP_REQUEST_TIMEOUT` are set. Only global rank 0 tracks. +- MLflow rejects `@` in metric names: the trainer's `test_map@10` is logged as `test/map_at_10`. Params MLflow rejects are counted in `kumo.dropped_params`. +- Snowflake keeps 200 metrics and 1000 params per run; the rest are counted in `kumo.dropped_metrics` / `kumo.dropped_params`. Snowflake has no tags (they are params of the same name), no child runs and no dataset inputs. A taken run name gets `_2`, `_3` and so on; an active run elsewhere in the process is never resumed. +- A fit's data are MLflow run inputs, each with its role as the context: the training table (`training examples`), the graph snapshot (`graph context`) and the snapshot's tables (`source table`). Databricks tables are Delta tables at the snapshot's version (so Unity Catalog links them to the registered model), files are URIs. diff --git a/packages/kumo-engine/kumo/engine/__init__.py b/packages/kumo-engine/kumo/engine/__init__.py index c12fc6a1f..34abb959d 100644 --- a/packages/kumo-engine/kumo/engine/__init__.py +++ b/packages/kumo-engine/kumo/engine/__init__.py @@ -21,10 +21,13 @@ TokenEmbeddingModel, ) from kumo_core.utils.nofile import ensure_nofile +from kumoapi.encoder import TextEmbedding from kumoapi.graph import Edge +from kumoapi.jobs import WriteMode from kumoapi.plugin import ModelRegistryLocation from kumoapi.table import Column +from kumo.engine import _quiet # noqa: F401 from kumo.engine._manifest import ( ModelManifest, SnapshotManifest, @@ -59,6 +62,7 @@ ModelOutputConfig, PredictionOutputConfig, PredictionResult, + ScoringDefaults, SnapshotOutputConfig, Trainer, TrainingResult, @@ -72,6 +76,8 @@ "configure", "EmbeddingModelStore", "GloVeVectors", + "TextEmbedding", + "WriteMode", "TokenEmbeddingModel", "ModelRegistryLocation", "Connector", @@ -105,6 +111,7 @@ "Trainer", "TrainingResult", "PredictionResult", + "ScoringDefaults", "ModelOutputConfig", "PredictionOutputConfig", "TrainingTableOutputConfig", diff --git a/packages/kumo-engine/kumo/engine/_data_summaries.py b/packages/kumo-engine/kumo/engine/_data_summaries.py new file mode 100644 index 000000000..a2cadc14f --- /dev/null +++ b/packages/kumo-engine/kumo/engine/_data_summaries.py @@ -0,0 +1,356 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +r"""Split and snapshot summaries, from the stats files builds wrote.""" + +import json +import logging +from pathlib import Path +from typing import Any, Dict, List, Mapping, Optional, Tuple + +from kumo_core.artifacts import PROVENANCE_FILE +from kumo_core.instrumentation import format_error +from kumo_core.summaries import Note, Summary, Table, Tag + +from kumo.engine._manifest import Details +from kumo.engine._text import entity + +logger = logging.getLogger(__name__) + +_SPLITS = ('train', 'val', 'test') +_COLUMN_STATS = ('count', 'null_rate', 'num_unique', 'min', 'max', 'mean', + 'std', 'q25', 'q50', 'q75', 'avg_tokens', 'max_tokens') +_TITLES = {'null_rate': 'null rate', 'num_unique': 'distinct≈', + 'q25': 'q25≈', 'q50': 'q50≈', 'q75': 'q75≈', + 'avg_tokens': 'avg tokens≈', 'max_tokens': 'max tokens≈', + 'top': 'top values≈'} +_TOP_CATEGORIES = 10 +_TARGET_BINS = 50 # as DiskGraph's node-regression histogram +_QUARTILES = ('target_min', 'target_percentile_25', 'target_median', + 'target_percentile_75', 'target_max') +_EDGE_STATS = {'FRAC_SRC_NONZERO_DEGREE': 'src_match_frac', + 'FRAC_DST_NONZERO_DEGREE': 'dst_match_frac'} +_HEALTH = ('table or link', 'rows', 'columns', 'dropped rows', + 'duplicate keys', 'rows without time', 'edges', + 'source matched', 'target matched', 'failed lookups') +_BAD_TIME_FRACTION = 0.05 +_INVALID_FRACTION = 0.10 +_MIN_MATCH_FRACTION = 0.20 +_MIN_ROWS = 10 + + +def split_summaries(staging_dir: Path) -> List[Summary]: + r"""Brief data/splits table and a note on a numerical target's bins.""" + stats = _read(staging_dir / 'split_stats.json') + tables = {split: stats.get(f'{split}_table_stats') for split in _SPLITS} + tables = {split: table for split, table in tables.items() + if isinstance(table, dict)} + targets = {split: table.get('target_distribution') or {} + for split, table in tables.items()} + if any(_number(target.get('target_mean')) is not None + for target in targets.values()): + columns, label = ('entities', 'target mean'), { + 'target mean': 'target_mean'} + elif any(set(target.get('count') or {}) <= {'true', 'false'} + and (target.get('count') or {}) for target in targets.values()): + columns, label = ('entities', 'positive rate'), {} + else: + columns, label = ('entities',), {} + rows = tuple( + (split, _whole(table.get('num_train_rows')), *( + _value(table, targets[split], name, label) for name in columns)) + for split, table in tables.items()) + query = _read(staging_dir / PROVENANCE_FILE).get('query') + entities = f"distinct {entity(query) or 'entities'} (approx.)" + summaries: List[Summary] = [] + if rows: + summaries.append(Table('data/splits', ('split', 'rows', *( + entities if name == 'entities' else name for name in columns)), + rows, brief=True)) + note = _bins_note((targets.get('train') or {})) + return summaries + ([Note('notes', note)] if note else []) + + +def _value(table: Mapping[str, Any], target: Mapping[str, Any], name: str, + label: Mapping[str, str]) -> Optional[float]: + if name == 'entities': + return _at_most(_whole(table.get('num_entities')), + _whole(table.get('num_train_rows'))) + if name == 'positive rate': + counts = target.get('count') or {} + total = sum(counts.values()) + return counts.get('true', 0) / total if total else None + return _number(target.get(label[name])) + + +def _target_bins(target: Mapping[str, Any]) -> Tuple[Tuple[float, ...], str]: + bins = (target.get('histogram') or {}).get('bins') or () + source = f"the training target's {len(bins) - 1} histogram bins" + low, high = (_number(target.get(key)) + for key in ('target_min', 'target_max')) + if not bins and low is not None and high is not None and high > low: + bins = [low + (high - low) * i / _TARGET_BINS + for i in range(_TARGET_BINS + 1)] + source = f'{_TARGET_BINS} equal bins' + edges = {*bins, *(target[key] for key in _QUARTILES + if _number(target.get(key)) is not None)} + parts = [source] if bins else [] + extra = len(edges) - len(set(bins)) + if extra: + parts.append(f"{extra} quartile edge{'s' * (extra > 1)}") + return ((tuple(sorted(edges)), ' + '.join(parts)) if len(edges) > 1 + else ((), '')) + + +def _bins_note(target: Mapping[str, Any]) -> Optional[str]: + edges, source = _target_bins(target) + if not edges: + return None + return f'Target histograms: {source}.' + + +def snapshot_summaries(index: Path, + graph_config: Mapping[str, Any]) -> List[Summary]: + r"""data/health, data/columns and a problems note per failed data check.""" + ingest = _read(index / 'ingest_stats.json').get('stats') or {} + columns: Dict[str, Dict[str, Any]] = _read(index / 'col_stats.json') + relations = _relations(_read(index / 'edge_stats.json'), ingest) + profiles = _profiles(columns) + problems = _problems(ingest, columns, relations, graph_config) + return [ + _health(ingest, columns, relations, graph_config), + *(Note('problems', problem) for problem in problems), + *([Tag('kumo.data_warnings', str(len(problems)))] if problems + else []), + Table('data/columns', + ('table', 'column', 'stype', + *(_TITLES.get(name, name) for name in _COLUMN_STATS), + _TITLES['top']), + tuple((table, column, *profile.values()) + for table, table_profiles in profiles.items() + for column, profile in table_profiles.items())), + ] + + +def _profiles(columns: Mapping[str, Mapping[str, Any]] + ) -> Dict[str, Dict[str, Any]]: + return { + table: {column: _profile(stats) for column, stats in table_stats.items()} + for table, table_stats in sorted(columns.items()) + } + + +def _relations(edge_stats: Mapping[str, Any], + ingest: Mapping[str, Any]) -> Dict[str, Dict[str, float]]: + relations: Dict[str, Dict[str, float]] = { + name: {'count': count} + for name, count in (ingest.get('edge_counts') or {}).items() + if not name.endswith('_reverse') + } + for key, entry in edge_stats.items(): + parts = key.split('__') + if len(parts) == 7 and isinstance(entry, dict): + src, _, dst, _, fkey, pkey, stat = parts + name = _EDGE_STATS.get(stat) + relation = relations.get(f'{src}_{fkey}_to_{dst}_{pkey}') + if name and relation is not None: + relation[name] = entry.get('value') + failures = ingest.get('key_failures_by_context') or {} + for name, relation in relations.items(): + missed = failures.get(f'dst_lookup:{name}') or {} + relation['fk_lookup_failures'] = (missed.get('lookup_failures', 0) + + missed.get('membership_failures', 0)) + return relations + + +def _links(table: Mapping[str, Any], + relations: Mapping[str, Mapping[str, float]] + ) -> List[Tuple[Mapping[str, Any], Mapping[str, float]]]: + return [(fkey, relations.get( + f"{table['table_name']}_{fkey['name']}_to_{fkey['dst_table']}_" + f"{fkey['dst_key']}") or {}) for fkey in table.get('fkeys') or ()] + + +def _health(ingest: Mapping[str, Any], columns: Mapping[str, Any], + relations: Mapping[str, Mapping[str, float]], + graph_config: Mapping[str, Any]) -> Table: + duplicates = _by_table(ingest.get('pk_duplicate_counts'), ':') + null_time = _by_table(ingest.get('pk_null_timestamp_drops'), ':shard-') + dropped = ingest.get('dropped_rows_counts') or {} + nodes = ingest.get('node_counts') or {} + rows: List[Tuple[Any, ...]] = [] + tables = graph_config.get('tables') or () + for table in tables: + name = table['table_name'] + rows.append((name, nodes.get(name), len(columns.get(name) or {}), + dropped.get(name, 0), duplicates.get(name, 0), + null_time.get(name, 0) if table.get('time_col') else '', + '', '', '', '')) + for table in tables: + for fkey, relation in _links(table, relations): + rows.append((_link(table, fkey), '', '', '', '', '', + *(relation.get(key) for key in ( + 'count', 'src_match_frac', 'dst_match_frac', + 'fk_lookup_failures')))) + return Table('data/health', _HEALTH, tuple(rows)) + + +def _link(table: Mapping[str, Any], fkey: Mapping[str, Any]) -> str: + return f"{table['table_name']}.{fkey['name']} → {fkey['dst_table']}" + + +def _problems(ingest: Mapping[str, Any], columns: Mapping[str, Any], + relations: Mapping[str, Mapping[str, float]], + graph_config: Mapping[str, Any]) -> List[str]: + duplicates = _by_table(ingest.get('pk_duplicate_counts'), ':') + null_time = _by_table(ingest.get('pk_null_timestamp_drops'), ':shard-') + sources = ingest.get('source_rows_by_table') or {} + failures = ingest.get('key_failures_by_context') or {} + problems: List[str] = [] + for table in graph_config.get('tables') or (): + name = table['table_name'] + stats = columns.get(name) or {} + nodes = (ingest.get('node_counts') or {}).get(name) + rows = sources.get(name) or nodes + pkey = (table.get('pkey') or {}).get('name') + bad_pks = (failures.get(f'pkey:{name}.{pkey}') + or {}).get('parse_failures', 0) + time = (table.get('time_col') or {}).get('name') + time_stats = stats.get(time) or {} + bad_times = (null_time.get(name, 0) + + time_stats.get('null_count', 0) + + time_stats.get('parse_fail_count', 0)) + invalid = [column for column, column_stats in stats.items() + if column_stats.get('count') and _invalid(column_stats) >= + _INVALID_FRACTION * column_stats['count']] + problems += [ + *([f'`{name}`: {bad_times / rows:.1%} of rows have a missing or ' + f'invalid time'] + if time and rows and bad_times >= _BAD_TIME_FRACTION * rows + else []), + *([f'`{name}`: {bad_pks} missing or invalid primary keys'] + if bad_pks else []), + *([f'`{name}`: {duplicates[name]} duplicate primary keys'] + if duplicates.get(name) else []), + *(f'`{name}.{column}`: at least {_INVALID_FRACTION:.0%} invalid ' + f'values' for column in sorted(invalid)), + *([f'`{name}`: only {nodes} rows'] if nodes is not None + and nodes < _MIN_ROWS else []), + ] + for fkey, relation in _links(table, relations): + for side, table_name in (('src', name), + ('dst', fkey['dst_table'])): + match = _number(relation.get(f'{side}_match_frac')) + if match is not None and match < _MIN_MATCH_FRACTION: + problems.append( + f'link `{_link(table, fkey)}`: {match:.1%} of ' + f'`{table_name}` rows match') + return problems + + +def table_details(index: Path, snapshot_id: Optional[str] = None) -> Details: + r"""Each source table's schema and profile, and the snapshot's.""" + rows = (_read(index / 'ingest_stats.json').get('stats') + or {}).get('node_counts') or {} + profiles = _profiles(_read(index / 'col_stats.json')) + tables = {table.get('kumo_source_table_name') or table.get('name'): table + for table in _read(index / 'diskgraph_config.json').get( + 'tables') or ()} + counts = {name: rows.get(table.get('name'), 0) + for name, table in tables.items()} + snapshot = {} if snapshot_id is None else {snapshot_id: ({}, { + 'snapshot': snapshot_id, 'tables': len(tables), + 'rows': sum(counts.values()), + **{f'{name} rows': count for name, count in counts.items()}})} + return {**snapshot, **{ + name: ( + {column['name']: {'stype': column.get('semantic_type'), + 'dtype': column.get('source_dtype')} + for column in table.get('columns') or ()}, + {'rows': rows.get(table.get('name')), + 'columns': profiles.get(table.get('name'), {})}) + for name, table in tables.items() + }} + + +def training_table_details(table: Any) -> Details: + r"""A training table's schema and profile from its Parquet metadata.""" + from pyarrow import parquet + + manifest = table.manifest + path = Path(table.train_table_path) + if manifest is None or not path.exists(): + return {} + parts = sorted(path.rglob('*.parquet')) if path.is_dir() else [path] + if not parts: + return {} + schema = parquet.read_schema(parts[0]) + return {manifest.id: ( + {field.name: {'stype': None, 'dtype': _dtype(field.type)} + for field in schema}, + {'rows': sum(parquet.ParquetFile(part).metadata.num_rows + for part in parts), + 'columns': len(schema)})} + + +def _dtype(arrow_type: Any) -> str: + text = str(arrow_type) + return {'double': 'float64', 'float': 'float32'}.get(text, ( + 'time' if text.startswith('timestamp') else + 'date' if text.startswith('date') else text)) + + +def _by_table(counts: Optional[Mapping[str, int]], + separator: str) -> Dict[str, int]: + tables: Dict[str, int] = {} + for key, count in (counts or {}).items(): + table = key.split(separator, 1)[0] + tables[table] = tables.get(table, 0) + count + return tables + + +def _invalid(stats: Mapping[str, Any]) -> int: + return stats.get('parse_fail_count', 0) + stats.get('invalid_count', 0) + + +def _profile(stats: Mapping[str, Any]) -> Dict[str, Any]: + count = _number(stats.get('count')) + profile: Dict[str, Any] = {'stype': stats.get('semantic_type')} + for name in _COLUMN_STATS: + value = (stats.get('null_count', 0) / count if count else None) \ + if name == 'null_rate' else _number(stats.get(name)) + profile[name] = value + profile['num_unique'] = _at_most(profile['num_unique'], count) + top = () if stats.get('semantic_type') == 'id' else ( + stats.get('category_counts') or ()) + profile['top'] = ', '.join(f'{value}: {n}' + for value, n in top[:_TOP_CATEGORIES]) or None + return profile + + +def _number(value: Any) -> Optional[float]: + if isinstance(value, bool) or not isinstance(value, (int, float)): + return None + return float(value) + + +def _at_most(estimate: Optional[float], + exact: Optional[float]) -> Optional[float]: + return None if estimate is None else ( + estimate if exact is None else min(estimate, exact)) + + +def _whole(value: Any) -> Optional[int]: + return None if _number(value) is None else int(value) + + +def _read(path: Path) -> Dict[str, Any]: + try: + value = json.loads(path.read_text()) + except FileNotFoundError: + return {} + except (OSError, ValueError) as error: + logger.warning('Could not read %s for the run summary: %s', path, + format_error(error)) + return {} + return value if isinstance(value, dict) else {} diff --git a/packages/kumo-engine/kumo/engine/_deps.py b/packages/kumo-engine/kumo/engine/_deps.py new file mode 100644 index 000000000..03230de81 --- /dev/null +++ b/packages/kumo-engine/kumo/engine/_deps.py @@ -0,0 +1,360 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +r"""A model's dependencies, recorded at fit and resolved at predict.""" + +import hashlib +import importlib +import os +import shutil +import tempfile +from contextlib import contextmanager +from contextvars import ContextVar +from pathlib import Path +from typing import ( + Any, + Callable, + Iterator, + List, + Mapping, + Optional, + Protocol, + Sequence, + Tuple, + Union, +) + +from kumo_core.artifacts import PROVENANCE_FILE + +from kumo.engine._manifest import ( + AssetManifest, + Dependency, + ModelManifest, + manifest_bytes, + read_manifest, + stamp, +) + +Location = Union[str, Tuple[Any, str]] +Hint = Callable[[str, str], str] +_READERS = (('@', 'SnowflakeConnector'), ('/Volumes/', 'DatabricksConnector'), + ('s3://', 'S3Connector'), ('gs://', 'GCSConnector')) + +_LOCATIONS: ContextVar[Tuple[Location, ...]] = ContextVar( + 'kumo_asset_locations', default=()) +_HINT: ContextVar[Optional[Hint]] = ContextVar('kumo_asset_hint', + default=None) + + +@contextmanager +def asset_locations(locations: Sequence[Location], + hint: Optional[Hint] = None) -> Iterator[None]: + r"""Asset roots to look in first, and how to hint at a missing reader.""" + token = _LOCATIONS.set((*locations, *_LOCATIONS.get())) + hinted = _HINT.set(hint or _HINT.get()) + try: + yield + finally: + _HINT.reset(hinted) + _LOCATIONS.reset(token) + + +class Asset(Protocol): + r"""A content-addressed file or directory a model predicts with.""" + + asset_kind: str + name: str + version: Optional[str] + + def fingerprint(self) -> str: + ... + + def files(self) -> Mapping[str, Path]: + ... + + +class MissingDependencyError(RuntimeError): + r"""A model's dependency is in none of the places looked in.""" + + +class FileAsset: + r"""A plain file, or a directory of files.""" + def __init__(self, asset_kind: str, name: str, path: Union[str, Path], + version: Optional[str] = None) -> None: + self.asset_kind, self.name, self.version = asset_kind, name, version + self.path = Path(path) + + def files(self) -> Mapping[str, Path]: + if self.path.is_file(): + return {self.path.name: self.path} + return { + path.relative_to(self.path).as_posix(): path + for path in sorted(self.path.rglob('*')) + if path.is_file() and path.name != PROVENANCE_FILE + } + + def fingerprint(self) -> str: + digest = hashlib.sha256() + for name, path in sorted(self.files().items()): + digest.update(f'{name}\0{path.stat().st_size}\0'.encode()) + with open(path, 'rb') as handle: + for block in iter(lambda: handle.read(1 << 20), b''): + digest.update(block) + return digest.hexdigest() + + @classmethod + def load(cls, directory: Path, dependency: Dependency) -> 'FileAsset': + files = [ + path for path in directory.iterdir() + if path.is_file() and path.name != PROVENANCE_FILE + ] + return cls(str(dependency.asset_kind), dependency.name, + files[0] if len(files) == 1 else directory, + dependency.version) + + +def predictor(entry_point: str) -> Dependency: + return Dependency(kind='predictor', name=entry_point) + + +def package(name: str, version: str, extras: Sequence[str]) -> Dependency: + return Dependency(kind='package', name=name, version=version, + extras=list(extras)) + + +def entry_point(value: Any) -> str: + cls = type(value) + return f'{cls.__module__}:{cls.__qualname__}' + + +def imported(entry: str) -> Any: + module, _, name = entry.partition(':') + value: Any = importlib.import_module(module) + for part in name.split('.'): + value = getattr(value, part) + return value + + +class Dependencies: + r"""Records a fit's dependencies and resolves a model's assets.""" + def __init__(self, records: Sequence[Dependency] = (), + settings: Any = None) -> None: + self.records: List[Dependency] = list(records) + self.errors: List[BaseException] = [] + self._settings = settings + self._locations = list(_LOCATIONS.get()) + self._hint = _HINT.get() + + @classmethod + def of(cls, manifest: Optional[ModelManifest], + settings: Any = None) -> 'Dependencies': + return cls((manifest and manifest.dependencies) or (), settings) + + def add(self, dependency: Dependency) -> Dependency: + self.records.append(dependency) + return dependency + + def require(self, asset: Asset) -> Dependency: + sha256 = asset.fingerprint() + location = None + settings = self._settings + if settings is not None and settings.wants('model'): + try: + location = _store(asset, sha256, settings) + except Exception as error: + self.errors.append(error) + return self.add( + Dependency(kind='asset', name=asset.name, version=asset.version, + asset_kind=asset.asset_kind, sha256=sha256, + entry_point=entry_point(asset), packing='external', + location=location)) + + def assets(self, asset_kind: str) -> List[Dependency]: + return [ + record for record in self.records + if record.kind == 'asset' and record.asset_kind == asset_kind + ] + + def resolve(self, asset_kind: str, name: str) -> Path: + r"""A local directory holding the asset, checked by sha256.""" + record = next((record for record in self.assets(asset_kind) + if record.name == name), None) + if record is None: + raise MissingDependencyError( + f'The model has no {asset_kind} dependency named {name!r}.') + tried = [] + for connector, path in self.candidates(record): + tried.append(str(path)) + try: + directory = _fetch(connector, path, str(record.sha256)) + except (FileNotFoundError, OSError, ValueError): + continue + if directory is not None and _verified(directory, record): + return directory + location = record.location or '' + reader = next((name for prefix, name in _READERS + if location.startswith(prefix)), None) + if reader is not None and self._hint is not None: + raise MissingDependencyError( + f"Cannot read {asset_kind} {name!r} (sha256 {record.sha256}) " + f"at {location}: {self._hint(reader, location)}.") + raise MissingDependencyError( + f"The model needs {asset_kind} {name!r} (sha256 " + f"{record.sha256}), which is in none of " + f"{tried or 'no configured locations'}; predict in a kumo.mlops " + f"block whose artifact root holds it.") + + def candidates(self, record: Dependency) -> List[Tuple[Any, str]]: + from kumo.engine.connector import connector_for_root + + relative = f'{record.asset_kind}/{record.sha256}' + found: List[Tuple[Any, str]] = [] + for location in self._locations: + if isinstance(location, str): + found.append((connector_for_root(location), relative)) + else: + connector, path = location + found.append((connector, f"{path.rstrip('/')}/{relative}")) + if record.location and not record.location.startswith('@'): + parent, _, leaf = record.location.rstrip('/').rpartition('/') + try: + found.append((connector_for_root(parent or '/'), leaf)) + except ValueError: + pass + settings = self._settings + if settings is not None: + found.append((settings.root.connector, + settings.root.location(f'assets/{relative}'))) + if record.location and reads(settings.root.connector, + record.location): + found.append((settings.root.connector, record.location)) + return found + + +def reads(connector: Any, location: str) -> bool: + r"""Whether ``connector`` reads the stored ``location`` itself.""" + from kumo.engine.connector import ( + FileConnector, + GCSConnector, + S3Connector, + is_file_destination, + ) + + for kind, prefix in ((S3Connector, 's3://'), (GCSConnector, 'gs://'), + (FileConnector, None)): + if isinstance(connector, kind): + return prefix is not None and location.startswith(prefix) + return (location.startswith(('@', '/Volumes/')) + and is_file_destination(connector, location)) + + +def connector_roots(manifest: Optional[ModelManifest], + connector: Any) -> List[Tuple[Any, str]]: + r"""The asset roots of ``manifest`` that ``connector`` reads.""" + if connector is None: + return [] + return [(connector, root) for root in sorted({ + record.location.rsplit('/', 2)[0] + for record in Dependencies.of(manifest).records + if record.location and reads(connector, record.location) + })] + + +@contextmanager +def assets_through(connector: Any, bundle: Union[str, Path], + hint: Hint) -> Iterator[None]: + r"""Resolves ``bundle``'s stored assets through ``connector`` too.""" + from kumo.engine.trainer._provenance import read_provenance + + with asset_locations( + connector_roots(read_provenance(Path(bundle)), connector), hint): + yield + + +def _store(asset: Asset, sha256: str, settings: Any) -> Optional[str]: + from kumo_core.object_store_sink import _prepared + + from kumo.engine.connector import file_store + + relative = f'assets/{asset.asset_kind}/{sha256}' + try: + filesystem, root = file_store(settings.root.connector, + settings.root.location(relative)) + root = root.rstrip('/') + marker = f'{root}/{PROVENANCE_FILE}' + if filesystem.get_file_info(marker).is_file: + return settings.root.uri(relative) + _prepared(filesystem, root) + for name, path in asset.files().items(): + with open(path, 'rb') as reader, \ + filesystem.open_output_stream(f'{root}/{name}') as writer: + shutil.copyfileobj(reader, writer, 1 << 20) + location = settings.root.uri(relative) + with filesystem.open_output_stream(marker) as writer: + writer.write(manifest_bytes(_manifest(asset, sha256, location))) + return location + except Exception as error: + raise RuntimeError( + f"Could not store {asset.asset_kind} {asset.name!r} under the " + f"artifact root: {error}") from error + + +def _manifest(asset: Asset, sha256: str, location: str) -> AssetManifest: + return AssetManifest(id=sha256, location=location, + asset_kind=asset.asset_kind, + name=asset.name, version=asset.version, + format=entry_point(asset), + source=str(getattr(asset, 'path', '') or '') or None, + **stamp()) + + +def _fetch(connector: Any, path: str, sha256: str) -> Optional[Path]: + from pyarrow.fs import FileSelector, FileType, LocalFileSystem + + from kumo.engine.connector import file_store + + store = file_store(connector, path) + if store is None: + return None + filesystem, root = store + root = root.rstrip('/') + if filesystem.get_file_info( + f'{root}/{PROVENANCE_FILE}').type != FileType.File: + return None + if isinstance(filesystem, LocalFileSystem): + return Path(root) + cache = os.environ.get('XDG_CACHE_HOME') or Path.home() / '.cache' + target = Path(cache) / 'kumo' / 'assets' / sha256 + if (target / PROVENANCE_FILE).is_file(): + return target + target.parent.mkdir(parents=True, exist_ok=True) + partial = Path(tempfile.mkdtemp(prefix=f'.{sha256}-', dir=target.parent)) + try: + for info in filesystem.get_file_info(FileSelector(root, + recursive=True)): + if info.type != FileType.File: + continue + local = partial / info.path[len(root) + 1:] + local.parent.mkdir(parents=True, exist_ok=True) + with filesystem.open_input_stream(info.path) as reader, \ + open(local, 'wb') as writer: + shutil.copyfileobj(reader, writer, 1 << 20) + try: + partial.rename(target) + except OSError: + if not (target / PROVENANCE_FILE).is_file(): + raise + finally: + shutil.rmtree(partial, ignore_errors=True) + return target + + +def _verified(directory: Path, record: Dependency) -> bool: + stored = read_manifest(directory / PROVENANCE_FILE, AssetManifest) + if stored is None or stored.id != record.sha256: + return False + return bool(loaded(directory, record).fingerprint() == record.sha256) + + +def loaded(directory: Path, record: Dependency) -> Any: + return imported(record.entry_point or '').load(directory, record) diff --git a/packages/kumo-engine/kumo/engine/_embeddings.py b/packages/kumo-engine/kumo/engine/_embeddings.py index 4916590d6..0591588b5 100644 --- a/packages/kumo-engine/kumo/engine/_embeddings.py +++ b/packages/kumo-engine/kumo/engine/_embeddings.py @@ -4,18 +4,24 @@ import logging from contextlib import contextmanager -from typing import Any, Iterator, Optional +from pathlib import Path +from typing import Any, Collection, Dict, Iterator, List, Optional from kumo_core import embedding_models from kumo_core.embedding_models import EmbeddingModelStore @contextmanager -def no_encoder_inference_warnings() -> Iterator[None]: +def no_encoder_inference_warnings( + overridden: Optional[Collection[str]] = None) -> Iterator[None]: # Only fit reports what inference chose: training tables do not use it, # and a trained model encodes columns as it was trained to. def drop(record: logging.LogRecord) -> bool: - return record.levelno < logging.WARNING + if overridden is None: + return record.levelno < logging.WARNING + return (record.levelno < logging.WARNING + or not isinstance(record.args, tuple) or not record.args + or record.args[0] not in overridden) inference = logging.getLogger('kumo_core.encoder_selection') inference.addFilter(drop) @@ -48,3 +54,86 @@ def require_trained_embedding_models(store: EmbeddingModelStore, # A trained model is matched by the identity of its models, not names. if encoders is not None: embedding_models.require_encoder_models(encoders, store, trained=True) + + +class TextEncoderAsset: + r"""The word vectors of a text encoder, as a model dependency.""" + asset_kind = 'text_encoder' + + def __init__(self, name: str, model: Any) -> None: + self.name, self.model = name, model + self.version = model.label + self.path = model.path + + def fingerprint(self) -> str: + return str(self.model.fingerprint()) + + def files(self) -> Dict[str, Path]: + vocab, vectors = self.model._files() + return {vocab.name: vocab, vectors.name: vectors} + + @classmethod + def load(cls, directory: Path, dependency: Any) -> 'TextEncoderAsset': + from kumo_core.embedding_models import GloVeVectors + + return cls(dependency.name, GloVeVectors.in_directory(directory)[0]) + + +def _text_encoders(encoders: Any) -> List[Any]: + from kumoapi.encoder import TextEmbedding + + found: Dict[str, Any] = {} + for table in (encoders.encoders.values() if encoders else ()): + for encoder in table.values(): + if isinstance(encoder, TextEmbedding) and encoder._identity: + found.setdefault(encoder._identity['fingerprint'], encoder) + return list(found.values()) + + +def require_text_encoders(deps: Any, store: EmbeddingModelStore, + encoders: Any) -> bool: + r"""Records the model's text-encoder vectors; returns whether any.""" + from kumo_core.embedding_models import GloVeVectors, resolve_encoder + + from kumo.engine._manifest import Dependency + + text = _text_encoders(encoders) + for encoder in text: + model = resolve_encoder(encoder, store) + if isinstance(model, GloVeVectors): + deps.require(TextEncoderAsset(encoder.model, model)) + else: + deps.add( + Dependency(kind='asset', name=encoder.model, + version=model.label, asset_kind='text_encoder', + sha256=encoder._identity['fingerprint'], + packing='external')) + return bool(text) + + +def with_text_encoders(store: EmbeddingModelStore, encoders: Any, + deps: Any) -> EmbeddingModelStore: + r"""``store`` plus the model's stored word vectors that ``deps`` finds.""" + from kumo.engine._deps import MissingDependencyError, loaded + + added: Dict[str, Any] = {} + for encoder in _text_encoders(encoders): + known = encoder._identity + record = next((record for record in deps.assets('text_encoder') + if record.sha256 == known['fingerprint']), None) + if record is None or not record.entry_point or not deps.candidates( + record): + continue + try: + directory = deps.resolve('text_encoder', record.name) + except MissingDependencyError: + if store._find(known, encoder.model) is not None: + continue + raise + added[f"{record.name}@{record.sha256[:12]}"] = loaded( + directory, record).model + if not added: + return store + extended = EmbeddingModelStore() + extended._models = {**store._models, **added} + return extended diff --git a/packages/kumo-engine/kumo/engine/_events.py b/packages/kumo-engine/kumo/engine/_events.py new file mode 100644 index 000000000..dbbef3b83 --- /dev/null +++ b/packages/kumo-engine/kumo/engine/_events.py @@ -0,0 +1,90 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +r"""Progress events reported by `Trainer.fit` and `Graph.snapshot`.""" + +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', 'snapshot_started', + 'table_staged', 'snapshot_ended'] + + +@dataclass(frozen=True) +class Event: + r"""One progress event. ``trial`` and ``epoch`` count from 1.""" + + name: EventName + 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 __str__(self) -> str: + fields = {'trial': self.trial, 'epoch': self.epoch, **self.metrics, + **self.params} + return ' '.join([self.name, *(f'{key}={value}' + for key, value in fields.items() + if value is not None)]) + + +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, callbacks: Sequence[Callable[[Event], None]]) -> None: + 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, 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 this call', callback, 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/_manifest.py b/packages/kumo-engine/kumo/engine/_manifest.py index 93391ce99..2f12ef9ca 100644 --- a/packages/kumo-engine/kumo/engine/_manifest.py +++ b/packages/kumo-engine/kumo/engine/_manifest.py @@ -2,18 +2,22 @@ # All rights reserved. # SPDX-License-Identifier: Apache-2.0 import json +import re import types import warnings from collections.abc import Mapping as AbcMapping -from dataclasses import dataclass, field, fields +from dataclasses import dataclass, field, fields, replace from datetime import datetime, timezone from pathlib import Path from typing import ( Any, + ClassVar, Dict, + Iterator, List, Mapping, Optional, + Tuple, Type, TypeVar, Union, @@ -23,6 +27,7 @@ ) from kumo_core.io.filesystem import atomic_write_bytes +from kumo_core.summaries import Dataset, Summary, Table, Tag SCHEMA_VERSION = 1 _OMIT_NONE = 'omit_none' @@ -35,6 +40,20 @@ def omit_none() -> Any: return field(default=None, metadata={_OMIT_NONE: True}) +@dataclass(frozen=True) +class Artifact: + kind: str + id: str + location: Optional[str] + is_dataset: bool + tag: Optional[str] = None + connector: Optional[str] = None + version: Optional[str] = None + version_kind: Optional[str] = None + context: Optional[str] = None + label: Optional[str] = None + + @dataclass(frozen=True, kw_only=True) class Record: """A JSON object that keeps unknown keys in ``extra``.""" @@ -42,6 +61,9 @@ class Record: extra: Mapping[str, Any] = field(default_factory=dict, repr=False, compare=False) + def as_artifact(self, key: str) -> Optional[Artifact]: + return None + def to_dict(self) -> Dict[str, Any]: data = dict(self.extra) for item in fields(self): @@ -121,6 +143,11 @@ class SourceTableEntry(Record): dialect: Optional[str] = omit_none() namespace: Optional[List[Any]] = omit_none() + def as_artifact(self, key: str) -> Artifact: + return Artifact('source_table', key, self.location, True, + connector=self.connector, version=self.version, + version_kind=self.version_kind) + KINDS = ('snapshot', 'training_table', 'model') @@ -133,6 +160,96 @@ class Manifest(Record): created_at: Optional[str] = None environment: Optional[Environment] = None reconstructed: Optional[bool] = omit_none() + location: Optional[str] = omit_none() + + EMBED_OMITS = () + IS_DATASET: ClassVar[bool] = False + URI_TAG: ClassVar[Optional[str]] = None + NOUN: ClassVar[Optional[str]] = None + + @property + def title(self) -> str: + short = self.created_at and ( + f"{self.created_at[:10]}-{re.split('[@-]', self.id)[-1][:6]}") + return f"{self.NOUN or self.kind.replace('_', ' ')} {short or self.id}" + + def as_artifact(self, key: str) -> Artifact: + return Artifact(self.kind, self.id, self.location, self.IS_DATASET, + self.URI_TAG, version=self.id, version_kind='exact', + label=self.title) + + +def artifacts(manifest: Manifest) -> List[Artifact]: + found: Dict[Any, Artifact] = {} + for artifact in _walk(manifest, manifest.kind, manifest): + found.setdefault((artifact.kind, artifact.id), artifact) + return list(found.values()) + + +def _walk(value: Any, key: str, root: Manifest) -> Iterator[Artifact]: + if isinstance(value, Record): + artifact = value.as_artifact(key) + if artifact is not None: + yield artifact + if isinstance(value, Manifest) and not value.IS_DATASET and \ + value is not root: + return + for item in fields(value): + yield from _walk(getattr(value, item.name), item.name, root) + elif isinstance(value, AbcMapping): + for name, item in value.items(): + yield from _walk(item, str(name), root) + elif isinstance(value, (list, tuple)): + for item in value: + yield from _walk(item, key, root) + + +_SCHEMES = {'s3': 'S3', 'gs': 'GCS', 'file': 'LOCAL_FILE', '': 'LOCAL_FILE'} + +Details = Mapping[str, Tuple[Mapping[str, Any], Mapping[str, Any]]] + + +FIT_ROLES = {'snapshot': 'graph context', 'training_table': + 'training examples', 'source_table': 'source table'} + + +def artifact_summaries(manifest: Manifest, roles: Mapping[str, str], + details: Optional[Details] = None) -> List[Summary]: + r"""Location tags, the locations table and datasets of a manifest.""" + located = [item for item in artifacts(manifest) if item.location] + tags: Dict[str, str] = {} + for item in located: + if item.tag: + tags.setdefault(item.tag, str(item.location)) + summaries: List[Summary] = [ + *(Tag(name, value) for name, value in tags.items()), + *([Table('locations', ('kind', 'id', 'location'), tuple( + (item.kind, item.id, str(item.location)) for item in located), + brief=True)] if located else []), + ] + for item in sorted(located, key=lambda item: (_DATASET_ORDER.get( + item.kind, len(_DATASET_ORDER)), item.label or item.id)): + if not item.is_dataset: + continue + location = str(item.location) + source = {'connector': item.connector or _connector(location), + 'location': location, 'version': item.version, + 'version_kind': item.version_kind} + schema, profile = (details or {}).get(item.id, ({}, {})) + summaries.append(Dataset( + item.label or item.id, source, schema, profile, item.context + or roles.get(item.kind, item.kind.replace('_', ' ')))) + return summaries + + +_DATASET_ORDER = {'training_table': 0, 'snapshot': 1, 'source_table': 2} + + +def _connector(location: str) -> Optional[str]: + from urllib.parse import urlparse + + return None if location.startswith('@') else _SCHEMES.get( + urlparse(location).scheme) def stamp() -> Dict[str, Any]: @@ -162,6 +279,13 @@ def stub(kind: str, id: str) -> 'Manifest': return _kinds()[kind](id=id, reconstructed=True) +def embedded(manifest: M) -> M: + return replace(manifest, **{ + name: None + for name in manifest.EMBED_OMITS + }) + + def _parent(hint: Type[M], value: Any) -> M: try: return hint.from_dict(value) @@ -182,6 +306,8 @@ class SnapshotManifest(Manifest): """The sources and engine build of a snapshot.""" kind: str = 'snapshot' + IS_DATASET: ClassVar[bool] = True + NOUN: ClassVar[Optional[str]] = 'graph snapshot' parents: Optional[SnapshotParents] = None record_sources: Optional[bool] = omit_none() @@ -197,6 +323,7 @@ class TrainingTableManifest(Manifest): """The snapshot, query and rows of a training table.""" kind: str = 'training_table' + IS_DATASET: ClassVar[bool] = True parents: Optional[TrainingTableParents] = None query: Optional[str] = None generation_plan: Optional[Dict[str, Any]] = None @@ -208,6 +335,7 @@ class TrainingTableManifest(Manifest): class ModelParents(Record): training_table: Optional[TrainingTableManifest] = None warm_start: Optional['ModelManifest'] = None + derived_from: Optional['ModelManifest'] = omit_none() @dataclass(frozen=True, kw_only=True) @@ -217,14 +345,63 @@ class FitTiming(Record): seconds: float +@dataclass(frozen=True, kw_only=True) +class TrackingRef(Record): + r"""A run in an experiment tracker.""" + + platform: str + run_id: str + url: Optional[str] = omit_none() + experiment: Optional[str] = omit_none() + run_name: Optional[str] = omit_none() + + +@dataclass(frozen=True, kw_only=True) +class Dependency(Record): + r"""A predictor, package or asset a model needs to predict.""" + + kind: str + name: str + version: Optional[str] = omit_none() + extras: Optional[List[str]] = omit_none() + asset_kind: Optional[str] = omit_none() + sha256: Optional[str] = omit_none() + entry_point: Optional[str] = omit_none() + packing: Optional[str] = omit_none() + location: Optional[str] = omit_none() + + def as_artifact(self, key: str) -> Optional[Artifact]: + kind = _kinds().get(self.kind) + return None if kind is None else Artifact( + self.kind, str(self.sha256), self.location, kind.IS_DATASET, + version=self.sha256, version_kind='exact') + + @dataclass(frozen=True, kw_only=True) class ModelManifest(Manifest): """How a model bundle was fitted; ``id`` is the bundle id.""" kind: str = 'model' + URI_TAG: ClassVar[Optional[str]] = 'kumo.bundle_uri' parents: Optional[ModelParents] = None fit: Optional[FitTiming] = None fit_mode: Optional[str] = None + tracking: Optional[TrackingRef] = omit_none() + dependencies: Optional[List[Dependency]] = omit_none() + + EMBED_OMITS = ('tracking', ) + + +@dataclass(frozen=True, kw_only=True) +class AssetManifest(Manifest): + """An asset a model depends on; its ``id`` is the sha256.""" + + kind: str = 'asset' + asset_kind: Optional[str] = None + name: Optional[str] = None + version: Optional[str] = None + format: Optional[str] = None + source: Optional[str] = None def manifest_bytes(manifest: Manifest) -> bytes: diff --git a/packages/kumo-engine/kumo/engine/_mlops.py b/packages/kumo-engine/kumo/engine/_mlops.py index 3c8b60fbb..b332ac101 100644 --- a/packages/kumo-engine/kumo/engine/_mlops.py +++ b/packages/kumo-engine/kumo/engine/_mlops.py @@ -5,8 +5,10 @@ import contextvars import os import posixpath +import re import warnings -from dataclasses import dataclass, field +from dataclasses import dataclass, field, replace +from pathlib import Path from typing import ( Any, FrozenSet, @@ -15,10 +17,13 @@ Literal, Optional, Tuple, + Type, Union, ) -from kumo.engine._manifest import KINDS +from kumoapi.plugin import ModelRegistryLocation, check_model_alias + +from kumo.engine._manifest import KINDS, M, read_manifest, write_manifest @dataclass(frozen=True) @@ -41,11 +46,14 @@ class MLOps: Args: artifact_root: a URI, local path or ``(connector, path)`` pair. + registry: where each fit registers its model, through the tracker. store: artifact kinds to store; ``True`` for all, ``False`` none. on_error: ``'warn'`` or ``'raise'`` when MLOps fails. """ artifact_root: Union[str, Tuple[Any, str]] + tracking: Any = field(default=None, kw_only=True) + registry: Any = field(default=None, kw_only=True) store: Union[bool, Iterable[str]] = field(default=True, kw_only=True) on_error: Literal['warn', 'raise'] = field(default='warn', kw_only=True) root: _Root = field(init=False, repr=False, compare=False) @@ -85,12 +93,57 @@ def __post_init__(self) -> None: f'{list(KINDS)}.') if kinds: _require_writable(root) + from kumo.engine._tracking import as_tracking + + tracker = as_tracking(self.tracking) + if self.registry is not None: + self._check_registry(tracker) + if tracker is not None: + tracker.check(self.registry) object.__setattr__(self, 'root', root) object.__setattr__(self, 'kinds', kinds) + def _check_registry(self, tracker: Any) -> None: + from kumo.engine.trainer.output_config import ( + ModelOutputConfig, + check_model_location, + ) + + registry = self.registry + if not isinstance(registry, ModelRegistryLocation): + raise TypeError( + f'registry takes a ModelRegistryLocation, not ' + f'{type(registry).__name__}; the tracker registers it.') + if registry.alias is not None: + check_model_alias(registry.alias) + if tracker is not None and tracker.log_model is not None: + if registry.version is not None: + raise ValueError('the tracker assigns registered model ' + 'versions; leave ModelRegistryLocation.version unset.') + return + connector = None if tracker is None else tracker.connector + if connector is None: + raise ValueError( + 'registry= needs tracking=MlflowTracking(...) or ' + 'SnowflakeTracking(connector, ...).') + check_model_location(ModelOutputConfig(connector, location=registry)) + + def snapshot_uri(self, snapshot: Any) -> str: + r"""Where a `GraphSnapshot`, or a snapshot id, is stored.""" + return self.root.uri(snapshot_path(snapshot)) + def wants(self, kind: str) -> bool: return kind in self.kinds + def locate(self, path: Path, kind: Type[M], relative: str) -> Optional[M]: + r"""Sets the manifest at ``path`` to say it is at ``relative``.""" + location = self.root.uri(relative) + manifest = read_manifest(path, kind) + if manifest is not None and manifest.location != location: + manifest = replace(manifest, location=location) + write_manifest(path, manifest) + return manifest + def _require_within_root(connector: Any, path: str) -> None: from kumo_core.sources.object_store_sessions import ObjectStoreSessions @@ -177,11 +230,27 @@ def raise_on_error(settings: Optional[MLOps], error: Optional[BaseException], call: str, stacklevel: int = 2) -> None: if error is None or settings is None: return - message = (f'{call} finished, but MLOps failed: ' - f'{type(error).__name__}: {error}') + message = f'{call} finished, but MLOps failed: {full_error(error)}' if settings.on_error == 'raise': raise RuntimeError(message) from error + warn_failure(message, error, stacklevel + 1) + + +def full_error(error: BaseException) -> str: + return f'{type(error).__name__}: {error}' + + +def warn_failure(message: str, error: BaseException, stacklevel: int = 2, + also: Optional[BaseException] = None) -> None: + r"""Warns ``message`` once per ``error`` (and ``also``).""" + if getattr(error, '_kumo_warned', False): + return warnings.warn(message, UserWarning, stacklevel=stacklevel + 1) + for reported in (error, also): + try: + setattr(reported, '_kumo_warned', True) + except AttributeError: + pass def on_root(source: Any, path: str) -> Tuple[Any, str]: @@ -193,3 +262,20 @@ def on_root(source: Any, path: str) -> Tuple[Any, str]: def snapshot_path(snapshot: Any) -> str: snapshot_id = snapshot if isinstance(snapshot, str) else snapshot.id return 'snapshots/' + snapshot_id.replace('@', '/') + + +_SNAPSHOT_ID = re.compile(r'[0-9a-f]+@[0-9a-f]+') + + +def snapshot_uri(snapshot: Any, + settings: Optional[MLOps] = None) -> Optional[str]: + r"""A snapshot location; an object or id is read from the root.""" + if snapshot is None or (isinstance(snapshot, str) + and not _SNAPSHOT_ID.fullmatch(snapshot)): + return snapshot + settings = resolve(settings) + if settings is None: + raise ValueError( + 'A snapshot object or id needs a kumo.mlops block; otherwise ' + 'pass the URI snapshot.export(...) returned.') + return settings.snapshot_uri(snapshot) diff --git a/packages/kumo-engine/kumo/engine/_predictor.py b/packages/kumo-engine/kumo/engine/_predictor.py new file mode 100644 index 000000000..f002968ea --- /dev/null +++ b/packages/kumo-engine/kumo/engine/_predictor.py @@ -0,0 +1,144 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +r"""The predictor a model bundle's manifest names.""" + +import logging +import tempfile +from pathlib import Path +from typing import Any, Mapping, Optional, Protocol, Tuple, Union + +import pandas as pd + +from kumo.engine._deps import imported + +TRAINER_PREDICTOR = 'kumo.engine._predictor:TrainerPredictor' +_ANCHOR = 'ANCHOR_TIMESTAMP' +logger = logging.getLogger(__name__) + + +class Predictor(Protocol): + r"""Scores ``rows`` (a frame or `PredictionTable`) on ``snapshot``.""" + def predict(self, rows: Any, snapshot: Any, + params: Mapping[str, Any]) -> Union[pd.DataFrame, + Tuple[str, ...]]: + ... + + +def load_predictor(bundle_dir: Union[str, Path]) -> Predictor: + r"""The predictor the bundle at ``bundle_dir`` names, loaded.""" + from kumo.engine.trainer._provenance import read_provenance + + manifest = read_provenance(Path(bundle_dir)) + entry = next((record.name + for record in (manifest and manifest.dependencies) or () + if record.kind == 'predictor'), TRAINER_PREDICTOR) + try: + predictor = imported(entry) + except (ImportError, AttributeError, ValueError) as error: + raise ValueError( + f"The model bundle at {bundle_dir} names predictor {entry!r}, " + f"which cannot be imported here: {error}") from error + return predictor.load(bundle_dir) + + +class TrainerPredictor: + r"""Scores with the `Trainer` the bundle restores.""" + def __init__(self, trainer: Any) -> None: + self.trainer = trainer + + @classmethod + def load(cls, bundle_dir: Union[str, Path]) -> 'TrainerPredictor': + from kumo.engine.trainer.trainer import Trainer + + return cls(Trainer.load(bundle_dir)) + + def predict(self, rows: Any, snapshot: Any, + params: Mapping[str, Any]) -> Union[pd.DataFrame, + Tuple[str, ...]]: + from kumo.engine.pquery.prediction_table import PredictionTable + + if isinstance(rows, PredictionTable): + return self.trainer.predict(snapshot, prediction_table=rows, + **params).files + return predict_frame(self.trainer, snapshot, rows, params) + + +def predict_frame(trainer: Any, snapshot: Any, 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'), + ) + _require_inputs(trainer, snapshot, rows, time) + trainer.scratch_dir.mkdir(parents=True, exist_ok=True) + with tempfile.TemporaryDirectory(dir=trainer.scratch_dir) as tmp, \ + scoped_scratch_dir(Path(tmp) / 'scratch'): + rows = rows.rename(columns={'ENTITY': 'ENTITY1'}) + if time is not None and 'TIMESTAMP' not in rows: + rows = rows.assign(TIMESTAMP=time) + (Path(tmp) / 'rows').mkdir() + rows.to_parquet(Path(tmp) / 'rows' / 'rows.parquet', index=False) + entities = FileConnector(root_dir=str(Path(tmp) / 'rows'))[ + 'rows.parquet'] + timed = 'TIMESTAMP' in rows and _temporal(trainer, snapshot) + scored = trainer._predict( + trainer._prediction_table(snapshot, entities, None), + snapshot=snapshot, extra_fields=[_ANCHOR] if timed else [], + **options).frame + return _in_input_order(rows, scored) + + +def _in_input_order(rows: pd.DataFrame, scored: pd.DataFrame) -> pd.DataFrame: + keys = ['ENTITY'] + left = pd.DataFrame({'ENTITY': rows['ENTITY1'].to_numpy()}) + right = scored[['ENTITY']].copy() + if _ANCHOR in scored: + keys.append(_ANCHOR) + left[_ANCHOR] = _seconds(rows['TIMESTAMP']) + right[_ANCHOR] = _seconds(scored[_ANCHOR]) + try: + left = left.astype({'ENTITY': right['ENTITY'].dtype}) + left['_n'] = left.groupby(keys, sort=False, dropna=False).cumcount() + right['_n'] = right.groupby(keys, sort=False, dropna=False).cumcount() + right['_row'] = range(len(right)) + order = left.merge(right, on=[*keys, '_n'])['_row'] + except (TypeError, ValueError): + order = None + scored = scored.drop(columns=_ANCHOR, errors='ignore') + if order is None or len(order) != len(scored): + logger.warning('Could not match the scores to the input rows; they ' + 'are returned in scoring order: join them on ENTITY') + return scored + return scored.iloc[order.to_numpy()].reset_index(drop=True) + + +def _temporal(trainer: Any, snapshot: Any) -> bool: + from kumo_core.pquery import query_is_temporal + + from kumo.engine.integrations._serving import parsed_query + + return query_is_temporal(parsed_query(trainer, snapshot), + snapshot._native_graph) + + +def _seconds(values: pd.Series) -> pd.Series: + stamps = pd.to_datetime(values) + if stamps.dt.tz is not None: + stamps = stamps.dt.tz_convert(None) + return stamps.astype('datetime64[s]') + + +def _require_inputs(trainer: Any, snapshot: Any, rows: pd.DataFrame, + time: Optional[pd.Timestamp]) -> None: + if len(rows) == 0 or time is not None or 'TIMESTAMP' in rows: + return + if _temporal(trainer, snapshot): + raise ValueError("temporal query: pass params={'prediction_time': " + "...} or a TIMESTAMP column") diff --git a/packages/kumo-engine/kumo/engine/_quiet.py b/packages/kumo-engine/kumo/engine/_quiet.py new file mode 100644 index 000000000..e498a9e2e --- /dev/null +++ b/packages/kumo-engine/kumo/engine/_quiet.py @@ -0,0 +1,11 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +r"""Imports PyTorch Geometric without its unactionable FutureWarning.""" + +import warnings + +with warnings.catch_warnings(): + warnings.filterwarnings('ignore', r'`torch\.jit\.script` is deprecated', + FutureWarning) + import torch_geometric.nn # noqa: F401 diff --git a/packages/kumo-engine/kumo/engine/_run_names.py b/packages/kumo-engine/kumo/engine/_run_names.py new file mode 100644 index 000000000..c4c8442ab --- /dev/null +++ b/packages/kumo-engine/kumo/engine/_run_names.py @@ -0,0 +1,67 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +r"""Run names, the same on every tracker.""" + +import re +from datetime import datetime, timezone +from typing import TYPE_CHECKING, Optional + +from kumo.engine._text import entity + +if TYPE_CHECKING: + from kumo.engine._tracking import RunRecord + +_KIND_WORDS = {'prediction': 'predict'} + + +def model_label(registered: Optional[str], query: Optional[str], + task_type: Optional[str]) -> str: + r"""The registered name's last part, else -.""" + if registered: + return registered.rsplit('.', 1)[-1] + return '-'.join(filter(None, (entity(query), task_type))) or 'model' + + +def run_name(kind: str, core: str, suffix: Optional[str] = None) -> str: + r"""-, the kind not repeated.""" + word = _KIND_WORDS.get(kind, kind) + named = re.match(rf'{word}([-_]|$)', core, re.IGNORECASE) + return (core if named else f'{word}-{core}') + (suffix or '') + + +def versioned(model: str, version: str) -> str: + r"""-v, or - for a named version.""" + return f'{model}-v{version}' if version.isdigit() else f'{model}-{version}' + + +def scored_model(model: str, version: Optional[str], bundle_id: str) -> str: + r"""-v, else -.""" + return versioned(model, version) if version else f'{model}-{bundle_id[:8]}' + + +def fit_name(model: str, started: datetime) -> str: + started = started.astimezone(timezone.utc) + return run_name('fit', f'{model}-{started:%Y%m%d-%H%M%S}') + + +def trial_name(trial: int, attempt: int, fit: str) -> str: + return f'trial-{trial}-{fit}' + (f'-attempt-{attempt}' if attempt > 1 + else '') + + +def record_name(record: 'RunRecord', named: Optional[str] = None) -> str: + r"""The run name of ``record``; ``named`` replaces -v.""" + tags = record.tags + kind = tags.get('kumo.kind') + model = record.model_label or 'model' + if kind == 'fit': + return run_name('fit', named) if named else fit_name( + model, record.started or datetime.now(timezone.utc)) + suffix = record.name_suffix or ( + f"-{tags['kumo.as_of']}" if tags.get('kumo.as_of') else '') + word = (kind if kind != 'monitoring' else + 'eval' if tags.get('kumo.monitor') == 'evaluate' else 'drift') + core = named or scored_model(model, tags.get('kumo.model_version'), + tags.get('kumo.bundle_id', '')) + return run_name(word, core, suffix) diff --git a/packages/kumo-engine/kumo/engine/_scratch.py b/packages/kumo-engine/kumo/engine/_scratch.py index 7d43b4828..008165c94 100644 --- a/packages/kumo-engine/kumo/engine/_scratch.py +++ b/packages/kumo-engine/kumo/engine/_scratch.py @@ -35,13 +35,17 @@ 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 from kumo_core.chunked_fs import set_spool_dir _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: @@ -67,7 +71,17 @@ 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 context's new scratch directories under ``path``.""" + token = _scoped_root.set(path) + try: + yield + finally: + _scoped_root.reset(token) set_spool_dir(lambda: str(get_scratch_dir() / "spool")) diff --git a/packages/kumo-engine/kumo/engine/_snapshot_io.py b/packages/kumo-engine/kumo/engine/_snapshot_io.py index 18ba8bbfb..7e7f72adb 100644 --- a/packages/kumo-engine/kumo/engine/_snapshot_io.py +++ b/packages/kumo-engine/kumo/engine/_snapshot_io.py @@ -201,7 +201,12 @@ def export_location(snapshot: Any, path: str) -> str: def store(snapshot: Any, settings: Any) -> None: + from kumo.engine._manifest import SnapshotManifest + from kumo.engine._mlops import snapshot_path + try: + settings.locate(Path(snapshot.index_path) / PROVENANCE_FILE, + SnapshotManifest, snapshot_path(snapshot)) connector = settings.root.connector path = settings.root.location('snapshots') filesystem, root = _location(connector, diff --git a/packages/kumo-engine/kumo/engine/_summary_tables.py b/packages/kumo-engine/kumo/engine/_summary_tables.py new file mode 100644 index 000000000..bbc7bf219 --- /dev/null +++ b/packages/kumo-engine/kumo/engine/_summary_tables.py @@ -0,0 +1,57 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +r"""Summaries routed for trackers, and tables as tracker JSON.""" + +import json +import math +from dataclasses import dataclass, field +from typing import Any, Dict, List, Sequence + +from kumo_core.summaries import Dataset, Note, Scalar, Summary, Table, Tag + +_DEFINITIONS = frozenset({'from', 'to', 'threshold'}) + + +@dataclass +class Routed: + r"""Summaries sorted by how a tracker writes them.""" + + scalars: List[Scalar] = field(default_factory=list) + tags: Dict[str, str] = field(default_factory=dict) + tables: List[Table] = field(default_factory=list) + card: List[Summary] = field(default_factory=list) + datasets: List[Dataset] = field(default_factory=list) + + +def route(summaries: Sequence[Summary]) -> Routed: + r"""Sorts summaries by how a tracker writes them; skips unknown types.""" + routed, tables = Routed(), [] + for summary in summaries: + if isinstance(summary, Scalar): + routed.scalars.append(summary) + elif isinstance(summary, Tag): + routed.tags[summary.name] = summary.value + elif isinstance(summary, Dataset): + routed.datasets.append(summary) + elif isinstance(summary, Note) or (isinstance(summary, Table) + and summary.brief): + routed.card.append(summary) + elif isinstance(summary, Table): + tables.append(summary) + routed.tables = tables + return routed + + +def table_json(table: Table) -> str: + r"""Table JSON for tracker table views.""" + digits = [10 if column in _DEFINITIONS else 6 for column in table.columns] + return json.dumps({'columns': list(table.columns), 'data': [ + [_rounded(value, places) for value, places in zip(row, digits)] + for row in table.rows]}, default=str) + + +def _rounded(value: Any, digits: int) -> Any: + if not isinstance(value, float) or not math.isfinite(value) or not value: + return value + return round(value, digits - 1 - math.floor(math.log10(abs(value)))) diff --git a/packages/kumo-engine/kumo/engine/_text.py b/packages/kumo-engine/kumo/engine/_text.py new file mode 100644 index 000000000..a660aebe2 --- /dev/null +++ b/packages/kumo-engine/kumo/engine/_text.py @@ -0,0 +1,17 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +r"""Text helpers shared by summaries and cards.""" + +import re +from typing import Optional + + +def entity(query: Optional[str]) -> Optional[str]: + r"""The entity table a predictive query predicts for.""" + found = re.search(r'FOR\s+EACH\s+(\w+)', query or '', re.IGNORECASE) + return found and found.group(1) + + +def metric_key(name: str) -> str: + return name.replace('@', '_at_') diff --git a/packages/kumo-engine/kumo/engine/_tracking.py b/packages/kumo-engine/kumo/engine/_tracking.py new file mode 100644 index 000000000..5131d62d9 --- /dev/null +++ b/packages/kumo-engine/kumo/engine/_tracking.py @@ -0,0 +1,577 @@ +# 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 datetime import datetime +from pathlib import Path +from typing import ( + TYPE_CHECKING, + Any, + Callable, + Dict, + List, + Mapping, + Optional, + Protocol, + Sequence, + Tuple, +) + +from kumo_core.instrumentation import format_error +from kumo_core.summaries import Summary + +from kumo.engine._events import Event, _finite_floats +from kumo.engine._mlops import full_error, warn_failure +from kumo.engine._version import __version__ + +if TYPE_CHECKING: + from kumoapi.plugin import ModelRegistryLocation, PublishedModel + + from kumo.engine._manifest import ModelManifest, TrackingRef + +logger = logging.getLogger(__name__) + +DROPPED_PARAMS = 'kumo.dropped_params' +DROPPED_METRICS = 'kumo.dropped_metrics' +QUEUE_SIZE = 1000 +FINISH_SECONDS = 30.0 +MODEL_SECONDS = 3600.0 + + +@dataclass(frozen=True) +class RunRecord: + r"""One tracked run; ``steps`` are ``(step, value)`` histories.""" + + 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) + summaries: Tuple[Summary, ...] = () + model_label: Optional[str] = None + started: Optional[datetime] = None + finished: Optional[datetime] = None + name_suffix: Optional[str] = None + + +@dataclass(frozen=True) +class RunSpec: + r"""What a run records; ``parents`` become ``kumo.`` tags.""" + + kind: str + name: Optional[str] = None + params: Mapping[str, Any] = field(default_factory=dict) + metrics: Mapping[str, Any] = field(default_factory=dict) + steps: Mapping[str, Sequence[Tuple[int, float]]] = field( + default_factory=dict) + tags: Mapping[str, Any] = field(default_factory=dict) + parents: Mapping[str, Optional[str]] = field(default_factory=dict) + artifacts: Mapping[str, str] = field(default_factory=dict) + summaries: Sequence[Summary] = () + status: Optional[str] = None + error: Optional[str] = None + model_label: Optional[str] = None + started: Optional[datetime] = None + finished: Optional[datetime] = None + name_suffix: Optional[str] = None + + +def metric_name(metric: str) -> str: + r"""The plan's acc is logged as accuracy.""" + return 'accuracy' if metric == 'acc' else metric + + +def logged_name(name: str) -> str: + r"""``val_auroc`` as ``val/auroc``.""" + best = name.startswith('best_') + split, _, metric = name.removeprefix('best_').partition('_') + if split not in ('train', 'val') or not metric: + return name + return f"{split}/{'best_' if best else ''}{metric_name(metric)}" + + +def logged_names(metrics: Mapping[str, Any], + refit_full: bool = False) -> Dict[str, Any]: + r"""Tracker names of metrics; test/ is the shipped model's holdout.""" + refit = any(name.startswith('test_refit_') for name in metrics) + named = {} + for name, value in metrics.items(): + if not name.startswith('test_'): + named[logged_name(name)] = value + continue + metric = name.removeprefix('test_') + latest = metric.startswith('refit_') == refit + metric = metric_name(metric.removeprefix('refit_')) + if not latest and refit_full: + continue + named[f'test/{metric}' if latest and not refit_full else + f'before_refit/{metric}'] = value + return named + + +def by_step(steps: Mapping[str, Sequence[Tuple[int, float]]] + ) -> List[Tuple[int, Dict[str, float]]]: + grouped: Dict[int, Dict[str, float]] = {} + for name, history in steps.items(): + for step, value in history: + grouped.setdefault(step, {})[name] = value + return sorted(grouped.items()) + + +_TRIAL_PARENTS = ('snapshot_id', 'task_type') + + +def run_record(spec: RunSpec) -> RunRecord: + params = {key: _param_value(value) for key, value in spec.params.items()} + return RunRecord( + name=spec.name, + params=params, + metrics=logged_names(_finite_floats(spec.metrics)), + steps={logged_name(name): tuple(history) + for name, history in spec.steps.items()}, + tags={ + **_user_tags(spec.tags, params), + 'kumo.kind': spec.kind, + 'kumo.engine_version': __version__, + **{f'kumo.{name}': str(value) + for name, value in spec.parents.items() if value}, + **({} if spec.status is None else {'kumo.status': spec.status}), + **({} if spec.error is None else {'kumo.error': spec.error}), + }, + artifacts=dict(spec.artifacts), + summaries=tuple(spec.summaries), + model_label=spec.model_label, + started=spec.started, + finished=spec.finished, + name_suffix=spec.name_suffix, + ) + + +def _user_tags(tags: Mapping[str, Any], + params: Mapping[str, str]) -> Dict[str, str]: + dropped = sorted(key for key in tags + 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.items() + if key not in dropped} + + +def trial_spec(events: Sequence[Event]) -> RunSpec: + started = next(event for event in events if event.name == 'trial_started') + trial, attempt = started.trial, int(started.params['attempt']) + metrics: Dict[str, float] = {} + steps: Dict[str, List[Tuple[int, float]]] = {} + ended: Dict[str, Any] = {} + 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 = dict(status=str(event.params.get('status')), + error=None if 'error' not in event.params else str( + event.params['error'])) + metrics.update({ + f'best_{name}': value + for name, value in event.metrics.items() if name != 'seconds' + }) + return RunSpec( + kind='trial', + name=f'trial-{trial}-attempt-{attempt}', + params=dict(started.params), + metrics=metrics, + steps=steps, + **ended, + ) + + +def fit_run_id(manifest: Any) -> Optional[str]: + r"""The tracker's run of the fit a `ModelManifest` records.""" + tracking = getattr(manifest, 'tracking', None) + return None if tracking is None else tracking.run_id + + +class Tracking(Protocol): + r"""Records one `Trainer.fit` in an experiment tracker.""" + 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 TrackingBase(Tracking): + r"""`Tracking` with defaults for its optional members.""" + + run_id: Optional[str] = None + run_url: Optional[str] = None + tags: Optional[Mapping[str, str]] = None + connector: Any = None + log_model: Optional[Callable[..., Optional['PublishedModel']]] = None + register: Optional[Callable[..., 'PublishedModel']] = None + load_model: Optional[Callable[..., Any]] = None + experiment: Optional[str] = None + logged_run_name: Optional[str] = None + + @property + def platform(self) -> str: + return type(self).__name__ + + def prepare(self) -> None: + pass + + def check(self, registry: Optional['ModelRegistryLocation'] = None + ) -> None: + pass + + def log_metrics(self, metrics: Mapping[str, float], step: int) -> None: + pass + + +class _Foreign(TrackingBase): + r"""A tracker that does not subclass `TrackingBase`, with its defaults.""" + + def __init__(self, tracker: Any) -> None: + self._tracker = tracker + + def __getattribute__(self, name: str) -> Any: + if name.startswith('_'): + return object.__getattribute__(self, name) + if name in ('start', 'on_trial', 'finish'): + return getattr(self._tracker, name) + value = getattr(self._tracker, name, _MISSING) + called = name in ('check', 'log_model', 'register') + if value is _MISSING or called and not callable(value): + return object.__getattribute__(self, name) + return value + + @property + def platform(self) -> str: + return type(self._tracker).__name__ + + def __repr__(self) -> str: + return repr(self._tracker) + + +def as_tracking(tracking: Any) -> Optional[TrackingBase]: + if tracking is None or isinstance(tracking, TrackingBase): + return tracking + return _Foreign(tracking) + + +@dataclass +class _Search: + r"""The fit run's series: one trial's epochs, or a search's best.""" + + single: bool + metric: str + lower: bool + best: Optional[Tuple[float, int]] = None + + @classmethod + def of(cls, params: Mapping[str, Any]) -> '_Search': + from kumo_core.evaluation.suite import MetricStore + + metric = f"val_{params.get('training_job.tune_metric', 'loss')}" + return cls(single=params.get('training_job.num_experiments') == 1, + metric=metric, + lower=MetricStore.tune_mode(metric) == 'min') + + def on(self, event: Event) -> Optional[Tuple[Dict[str, float], int]]: + if event.name == 'epoch_ended' and self.single: + return dict(event.metrics), event.epoch or 0 + value = event.metrics.get(self.metric) + if (event.name != 'trial_ended' or self.single or value is None + or event.params.get('status') != 'succeeded'): + return None + if self.best is None or (value < self.best[0] if self.lower else + value > self.best[0]): + self.best = (value, event.trial or 0) + return None + + +class RegistrationError(Exception): + r"""Registering the logged model failed.""" + + +class TrackingSession: + def __init__(self, tracking: Any, warn: bool = True) -> None: + self._tracker = as_tracking(tracking) if _rank() == 0 else None + self._warn = warn + self.spec: Optional[RunSpec] = None + self._trials: Dict[int, List[Event]] = {} + self._search: Optional[_Search] = None + self._jobs: 'queue.Queue[Optional[Callable[[], Any]]]' = queue.Queue( + QUEUE_SIZE) + self._worker: Optional[threading.Thread] = None + self._behind = False + self._started = threading.Event() + self.error: Optional[BaseException] = None + + @property + def enabled(self) -> bool: + return self._tracker is not None + + def start(self, spec: Callable[[], RunSpec]) -> None: + if self._tracker is None: + return + try: + value = spec() + except Exception as error: + logger.warning('Experiment tracking failed: %s; the fit continues', + full_error(error), exc_info=_debug()) + self._fail(error) + self._tracker = None + return + self.spec = value + self._search = _Search.of(value.params) + self._call(lambda tracker: tracker.prepare()) + self._worker = threading.Thread(target=self._work, + name='kumo-tracking', daemon=True) + self._worker.start() + + def begin(tracker: TrackingBase) -> None: + tracker.start(run_record(dataclasses.replace(value, tags={ + **value.tags, **(tracker.tags or {}) + }))) + + started = self._submit(lambda: self._call(begin)) + if not (started and self._submit(self._started.set)): + self._started.set() + + def pointer(self) -> Optional['TrackingRef']: + r"""The tracker's own run, once it started (`tracking_ref`).""" + if self.spec is None or not self._started.wait(FINISH_SECONDS): + return None + return tracking_ref(self._tracker) + + def on_event(self, event: Event) -> None: + if self.spec is None or 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) + if self._search is not None and not self._search.single: + trial = trial_spec(events) + record = run_record(dataclasses.replace( + trial, model_label=self.spec.model_label, parents={ + **{name: self.spec.parents.get(name) + for name in _TRIAL_PARENTS}, + **trial.parents + })) + logged = dataclasses.replace(event, + metrics=logged_names(event.metrics)) + self._submit(lambda: self._call( + lambda tracker: tracker.on_trial(record, logged))) + series = None if self._search is None else self._search.on(event) + if series is not None: + metrics, step = logged_names(series[0]), series[1] + self._submit(lambda: self._call( + lambda tracker: tracker.log_metrics(metrics, step))) + + def log_model( + self, + bundle_dir: Path, + manifest: Optional['ModelManifest'], + snapshot: Any, + registry: Optional['ModelRegistryLocation'] = None, + ) -> Tuple[Optional['PublishedModel'], Optional[BaseException]]: + if self._worker is None or (self._tracker is not None + and self._tracker.log_model is None): + return None, None + done = threading.Event() + failed: List[BaseException] = [] + registered: List[Any] = [] + + def log() -> None: + try: + tracker = self._tracker + if tracker is None: + raise RuntimeError( + 'tracking stopped after an earlier failure' + ( + '' if self.error is None else + f' ({full_error(self.error)})')) + registered.append(tracker.log_model(bundle_dir, manifest, + snapshot, registry)) + 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: + error = failed[0] + verb = 'log' + if isinstance(error, RegistrationError) and error.__cause__: + verb, error = 'register', error.__cause__ + if self._warn: + warn_failure( + f'Could not {verb} the model: {full_error(error)}', + error, 3, also=self.error) + self._fail(error) + return None, error + return (registered or [None])[0], None + + def finish(self, spec: Optional[Callable[[], RunSpec]], + error: Optional[BaseException]) -> None: + r"""Ends the run, failed if error.""" + worker = self._worker + if self.spec is None or worker is None: + return + try: + final = self.spec if spec is None else spec() + selected = None if self._search is None else self._search.best + if error is None and selected is not None: + final = dataclasses.replace( + final, steps=trial_spec(self._trials[selected[1]]).steps, + parents={**final.parents, 'best_trial': selected[1]}) + if error is not None and final.error is None: + final = dataclasses.replace(final, error=format_error(error)) + record: Optional[RunRecord] = run_record(final) + except Exception as failure: + logger.warning('Could not build the tracking record: %s', + full_error(failure), exc_info=_debug()) + self._fail(failure) + record = None + + def finish_tracker() -> None: + tracker, self._tracker = self._tracker, None + if tracker is not None: + try: + tracker.finish(record, error) + except Exception as failure: + self._fail(failure) + raise + + 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(): + self._fail(TimeoutError( + f'experiment tracking did not finish within ' + f'{FINISH_SECONDS:g}s; it goes on in the background')) + + 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 as error: + logger.warning('Experiment tracking failed: %s', + full_error(error), exc_info=_debug()) + + def _call(self, call: Callable[[TrackingBase], Any]) -> None: + tracker = self._tracker + if tracker is None: + return + try: + call(tracker) + except Exception as error: + logger.warning( + 'Experiment tracking failed: %s; it is off for the rest of ' + 'this run', full_error(error), exc_info=_debug()) + self._fail(error) + self._tracker = None + _abandon(tracker) + + def _fail(self, error: BaseException) -> None: + if self.error is None: + self.error = error + + +def tracking_ref(tracking: Any) -> Optional['TrackingRef']: + r"""The tracker's latest run.""" + from kumo.engine._manifest import TrackingRef + + tracker = as_tracking(tracking) + if tracker is None or not tracker.run_id: + return None + return TrackingRef( + platform=str(tracker.platform), run_id=str(tracker.run_id), + url=tracker.run_url, experiment=tracker.experiment, + run_name=tracker.logged_run_name) + + +def _debug() -> bool: + return logger.isEnabledFor(logging.DEBUG) + + +_MISSING = object() + + +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 _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/snowflake.py b/packages/kumo-engine/kumo/engine/connector/snowflake.py index 0e75a9947..48c1b7526 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,11 @@ def sink_backend(self) -> ArtifactSink: "snowflake").SnowflakeSinkBackend return SnowflakeSinkBackend(identity=self.identity) + def session(self) -> Any: + r"""The Snowpark session this connector reads through.""" + 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 4c80a5191..f125a498b 100644 --- a/packages/kumo-engine/kumo/engine/graph.py +++ b/packages/kumo-engine/kumo/engine/graph.py @@ -60,18 +60,30 @@ """ import json +import logging import threading +import time import warnings from dataclasses import dataclass, field, replace from functools import cached_property, partial from pathlib import Path -from typing import TYPE_CHECKING, Any, Dict, List, Mapping, Optional, Sequence +from typing import ( + TYPE_CHECKING, + Any, + Callable, + Dict, + List, + Mapping, + Optional, + Sequence, +) from kumo_core.artifacts import PROVENANCE_FILE from kumo_core.pquery import NativeGraph, normalize_graph from kumoapi.graph import Edge from kumoapi.typing import TimestampPrecision, validate_timestamp_precision +from kumo.engine._events import Event, EventEmitter from kumo.engine._manifest import ( RECONSTRUCTED, SnapshotManifest, @@ -95,6 +107,7 @@ if TYPE_CHECKING: from kumo.engine.trainer.output_config import SnapshotOutputConfig +logger = logging.getLogger(__name__) _warned_scratch: set[Path] = set() _warned_backfill: set[Path] = set() @@ -269,6 +282,19 @@ def report_other_graph(message: str, before: Optional[Mapping[str, Any]], stacklevel=stacklevel + 1) +class _SnapshotEvents(EventEmitter): + + def __init__(self, callbacks: Sequence[Callable[[Event], None]]): + super().__init__(callbacks) + self._started = time.monotonic() + + def emit(self, name: Any, **fields: Any) -> None: + params = {**fields.pop('params', {}), + 'seconds': round(time.monotonic() - self._started, 1)} + logger.info('Graph.snapshot: %s %s', name, params) + super().emit(name, params=params, **fields) + + def _store() -> Any: from kumo_core.diskgraph.snapshot import ( DiskgraphLocalConfig, @@ -734,7 +760,9 @@ def visualize(self, path: str | Path | None = None, *, def snapshot(self, timestamp_precision: TimestampPrecision = 'seconds', *, record_sources: bool = True, - mlops: Optional[Any] = None) -> GraphSnapshot: + mlops: Optional[Any] = None, + callbacks: Sequence[Callable[[Event], None]] = () + ) -> 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` @@ -754,6 +782,7 @@ def snapshot(self, mlops: also export the snapshot under this `kumo.MLOps` root. record_sources: if False, record sources without locations or SQL text. + callbacks: called with each progress `Event`. """ validate_timestamp_precision(timestamp_precision) self.validate() @@ -783,10 +812,13 @@ def snapshot(self, from kumo.engine._mlops import resolve settings = resolve(mlops) + events = _SnapshotEvents(callbacks) + events.emit('snapshot_started', params={ + 'graph': graph_key, 'tables': list(self._named_tables)}) with _lock(graph_key): snapshot = self._snapshot_no_lock(graph_key, persist_dir, timestamp_precision, - record_sources) + record_sources, events) if settings is not None and settings.wants('snapshot'): from kumo.engine._mlops import raise_on_error from kumo.engine._snapshot_io import store @@ -800,7 +832,8 @@ def snapshot(self, def _snapshot_no_lock( self, graph_key: str, persist_dir: Path, timestamp_precision: TimestampPrecision = 'seconds', - record_sources: bool = True) -> GraphSnapshot: + record_sources: bool = True, + events: Optional['_SnapshotEvents'] = None) -> 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, @@ -862,6 +895,8 @@ def _snapshot_no_lock( "limitation in kumoai/backend/local/graph.py's " "_stage_sources, which raises the same way.") table_data_versions[alias] = prepared.version + if events is not None: + events.emit('table_staged', params={'table': alias}) source_tables[alias] = source_table( connector, name, sql, prepared.version, record_sources=record_sources) @@ -923,6 +958,9 @@ def _snapshot_no_lock( record_sources, fresh=existing is None) _warn_ingest_data_loss(result, graph_config) + if events is not None: + events.emit('snapshot_ended', params={ + 'snapshot_id': snapshot_id, 'reused': existing is not None}) return GraphSnapshot( id=snapshot_id, index_path=Path(result.index_path), 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..8233a4094 --- /dev/null +++ b/packages/kumo-engine/kumo/engine/integrations/_mlflow_model.py @@ -0,0 +1,133 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +import os +from typing import Any, Dict, Optional, Tuple + +import pandas as pd + +from kumo.engine.connector._extras import import_provider +from kumo.engine.graph import GraphSnapshot +from kumo.engine.integrations._serving import snapshot_location + +_pyfunc = import_provider('mlflow.pyfunc', 'mlflow') +_schema = import_provider('mlflow.types.schema', 'mlflow') +_signature = import_provider('mlflow.models.signature', 'mlflow') + +_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._deps import assets_through + from kumo.engine._predictor import load_predictor + + self._reader = context.model_config.get('assets_from') + bundle = context.artifacts['bundle'] + with assets_through( + self._reader, bundle, lambda reader, _: + f'{_assets_from(reader)}, with an identity that reads it'): + self._predictor = load_predictor(bundle) + self._columns = list(context.model_config['columns']) + self._snapshots: Dict[str, GraphSnapshot] = {} + self._default = context.model_config.get('snapshot') + 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: + return pd.DataFrame(columns=self._columns) + uri = params.get('snapshot') or self._default + if not uri: + raise ValueError("Pass params={'snapshot': }.") + return self._predictor.predict(model_input, self._snapshot(uri), + params) + + def _snapshot(self, uri: str) -> GraphSnapshot: + from kumo.engine._deps import reads + + if uri not in self._snapshots: + connector, path = ((self._reader, uri) + if self._reader is not None + and reads(self._reader, uri) else + snapshot_location(uri)) + try: + self._snapshots[uri] = GraphSnapshot.load_from_connector( + connector, path=path) + except FileNotFoundError as error: + reader = _remote_reader(uri) + if reader is None or connector is self._reader: + raise + raise FileNotFoundError( + f'{error} To read {uri} from here, ' + f'{_assets_from(reader)}.') from error + return self._snapshots[uri] + + +def _assets_from(reader: str) -> str: + return (f"pass model_config={{'assets_from': {reader}(...)}} to " + f"mlflow.pyfunc.load_model(model_uri, ...)") + + +def _remote_reader(uri: str) -> Optional[str]: + if uri.startswith('@'): + return 'SnowflakeConnector' + if uri.startswith('/Volumes/') and not os.path.isdir('/Volumes'): + return 'DatabricksConnector' + return None + + +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), + ]) + 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_target.py b/packages/kumo-engine/kumo/engine/integrations/_mlflow_target.py new file mode 100644 index 000000000..47b25047c --- /dev/null +++ b/packages/kumo-engine/kumo/engine/integrations/_mlflow_target.py @@ -0,0 +1,271 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +r"""Where `MlflowTracking` records, and the credentials that reach it.""" + +import logging +import os +import re +import threading +from contextlib import contextmanager +from dataclasses import dataclass +from typing import Any, Callable, ContextManager, Iterator, Mapping, Optional + +from kumo.engine.connector._extras import import_provider + +SAGEMAKER_ARN = re.compile(r'arn:aws[a-z-]*:sagemaker:([^:]*):[^:]*:' + r'(mlflow-tracking-server|mlflow-app)/(.+)') +_LOCK = threading.RLock() +_NO_LINKS = 'MLFLOW_SUPPRESS_PRINTING_URL_TO_STDOUT' +_HEADERS = logging.getLogger('mlflow.tracking.request_header.registry') +_quiet = 0 +_QUIET_LOCK = threading.Lock() + + +def _no_header_error(record: logging.LogRecord) -> bool: + return not record.getMessage().startswith( + 'Encountered unexpected error during resolving request headers') + + +@contextmanager +def _without_header_warnings() -> Iterator[None]: + global _quiet + with _QUIET_LOCK: + _quiet += 1 + if _quiet == 1: + _HEADERS.addFilter(_no_header_error) + try: + yield + finally: + with _QUIET_LOCK: + _quiet -= 1 + if not _quiet: + _HEADERS.removeFilter(_no_header_error) + + +@dataclass(frozen=True) +class SageMakerMLflow: + r"""A SageMaker managed MLflow server, for ``MlflowTracking(target=)``. + + Args: + identity: an ``S3Identity``; `None` uses the ambient AWS credentials. + """ + + tracking_server_arn: str + identity: Any = None + + def __post_init__(self) -> None: + if not SAGEMAKER_ARN.fullmatch(self.tracking_server_arn): + raise ValueError( + f'{self.tracking_server_arn!r} is not a SageMaker MLflow ' + f'tracking server or app ARN.') + + +class Target: + r"""A tracking URI, its registry URI and how to reach them.""" + + def __init__(self, tracking_uri: str, registry_uri: Optional[str], + swap: Optional[Callable[[], ContextManager[Any]]] = None, + aws_session: Any = None, host: Optional[str] = None, + assets: Any = None) -> None: + self.tracking_uri = tracking_uri + self.assets = assets + self.registry_uri = registry_uri + self._swap = swap + self._aws_session = aws_session + self._host = host + + @property + def platform(self) -> str: + if SAGEMAKER_ARN.fullmatch(self.tracking_uri or ''): + return 'sagemaker' + if (self.tracking_uri or '').partition(':')[0] == 'databricks': + return 'databricks' + return 'mlflow' + + def run_url(self, experiment_id: str, run_id: str) -> Optional[str]: + r"""The run's page in the MLflow UI, where there is one.""" + if self.platform == 'databricks': + host = self._host or _workspace_host() + return None if host is None else ( + f"{host.rstrip('/')}/ml/experiments/{experiment_id}/runs/" + f'{run_id}') + if self.platform == 'mlflow' and self.tracking_uri.startswith( + ('http://', 'https://')): + return (f"{self.tracking_uri.rstrip('/')}/#/experiments/" + f'{experiment_id}/runs/{run_id}') + return None + + @property + def explicit(self) -> bool: + r"""Whether calls run with process-wide state set around them.""" + return self._swap is not None + + @contextmanager + def scope(self) -> Iterator[None]: + r"""Runs the block with this target's credentials.""" + links = {_NO_LINKS: 'true'} if self.platform == 'sagemaker' else {} + with _without_header_warnings(), environ(links): + if self._swap is None: + yield + else: + with _LOCK, self._swap(): + yield + + def client(self) -> Any: + mlflow = import_provider('mlflow', 'mlflow') + return mlflow.MlflowClient(self.tracking_uri, self.registry_uri) + + def artifact_repository(self, uri: str) -> Any: + r"""An S3 repository signed as this target, or `None`.""" + if self._aws_session is None or not uri.startswith('s3://'): + return None + s3 = import_provider('mlflow.store.artifact.s3_artifact_repo', + 'mlflow') + frozen = self._aws_session.get_credentials().get_frozen_credentials() + return s3.S3ArtifactRepository(uri, access_key_id=frozen.access_key, + secret_access_key=frozen.secret_key, + session_token=frozen.token) + + +def resolve(target: Any) -> Target: + mlflow = import_provider('mlflow', 'mlflow') + if isinstance(target, Target): + return target + if target is None: + return Target(mlflow.get_tracking_uri(), mlflow.get_registry_uri()) + if isinstance(target, str): + if SAGEMAKER_ARN.fullmatch(target): + return _sagemaker(SageMakerMLflow(target)) + scheme, colon, rest = target.partition(':') + return Target(target, f'databricks-uc{colon}{rest}' + if scheme == 'databricks' else target) + if isinstance(target, SageMakerMLflow): + return _sagemaker(target) + from kumo.engine.connector.databricks import DatabricksConnector + + if isinstance(target, DatabricksConnector): + return _databricks(target) + raise TypeError(f'MlflowTracking(target=...) takes a DatabricksConnector, ' + f'SageMakerMLflow or tracking URI, not ' + f'{type(target).__name__}.') + + +def _workspace_host() -> Optional[str]: + try: + utils = import_provider('mlflow.utils.databricks_utils', 'mlflow') + host, _ = utils.get_workspace_info_from_dbutils() + if host: + return host + except Exception: + pass + try: + return import_provider('databricks.sdk.config', + 'databricks').Config().host + except Exception: + return None + + +def _principal(identity: Any) -> Any: + from kumo_core.sources.principals import store + + return None if identity is None else store.get(identity.slug) + + +def _databricks(connector: Any) -> Target: + from kumoapi.data_source import DatabricksAuthenticationKind + + principal = _principal(connector.identity) + if (principal is None or principal.config.authentication + is DatabricksAuthenticationKind.AMBIENT): + return Target('databricks', 'databricks-uc', + swap=lambda: global_tracking_uri('databricks'), + assets=connector) + client = import_provider('kumo_databricks.session', + 'databricks')._client_for(principal) + return Target('databricks', 'databricks-uc', + swap=lambda: _databricks_token(client), + host=client.config.host, assets=connector) + + +@contextmanager +def _databricks_token(client: Any) -> Iterator[None]: + config_type = import_provider('databricks.sdk.config', + 'databricks').Config + authorization = client.config.authenticate()['Authorization'] + names = { + attribute.env + for attribute in config_type.attributes() + if (attribute.env or '').startswith('DATABRICKS_') + } + with environ({ + **{name: None + for name in names}, 'DATABRICKS_HOST': client.config.host, + 'DATABRICKS_TOKEN': authorization.partition(' ')[2] + }), global_tracking_uri('databricks'): + yield + + +def _sagemaker(target: SageMakerMLflow) -> Target: + from kumoapi.data_source import S3AmbientAuthentication + + from kumo.engine.connector import S3Connector + + arn = target.tracking_server_arn + assets = S3Connector(identity=target.identity) + principal = _principal(target.identity) + if principal is None or isinstance(principal.config.authentication, + S3AmbientAuthentication): + return Target(arn, arn, swap=lambda: global_tracking_uri(arn), + assets=assets) + session = import_provider('kumo_s3.backend', + 's3')._session_for(principal) + return Target(arn, arn, swap=lambda: _sagemaker_session(arn, session), + aws_session=session, assets=assets) + + +@contextmanager +def _sagemaker_session(arn: str, session: Any) -> Iterator[None]: + plugin = import_provider('sagemaker_mlflow.session', 'sagemaker') + with environ({'SAGEMAKER_MLFLOW_ASSUME_ROLE_ARN': None}), \ + global_tracking_uri(arn), plugin.use_session(session): + yield + + +@contextmanager +def global_tracking_uri(uri: str) -> Iterator[None]: + r"""``mlflow.set_tracking_uri(uri)`` for the block.""" + mlflow = import_provider('mlflow', 'mlflow') + utils = import_provider('mlflow.tracking._tracking_service.utils', + 'mlflow') + if mlflow.get_tracking_uri() == uri: + yield + return + saved, saved_env = utils._tracking_uri, os.environ.get( + 'MLFLOW_TRACKING_URI') + mlflow.set_tracking_uri(uri) + try: + yield + finally: + mlflow.set_tracking_uri(saved) + if saved_env is not None: + os.environ['MLFLOW_TRACKING_URI'] = saved_env + + +@contextmanager +def environ(values: Mapping[str, Optional[str]]) -> Iterator[None]: + r"""Sets (or, for `None`, unsets) environment variables for the block.""" + saved = {name: os.environ.get(name) for name in values} + try: + _apply(values) + yield + finally: + _apply(saved) + + +def _apply(values: Mapping[str, Optional[str]]) -> None: + for name, value in values.items(): + if value is None: + os.environ.pop(name, None) + else: + os.environ[name] = value diff --git a/packages/kumo-engine/kumo/engine/integrations/_serving.py b/packages/kumo-engine/kumo/engine/integrations/_serving.py new file mode 100644 index 000000000..b457ce35f --- /dev/null +++ b/packages/kumo-engine/kumo/engine/integrations/_serving.py @@ -0,0 +1,23 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +r"""What MLflow and Snowflake serving share: the query and snapshot.""" + +from typing import Any, Tuple + +from kumo.engine.connector.base import Connector +from kumo.engine.graph import GraphSnapshot + + +def parsed_query(trainer: Any, snapshot: GraphSnapshot) -> Any: + from kumo.engine.pquery import PredictiveQuery + + return PredictiveQuery._from_snapshot(snapshot, + trainer._training_query)._parsed + + +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 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..56cc046dc --- /dev/null +++ b/packages/kumo-engine/kumo/engine/integrations/mlflow.py @@ -0,0 +1,1083 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +r"""MLflow tracking and model logging for `Trainer.fit`.""" + +import hashlib +import json +import logging +import os +import posixpath +import re +import tempfile +import threading +import time +import warnings +from contextlib import ExitStack, contextmanager +from datetime import datetime +from pathlib import Path +from types import SimpleNamespace +from typing import ( + TYPE_CHECKING, + Any, + Callable, + Dict, + FrozenSet, + Iterator, + List, + Literal, + Mapping, + Optional, + Sequence, + Set, + Tuple, + cast, +) +from urllib.parse import quote, urlparse + +from kumo_core.instrumentation import format_error +from kumo_core.summaries import Dataset, Summary, Table +from kumoapi.plugin import BUNDLE_ID_KEY, ModelRegistryLocation, PublishedModel + +from kumo.engine._events import Event +from kumo.engine._mlops import full_error +from kumo.engine._run_names import model_label, record_name, trial_name +from kumo.engine._summary_tables import Routed, route, table_json +from kumo.engine._text import metric_key +from kumo.engine._tracking import ( + DROPPED_PARAMS, + RegistrationError, + RunRecord, + TrackingBase, + by_step, + fit_run_id, +) +from kumo.engine._version import __version__ +from kumo.engine.connector._extras import import_provider +from kumo.engine.integrations._mlflow_target import ( + SAGEMAKER_ARN, + SageMakerMLflow, + environ, + global_tracking_uri, +) +from kumo.engine.integrations._mlflow_target import resolve as resolve_target + +__all__ = ['MlflowTracking', 'SageMakerMLflow', 'load_model', 'log_model'] + +if TYPE_CHECKING: + from kumo_core.embedding_models import EmbeddingModelStore + + from kumo.engine._manifest import ModelManifest + from kumo.engine.graph import GraphSnapshot + from kumo.engine.trainer.trainer import Trainer + +logger = logging.getLogger(__name__) + +_MAX_PARAM_KEY = 250 +_MAX_PARAM_VALUE = 6000 +_PARAM_KEY = re.compile(r'[/\w.\- ]+') +_HTTP_LIMITS = { + 'MLFLOW_HTTP_REQUEST_MAX_RETRIES': '2', + 'MLFLOW_HTTP_REQUEST_TIMEOUT': '30', +} +_UPLOAD_HTTP = {'MLFLOW_HTTP_REQUEST_TIMEOUT': '600'} +_PATH_CONNECTORS = ('LOCAL_FILE', 'S3', 'GCS') +_SAGEMAKER_MODEL = re.compile(r'[a-zA-Z0-9](-*[a-zA-Z0-9]){0,56}') +_UC_FORBIDDEN = re.compile(r'[ /\x00-\x1f\x7f]') +_HINTED: set = set() +_BATCH_METRICS = 1000 +_SYSTEM_METRICS_FAILURES = 6 +_COLUMN_TYPES = {'int64': 'long', 'int32': 'integer', 'float64': 'double', + 'float32': 'float', 'bool': 'boolean', 'date': 'datetime', + 'time': 'datetime'} +_Reading = Tuple[Mapping[str, float], int, int] # values, second, time +_FLUSH_SECONDS = 10.0 +_MIN_READINGS = 2 + + +class MlflowTracking(TrackingBase): + r"""Records `Trainer.fit` in MLflow, one run per fit. + + Args: + target: a ``DatabricksConnector``, `SageMakerMLflow` or tracking URI. + """ + + 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, + *, + target: Any = None, + ) -> None: + self._mlflow = import_provider('mlflow', 'mlflow') + self.target = resolve_target(target) + self._named = target is not None + self._default_experiment: Optional[str] = ( + None if experiment_name is not None or self._named else + import_provider('mlflow.tracking.fluent', + 'mlflow')._get_experiment_id()) + self.experiment_name = experiment_name + self.run_name = run_name + self.log_system_metrics = log_system_metrics + self.tags = dict(tags or {}) + self.logs_model = log_model + self._outer: Optional[str] = None + self.logged_run_name: Optional[str] = None + self.experiment: Optional[str] = None + self._experiment_id: Optional[str] = None + self._reset() + + def _reset(self) -> None: + self._client: Any = None + self._parent: 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() + + @property + def platform(self) -> str: + return self.target.platform + + @property + def run_url(self) -> Optional[str]: + if self.run_id is None or self._experiment_id is None: + return None + return self.target.run_url(self._experiment_id, self.run_id) + + @contextmanager + def client(self) -> Iterator[Any]: + r"""An ``MlflowClient`` with this tracker's credentials.""" + with self.target.scope(), global_tracking_uri( + self.target.tracking_uri): + yield self.target.client() + + def check_model_name(self, name: str) -> None: + r"""Raises :class:`ValueError` unless the registry accepts ``name``.""" + check_model_name(name, self.target.registry_uri or '') + + def check(self, registry: Optional[ModelRegistryLocation] = None) -> None: + r"""Raises unless the server (and ``registry``) is usable.""" + self._require_experiment() + if registry is not None: + self.check_model_name(registry.model_name) + uri = self.target.tracking_uri + missing = False + try: + with _limited_http(), self.target.scope(): + client = self.target.client() + if self.experiment_name is not None: + missing = client.get_experiment_by_name( + self.experiment_name) is None + else: + client.search_experiments(max_results=1) + if registry is not None: + client.search_registered_models(max_results=1) + check_registry_schema(registry.model_name, + self.target.registry_uri or '') + except ValueError: + raise + except Exception as error: + raise _unusable(uri, error) from error + if not missing: + return + with _limited_http(), self.target.scope(): + try: + client.create_experiment(self.experiment_name) + except Exception as error: + if client.get_experiment_by_name(self.experiment_name): + return + raise RuntimeError( + f'Could not create the MLflow experiment ' + f'{self.experiment_name!r}: {full_error(error)}; create ' + f'its parent folder first.') from error + + def prepare(self) -> None: + active = self._mlflow.active_run() + same = self._mlflow.get_tracking_uri() == self.target.tracking_uri + self._outer = (active.info.run_id + if active is not None and same else None) + + def _experiment_for(self, client: Any, + name: Optional[str] = None) -> str: + name = name or self.experiment_name + if name is None: + return self._default_experiment or '0' + experiment = client.get_experiment_by_name(name) + return (experiment.experiment_id + if experiment else client.create_experiment(name)) + + def _require_experiment(self) -> None: + if self._named and self.experiment_name is None: + raise ValueError('MlflowTracking(target=...) needs ' + 'experiment_name.') + + def load_model(self, location: ModelRegistryLocation, + **kwargs: Any) -> 'Trainer': + r"""The `Trainer` of ``location``'s alias or version.""" + if (location.alias is None) == (location.version is None): + raise ValueError('Pass a ModelRegistryLocation with either an ' + 'alias or a version.') + name, alias, version = (location.model_name, location.alias, + location.version) + if alias: + with self.target.scope(): + version = self.target.client().get_model_version_by_alias( + name, alias).version + return load_model(f'models:/{name}/{version}', target=self.target, + **kwargs) + + def register(self, result: Any, + location: ModelRegistryLocation) -> PublishedModel: + r"""Registers a `TrainingResult` at ``location``.""" + from kumo.engine.graph import GraphSnapshot + from kumo.engine.trainer._provenance import bundle_manifest + from kumo.engine.trainer._run_spec import registry_metrics + from kumo.engine.trainer._spec import read_training_spec + from kumo.engine.trainer.trainer import registry_pointers + + self._require_experiment() + self._client = client = self.target.client() + pointers = registry_pointers(result.id) + metrics = registry_metrics(result) + with _limited_http(): + with self.target.scope(): + held = _registered_version(client, location.model_name, + pointers) + fit, model = self._fit_of(fit_run_id(result.manifest)) + info = None if model is None else SimpleNamespace( + model_uri=f'models:/{model.model_id}', run_id=fit.run_id, + model_id=model.model_id) + if held is None and info is None: + spec = read_training_spec(result.bundle_dir) + snapshot = GraphSnapshot(spec.snapshot_id, Path(), + spec.graph_config or {}) + info = _log_model( + self.target, client, result.bundle_dir, + bundle_manifest(result.bundle_dir, result.id), snapshot, + name=_logged_model_name(location.model_name, + result.manifest), + run_id=None if fit is None else fit.run_id, + experiment_id=(self._registry_experiment( + location.model_name) if fit is None else + fit.experiment_id), serving_snapshot=None) + if fit is not None: + self._log_metrics(fit.run_id, metrics, + model_id=info.model_id) + return _register(self.target, client, info, location, pointers, + held) + + def _logged_model(self, fit: Any) -> Any: + try: + with self.target.scope(): + models = self._client.search_logged_models( + [fit.experiment_id], + filter_string=f"source_run_id = '{fit.run_id}'", + max_results=1) + except Exception: + logger.debug('No logged model of %s', fit.run_id, exc_info=True) + return None + return models[0] if models else None + + def _registry_experiment(self, model_name: str) -> str: + r"""The experiment for a model logged in no run, never Default.""" + name = self.experiment_name or model_label(model_name, None, None) + if self.experiment_name is None and self.platform == 'databricks': + name = f'/Shared/kumo/{name}' + with self.target.scope(): + return self._experiment_for(self._client, name) + + def start(self, record: RunRecord) -> None: + self._require_experiment() + self._reset() + self.run_id = None + self._http_owned = self._guard.enter_context(_limited_http()) + try: + autologging = import_provider('mlflow.utils.autologging_utils', + 'mlflow') + if not getattr(autologging, '_AUTOLOGGING_GLOBALLY_DISABLED', + False): + self._guard.enter_context(autologging.disable_autologging()) + except Exception as error: + logger.warning('Could not suspend MLflow autologging: %s', + format_error(error)) + self._client = client = self.target.client() + context = import_provider('mlflow.tracking.context.registry', + 'mlflow') + with self.target.scope(): + experiment = (client.get_experiment_by_name(self.experiment_name) + if self.experiment_name is not None else + client.get_experiment(self._default_experiment + or '0')) + if experiment is None: + experiment = client.get_experiment( + client.create_experiment(self.experiment_name)) + experiment_id, self.experiment = (experiment.experiment_id, + experiment.name) + tags = dict(record.tags) + if self._outer is not None: + tags['mlflow.parentRunId'] = self._outer + kind = record.tags.get('kumo.kind') + name = (record_name(record, self.run_name) + if self.run_name and kind in ('fit', 'prediction') + else record.name) + run = client.create_run( + experiment_id, run_name=name, + start_time=_millis(record.started), + tags=context.resolve_tags(tags)) + self._parent = self.run_id = run.info.run_id + self.logged_run_name = run.info.run_name + self._experiment_id = run.info.experiment_id + if self.log_system_metrics and record.tags.get('kumo.kind') == 'fit': + self._guard.enter_context(_SystemMetrics( + lambda readings: self._log_readings(client, run.info.run_id, + readings))) + self._log_params(self._parent, record.params) + self._log_artifacts(record.artifacts) + + def on_trial(self, record: RunRecord, event: Event) -> None: + key = (event.trial or 0, int(event.params['attempt'])) + if event.name == 'trial_started': + with self.target.scope(): + run = self._client.create_run( + self._experiment_id, run_name=trial_name( + *key, cast(str, self.logged_run_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' + with self.target.scope(): + self._client.set_terminated( + run_id, 'FINISHED' if succeeded else 'FAILED') + + def log_model( + self, + bundle_dir: Path, + manifest: Optional['ModelManifest'], + snapshot: 'GraphSnapshot', + registry: Optional[ModelRegistryLocation] = None, + ) -> Optional[PublishedModel]: + from kumo.engine.trainer.trainer import registry_pointers + + run_id = self._parent + if not (self.logs_model or registry) or run_id is None: + return None + 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.target, self._client, bundle_dir, manifest, + snapshot, name=_logged_model_name( + registry and registry.model_name, manifest), + run_id=run_id, + experiment_id=self._experiment_id, + serving_snapshot=None) + 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 None + self._model_id = info.model_id + if registry is None: + return None + pointers = registry_pointers(_id(manifest)) + try: + with self.target.scope(): + held = _registered_version(self._client, + registry.model_name, pointers) + return _register(self.target, self._client, info, registry, + pointers, held) + except Exception as error: + raise RegistrationError(full_error(error)) from error + + def finish(self, record: Optional[RunRecord], + error: Optional[BaseException]) -> None: + status = 'FAILED' + kind = None if record is None else record.tags.get('kumo.kind') + started = self._parent is not None + ended = None if record is None else _millis(record.finished) + try: + if record is not None and self._parent is not None: + for epoch, metrics in by_step(record.steps): + self._log_metrics(self._parent, metrics, epoch) + self._log_summaries(record.summaries) + 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: + self._guard.close() + with self.target.scope(): + for run_id in self._children.values(): + _quietly(self._client.set_terminated, run_id, 'FAILED') + if self._parent is not None: + _quietly(self._client.set_terminated, self._parent, + status, ended) + self._reset() + hint = _sagemaker_ui_hint(self.target.tracking_uri) + if (hint and started and kind in ('fit', 'prediction') + and self.target.tracking_uri not in _HINTED): + _HINTED.add(self.target.tracking_uri) + warnings.warn(f'MLflow run {self.logged_run_name!r} in experiment ' + f'{self.experiment!r}: {hint}', UserWarning, + stacklevel=2) + + def _fit_of(self, fit_id: Optional[str]) -> Tuple[Any, Any]: + r"""The fit run's info and its logged model, each None if missing.""" + if not fit_id: + return None, None + try: + with self.target.scope(): + fit = self._client.get_run(fit_id).info + except Exception: + logger.debug('No fit run %s here', fit_id, exc_info=True) + return None, None + return fit, self._logged_model(fit) + + def log_metrics(self, metrics: Mapping[str, float], step: int) -> None: + self._log_metrics(self._parent, metrics, step) + + def _log_readings(self, client: Any, run_id: str, + readings: Sequence['_Reading']) -> None: + entities = self._mlflow.entities + self._log_batches(client, run_id, [ + entities.Metric(metric_key(name), float(value), timestamp, step) + for values, step, timestamp in readings + for name, value in values.items() + ]) + + def _log_batches(self, client: Any, run_id: str, + metrics: Sequence[Any]) -> None: + with self.target.scope(): + for start in range(0, len(metrics), _BATCH_METRICS): + client.log_batch( + run_id, metrics=metrics[start:start + _BATCH_METRICS]) + + def _log_summaries(self, summaries: Sequence[Summary]) -> None: + self._log_routed(route(summaries), self._parent) + + def _log_routed(self, routed: Routed, run_id: str) -> None: + self._log_inputs(run_id, routed.datasets) + scalars = routed.scalars + self._log_metrics(run_id, {s.name: s.value for s in scalars + if s.role == 'model'}, + model_id=self._model_id) + self._log_metrics(run_id, {s.name: s.value for s in scalars + if s.role != 'model'}) + if routed.tags: + self._log_tags(run_id, routed.tags) + self._write_tables(run_id, routed.tables) + + def _log_inputs(self, run_id: str, summaries: Sequence[Dataset]) -> None: + datasets = [d for d in map(self._dataset_input, summaries) if d] + if datasets: + with self.target.scope(): + self._client.log_inputs(run_id, datasets=datasets) + + def _dataset_input(self, summary: Dataset) -> Any: + source = _dataset_source(summary.source) + if source is None: + return None + entities = self._mlflow.entities + dataset = entities.Dataset( + name=summary.name, + digest=_digest(summary.name, summary.source, summary.schema), + source_type=source._get_source_type(), + source=source.to_json(), + schema=json.dumps({'mlflow_colspec': [ + {'name': name, 'type': _COLUMN_TYPES.get(column['dtype'], + 'string')} + for name, column in summary.schema.items()]}), + profile=json.dumps({ + key: value for key, value in summary.profile.items() + if not isinstance(value, (Mapping, list, tuple)) + }, default=str)) + return entities.DatasetInput( + dataset, tags=[entities.InputTag('mlflow.data.context', + summary.context)]) + + def _write_tables(self, run_id: str, tables: Sequence[Table]) -> None: + if not tables: + return + for table in tables: + self._write_text(run_id, f'{table.name}.json', table_json(table)) + with self.target.scope(): + tags = self._client.get_run(run_id).data.tags + logged = json.loads(tags.get('mlflow.loggedArtifacts', '[]')) + entries = [{'path': f'{table.name}.json', 'type': 'table'} + for table in tables] + added = [entry for entry in entries if entry not in logged] + if added: + self._client.set_tag(run_id, 'mlflow.loggedArtifacts', + json.dumps([*logged, *added])) + + def _log_artifacts(self, artifacts: Mapping[str, str]) -> None: + for name, text in artifacts.items(): + if name not in self._artifacts: + self._write_text(self._parent, name, text) + self._artifacts.add(name) + + def _write_text(self, run_id: str, name: str, text: str) -> None: + with self.target.scope(): + repository = self.target.artifact_repository( + self._client.get_run(run_id).info.artifact_uri + ) if self.target.explicit else None + if repository is None: + self._client.log_text(run_id, text, name) + else: + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / posixpath.basename(name) + path.write_text(text) + repository.log_artifact(str(path), + posixpath.dirname(name) or None) + + def _log_tags(self, run_id: str, tags: Mapping[str, str]) -> None: + entities = self._mlflow.entities + with self.target.scope(): + 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 [] + with self.target.scope(): + 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._log_batches(self._client, run_id, [ + entities.Metric(metric_key(key), float(value), timestamp, step, + model_id=model_id) + for key, value in metrics.items() + ]) + + +def _digest(name: str, source: Mapping[str, Optional[str]], + schema: Mapping[str, Any]) -> str: + return hashlib.sha256(json.dumps( + [name, dict(source), dict(schema)], sort_keys=True, + default=str).encode()).hexdigest()[:8] + + +def _millis(moment: Optional[datetime]) -> Optional[int]: + return None if moment is None else int(moment.timestamp() * 1000) + + +class _SystemMetrics: + def __init__(self, log: Callable[[List[_Reading]], None], + interval: Optional[float] = None) -> None: + fixed = os.environ.get('MLFLOW_SYSTEM_METRICS_SAMPLING_INTERVAL') + self._log = log + self._interval = interval or float(fixed or 1) + self._monitors = _metric_monitors() + self._failed = False + self._stop = threading.Event() + self._thread = threading.Thread(target=self._run, daemon=True, + name='kumo-system-metrics') + + def __enter__(self) -> '_SystemMetrics': + if self._monitors: + self._thread.start() + return self + + def __exit__(self, *exc: Any) -> None: + self._stop.set() + if self._thread.is_alive(): + self._thread.join(_FLUSH_SECONDS + 5) + + def _run(self) -> None: + began = flushed = time.monotonic() + readings: List[_Reading] = [] + taken = failures = 0 + step = -1 + stopped = False + while not stopped: + stopped = self._stop.wait(self._interval) + metrics: Dict[str, float] = {} + for monitor in self._monitors: + monitor.collect_metrics() + metrics.update(monitor.aggregate_metrics()) + monitor.clear_metrics() + step = max(round(time.monotonic() - began), step + 1) + readings.append(({f'system/{k}': v for k, v in metrics.items()}, + step, int(time.time() * 1000))) + taken += 1 + if taken < _MIN_READINGS or not ( + stopped or time.monotonic() - flushed >= _FLUSH_SECONDS): + continue + batch, readings, flushed = readings, [], time.monotonic() + try: + self._log(batch) + failures = 0 + except Exception as error: + if not self._failed: + logger.info('Could not log system metrics, skipping those ' + 'readings: %s', format_error(error)) + self._failed = True + failures += 1 + if failures == _SYSTEM_METRICS_FAILURES: + logger.info('System metrics stopped after %d failed logs ' + 'in a row', failures) + return + + +def _metric_monitors() -> List[Any]: + monitors = [] + for module, name in (('cpu_monitor', 'CPUMonitor'), + ('disk_monitor', 'DiskMonitor'), + ('network_monitor', 'NetworkMonitor'), + ('gpu_monitor', 'GPUMonitor')): + try: + monitors.append( + getattr( + import_provider(f'mlflow.system_metrics.metrics.{module}', + 'mlflow'), name)()) + except Exception: + logger.debug('No %s system metrics', name, exc_info=True) + return monitors + + +def check_model_name(name: str, registry_uri: str) -> None: + r"""Raises :class:`ValueError` unless ``registry_uri`` accepts ``name``.""" + is_databricks_unity_catalog_uri = import_provider( + 'mlflow.utils.uri', 'mlflow').is_databricks_unity_catalog_uri + + if SAGEMAKER_ARN.fullmatch(registry_uri or ''): + if not _SAGEMAKER_MODEL.fullmatch(name): + raise ValueError( + f'SageMaker MLflow cannot register a model named {name!r}: ' + f'use up to 57 letters and digits joined by hyphens, e.g. ' + f"'users-churn'.") + elif is_databricks_unity_catalog_uri(registry_uri or ''): + parts = name.split('.') + if len(parts) > 3 or not all(parts) or any( + len(part) > 255 or _UC_FORBIDDEN.search(part) + for part in parts): + raise ValueError( + f'Unity Catalog cannot register a model named {name!r}: use ' + f'catalog.schema.model without spaces or "/", e.g. ' + f"'main.ml.users_churn'.") + + +def _unusable(uri: str, error: Exception) -> RuntimeError: + if (urlparse(uri).scheme in ('', 'file') + and 'MLFLOW_ALLOW_FILE_STORE' in str(error)): + return RuntimeError( + f'MLflow refused the local folder store at {uri}: use its default ' + f'sqlite:///mlflow.db or another tracking backend.') + return RuntimeError(f'MLflow at {uri} is not usable: {full_error(error)}') + + +def _sagemaker_ui_hint(tracking_uri: str) -> Optional[str]: + match = SAGEMAKER_ARN.fullmatch(tracking_uri or '') + if match is None: + return None + region, kind, name = match.groups() + command = (f'aws sagemaker create-presigned-mlflow-tracking-server-url ' + f'--tracking-server-name {name} --region {region}' + if kind == 'mlflow-tracking-server' else + f'aws sagemaker create-presigned-mlflow-app-url ' + f'--arn {tracking_uri} --region {region}') + return ("MLflow's run links do not open on SageMaker: open the MLflow UI " + f'from SageMaker Studio or at the URL `{command}` prints.') + + +def log_model( + trainer: 'Trainer', + *, + snapshot: 'GraphSnapshot', + name: Optional[str] = None, + run_id: Optional[str] = None, + serving_snapshot: Any = None, + target: Any = None, + experiment_name: Optional[str] = None, +) -> Any: + r"""Logs a trainer as an MLflow pyfunc model; returns its `ModelInfo`. + + Args: + serving_snapshot: snapshot to score against when the request names + none. + target: as in `MlflowTracking`. + """ + from kumo.engine._manifest import FIT_ROLES, artifact_summaries + from kumo.engine._mlops import snapshot_uri + from kumo.engine.trainer._provenance import bundle_manifest + from kumo.engine.trainer._run_spec import registry_metrics + + mlflow = import_provider('mlflow', 'mlflow') + result = trainer.result + if result is None: + raise ValueError('log_model requires a fitted or loaded Trainer') + tracking = MlflowTracking(experiment_name, target=target) + target = tracking.target + if run_id is None: + tracking._require_experiment() + client = tracking._client = target.client() + active = mlflow.active_run() + if run_id is None and active is not None and ( + mlflow.get_tracking_uri() == target.tracking_uri): + run_id = active.info.run_id + created = run_id is None + with target.scope(): + if created: + run_id = client.create_run( + tracking._experiment_for(client)).info.run_id + experiment_id = client.get_run(run_id).info.experiment_id + summaries = [] if result.manifest is None else artifact_summaries( + result.manifest, FIT_ROLES) + try: + tracking._log_inputs( + run_id, [s for s in summaries if isinstance(s, Dataset)]) + except Exception as error: + logger.warning('Could not log the fit\'s source tables to MLflow: %s', + format_error(error), exc_info=True) + registered = getattr(result.registered_model, 'name', None) + info = _log_model(target, client, result.bundle_dir, + bundle_manifest(result.bundle_dir, result.id), + snapshot, name=name or _logged_model_name( + registered, result.manifest), run_id=run_id, + experiment_id=experiment_id, + serving_snapshot=snapshot_uri(serving_snapshot)) + tracking._log_metrics(run_id, registry_metrics(result), + model_id=info.model_id) + if created: + with target.scope(): + client.set_terminated(run_id, 'FINISHED') + return info + + +def load_model( + model_uri: str, + embedding_models: Optional['EmbeddingModelStore'] = None, + *, + target: Any = None, + accelerator: Literal['cpu', 'gpu', 'mps', 'auto'] = 'auto', +) -> 'Trainer': + r"""Restores the `Trainer` of a model `log_model` logged. + + Args: + target: as in `MlflowTracking`. + """ + from kumo.engine._deps import assets_through + 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') + target = resolve_target(target) + destination = trainer_scratch_dir(new_run_id()) / 'mlflow-model' + destination.mkdir(parents=True) + with target.scope(): + local = Path(_download(mlflow, target, model_uri, destination)) + model = mlflow.models.Model.load(str(local)) + bundle = model.flavors['python_function']['artifacts']['bundle']['path'] + with assets_through(target.assets, local / bundle, _target_hint): + trainer = Trainer.load(local / bundle, embedding_models, + accelerator=accelerator) + warn_if_other_bundle(model.metadata, trainer) + return trainer + + +def _target_hint(reader: str, location: str) -> str: + target = { + 'DatabricksConnector': 'DatabricksConnector(identity=...)', + 'S3Connector': 'SageMakerMLflow(arn, identity=S3Identity(...))', + }.get(reader, f'a {reader}') + return (f'pass target={target} to load_model(model_uri, ...), with an ' + f'identity that reads {location.rsplit("/", 2)[0]}') + + +def _download(mlflow: Any, target: Any, model_uri: str, + destination: Path) -> str: + if target.explicit and model_uri.startswith('models:/'): + client = target.client() + source = model_uri + reference = model_uri[len('models:/'):] + name, at, alias = reference.partition('@') + if at: + source = client.get_model_version_by_alias(name, alias).source + elif '/' in reference: + name, _, version = reference.partition('/') + source = client.get_model_version(name, version).source + if source.startswith('models:/') and '/' not in source[8:]: + location = client.get_logged_model(source[8:]).artifact_location + repository = target.artifact_repository(location) + if repository is not None: + return repository.download_artifacts('', str(destination)) + return mlflow.artifacts.download_artifacts( + artifact_uri=model_uri, dst_path=str(destination), + tracking_uri=target.tracking_uri, registry_uri=target.registry_uri) + + +def _logged_model_name(registered: Optional[str], manifest: Any) -> str: + r"""model_label without the characters MLflow rejects.""" + table = getattr(getattr(manifest, 'parents', None), 'training_table', + None) + return re.sub(r'[/:.%"\']', '-', model_label( + registered, getattr(table, 'query', None), + getattr(table, 'task_type', None))) + + +def _log_model(target: Any, client: Any, bundle_dir: Path, + manifest: Optional['ModelManifest'], + snapshot: Optional['GraphSnapshot'], *, name: str, + run_id: Optional[str], experiment_id: Optional[str], + serving_snapshot: Optional[str]) -> Any: + from kumo_core.artifacts import PROVENANCE_FILE + + from kumo.engine.integrations import _mlflow_model, _mlflow_model_code + from kumo.engine.trainer._spec import SPEC_FILE, read_training_spec + from kumo.engine.trainer.trainer import _encoders, registry_pointers + + mlflow = import_provider('mlflow', 'mlflow') + if snapshot is None: + raise ValueError('log_model needs the snapshot the model trained on') + bundle_dir = Path(bundle_dir) + signature, example = None, None + if (bundle_dir / SPEC_FILE).is_file(): + signature, example = _mlflow_model.signature( + snapshot, + read_training_spec(bundle_dir).predictive_query) + pointers = registry_pointers(_id(manifest)) + provenance = bundle_dir / PROVENANCE_FILE + with target.scope(): + run = None if run_id is None else client.get_run(run_id) + model = client.create_logged_model( + experiment_id or run.info.experiment_id, name=name, + source_run_id=run_id, + params=None if run is None else dict(run.data.params)) + mlflow_model = mlflow.models.Model( + artifact_path=model.artifact_location, model_uuid=model.model_id, + run_id=run_id, metadata=pointers, model_id=model.model_id) + model_logger = logging.getLogger('mlflow.models.model') + model_logger.addFilter(_no_serving_check) + try: + with tempfile.TemporaryDirectory() as tmp: + local = str(Path(tmp) / 'model') + mlflow.pyfunc.save_model( + local, mlflow_model=mlflow_model, + python_model=_mlflow_model_code.__file__, + model_config={ + 'columns': [] if signature is None else + signature.outputs.input_names(), + 'snapshot': serving_snapshot, + 'assets_from': None, + }, + artifacts={ + 'bundle': str(bundle_dir), + **({'provenance': str(provenance)} + if provenance.is_file() else {}), + }, + signature=signature, + input_example=example, + metadata=pointers, + pip_requirements=[ + _requirement(_encoders(bundle_dir)) + ], + ) + with target.scope(): + repository = target.artifact_repository( + model.artifact_location) + if repository is None: + client.log_model_artifacts(model.model_id, local) + else: + repository.log_artifacts(local) + logged = client.finalize_logged_model(model.model_id, + 'READY') + if run_id is not None: + client.log_outputs(run_id, [ + mlflow.entities.LoggedModelOutput(model.model_id, 0) + ]) + except BaseException: + with target.scope(): + _quietly(client.finalize_logged_model, model.model_id, 'FAILED') + raise + finally: + model_logger.removeFilter(_no_serving_check) + return mlflow_model.get_model_info(logged) + + +def _id(manifest: Optional['ModelManifest']) -> Optional[str]: + return None if manifest is None else manifest.id + + +def _requirement(config: Any) -> str: + from kumoapi.encoder import TextEmbedding + + extras = ['s3', 'gcs'] + if config is not None and any( + isinstance(encoder, TextEmbedding) + for table in config.encoders.values() + for encoder in table.values()): + extras.append('text') + return f"kumo-engine[{','.join(extras)}]=={__version__}" + + +def _register(target: Any, client: Any, info: Any, + location: ModelRegistryLocation, + metadata: Mapping[str, str], + version: Optional[str]) -> PublishedModel: + r"""Registers ``info`` as a new version unless ``version`` holds it.""" + name = location.model_name + with target.scope(): + if version is None: + try: + client.create_registered_model(name) + except Exception as error: + if _code(error) != 'RESOURCE_ALREADY_EXISTS': + raise + version = str( + client.create_model_version(name, info.model_uri, + run_id=info.run_id, tags=metadata, + model_id=info.model_id).version) + else: + logger.info('%s version %s already holds bundle %s; reused it', + name, version, metadata.get(BUNDLE_ID_KEY)) + if location.comment: + client.update_model_version(name, version, + description=location.comment) + if location.alias: + client.set_registered_model_alias(name, location.alias, version) + return PublishedModel(name=name, version=version, + uri=f'models:/{name}/{version}', + metadata=dict(metadata)) + + +def _registered_version(client: Any, name: str, + metadata: Mapping[str, str]) -> Optional[str]: + bundle = metadata.get(BUNDLE_ID_KEY) + if not bundle: + return None + token = None + while True: + try: + page = client.search_model_versions(f"name = '{name}'", + page_token=token) + except Exception as error: + if _code(error) == 'RESOURCE_DOES_NOT_EXIST': + return None + raise + for version in page: + # Unity Catalog returns tags as a method. + tags = version.tags + if not isinstance(tags, Mapping) or not tags: + tags = client.get_model_version(name, version.version).tags + if (tags or {}).get(BUNDLE_ID_KEY) == bundle: + return str(version.version) + token = page.token + if not token: + return None + + +def _code(error: BaseException) -> str: + return str(getattr(error, 'error_code', '') or '') + + +def check_registry_schema(name: str, registry_uri: str) -> None: + r"""On Unity Catalog, raises unless the model's schema is usable.""" + uri = import_provider('mlflow.utils.uri', 'mlflow') + parts = name.split('.') + if not uri.is_databricks_unity_catalog_uri(registry_uri) or len( + parts) != 3: + return + databricks = import_provider('mlflow.utils.databricks_utils', 'mlflow') + rest = import_provider('mlflow.utils.rest_utils', 'mlflow') + schema = f'{parts[0]}.{parts[1]}' + response = rest.http_request( + databricks.get_databricks_host_creds(registry_uri), + f'/api/2.1/unity-catalog/schemas/{quote(schema, safe="")}', 'GET', + max_retries=1, timeout=30) + if response.status_code in (403, 404): + try: + body = response.json() + except ValueError: + body = {} + raise ValueError( + f"Unity Catalog cannot register {name!r}: " + f"{body.get('error_code', response.status_code)}: " + f"{body.get('message', response.reason)}; create the schema or " + f"grant USE CATALOG, USE SCHEMA and CREATE MODEL on it.") + response.raise_for_status() + + +def _no_serving_check(record: logging.LogRecord) -> bool: + return not record.getMessage().startswith( + 'Failed to validate serving input example') + + +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') + parts = re.findall(r'`((?:[^`]|``)*)`', location) + head = (entry.get('version') or '').split('_', 1)[0] + return delta.DeltaDatasetSource( + delta_table_name='.'.join(part.replace('``', '`') + for part in parts) or location, + delta_table_version=int(head) if entry.get('version_kind') + == 'exact' and head.isdigit() else None) + 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 _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() -> Iterator[FrozenSet[str]]: + r"""MLflow's HTTP limits where unset; yields the names it set.""" + owned = frozenset(name for name in _HTTP_LIMITS if name not in os.environ) + with environ({name: _HTTP_LIMITS[name] for name in owned}): + yield owned 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..62f7882f3 --- /dev/null +++ b/packages/kumo-engine/kumo/engine/integrations/snowflake.py @@ -0,0 +1,292 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +r"""Snowflake Experiment Tracking for `Trainer.fit`.""" + +import logging +import posixpath +import re +import tempfile +from contextlib import ExitStack +from datetime import datetime, timezone +from pathlib import Path +from typing import Any, List, Mapping, Optional, Sequence, Set, Tuple + +from kumo_core.instrumentation import format_error +from kumo_core.sources.sql import resolve_snowflake_identifier +from kumo_core.summaries import Summary + +from kumo.engine._events import Event +from kumo.engine._mlops import full_error +from kumo.engine._run_names import record_name +from kumo.engine._summary_tables import route, table_json +from kumo.engine._tracking import ( + DROPPED_METRICS, + DROPPED_PARAMS, + RunRecord, + TrackingBase, + by_step, +) +from kumo.engine.connector._extras import import_provider + +logger = logging.getLogger(__name__) + +_MAX_METRICS = 200 +_MAX_PARAMS = 1000 +_EPOCH_METRICS = _MAX_METRICS - 60 +_SUMMARY_METRICS = _MAX_METRICS - 30 +_FINISH_PARAMS = 20 + + +class SnowflakeTracking(TrackingBase): + r"""Records `Trainer.fit` in Snowflake Experiment Tracking as one run. + + Args: + source: a ``SnowflakeConnector``, or a Snowpark session (tracking only). + native_model_lineage: link models registered during the fit to its run. + """ + + platform = 'snowflake' + + def __init__( + self, + source: Any, + *, + experiment_name: str, + database: Optional[str] = None, + schema_name: Optional[str] = None, + run_name: Optional[str] = None, + tags: Optional[Mapping[str, str]] = None, + native_model_lineage: bool = False, + ) -> None: + from kumo.engine.connector.base import Connector + + self._module = import_provider('snowflake.ml.experiment', 'snowflake') + if source is None: + raise ValueError('SnowflakeTracking takes a SnowflakeConnector ' + 'or a Snowpark session.') + self.connector = source if isinstance(source, Connector) else None + self.session = (source.session() + if self.connector is not None else source) + self.experiment_name = experiment_name + self.database = database + self.schema = schema_name + self.run_name = run_name + self.tags = dict(tags or {}) + self.native_model_lineage = native_model_lineage + self.logged_run_name: Optional[str] = None + self._reset() + + def _reset(self) -> None: + self._tracking: Any = None + self._run: Any = None + self._run_context = ExitStack() + self._tune = 'loss' + self._metrics: Set[str] = set() + self._started = _now() + 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 load_model(self, location: Any, **kwargs: Any) -> Any: + r"""The `Trainer` registered at ``location``.""" + from kumo.engine.trainer.trainer import Trainer + + if self.connector is None: + raise ValueError('Loading a registered model needs ' + 'SnowflakeTracking(connector, ...).') + return Trainer.load_from_connector(self.connector, location=location, + **kwargs) + + def _location(self) -> Tuple[Optional[str], str]: + database = self.database or self.session.get_current_database() + schema = self.schema or (None if self.database else + self.session.get_current_schema()) + return database, schema or 'PUBLIC' + + def check(self, registry: Any = None) -> None: + r"""Raises unless the session can see the experiment's schema.""" + database, schema = self._location() + if database is None: + return + name = resolve_snowflake_identifier(schema) + try: + rows = self.session.sql( + f"SHOW SCHEMAS LIKE '{name.replace(chr(39), chr(39) * 2)}' " + f"IN DATABASE {database}").collect() + except Exception as error: + raise RuntimeError( + f'Snowflake Experiment Tracking cannot use {database}: ' + f'{full_error(error)}') from error + if not rows: + raise ValueError( + f'Snowflake schema {database}.{schema} does not exist or this ' + f'role cannot see it.') + + def start(self, record: RunRecord) -> None: + self._reset() + self.run_id = self.logged_run_name = self.experiment = 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): + raise RuntimeError( + f'Snowflake ExperimentTracking is already in use in this ' + f'process, so this {record.tags.get("kumo.kind")} run is not ' + f'tracked; end that run or pass its session.') + self._tracking = tracking + database, schema = self._location() + experiment = tracking.set_experiment(self.experiment_name, + database_name=database, + schema_name=schema) + self.experiment = str(experiment.name) + existing = [ + str(row['name']) + for row in tracking._sql_client.show_runs_in_experiment( + experiment_name=experiment.name) + ] + self._started = record.started or _now() + self._run = tracking.start_run(self._new_run_name( + existing, _identifier(record_name(record, self.run_name)))) + self.run_id = self.logged_run_name = str(self._run.name) + if self.native_model_lineage: + self._run_context.enter_context(self._run) + self._start_tags = set(record.tags) + self._tune = record.tags.get('kumo.tune_metric', 'loss') + 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 + attempt = int(event.params['attempt']) + prefix = f'TRIAL_{event.trial or 0}' + ( + f'_ATTEMPT_{attempt}/' if attempt > 1 else '/') + if event.name == 'trial_started': + self._log_params({prefix + k: v for k, v in record.params.items()}) + elif event.name == 'trial_ended': + best = f'val/best_{self._tune}' + self._log_metrics({prefix + best: value + for name, value in record.metrics.items() + if name == best}, _SUMMARY_METRICS) + self._log_params({prefix + k: v for k, v in record.tags.items() + if k in ('kumo.status', 'kumo.error')}) + + def finish(self, record: Optional[RunRecord], + error: Optional[BaseException]) -> None: + status = 'FAILED' + try: + if self._run is not None and record is not None: + for step, metrics in by_step(record.steps): + self.log_metrics(metrics, step) + 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_summaries(record.summaries) + finished = record.finished or _now() + self._log_metrics({'duration_seconds': ( + finished - self._started).total_seconds()}, _MAX_METRICS) + 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 as error: + logger.warning('Could not end the Snowflake run %s: %s', + status, format_error(error), exc_info=True) + finally: + self._run_context.close() + self._reset() + + def log_metrics(self, metrics: Mapping[str, float], step: int) -> None: + self._log_metrics(metrics, _EPOCH_METRICS, step=step) + + def _new_run_name(self, existing: List[str], run_name: str) -> str: + # start_run matches name with LIKE. + name, number = run_name, 1 + while _taken(name, existing): + number += 1 + name = f'{run_name}_{number}' + if number > 1: + logger.warning( + 'Snowflake run %s already exists in experiment %s; recording ' + 'this run as %s', run_name, self.experiment_name, name) + return name + + def _log_summaries(self, summaries: Sequence[Summary]) -> None: + routed = route(summaries) + self._log_metrics({s.name: s.value for s in routed.scalars}, + _MAX_METRICS) + self._log_params(routed.tags) + self._log_artifacts({f'{table.name}.json': table_json(table) + for table in routed.tables}) + + 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) / posixpath.basename(name) + path.write_text(text, encoding='utf-8') + self._tracking.log_artifact( + str(path), posixpath.dirname(name) or None) + 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 _now() -> datetime: + return datetime.now(timezone.utc) + + +def _identifier(name: str) -> str: + return re.sub(r'[^A-Za-z0-9_]+', '_', name).upper() + + +def _taken(name: str, existing: List[str]) -> bool: + regex = ''.join('.' if c == '_' else re.escape(c) for c in name) + return any(re.fullmatch(regex, other, re.IGNORECASE | re.DOTALL) + for other in existing) diff --git a/packages/kumo-engine/kumo/engine/pquery/training_table.py b/packages/kumo-engine/kumo/engine/pquery/training_table.py index 73e838bee..0ac980533 100644 --- a/packages/kumo-engine/kumo/engine/pquery/training_table.py +++ b/packages/kumo-engine/kumo/engine/pquery/training_table.py @@ -24,6 +24,7 @@ SourceTableEntry, TrainingTableManifest, TrainingTableParents, + embedded, read_manifest, stamp, stub, @@ -144,6 +145,8 @@ def _store(self, settings: Any) -> None: path = settings.root.location(f'training_tables/{self.id}/') try: + settings.locate(self.staging_dir / PROVENANCE_FILE, + TrainingTableManifest, f'training_tables/{self.id}') filesystem, root = file_store(settings.root.connector, path) if filesystem.get_file_info( f'{root}/_SUCCESS').type is not FileType.File: @@ -395,7 +398,7 @@ def _record_manifest(table: TrainingTable, snapshot: GraphSnapshot, TrainingTableManifest( id=table.id, parents=TrainingTableParents( - snapshot=parent if parent is not None else + snapshot=embedded(parent) if parent is not None else stub('snapshot', snapshot.id), source_table=source_table), query=table.query, diff --git a/packages/kumo-engine/kumo/engine/tracking.py b/packages/kumo-engine/kumo/engine/tracking.py new file mode 100644 index 000000000..4a7dc3341 --- /dev/null +++ b/packages/kumo-engine/kumo/engine/tracking.py @@ -0,0 +1,13 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +r"""API for writing an experiment tracker for ``kumo.MLOps``.""" + +from kumo.engine._events import Event +from kumo.engine._tracking import RunRecord, Tracking + +__all__ = [ + 'Event', + 'Tracking', + 'RunRecord', +] diff --git a/packages/kumo-engine/kumo/engine/trainer/__init__.py b/packages/kumo-engine/kumo/engine/trainer/__init__.py index 0f4c1cfff..abc11c7e3 100644 --- a/packages/kumo-engine/kumo/engine/trainer/__init__.py +++ b/packages/kumo-engine/kumo/engine/trainer/__init__.py @@ -2,6 +2,7 @@ # 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, @@ -18,6 +19,7 @@ "Trainer", "TrainingResult", "PredictionResult", + "ScoringDefaults", "ModelOutputConfig", "PredictionOutputConfig", "SnapshotOutputConfig", diff --git a/packages/kumo-engine/kumo/engine/trainer/_provenance.py b/packages/kumo-engine/kumo/engine/trainer/_provenance.py index 9a92656b0..ec282d74e 100644 --- a/packages/kumo-engine/kumo/engine/trainer/_provenance.py +++ b/packages/kumo-engine/kumo/engine/trainer/_provenance.py @@ -4,12 +4,15 @@ from dataclasses import dataclass, replace from datetime import datetime, timezone from pathlib import Path -from typing import Optional +from typing import List, Optional from kumo.engine._manifest import ( + Dependency, FitTiming, ModelManifest, ModelParents, + TrackingRef, + embedded, read_manifest, stamp, stub, @@ -25,6 +28,8 @@ class FitRecord: warm_start: Optional[ModelManifest] started_at: datetime fit_mode: Optional[str] + tracking: Optional[TrackingRef] = None + dependencies: Optional[List[Dependency]] = None def write(self, bundle_dir: Path, bundle_id: str) -> None: from kumo_core.artifacts import PROVENANCE_FILE, add_bundle_file @@ -37,7 +42,7 @@ def write(self, bundle_dir: Path, bundle_id: str) -> None: ModelManifest( id=bundle_id, parents=ModelParents( - training_table=table if table is not None else + training_table=embedded(table) if table is not None else stub('training_table', train_table.id or train_table.tt_id), warm_start=None if self.warm_start is None else @@ -49,21 +54,25 @@ def write(self, bundle_dir: Path, bundle_id: str) -> None: (finished_at - started_at).total_seconds(), 3), ), fit_mode=self.fit_mode, + tracking=self.tracking, + dependencies=self.dependencies, **stamp(), )) add_bundle_file(bundle_dir, PROVENANCE_ROLE, PROVENANCE_FILE) def embedded_model(manifest: ModelManifest) -> ModelManifest: - parents = manifest.parents + copy = embedded(manifest) + parents = copy.parents if parents is None: - return manifest + return copy def cut(model: Optional[ModelManifest]) -> Optional[ModelManifest]: return None if model is None else stub('model', model.id) - return replace(manifest, parents=replace( - parents, warm_start=cut(parents.warm_start))) + return replace(copy, parents=replace( + parents, warm_start=cut(parents.warm_start), + derived_from=cut(parents.derived_from))) def bundle_manifest(bundle_dir: Path, @@ -77,3 +86,26 @@ def bundle_manifest(bundle_dir: Path, def read_provenance(bundle_dir: Path) -> Optional[ModelManifest]: path = listed_file(bundle_dir, PROVENANCE_ROLE) return None if path is None else read_manifest(path, ModelManifest) + + +def record_derivation(bundle_dir: Path, parent_dir: Path, + parent_id: Optional[str]) -> None: + r"""Writes the manifest of a bundle derived from ``parent_dir``.""" + from kumo_core.artifacts import ( + PROVENANCE_FILE, + add_bundle_file, + bundle_id_from_manifest, + ) + + parent = bundle_manifest(parent_dir, parent_id) + write_manifest( + bundle_dir / PROVENANCE_FILE, + ModelManifest( + id=bundle_id_from_manifest(bundle_dir), + parents=ModelParents( + derived_from=None if parent is None else embedded_model( + parent)), + dependencies=None if parent is None else parent.dependencies, + **stamp(), + )) + add_bundle_file(bundle_dir, PROVENANCE_ROLE, PROVENANCE_FILE) diff --git a/packages/kumo-engine/kumo/engine/trainer/_run_spec.py b/packages/kumo-engine/kumo/engine/trainer/_run_spec.py new file mode 100644 index 000000000..bb044bf1f --- /dev/null +++ b/packages/kumo-engine/kumo/engine/trainer/_run_spec.py @@ -0,0 +1,102 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +r"""A `Trainer.fit` as a tracked run's `RunSpec`.""" + +import json +from dataclasses import dataclass, replace +from datetime import datetime, timezone +from pathlib import Path +from typing import Any, Dict, Iterator, Mapping, Optional, Sequence, Tuple + +from kumo_core.artifacts import PROVENANCE_FILE +from kumo_core.summaries import Scalar, Summary + +from kumo.engine._run_names import fit_name, model_label +from kumo.engine._text import entity, metric_key +from kumo.engine._tracking import RunSpec, logged_names, metric_name + + +@dataclass(frozen=True) +class FitContext: + r"""What a fit trains, known once its model plan is resolved.""" + + query: str + model_plan: Mapping[str, Any] + started: Optional[datetime] = None + snapshot_id: Optional[str] = None + task_type: Optional[str] = None + model_name: Optional[str] = None + + +_UNSET = '???' + + +def fit_spec(context: FitContext) -> RunSpec: + started = context.started or datetime.now(timezone.utc) + model = model_label(context.model_name, context.query, context.task_type) + tune = (context.model_plan.get('training_job') or {}).get('tune_metric') + return RunSpec( + kind='fit', + name=fit_name(model, started), + model_label=model, + started=started, + params={key: value for key, value in _flatten(context.model_plan) + if value is not None and value != _UNSET}, + parents={'snapshot_id': context.snapshot_id, + 'task_type': context.task_type, + 'entity': entity(context.query), + 'tune_metric': tune and metric_name(tune)}, + artifacts={ + 'query.pql': context.query, + 'model_plan.json': json.dumps(dict(context.model_plan), indent=2, + default=str), + }, + ) + + +def fitted_spec(spec: RunSpec, result: Any, + summaries: Sequence[Summary] = ()) -> RunSpec: + r"""The fit run's spec once the fit returned.""" + artifacts = dict(spec.artifacts) + provenance = Path(result.bundle_dir) / PROVENANCE_FILE + if provenance.is_file(): + artifacts[PROVENANCE_FILE] = provenance.read_text(encoding='utf-8') + return replace( + spec, metrics=holdout_metrics(result), + parents={**spec.parents, 'bundle_id': result.id, + **_registered(result.registered_model)}, + artifacts=artifacts, summaries=summaries) + + +def holdout_metrics(result: Any) -> Dict[str, float]: + r"""The trainer's test metrics under their logged names.""" + return logged_names({ + name: value for name, value in result.metrics.items() + if not name.startswith(('train_', 'val_')) + }, refit_full=getattr(getattr(result, 'manifest', None), 'fit_mode', + None) == 'refit_full') + + +def registry_metrics(result: Any) -> Dict[str, float]: + r"""Holdout metrics and model scalars, as registries record them.""" + return {metric_key(name): value for name, value in { + **holdout_metrics(result), + **{summary.name: summary.value + for summary in getattr(result, '_summaries', ()) + if isinstance(summary, Scalar) and summary.role == 'model'} + }.items()} + + +def _registered(model: Any) -> Mapping[str, Any]: + if model is None: + return {} + return {'model_name': model.name, 'model_version': model.version} + + +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 diff --git a/packages/kumo-engine/kumo/engine/trainer/_runtime.py b/packages/kumo-engine/kumo/engine/trainer/_runtime.py index f258d6ac1..564bdcc6d 100644 --- a/packages/kumo-engine/kumo/engine/trainer/_runtime.py +++ b/packages/kumo-engine/kumo/engine/trainer/_runtime.py @@ -11,8 +11,10 @@ def local_training_bindings(training_input: Any, run_name: Optional[str] = None, embedding_models: Any = None, - accelerator: str = 'auto') -> Any: + accelerator: str = 'auto', + 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, accelerator=accelerator) + embedding_models=embedding_models, accelerator=accelerator, + 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..c6a1f0655 --- /dev/null +++ b/packages/kumo-engine/kumo/engine/trainer/_scoring.py @@ -0,0 +1,81 @@ +# 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.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 + from kumo_core.io.filesystem import atomic_write_bytes + + if defaults == ScoringDefaults(): + remove_bundle_file(bundle_dir, SCORING_DEFAULTS_ROLE) + return + atomic_write_bytes( + bundle_dir / SCORING_DEFAULTS_FILE, + (json.dumps(asdict(defaults), indent=2, sort_keys=True) + + '\n').encode('utf-8')) + 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 5420d9f11..8e5b712b0 100644 --- a/packages/kumo-engine/kumo/engine/trainer/_spec.py +++ b/packages/kumo-engine/kumo/engine/trainer/_spec.py @@ -24,6 +24,7 @@ MANIFEST_FILE = 'publication-manifest.json' SPEC_ROLE = 'training_spec' PROVENANCE_ROLE = 'provenance' +SCORING_DEFAULTS_ROLE = 'scoring_defaults' HOLDOUT_ROLE = 'holdout' HOLDOUT_META_ROLE = 'holdout_meta' @@ -151,6 +152,7 @@ def read_bundle_artifact_paths( if entry['role'] not in ( SPEC_ROLE, PROVENANCE_ROLE, + SCORING_DEFAULTS_ROLE, HOLDOUT_ROLE, HOLDOUT_META_ROLE, ) diff --git a/packages/kumo-engine/kumo/engine/trainer/output_config.py b/packages/kumo-engine/kumo/engine/trainer/output_config.py index 54ed5a68b..2e75163f0 100644 --- a/packages/kumo-engine/kumo/engine/trainer/output_config.py +++ b/packages/kumo-engine/kumo/engine/trainer/output_config.py @@ -12,6 +12,7 @@ FileArtifactSink, ModelRegistryLocation, SourceConnector, + check_model_alias, ) from kumoapi.prediction_outputs import resolve_prediction_output_paths @@ -113,3 +114,16 @@ class ModelOutputConfig: def __post_init__(self) -> None: if (self.path is None) == (self.location is None): raise ValueError("Exactly one of path or location must be set.") + + +def check_model_location(config: ModelOutputConfig) -> None: + """Raises ``ValueError`` unless ``config.location`` is registrable.""" + location = config.location + if location is None: + return + if location.alias: + check_model_alias(location.alias) + check = getattr(config.connector.sink_backend, 'check_model_location', + None) + if callable(check): + check(config.connector, location) diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index 979964962..96107524e 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -8,6 +8,7 @@ import hashlib import json import logging +import os import shutil import time import warnings @@ -18,9 +19,11 @@ Any, Callable, Dict, + List, Literal, Mapping, Optional, + Sequence, Tuple, TypeVar, Union, @@ -29,6 +32,7 @@ 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 @@ -43,6 +47,8 @@ from kumo_core.types import ExplainMethod, PredictionArtifactType from kumoapi.model_plan import ModelPlan, TrainingTableGenerationPlan from kumoapi.plugin import ( + BUNDLE_ID_KEY, + ModelAliasingSink, ModelFetchingSink, ModelPublishingSink, ModelRegistryLocation, @@ -53,13 +59,33 @@ from kumoapi.typing import check_prediction_precision from typing_extensions import Self +from kumo.engine import _deps as dependency +from kumo.engine._data_summaries import ( + snapshot_summaries, + split_summaries, + table_details, + training_table_details, +) +from kumo.engine._deps import Dependencies from kumo.engine._embeddings import ( no_encoder_inference_warnings, require_embedding_models, + require_text_encoders, require_trained_embedding_models, + with_text_encoders, +) +from kumo.engine._events import Event, EventEmitter, query_sha256 +from kumo.engine._manifest import ( + FIT_ROLES, + Dependency, + ModelManifest, + TrackingRef, + artifact_summaries, ) -from kumo.engine._manifest import ModelManifest +from kumo.engine._mlops import full_error, warn_failure +from kumo.engine._predictor import TRAINER_PREDICTOR from kumo.engine._scratch import new_run_id, trainer_scratch_dir +from kumo.engine._tracking import RunSpec, TrackingSession, as_tracking from kumo.engine.connector.base import Connector from kumo.engine.connector.source_table import SourceTable from kumo.engine.graph import ( @@ -73,6 +99,19 @@ FitRecord, bundle_manifest, read_provenance, + record_derivation, +) +from kumo.engine.trainer._run_spec import ( + FitContext, + fit_spec, + fitted_spec, + registry_metrics, +) +from kumo.engine.trainer._scoring import ( + ScoringDefaults, + normalized_scoring_defaults, + read_scoring_defaults, + write_scoring_defaults, ) from kumo.engine.trainer._spec import ( HOLDOUT_META_ROLE, @@ -86,6 +125,7 @@ from kumo.engine.trainer.output_config import ( ModelOutputConfig, PredictionOutputConfig, + check_model_location, ) logger = logging.getLogger(__name__) @@ -97,6 +137,7 @@ class TrainingResult: `published_model` records the destination's assigned name, version and URI when `fit()` exports the bundle. It is `None` without an export. + `registered_model` is the version a `kumo.MLOps` ``registry`` registered. `holdout_path` is the published model's holdout Parquet file written by `fit()`, or `None` when it has none (no holdout task, or a model refit on @@ -114,6 +155,8 @@ class TrainingResult: holdout_path: Optional[Path] = None holdout_truncated: Optional[bool] = False bundle_id: Optional[str] = None + registered_model: Optional[PublishedModel] = None + scoring_defaults: ScoringDefaults = field(default_factory=ScoringDefaults) _trainer: Optional['Trainer'] = field(default=None, repr=False, compare=False) @@ -134,11 +177,63 @@ def snapshot_id(self) -> Optional[str]: return None return read_training_spec(self.bundle_dir).snapshot_id + @property + def tracking(self) -> Optional[TrackingRef]: + """The experiment tracker's run of the fit, if tracked.""" + manifest = self.manifest + return None if manifest is None else manifest.tracking + def predict(self, *args: Any, **kwargs: Any) -> 'PredictionResult': """Score with this model, as :meth:`Trainer.predict`.""" + return self._model('predict').predict(*args, **kwargs) + + def export(self, config: Any) -> PublishedModel: + """Publishes the bundle to a ``ModelOutputConfig`` or ``MLOps``.""" + from kumo.engine._mlops import MLOps + + if isinstance(config, MLOps): + if config.registry is None: + raise ValueError('export(cfg) registers at cfg.registry, ' + 'which is not set.') + tracker = as_tracking(config.tracking) + if tracker is not None and tracker.register is not None: + return cast(PublishedModel, _with_metrics( + tracker.register(self, config.registry), + registry_metrics(self))) + config = _connector_registry(config) + _require_publisher(config) + return _publish(self.bundle_dir, self.bundle_id, + registry_metrics(self), config)[0] + + def with_scoring_defaults(self, + defaults: ScoringDefaults) -> 'TrainingResult': + """A copy of this model whose ``predict`` applies ``defaults``.""" + from kumo_core.artifacts import bundle_task_type + + if self._trainer is None: + raise RuntimeError('This result has no model; use Trainer.load on ' + 'its bundle_dir.') + defaults = normalized_scoring_defaults( + defaults, bundle_task_type(self.artifact_paths)) + bundle_dir = trainer_scratch_dir(new_run_id()) / 'derived-model' + shutil.copytree(self.bundle_dir, bundle_dir, + copy_function=_link_or_copy) + write_scoring_defaults(bundle_dir, defaults) + try: + record_derivation(bundle_dir, self.bundle_dir, self.bundle_id) + except Exception as error: + warnings.warn( + f"Could not record the parent model in _provenance.json: " + f"{error}", UserWarning, stacklevel=2) + result = Trainer._load_bundle(bundle_dir, + self._trainer._embedding_models).result + assert result is not None + return result + + def _model(self, call: str) -> 'Trainer': trainer = self._trainer if trainer is None: - raise RuntimeError('This result has no model to predict with; ' + raise RuntimeError(f'This result has no model to {call} with; ' 'use Trainer.load on its bundle_dir.') current = trainer._training_result if current is None or current.bundle_dir != self.bundle_dir: @@ -146,7 +241,7 @@ def predict(self, *args: Any, **kwargs: Any) -> 'PredictionResult': trainer._embedding_models, accelerator=trainer._accelerator) self._trainer = trainer - return trainer.predict(*args, **kwargs) + return trainer def holdout_df(self) -> pd.DataFrame: r"""Read the holdout predictions as a pandas DataFrame. @@ -278,16 +373,71 @@ def _listed_holdout(bundle_dir: Path) -> tuple[Optional[Path], Optional[bool]]: return path, truncated +def _store_and_register( + bundle_dir: Path, + manifest: Optional[ModelManifest], + settings: Any, + *, + session: Optional[TrackingSession] = None, + snapshot: Any = None, + metrics: Mapping[str, float] = {}, + warn: bool = True, + stored: Optional[BaseException] = None, +) -> Tuple[Optional[PublishedModel], Optional[BaseException]]: + if stored is None and settings is not None and settings.wants('model'): + try: + _store_model(bundle_dir, manifest, settings) + except Exception as error: + stored = error + if stored is not None and warn: + warn_failure(f'{stored}; the model is not registered', stored, 3) + registered, error = (None, None) if session is None else session.log_model( + bundle_dir, manifest, snapshot, + None if stored or not _logs_models(settings) else settings.registry) + registered = _with_metrics(registered, metrics) + error = stored or error + registry = None if stored else _connector_registry(settings) + if registry is None: + return registered, error + try: + registered, alias_error = _publish( + bundle_dir, None if manifest is None else manifest.id, metrics, + registry, warn) + except Exception as failure: + if warn: + warn_failure( + f"Could not register the model: {full_error(failure)}", + failure, 3) + return registered, error or failure + return registered, error or alias_error + + +def _fit_identity(snapshot: GraphSnapshot, + train_table: TrainingTable) -> Dict[str, Any]: + try: + return { + 'snapshot_id': snapshot.id, + 'query_sha256': query_sha256(train_table.query), + } + except Exception: + return {} + + def _store_model(bundle_dir: Path, manifest: Optional[ModelManifest], settings: Any) -> None: + from kumo_core.artifacts import PROVENANCE_FILE, add_bundle_file from kumo_core.object_store_sink import holds_bundle, upload_bundle from pyarrow.fs import FileSelector from kumo.engine.connector import file_store + from kumo.engine.trainer._spec import PROVENANCE_ROLE if manifest is None: return try: + if settings.locate(bundle_dir / PROVENANCE_FILE, ModelManifest, + f'models/{manifest.id}') is not None: + add_bundle_file(bundle_dir, PROVENANCE_ROLE, PROVENANCE_FILE) filesystem, root = file_store( settings.root.connector, settings.root.location(f'models/{manifest.id}')) @@ -328,12 +478,40 @@ 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 isinstance(export_config.connector.sink_backend, ModelPublishingSink): raise NotImplementedError( f'{type(export_config.connector).__name__} cannot publish ' f'models') + check_model_location(export_config) + + +def _with_metrics(model: Optional[PublishedModel], + metrics: Mapping[str, float]) -> Optional[PublishedModel]: + if model is None or model.metrics: + return model + return replace(model, metrics=dict(metrics)) + + +def _connector_registry(settings: Any) -> Optional[ModelOutputConfig]: + if (settings is None or settings.registry is None + or _logs_models(settings)): + return None + return ModelOutputConfig(as_tracking(settings.tracking).connector, + location=settings.registry) + + +def _logs_models(settings: Any) -> bool: + tracker = None if settings is None else as_tracking(settings.tracking) + return tracker is not None and tracker.log_model is not None def _preflight(output_config: PredictionOutputConfig) -> None: @@ -344,20 +522,50 @@ def _preflight(output_config: PredictionOutputConfig) -> None: write_mode=output_config.write_mode) -def _publish(bundle_dir: Path, metrics: Mapping[str, float], - export_config: ModelOutputConfig) -> PublishedModel: - spec = read_training_spec(bundle_dir) - return export_config.connector.sink_backend.publish_model( +def _publish(bundle_dir: Path, bundle_id: Optional[str], + metrics: Mapping[str, float], export_config: ModelOutputConfig, + warn: bool = True + ) -> Tuple[PublishedModel, Optional[BaseException]]: + r"""The published model, and the error setting its alias, if any.""" + metadata = {} + if (bundle_dir / SPEC_FILE).is_file(): + spec = read_training_spec(bundle_dir) + metadata = { + name: str(getattr(spec, name)) + for name in ('spec_version', 'predictive_query', 'snapshot_id') + } + metadata.update(registry_pointers(bundle_id)) + location = export_config.location + alias = location.alias if location else None + if alias: + location = replace(location, alias=None) + published = export_config.connector.sink_backend.publish_model( export_config.connector, str(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') - }, + location=location, + metadata=metadata, metrics=metrics, ) + if not alias: + return published, None + try: + aliasing = export_config.connector.sink_backend + if not isinstance(aliasing, ModelAliasingSink): + raise NotImplementedError( + f'{type(export_config.connector).__name__} cannot set ' + f'model aliases') + aliasing.set_model_alias( + cast(SourceConnector, export_config.connector), + replace(location, version=published.version), alias) + except Exception as error: + if warn: + warn_failure( + f"Published {published.name} version " + f"{published.version} without alias {alias}: " + f"{full_error(error)}", error, 4) + return published, error + return published, None def _warm_start_plan( @@ -448,28 +656,35 @@ def _stage_model_bundle(artifacts: Mapping[str, str], dest: Path) -> Path: _F = TypeVar('_F', bound=Callable[..., Any]) -def _without_torch_scatter_hint(method: _F) -> _F: - r"""Runs ``method`` without PyTorch Geometric's GPU hint that min/max - aggregation "can be accelerated via the 'torch-scatter' package". +_UNACTIONABLE = ( + (r"The usage of `scatter\(reduce='\w+'\)` can be accelerated via the " + r"'torch-scatter' package", UserWarning), + (r'Your `IterableDataset` has `__len__`', UserWarning), + (r'This process .* is multi-threaded, use of fork\(\) may lead to ' + r'deadlocks', DeprecationWarning), +) - kumo-engine uses PyTorch's own scatter instead of that optional extension - (see ``test_wheel_uses_torch_fallbacks_without_native_extensions``), so - the hint, which PyG raises for every such aggregation, is not actionable. - """ + +def _without_unactionable_warnings(method: _F) -> _F: @functools.wraps(method) def wrapper(*args: Any, **kwargs: Any) -> Any: - with warnings.catch_warnings(): - warnings.filterwarnings( - 'ignore', - message=r"The usage of `scatter\(reduce='\w+'\)` can be " - r"accelerated via the 'torch-scatter' package", - category=UserWarning, - ) - return method(*args, **kwargs) + tips = logging.getLogger('pytorch_lightning.utilities.rank_zero') + tips.addFilter(_no_lightning_tip) + try: + with warnings.catch_warnings(): + for message, category in _UNACTIONABLE: + warnings.filterwarnings('ignore', message, category) + return method(*args, **kwargs) + finally: + tips.removeFilter(_no_lightning_tip) return cast(_F, wrapper) +def _no_lightning_tip(record: logging.LogRecord) -> bool: + return not str(record.msg).startswith('💡 Tip') + + class Trainer: r"""Trains and runs a GNN model in-process. @@ -530,15 +745,28 @@ def result(self) -> Optional[TrainingResult]: """The last fit's result, or the loaded bundle's.""" return self._training_result - @_without_torch_scatter_hint + 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 export(self, config: Any) -> PublishedModel: + """Publishes the model, as :meth:`TrainingResult.export`.""" + return self._fitted('export').export(config) + + @_without_unactionable_warnings def fit( self, train_table: TrainingTable, *, snapshot: GraphSnapshot, export_config: Optional[ModelOutputConfig] = None, + scoring_defaults: Optional[ScoringDefaults] = None, warm_start: Optional[Self] = None, mlops: Optional[Any] = None, + callbacks: Sequence[Callable[[Event], None]] = (), ) -> TrainingResult: r"""Train a model against `train_table`. @@ -559,6 +787,7 @@ def fit( generate_training_table(snapshot)`. snapshot: graph data and features to train against. export_config: if given, publishes the trained model bundle. + scoring_defaults: arguments ``predict`` applies when omitted. warm_start: a fitted trainer or one restored with ``load`` or ``load_from_connector``. Reuses its selected weights, optimizer state, fitted encoders, architecture and neighbor sampling; @@ -568,7 +797,8 @@ def fit( architecture; they support prediction but are unreliable warm-start sources. mlops: a `kumo.MLOps`; with MLOps on (this or a `kumo.mlops` - block), the bundle is also stored under the artifact root. + block), the bundle is also stored and the fit tracked. + callbacks: called with each progress `Event`. Each call returns a new `TrainingResult`; `Trainer.result` holds the latest. @@ -578,19 +808,59 @@ def fit( from kumo.engine._mlops import raise_on_error, resolve settings = resolve(mlops) + registry = _connector_registry(settings) + if registry is not None: + _require_publisher(registry) run_id = f'{self._id}-fit-{new_run_id()}' - result = self._fit(train_table, snapshot=snapshot, - export_config=export_config, - warm_start=warm_start, run_id=run_id) - if settings is not None and settings.wants('model'): - try: - _store_model( - result.bundle_dir, - bundle_manifest(result.bundle_dir, result.bundle_id), - settings) - except Exception as error: - raise_on_error(settings, error, - 'Trainer.fit (its model is in trainer.result)') + warn = settings is None or settings.on_error == 'warn' + session = TrackingSession( + None if settings is None else settings.tracking, warn=warn) + events = EventEmitter( + [*callbacks, session.on_event] if session.enabled else callbacks) + deps = Dependencies(settings=settings) + fit_start = time.monotonic() + events.emit('fit_started', params=_fit_identity(snapshot, train_table)) + try: + result = self._fit(train_table, snapshot=snapshot, + export_config=export_config, + scoring_defaults=scoring_defaults, + warm_start=warm_start, run_id=run_id, + events=events, tracking=session, + deps=deps, model_name=None if ( + settings is None or settings.registry is + None) else settings.registry.model_name) + 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'}) + registered, error = _store_and_register( + result.bundle_dir, + bundle_manifest(result.bundle_dir, result.bundle_id), settings, + session=session, snapshot=snapshot, + metrics=registry_metrics(result), warn=warn, + stored=(deps.errors or [None])[0]) + if registered is not None: + result = replace(result, registered_model=registered) + self._training_result = result + session.finish( + lambda: fitted_spec( + cast(RunSpec, session.spec), result, ( + *snapshot_summaries(Path(snapshot.index_path), + snapshot.graph_config), + *([] if result.manifest is None else artifact_summaries( + result.manifest, FIT_ROLES, + {**table_details(Path(snapshot.index_path), + snapshot.id), + **training_table_details(train_table)})), + *split_summaries(train_table.staging_dir))), error) + raise_on_error(settings, session.error or error, + 'Trainer.fit (its model is in trainer.result)') return result def _fit( @@ -599,10 +869,17 @@ def _fit( *, snapshot: GraphSnapshot, export_config: Optional[ModelOutputConfig], + scoring_defaults: Optional[ScoringDefaults], warm_start: Optional[Self], run_id: str, + events: EventEmitter, + tracking: TrackingSession, + deps: Dependencies, + model_name: Optional[str] = None, ) -> TrainingResult: + from kumo_core.artifacts import bundle_task_type from kumo_core.automl.runtime import validate_warm_start_task_type + from kumo_core.io.json import to_json_dict from kumo_core.training.search import AutoMLSearchOptions from kumo_core.training.workflow import PreparedGraph from kumo_core.types import TaskType @@ -625,18 +902,27 @@ def _fit( training_table_generation_plan=train_table.generation_plan, search=AutoMLSearchOptions(), ) + if scoring_defaults is not None: + scoring_defaults = normalized_scoring_defaults( + scoring_defaults, TaskType(workflow.query.task_type)) if warm_start_artifacts is not None: - source_schema = json.loads( - Path(warm_start_artifacts['prediction_schema']).read_text()) validate_warm_start_task_type( TaskType(workflow.query.task_type), - TaskType(source_schema['task_type']), + bundle_task_type(warm_start_artifacts), ) _validate_model_plan(model_plan, workflow, train_table) materialized = self._materialize_training(workflow, train_table, run_dir, effective_plan) + tracking.start(lambda: fit_spec(FitContext( + query=train_table.query, + model_plan=to_json_dict(materialized.model_plan), + started=started_at, + snapshot_id=snapshot.id, + task_type=materialized.task.task_type.value, + model_name=model_name, + ))) training_result, fit_mode = self._train(materialized, run_dir, run_id, - warm_start_artifacts) + warm_start_artifacts, events) _collect_garbage('fit') return self._finish_fit( training_result, @@ -646,13 +932,16 @@ def _fit( graph_config=snapshot.graph_config, training_table_generation=train_table.generation_plan, export_config=export_config, + scoring_defaults=scoring_defaults, model_plan=effective_plan, provenance=FitRecord( train_table=train_table, warm_start=None if warm_start is None else bundle_manifest( warm_start._training_result.bundle_dir, warm_start_bundle_id), - started_at=started_at, fit_mode=fit_mode), + started_at=started_at, fit_mode=fit_mode, + tracking=tracking.pointer()), + deps=deps, ) def _materialize_training(self, workflow: TrainingWorkflowRequest, @@ -668,18 +957,20 @@ def _materialize_training(self, workflow: TrainingWorkflowRequest, train_table.generation_plan.lag_timesteps) year_over_year = inferred_or_missing_type_to_none( train_table.generation_plan.year_over_year) - task_dataset = build_generated_training_task_dataset( - query=workflow.query, - graph=workflow.graph, - graph_index_path=workflow.prepared_graph.index_path, - run_id=workflow.run_id, - artifact_dir=str(run_dir), - tt_id=ingest_tt_id, - required_stages=train_table.required_stages, - lag_timesteps=(None - if lag_timesteps is None else int(lag_timesteps)), - year_over_year=bool(year_over_year), - ) + with no_encoder_inference_warnings( + effective_plan.column_processing.encoder_overrides or {}): + task_dataset = build_generated_training_task_dataset( + query=workflow.query, + graph=workflow.graph, + graph_index_path=workflow.prepared_graph.index_path, + run_id=workflow.run_id, + artifact_dir=str(run_dir), + tt_id=ingest_tt_id, + required_stages=train_table.required_stages, + lag_timesteps=(None if lag_timesteps is None else + int(lag_timesteps)), + year_over_year=bool(year_over_year), + ) require_embedding_models( self._embedding_models, task_dataset.feature_schema, effective_plan.column_processing.encoder_overrides) @@ -702,7 +993,8 @@ def _materialize_training(self, workflow: TrainingWorkflowRequest, def _train( self, materialized: Any, run_dir: Path, run_id: str, - warm_start_artifacts: Optional[Mapping[str, str]] + warm_start_artifacts: Optional[Mapping[str, str]], + events: EventEmitter ) -> Tuple[TrainingResult, Optional[str]]: from kumo_core.training.orchestration import ( CallableTrainingOptions, @@ -724,7 +1016,7 @@ def _train( bindings=local_training_bindings( training_input, run_name=run_id, embedding_models=self._embedding_models, - accelerator=self._accelerator), + accelerator=self._accelerator, instrumentation=events), options=CallableTrainingOptions(), ) if not result.explain_disabled: @@ -754,6 +1046,8 @@ def _finish_fit( graph_config: Optional[Dict[str, Any]] = None, model_plan: Optional[ModelPlan] = None, provenance: Optional[FitRecord] = None, + deps: Optional[Dependencies] = None, + scoring_defaults: Optional[ScoringDefaults] = None, ) -> TrainingResult: """Commit durable fit outputs before making this trainer usable.""" from kumo_core.artifacts import bundle_id_from_manifest @@ -768,6 +1062,11 @@ def _finish_fit( graph_config=graph_config, timestamp_precision=timestamp_precision, ) + if scoring_defaults is not None: + write_scoring_defaults(training_result.bundle_dir, + scoring_defaults) + training_result = replace(training_result, + scoring_defaults=scoring_defaults) _bundle_holdout(training_result) try: training_result = replace( @@ -778,8 +1077,9 @@ def _finish_fit( UserWarning, stacklevel=3) if provenance is not None and training_result.bundle_id is not None: try: - provenance.write(training_result.bundle_dir, - training_result.bundle_id) + replace(provenance, dependencies=self._dependencies( + training_result, deps or Dependencies())).write( + training_result.bundle_dir, training_result.bundle_id) except Exception as error: warnings.warn( f"Publishing the model bundle without _provenance.json: " @@ -787,8 +1087,8 @@ def _finish_fit( if export_config is not None: training_result = replace( training_result, published_model=_publish( - training_result.bundle_dir, training_result.metrics, - export_config)) + training_result.bundle_dir, training_result.bundle_id, + registry_metrics(training_result), export_config)[0]) training_result = replace(training_result, _trainer=self) self._training_result = training_result @@ -800,7 +1100,19 @@ def _finish_fit( return training_result - @_without_torch_scatter_hint + def _dependencies(self, result: TrainingResult, + deps: Dependencies) -> List[Dependency]: + from kumo.engine._version import __version__ + + deps.add(dependency.predictor(TRAINER_PREDICTOR)) + text = require_text_encoders(deps, self._embedding_models, + _encoders(result.bundle_dir)) + deps.add( + dependency.package('kumo-engine', __version__, + ['text'] if text else [])) + return deps.records + + @_without_unactionable_warnings def predict( self, source: Union[GraphSnapshot, PredictionTable, None] = None, @@ -833,9 +1145,12 @@ def predict( Link predictions require an explicit count of at most 500. Embeddings alone do not require a count. - 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. + Binary classification always returns ``True_PROB``; ``TARGET_PRED`` + is set only with a threshold, else null. 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, @@ -858,10 +1173,7 @@ def predict( raise ValueError('A prediction table already holds its entities ' 'and times; pass it or entities=/' 'prediction_time=, not both.') - if self._training_result is None: - raise RuntimeError( - "Trainer.predict() requires a prior Trainer.fit() call on " - "this same Trainer instance.") + self._fitted('predict') check_prediction_precision( self._timestamp_precision, index_timestamp_precision(snapshot.index_path)) @@ -909,7 +1221,12 @@ def _predict( num_classes_to_return: Optional[int] = None, num_workers: int = 1, binary_classification_threshold: Optional[float] = None, + extra_fields: Sequence[str] = (), ) -> PredictionResult: + training_result = self._fitted('predict') + if binary_classification_threshold is None: + binary_classification_threshold = ( + training_result.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, @@ -927,8 +1244,10 @@ def _predict( 'Prediction query target/entity contract differs from the trained model' ) task_type = TaskType(contract[1]) - threshold = resolve_classification_threshold( - task_type, binary_classification_threshold, binary_required=True) + threshold = (None if binary_classification_threshold is None + and task_type is TaskType.BINARY_CLASSIFICATION else + resolve_classification_threshold( + task_type, binary_classification_threshold)) destinations = (output_config.destinations() if output_config is not None else {}) artifacts = frozenset( @@ -943,6 +1262,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 = training_result.scoring_defaults.num_classes_to_return num_classes_to_return = resolve_num_classes_to_return( task_type, num_classes_to_return, @@ -957,7 +1278,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) @@ -1021,12 +1342,15 @@ 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 []) + extra_fields = [ + *(f.value for f in (output_config.extra_fields + if output_config is not None else [])), + *extra_fields + ] source_identity = [ snapshot.id, prediction_table.query, timestamps, sorted(kind.value for kind in artifacts), num_classes_to_return, @@ -1103,12 +1427,20 @@ def load_from_connector( from trusted publishers; checkpoint deserialization happens at predict. ``accelerator`` selects the device for this instance (default: ``auto``). A `kumo.MLOps` ``model_registry`` reads ``path`` under its artifact - root (``models/``). + root (``models/``) and ``location`` in its tracker's registry. """ from kumo.engine._mlops import MLOps, on_root - if isinstance(model_registry, MLOps) and path is None: - raise TypeError('Loading from a kumo.MLOps takes path=.') + if isinstance(model_registry, MLOps) and location is not None: + tracker = as_tracking(model_registry.tracking) + load = None if tracker is None else tracker.load_model + if path is not None or load is None: + raise TypeError( + 'Loading a registered model from a kumo.MLOps takes ' + 'location= and its tracking=MlflowTracking(...) or ' + 'SnowflakeTracking(connector, ...).') + return load(location, embedding_models=embedding_models, + accelerator=accelerator) if path is not None: model_registry, path = on_root(model_registry, path) address = ModelOutputConfig(model_registry, path=path, @@ -1118,10 +1450,22 @@ 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, - accelerator=accelerator) + 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) + with dependency.assets_through( + model_registry, bundle_dir, lambda reader, _: ( + f'pass a {reader}(...) with an identity that reads it ' + f'as Trainer.load_from_connector(model_registry, ...)')): + trainer = cls._load_bundle(bundle_dir, embedding_models, + accelerator=accelerator) + warn_if_other_bundle(published.metadata, trainer) + return trainer @classmethod def _load_bundle( @@ -1146,9 +1490,39 @@ def _load_bundle( trainer._training_query = spec.predictive_query trainer._training_table_generation_plan = from_json_dict( spec.training_table_generation, TrainingTableGenerationPlan) + from kumo.engine._mlops import resolve + + trainer._embedding_models = with_text_encoders( + trainer._embedding_models, _encoders(bundle_dir), + Dependencies.of(read_provenance(bundle_dir), resolve())) holdout_path, holdout_truncated = _listed_holdout(bundle_dir) trainer._training_result = TrainingResult( metrics, paths, bundle_dir, holdout_path=holdout_path, holdout_truncated=holdout_truncated, bundle_id=bundle_id, + scoring_defaults=read_scoring_defaults(bundle_dir, paths), _trainer=trainer) return trainer + + +def _encoders(bundle_dir: Path) -> Any: + from kumo_core.encoders.config import EncodersConfig + from kumo_core.io.json import load_json + + path = listed_file(bundle_dir, 'encoders') + return (load_json(str(path), EncodersConfig) + if path is not None and path.is_file() else None) + + +def warn_if_other_bundle(metadata: Optional[Mapping[str, Any]], + trainer: Trainer) -> None: + recorded = (metadata or {}).get(BUNDLE_ID_KEY) + current = None if trainer.result is None else trainer.result.id + if recorded and current and recorded != current: + warnings.warn( + f"The registry records bundle {recorded} for this version, but " + f"the fetched bundle is {current}.", UserWarning, + stacklevel=3) + + +def registry_pointers(bundle_id: Optional[str]) -> Dict[str, str]: + return {BUNDLE_ID_KEY: bundle_id} if bundle_id else {} diff --git a/packages/kumo-engine/pyproject.toml b/packages/kumo-engine/pyproject.toml index a63d7b6f9..2387d76cb 100644 --- a/packages/kumo-engine/pyproject.toml +++ b/packages/kumo-engine/pyproject.toml @@ -74,6 +74,18 @@ adls=[ gcs=[ "kumo-gcs>=0.1.0.dev1", ] +# mlflow-skinny 3.5 added download_artifacts(registry_uri=), which loading +# by alias uses; nvidia-ml-py adds GPU system metrics. +mlflow=[ + "mlflow-skinny>=3.5,<4", + "nvidia-ml-py", +] +# SageMaker managed MLflow: 0.5.0 added the per-thread boto3 session +# (sagemaker_mlflow.session.use_session) that explicit AWS identities use. +sagemaker=[ + "kumo-engine[mlflow,s3]", + "sagemaker-mlflow>=0.5.0", +] # 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/test/_snowflake_experiment_fake.py b/packages/kumo-engine/test/_snowflake_experiment_fake.py new file mode 100644 index 000000000..ac25b360c --- /dev/null +++ b/packages/kumo-engine/test/_snowflake_experiment_fake.py @@ -0,0 +1,195 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import os +import re +import sys +import types +from typing import Any, Dict, List, Optional, Sequence, 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) + if not re.fullmatch(r'[A-Z_][A-Z0-9_$]*', self._run.name): + raise RuntimeError(f'Run "{self._run.name}" does not exist in ' + f'Experiment {self.experiment}.') + with open(local_path, encoding='utf-8') as file: + self._active()['artifacts'][os.path.join( + artifact_path or '', os.path.basename(local_path))] = ( + file.read()) + + +class FakeSession: + def __init__(self, database: Optional[str] = 'DB', + schema: Optional[str] = 'PUBLIC', + schemas: Sequence[Tuple[str]] = ()) -> None: + self.database, self.schema = database, schema + self.schemas = list(schemas) + self.statements: List[str] = [] + + def sql(self, query: str) -> Any: + self.statements.append(query) + return types.SimpleNamespace(collect=lambda: self.schemas) + + 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: + 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 4b60addea..28e0791fb 100644 --- a/packages/kumo-engine/test/conftest.py +++ b/packages/kumo-engine/test/conftest.py @@ -164,3 +164,53 @@ def fitted(tmp_path_factory): yield trainer, pquery.generate_prediction_table(snapshot), snapshot finally: kumo.configure(scratch_dir=previous) + + +@pytest.fixture +def tracking_uri(tmp_path, monkeypatch): + mlflow = pytest.importorskip('mlflow') + monkeypatch.chdir(tmp_path) + 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): + 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='PREDICT ents.Y FOR EACH ents.ID') + yield query.generate_training_table(snapshot), snapshot + finally: + kumo.configure(scratch_dir=previous_scratch) + + +@pytest.fixture +def quick_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 diff --git a/packages/kumo-engine/test/test_adls_connector.py b/packages/kumo-engine/test/test_adls_connector.py index b356a75f5..cf8fc8af5 100644 --- a/packages/kumo-engine/test/test_adls_connector.py +++ b/packages/kumo-engine/test/test_adls_connector.py @@ -20,6 +20,7 @@ from kumo.engine import ADLSConnector, ADLSIdentity, Table from kumo.engine._portable import _source_table_name, portable_graph_config +from kumo.engine.trainer._scoring import ScoringDefaults ROOT = 'abfss://curated@companylake.dfs.core.windows.net/orders/' @@ -159,8 +160,9 @@ def test_invalid_prediction_output_fails_before_scoring(monkeypatch, tmp_path, table_name=destination, write_mode=WriteMode.APPEND if case == 'append' else WriteMode.OVERWRITE) trainer = Trainer(ModelPlan(), accelerator='cpu') - monkeypatch.setattr(trainer, '_training_result', - SimpleNamespace(artifact_paths={})) + monkeypatch.setattr( + trainer, '_training_result', + SimpleNamespace(artifact_paths={}, scoring_defaults=ScoringDefaults())) trainer._training_query = 'trained query' monkeypatch.setattr(Trainer, 'scratch_dir', property(lambda _: tmp_path)) # Isolate output validation from query parsing and model execution. diff --git a/packages/kumo-engine/test/test_artifacts.py b/packages/kumo-engine/test/test_artifacts.py new file mode 100644 index 000000000..01566f2e8 --- /dev/null +++ b/packages/kumo-engine/test/test_artifacts.py @@ -0,0 +1,159 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from dataclasses import dataclass +from typing import List, Optional + +from kumo_core.artifacts import bundle_id_from_manifest +from kumo_core.summaries import Dataset, Table, Tag + +import kumo.engine as kumo +from kumo.engine._deps import Dependencies, FileAsset +from kumo.engine._manifest import ( + FIT_ROLES, + Artifact, + Dependency, + Manifest, + ModelManifest, + ModelParents, + Record, + SnapshotManifest, + SnapshotParents, + SourceTableEntry, + TrainingTableManifest, + TrainingTableParents, + artifact_summaries, + artifacts, + read_manifest, +) +from kumo.engine._snapshot_io import store +from kumo.engine.trainer.trainer import _store_and_register + + +@dataclass(frozen=True, kw_only=True) +class ReportParents(Record): + snapshot: SnapshotManifest + model: ModelManifest + + +@dataclass(frozen=True, kw_only=True) +class ReportManifest(Manifest): + kind: str = 'report' + parents: ReportParents + pages: Optional[List[SnapshotManifest]] = None + dependencies: Optional[List[Dependency]] = None + + IS_DATASET = True + + +def test_a_new_kind_names_its_artifacts_without_other_changes(): + orders = SourceTableEntry(connector='S3', location='s3://data/orders', + version='7', version_kind='exact') + snapshot = SnapshotManifest( + id='g@d', location='s3://root/snapshots/g/d', + parents=SnapshotParents(source_tables={'orders': orders})) + model = ModelManifest( + id='m', location='s3://root/models/m', parents=ModelParents( + training_table=TrainingTableManifest(id='t'))) + report = ReportManifest( + id='r', location='s3://root/reports/r', + parents=ReportParents(snapshot=snapshot, model=model), + pages=[snapshot, SnapshotManifest(id='p', reconstructed=True)]) + + assert artifacts(report) == [ + Artifact('report', 'r', 's3://root/reports/r', True, + version='r', version_kind='exact', label='report r'), + Artifact('snapshot', 'g@d', 's3://root/snapshots/g/d', True, + version='g@d', version_kind='exact', + label='graph snapshot g@d'), + Artifact('source_table', 'orders', 's3://data/orders', True, + connector='S3', version='7', version_kind='exact'), + Artifact('model', 'm', 's3://root/models/m', False, + 'kumo.bundle_uri', version='m', version_kind='exact', + label='model m'), + Artifact('snapshot', 'p', None, True, version='p', + version_kind='exact', label='graph snapshot p'), + ] + + +def test_each_store_records_its_location_and_keeps_the_ids(fitted, tmp_path): + trainer, _, snapshot = fitted + bundle = trainer.result.bundle_dir + ids = snapshot.id, trainer.result.id, bundle_id_from_manifest(bundle) + for root in (tmp_path / 'a', tmp_path / 'b'): + settings = kumo.MLOps(str(root)) + store(snapshot, settings) + _store_and_register(bundle, trainer.result.manifest, settings) + assert snapshot.manifest.location == settings.snapshot_uri(snapshot) + assert trainer.result.manifest.location == f'{root}/models/{ids[1]}' + assert read_manifest(root / 'models' / ids[1] / '_provenance.json', + ModelManifest).location == f'{root}/models/{ids[1]}' + assert (snapshot.id, trainer.result.id, + bundle_id_from_manifest(bundle)) == ids + + +def test_a_new_kind_and_asset_reach_tags_and_datasets(tmp_path): + settings = kumo.MLOps(str(tmp_path / 'root')) + (tmp_path / 'notes.txt').write_text('notes') + asset = FileAsset('dummy', 'notes', tmp_path / 'notes.txt') + dependency = Dependencies(settings=settings).require(asset) + root = f"{tmp_path / 'root'}" + orders = SourceTableEntry(connector='S3', location='s3://data/orders', + version='7', version_kind='exact') + report = ReportManifest( + id='r', location=f'{root}/reports/r', + parents=ReportParents( + snapshot=SnapshotManifest( + id='g@d', location=f'{root}/snapshots/g/d', + parents=SnapshotParents(source_tables={'orders': orders})), + model=ModelManifest(id='m', location=f'{root}/models/m')), + dependencies=[dependency]) + + assert dependency.location == f'{root}/assets/dummy/{dependency.sha256}' + assert ReportManifest.from_dict(report.to_dict()) == report + assert artifact_summaries(report, FIT_ROLES, { + 'orders': ({'ID': {'stype': 'ID', 'dtype': 'int64'}}, {'rows': 3}) + }) == [ + Tag('kumo.bundle_uri', f'{root}/models/m'), + Table('locations', ('kind', 'id', 'location'), ( + ('report', 'r', f'{root}/reports/r'), + ('snapshot', 'g@d', f'{root}/snapshots/g/d'), + ('source_table', 'orders', 's3://data/orders'), + ('model', 'm', f'{root}/models/m'), + ('asset', dependency.sha256, dependency.location)), brief=True), + Dataset('graph snapshot g@d', { + 'connector': 'LOCAL_FILE', 'location': f'{root}/snapshots/g/d', + 'version': 'g@d', 'version_kind': 'exact'}, {}, {}, + 'graph context'), + Dataset('orders', {'connector': 'S3', 'location': 's3://data/orders', + 'version': '7', 'version_kind': 'exact'}, + {'ID': {'stype': 'ID', 'dtype': 'int64'}}, {'rows': 3}, + 'source table'), + Dataset('report r', {'connector': 'LOCAL_FILE', + 'location': f'{root}/reports/r', 'version': 'r', + 'version_kind': 'exact'}, {}, {}, 'report'), + ] + + +def test_a_fits_datasets_are_named_and_labelled_by_role(): + users = SourceTableEntry(connector='S3', location='s3://data/users', + version='3', version_kind='exact') + snapshot = SnapshotManifest( + id='g@5ac245', created_at='2026-11-01T09:00:00+00:00', + location='s3://root/snapshots/g/5ac245', + parents=SnapshotParents(source_tables={'users': users})) + model = ModelManifest(id='m', parents=ModelParents( + training_table=TrainingTableManifest( + id='tt-fadf04', created_at='2026-11-01T10:00:00+00:00', + location='s3://root/training_tables/tt-fadf04', + parents=TrainingTableParents(snapshot=snapshot)))) + + assert [(s.name, s.context, s.source['version']) + for s in artifact_summaries(model, FIT_ROLES) + if isinstance(s, Dataset)] == [ + ('training table 2026-11-01-fadf04', 'training examples', + 'tt-fadf04'), + ('graph snapshot 2026-11-01-5ac245', 'graph context', 'g@5ac245'), + ('users', 'source table', '3'), + ] diff --git a/packages/kumo-engine/test/test_bundle_provenance.py b/packages/kumo-engine/test/test_bundle_provenance.py index 09f3f5567..2723b7d9c 100644 --- a/packages/kumo-engine/test/test_bundle_provenance.py +++ b/packages/kumo-engine/test/test_bundle_provenance.py @@ -22,7 +22,7 @@ ) from kumo.engine.trainer._provenance import FitRecord, embedded_model from kumo.engine.trainer._spec import MANIFEST_FILE -from kumo.engine.trainer.trainer import TrainingResult, _store_model +from kumo.engine.trainer.trainer import TrainingResult, _store_and_register def test_fit_records_a_self_contained_manifest(fitted): @@ -114,8 +114,7 @@ def _finish_stub_fit(bundle, trainer=None, warm_start=None, settings=None): training_table_generation=TrainingTableGenerationPlan(), provenance=FitRecord(table, warm_start, datetime.now(timezone.utc), None)) - if settings is not None and settings.wants('model'): - _store_model(result.bundle_dir, result.manifest, settings) + _store_and_register(result.bundle_dir, result.manifest, settings) return result diff --git a/packages/kumo-engine/test/test_data_summaries.py b/packages/kumo-engine/test/test_data_summaries.py new file mode 100644 index 000000000..0aec16d69 --- /dev/null +++ b/packages/kumo-engine/test/test_data_summaries.py @@ -0,0 +1,179 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +import json +from pathlib import Path + +import numpy as np +import pandas as pd +import pytest +from kumo_core.summaries import Dataset, Note, Table, Tag + +import kumo.engine as kumo +from kumo.engine._data_summaries import ( + snapshot_summaries, + split_summaries, + table_details, +) +from kumo.engine._manifest import FIT_ROLES, artifact_summaries +from kumo.engine._scratch import get_scratch_dir +from kumo.engine.pquery import PredictiveQuery + + +@pytest.fixture +def scratch(tmp_path): + previous = get_scratch_dir() + kumo.configure(scratch_dir=tmp_path / 'scratch') + yield tmp_path + kumo.configure(scratch_dir=previous) + + +def test_summaries_match_the_snapshot_and_training_table(scratch): + rng = np.random.default_rng(0) + data = scratch / 'data' + data.mkdir() + users = pd.DataFrame({'USER_ID': np.arange(300), + 'AGE': rng.normal(40, 10, 300)}) + users.loc[0, 'AGE'] = None + orders = pd.DataFrame({ + 'ORDER_ID': np.arange(6000), + 'USER_ID': rng.integers(0, 360, 6000), + 'ORDER_TIME': pd.Timestamp('2024-01-01') + pd.to_timedelta( + rng.integers(0, 120, 6000), unit='D'), + }) + users.to_parquet(data / 'users.parquet') + orders.to_parquet(data / 'orders.parquet') + connector = kumo.FileConnector(root_dir=str(data)) + graph = kumo.Graph( + tables={ + 'users': kumo.Table.from_source_table( + connector['users'], primary_key='USER_ID'), + 'orders': kumo.Table.from_source_table( + connector['orders'], primary_key='ORDER_ID', + time_column='ORDER_TIME').with_column('USER_ID', stype='ID'), + }, + edges=[kumo.Edge(src_table='orders', fkey='USER_ID', + dst_table='users')]) + snapshot = graph.snapshot() + table = PredictiveQuery( + graph, 'PREDICT COUNT(orders.*, 0, 7, days) > 0 FOR EACH users.USER_ID' + ).generate_training_table(snapshot) + + splits, = split_summaries(table.staging_dir) + assert splits.brief and splits.columns == ( + 'split', 'rows', 'distinct users (approx.)', 'positive rate') + for split, row in zip(('train', 'val', 'test'), splits.rows): + seeds = pd.read_parquet(table.staging_dir / f'{split}_seeds') + assert row[:2] == (split, len(seeds)) + assert row[3] == pytest.approx(seeds['TARGET1'].mean()) + + summaries = {s.name: s for s in snapshot_summaries( + Path(snapshot.index_path), snapshot.graph_config)} + assert set(summaries) == {'data/health', 'data/columns'} + matched = orders['USER_ID'] < 300 + health = summaries['data/health'] + assert [dict(zip(health.columns, row)) for row in health.rows] == [ + {'table or link': 'users', 'rows': 300, 'columns': 2, + 'dropped rows': 0, 'duplicate keys': 0, 'rows without time': '', + 'edges': '', 'source matched': '', + 'target matched': '', 'failed lookups': ''}, + {'table or link': 'orders', 'rows': 6000, 'columns': 3, + 'dropped rows': 0, 'duplicate keys': 0, 'rows without time': 0, + 'edges': '', 'source matched': '', + 'target matched': '', 'failed lookups': ''}, + {'table or link': 'orders.USER_ID → users', 'rows': '', + 'columns': '', 'dropped rows': '', 'duplicate keys': '', + 'rows without time': '', 'edges': matched.sum(), + 'source matched': pytest.approx(matched.mean()), + 'target matched': pytest.approx( + orders.loc[matched, 'USER_ID'].nunique() / 300), + 'failed lookups': (~matched).sum()}] + stats = summaries['data/columns'] + assert stats.columns[3:] == ( + 'count', 'null rate', 'distinct≈', 'min', 'max', 'mean', 'std', + 'q25≈', 'q50≈', 'q75≈', 'avg tokens≈', 'max tokens≈', 'top values≈') + age = next(dict(zip(stats.columns, row)) for row in stats.rows + if row[:2] == ('users', 'AGE')) + assert (age['stype'], age['count'], age['null rate'], age['min'], + age['max']) == ('numerical', 300, 1 / 300, users['AGE'].min(), + users['AGE'].max()) + assert age['mean'] == pytest.approx(users['AGE'].mean()) + datasets = {s.name: s for s in artifact_summaries( + snapshot.manifest, FIT_ROLES, + details=table_details(Path(snapshot.index_path))) + if isinstance(s, Dataset)} + assert set(datasets) == {'users', 'orders'} + assert datasets['orders'].schema == { + 'ORDER_ID': {'stype': 'id', 'dtype': 'int64'}, + 'USER_ID': {'stype': 'id', 'dtype': 'int64'}, + 'ORDER_TIME': {'stype': 'timestamp', 'dtype': 'date'}} + assert datasets['orders'].profile['rows'] == 6000 + assert datasets['users'].source['location'] == str( + (data / 'users.parquet').resolve()) + assert datasets['users'].context == 'source table' + + +def test_data_warnings_name_each_table_that_breaks_a_check(tmp_path): + def write(name, value): + (tmp_path / name).write_text(json.dumps(value)) + + write('ingest_stats.json', {'stats': { + 'node_counts': {'users': 9, 'orders': 100}, + 'source_rows_by_table': {'users': 12, 'orders': 100}, + 'edge_counts': {'orders_USER_ID_to_users_USER_ID': 10}, + 'pk_duplicate_counts': {'users:shard-0': 2}, + 'pk_null_timestamp_drops': {'orders:shard-0': 3}, + 'dropped_rows_counts': {'users': 1, 'orders': 3}, + 'key_failures_by_context': {'pkey:users.USER_ID': { + 'parse_failures': 1}}, + }}) + write('col_stats.json', {'orders': { + 'ORDER_TIME': {'semantic_type': 'timestamp', 'count': 97, + 'null_count': 2, 'parse_fail_count': 0}, + 'AMOUNT': {'semantic_type': 'numerical', 'count': 100, + 'null_count': 0, 'parse_fail_count': 4, + 'invalid_count': 6}, + }}) + write('edge_stats.json', { + f'orders__c__users__c__USER_ID__USER_ID__{stat}': {'value': value} + for stat, value in (('FRAC_SRC_NONZERO_DEGREE', 0.1), + ('FRAC_DST_NONZERO_DEGREE', 0.5))}) + config = {'tables': [ + {'table_name': 'users', 'pkey': {'name': 'USER_ID'}, 'fkeys': []}, + {'table_name': 'orders', 'pkey': {'name': 'ORDER_ID'}, + 'time_col': {'name': 'ORDER_TIME'}, + 'fkeys': [{'name': 'USER_ID', 'dst_table': 'users', + 'dst_key': 'USER_ID'}]}, + ]} + + summaries = snapshot_summaries(tmp_path, config) + assert [s for s in summaries if isinstance(s, (Note, Tag))] == [ + Note('problems', '`users`: 1 missing or invalid primary keys'), + Note('problems', '`users`: 2 duplicate primary keys'), + Note('problems', '`users`: only 9 rows'), + Note('problems', + '`orders`: 5.0% of rows have a missing or invalid time'), + Note('problems', '`orders.AMOUNT`: at least 10% invalid values'), + Note('problems', + 'link `orders.USER_ID → users`: 10.0% of `orders` rows match'), + Tag('kumo.data_warnings', '6'), + ] + + +def test_approximate_counts_never_exceed_the_exact_ones(tmp_path): + (tmp_path / 'split_stats.json').write_text(json.dumps({ + 'val_table_stats': {'num_train_rows': 947, 'num_entities': 944}, + 'test_table_stats': {'num_train_rows': 976, 'num_entities': 985}})) + (tmp_path / 'col_stats.json').write_text(json.dumps({'users': { + 'BIO': {'semantic_type': 'text', 'count': 2000, 'num_unique': 2007}, + 'AGE': {'semantic_type': 'numerical', 'count': 2000, + 'num_unique': 61}}})) + + assert split_summaries(tmp_path) == [Table( + 'data/splits', ('split', 'rows', 'distinct entities (approx.)'), + (('val', 947, 944), ('test', 976, 976)), brief=True)] + columns = {s.name: s for s in snapshot_summaries(tmp_path, {})}[ + 'data/columns'] + distinct = columns.columns.index('distinct≈') + assert [(row[1], row[distinct]) for row in columns.rows] == [ + ('BIO', 2000), ('AGE', 61)] diff --git a/packages/kumo-engine/test/test_embedding_models.py b/packages/kumo-engine/test/test_embedding_models.py index ee2eb68da..681fef6c9 100644 --- a/packages/kumo-engine/test/test_embedding_models.py +++ b/packages/kumo-engine/test/test_embedding_models.py @@ -14,7 +14,7 @@ from kumo_core import embedding_models from kumo_core.embedding_models import EmbeddingModelNotFoundError, Release from kumo_core.glove import TokenEmbeddingTable -from kumoapi.encoder import TextEmbedding +from kumoapi.encoder import Null from kumoapi.typing import Stype from kumo.engine import ( @@ -24,6 +24,7 @@ Graph, PredictiveQuery, Table, + TextEmbedding, Trainer, ) @@ -70,7 +71,8 @@ def fitted(tmp_path_factory): warnings['failed fit'] = warned.take() plan.column_processing.encoder_overrides = { - 'users.bio': TextEmbedding('words') + 'users.bio': TextEmbedding('words'), + 'users.code': Null(), } store = EmbeddingModelStore() store.add('words', GloVeVectors(archive, cache_dir=root / 'emb')) @@ -148,11 +150,14 @@ def test_a_missing_model_fails_fit_before_any_work(fitted, monkeypatch): def test_each_fit_reports_what_inference_leaves_unencoded(fitted): unencoded = ("Text column 'users.code' will not be encoded by default: its " "values are too short or have too few known words. To " - "encode it anyway, set its encoder to TextEmbedding().") + "encode it anyway, set its encoder to TextEmbedding(). No " + "word-vector model is installed for it either: pip install " + "\"kumo-engine[text]\", or pass embedding_models= to " + "Trainer.") assert fitted[-1] == { 'tables': [], 'failed fit': [unencoded], - 'fit': [unencoded], + 'fit': [], } diff --git a/packages/kumo-engine/test/test_filestore_sink.py b/packages/kumo-engine/test/test_filestore_sink.py index ca8a26311..3cc95f40a 100644 --- a/packages/kumo-engine/test/test_filestore_sink.py +++ b/packages/kumo-engine/test/test_filestore_sink.py @@ -132,24 +132,19 @@ def test_a_trainer_publishes_and_loads_through_a_filesystem_only_sink( trainer, *_ = fitted connector = SimpleNamespace(source_type='SNOWFLAKE', sink_backend=StageSink()) - result = trainer._finish_fit( - trainer._training_result, predictive_query=trainer._training_query, - snapshot_id=trainer._snapshot_id, - training_table_generation=trainer.training_table_generation_plan, - export_config=kumo.ModelOutputConfig(connector, - path='@stage/models/m')) - assert result.published_model.uri == '@stage/models/m' + published = trainer.result.export( + kumo.ModelOutputConfig(connector, path='@stage/models/m')) + assert published.uri == '@stage/models/m' loaded = kumo.Trainer.load_from_connector(connector, path='@stage/models/m') - assert loaded._snapshot_id == trainer._snapshot_id + assert loaded.result.id == trainer.result.id def test_predictions_export_through_a_filesystem_only_sink(fitted, connector): trainer, table, snapshot = fitted output = PredictionOutputConfig( connector, output_paths={'predictions': '@stage/preds/'}) - result = trainer.predict(table, snapshot=snapshot, output_config=output, - binary_classification_threshold=0.5) + result = trainer.predict(table, snapshot=snapshot, output_config=output) store = connector.sink_backend.store assert any(name.startswith('@stage/preds/') for name in store.objects) result.export( 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..ed70a3a36 --- /dev/null +++ b/packages/kumo-engine/test/test_fit_events.py @@ -0,0 +1,215 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import json +import logging +import threading + +import pytest +from kumoapi.model_plan import ModelPlan + +import kumo.engine as kumo +from kumo.engine._scratch import get_scratch_dir +from kumo.engine.tracking import Event + + +def test_fit_reports_progress_events(table_and_snapshot, quick_plan, + 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 = quick_plan + plan.training_job.num_experiments = 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 + assert all(isinstance(event, Event) for event in events) + 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 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[2].metrics['train_epoch_seconds'] > 0 + 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' + } + + +@pytest.mark.parametrize('on_error', ['warn', 'raise']) +def test_an_unfinished_tracker_is_reported_by_the_error_policy( + table_and_snapshot, quick_plan, tmp_path, monkeypatch, on_error): + import kumo.engine._tracking as tracking + + monkeypatch.setattr(tracking, 'FINISH_SECONDS', 0.5) + release = threading.Event() + + class Blocked: + def start(self, record): + pass + + def on_trial(self, record, event): + pass + + def finish(self, record, error): + release.wait(30) + + table, snapshot = table_and_snapshot + trainer = kumo.Trainer(quick_plan) + cfg = kumo.MLOps(str(tmp_path / 'root'), tracking=Blocked(), store=False, + on_error=on_error) + try: + if on_error == 'warn': + with pytest.warns(UserWarning, match='TimeoutError: experiment ' + 'tracking did not finish within 0.5s'): + assert trainer.fit(table, snapshot=snapshot, + mlops=cfg) is trainer.result + else: + with pytest.raises(RuntimeError) as raised: + trainer.fit(table, snapshot=snapshot, mlops=cfg) + assert isinstance(raised.value.__cause__, TimeoutError) + finally: + release.set() + assert trainer.result.bundle_dir is not None + + +def test_event_metrics_keep_finite_numbers(): + from kumo.engine._events import EventEmitter + + received = [] + EventEmitter([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 == [] + + +def test_a_snapshot_reports_progress_events(graph, tmp_path): + previous = get_scratch_dir() + kumo.configure(scratch_dir=tmp_path / 'scratch') + try: + events, again = [], [] + snapshot = graph.snapshot(callbacks=[events.append]) + graph.snapshot(callbacks=[again.append]) + finally: + kumo.configure(scratch_dir=previous) + + assert [event.name for event in events] == [ + 'snapshot_started', 'table_staged', 'table_staged', 'snapshot_ended'] + assert [event.params['table'] for event in events[1:3]] == [ + 'users', 'events'] + assert events[-1].params['snapshot_id'] == snapshot.id + assert events[-1].params['reused'] is False + assert again[-1].params['reused'] is True + assert all(event.params['seconds'] >= 0 for event in events) + + +def test_kumo_calls_drop_lightning_tips_and_only_those(caplog): + from pytorch_lightning.utilities.rank_zero import rank_zero_info + + from kumo.engine.trainer.trainer import _without_unactionable_warnings + + @_without_unactionable_warnings + def fit(): + rank_zero_info('💡 Tip: inside') + rank_zero_info('GPU available: False') + + lightning = logging.getLogger('pytorch_lightning') + lightning.addHandler(caplog.handler) + try: + with caplog.at_level('INFO', logger='pytorch_lightning'): + fit() + rank_zero_info('💡 Tip: outside') + finally: + lightning.removeHandler(caplog.handler) + assert 'inside' not in caplog.text + assert 'GPU available' in caplog.text and 'outside' in caplog.text + + +def test_an_event_prints_as_its_name_and_fields(): + event = Event(name='table_staged', time=1.0, + params={'table': 'users', 'seconds': 12.0}) + assert str(event) == 'table_staged table=users seconds=12.0' + assert str(Event(name='epoch_ended', time=1.0, trial=1, epoch=2, + metrics={'val_auroc': 0.5})) == ( + 'epoch_ended trial=1 epoch=2 val_auroc=0.5') 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..e81734f95 --- /dev/null +++ b/packages/kumo-engine/test/test_mlflow_model.py @@ -0,0 +1,529 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import dataclasses +import json +import warnings +from types import SimpleNamespace + +import pandas as pd +import pytest +from kumo_core.artifacts import PROVENANCE_FILE +from kumoapi.plugin import PublishedModel +from test_mlflow_tracking import _fit_run + +import kumo.engine as kumo +from kumo.engine._scratch import get_scratch_dir +from kumo.engine._tracking import logged_names +from kumo.engine.trainer._run_spec import registry_metrics + +mlflow = pytest.importorskip('mlflow') +pytestmark = pytest.mark.usefixtures('tracking_uri') + + +def _sorted(frame): + return frame.sort_values('ENTITY').reset_index(drop=True) + + +class _Tracker: + def __init__(self): + self.calls = [] + + def start(self, record): + self.calls.append('start') + + def on_trial(self, record, event): + pass + + def finish(self, record, error): + self.calls.append(('finish', record, error)) + + +def _started(tracker): + from kumo.engine._tracking import TrackingSession + from kumo.engine.trainer._run_spec import FitContext, fit_spec + + session = TrackingSession(tracker) + session.start(lambda: fit_spec(FitContext('q', {}))) + return session + + +def test_fit_logs_registers_and_serves_the_model(table_and_snapshot, + quick_plan, tmp_path, caplog): + from kumo.engine.integrations.mlflow import MlflowTracking, load_model + + table, snapshot = table_and_snapshot + cfg = kumo.MLOps(str(tmp_path / 'root'), + tracking=MlflowTracking(log_system_metrics=False), + registry=kumo.ModelRegistryLocation('m', + alias='champion')) + trainer = kumo.Trainer(quick_plan) + result = trainer.fit(table, snapshot=snapshot, mlops=cfg) + + assert 'MLflow tracking failed' not in caplog.text + parent = _fit_run() + model, = mlflow.search_logged_models(output_format='list') + assert (model.name, model.source_run_id) == ('m', parent.info.run_id) + version = mlflow.MlflowClient().get_model_version_by_alias('m', 'champion') + metrics = {metric.key: metric.value for metric in model.metrics} + holdout = {name: value for name, value in logged_names( + result.metrics).items() if not name.startswith(('train/', 'val/'))} + assert metrics == holdout + assert version.source == model.model_uri + assert parent.data.tags.items() >= { + 'kumo.kind': 'fit', 'kumo.bundle_id': result.id, + 'kumo.model_name': 'm', + 'kumo.model_version': str(version.version)}.items() + assert result.registered_model == PublishedModel( + name='m', version=str(version.version), + uri=f'models:/m/{version.version}', + metadata={'kumo_bundle_id': result.id}, + metrics=registry_metrics(result)) + assert version.tags == {'kumo_bundle_id': result.id} + assert trainer.result is result + + info = mlflow.models.get_model_info('models:/m@champion') + assert info.metadata == {'kumo_bundle_id': result.id} + provenance = json.loads(mlflow.artifacts.load_text( + f'{model.artifact_location}/artifacts/{PROVENANCE_FILE}')) + assert provenance == result.manifest.to_dict() + 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.SnapshotOutputConfig( + 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').result.id == result.id + champion = kumo.Trainer.load_from_connector( + cfg, location=kumo.ModelRegistryLocation('m', alias='champion')) + assert champion.result.id == result.id + assert champion.export(cfg).version == str(version.version) + client = mlflow.MlflowClient() + assert len(client.search_model_versions("name = 'm'")) == 1 + other = kumo.MLOps(str(tmp_path / 'root'), tracking=cfg.tracking, + registry=kumo.ModelRegistryLocation('m2')) + first = champion.export(other) + assert first.version == '1' and champion.export(other) == first + assert len(mlflow.search_logged_models(output_format='list')) == 1 + assert load_model('models:/m2/1').result.id == result.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_given_rows(fitted, tmp_path): + from kumo.engine.integrations.mlflow import log_model + + trainer, table, snapshot = fitted + location = snapshot.export( + kumo.SnapshotOutputConfig( + kumo.FileConnector(root_dir=str(tmp_path / 'exports')), + 'snapshots')) + + info = log_model(trainer, snapshot=snapshot, + serving_snapshot=str(tmp_path / 'exports' / location)) + + run = mlflow.get_run(info.run_id) + assert run.info.status == 'FINISHED' + assert 'ents' in {item.dataset.name for item in run.inputs.dataset_inputs} + 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: metric.value + for metric in mlflow.get_logged_model(info.model_id).metrics + } == registry_metrics(trainer.result) + assert mlflow.get_logged_model(info.model_id).name == ( + 'ents-binary_classification') + 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('*')) + + assert pyfunc.predict(empty, params=params).empty + served = pyfunc.unwrap_python_model()._snapshots + assert list(served) == [params['snapshot']] + rows = pd.DataFrame({'ENTITY': range(50)}) + for snapshot_param in (None, params['snapshot']): + pd.testing.assert_frame_equal( + _sorted( + pyfunc.predict(rows, + params={**params, 'snapshot': snapshot_param})), + expected) + scratch = pyfunc.unwrap_python_model()._predictor.trainer.scratch_dir + assert list(scratch.iterdir()) == [] + assert set(tables.glob('*')) == before + + +def test_pyfunc_names_a_missing_prediction_time_and_needs_no_threshold( + fitted): + from kumo.engine._predictor import predict_frame + + trainer, table, snapshot = fitted + rows = pd.DataFrame({'ENTITY': range(5)}) + with pytest.raises(ValueError) as error: + predict_frame(trainer, snapshot, rows, + {'binary_classification_threshold': .5}) + assert str(error.value) == ("temporal query: pass params={" + "'prediction_time': ...} or a TIMESTAMP " + "column") + + time = str(pd.Timestamp(table.timestamps[0])) + scores = predict_frame(trainer, snapshot, rows, {'prediction_time': time}) + assert len(scores) == 5 + assert scores['TARGET_PRED'].isna().all() + assert scores['True_PROB'].between(0, 1).all() + + +def test_pyfunc_returns_rows_in_input_order(fitted): + from kumo.engine._predictor import predict_frame + + trainer, table, snapshot = fitted + time = pd.Timestamp(table.timestamps[0]) + entities = [7, 3, 9, 3, 999, 1, 40, 2, 33, 11, 3, 25] * 10 + rows = pd.DataFrame({ + 'ENTITY': entities, + 'TIMESTAMP': [time - pd.Timedelta(days=i % 3) + for i in range(len(entities))], + }) + scored = predict_frame(trainer, snapshot, rows, + {'binary_classification_threshold': .5}) + assert scored['ENTITY'].tolist() == [e for e in entities if e != 999] + assert 'ANCHOR_TIMESTAMP' not in scored + + +def test_scores_pair_with_their_rows_when_scoring_reorders(): + from kumo.engine._predictor import _in_input_order + + day = pd.Timestamp('2024-05-01 12:00:00.250') + rows = pd.DataFrame({ + 'ENTITY1': ['b', 'a', 'b', 'c', 'a', 'gone'], + 'TIMESTAMP': [day, day, day - pd.Timedelta(days=1), day, day, day], + }) + seconds = day.floor('s') + scored = pd.DataFrame({ + 'ENTITY': ['a', 'c', 'b', 'a', 'b'], + 'ANCHOR_TIMESTAMP': [seconds, seconds, seconds - pd.Timedelta(days=1), + seconds, seconds], + 'SCORE': [1.0, 3.0, 2.5, 1.0, 2.0], + }) + + ordered = _in_input_order(rows, scored) + + assert ordered.to_dict('list') == { + 'ENTITY': ['b', 'a', 'b', 'c', 'a'], + 'SCORE': [2.0, 1.0, 2.5, 3.0, 1.0], + } + + +@pytest.mark.parametrize('tracking,registry', [ + (None, kumo.ModelRegistryLocation('m')), + (SimpleNamespace(start=None), kumo.ModelRegistryLocation('m')), + ('mlflow', kumo.ModelRegistryLocation('m', alias='SageMakerLifecycle-p')), + ('mlflow', kumo.ModelRegistryLocation('m', version='3')), +]) +def test_invalid_mlflow_registry_is_refused(tmp_path, tracking, registry): + from kumo.engine.integrations.mlflow import MlflowTracking + + if tracking == 'mlflow': + tracking = MlflowTracking() + with pytest.raises(ValueError): + kumo.MLOps(str(tmp_path), tracking=tracking, registry=registry) + + +def test_fit_waits_for_model_logging_and_fails_the_run_if_it_fails( + monkeypatch): + import time + + from kumo.engine import _tracking + + monkeypatch.setattr(_tracking, 'FINISH_SECONDS', 0.01) + + class Tracker(_Tracker): + def __init__(self, fail): + super().__init__() + self.fail = fail + + def log_model(self, bundle_dir, manifest, snapshot, registry): + time.sleep(0.2) + self.calls.append('log_model') + if self.fail: + raise RuntimeError('upload failed') + + for fail in (False, True): + tracker = Tracker(fail) + session = _started(tracker) + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter('always') + _, error = session.log_model('missing', None, None) + session.finish(None, error) + session._worker.join(5) + started, logged, (finished, _, ended) = tracker.calls + assert (started, logged, finished, ended) == ('start', 'log_model', + '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 + + monkeypatch.setattr(_tracking, 'MODEL_SECONDS', 0.5) + release = threading.Event() + + class Tracker(_Tracker): + def log_model(self, bundle_dir, manifest, snapshot, registry): + release.wait(30) + + tracker = Tracker() + session = _started(tracker) + 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(None, None, None) + release.set() + session.finish(None, error) + + assert type(error) is (TimeoutError + if stall == 'slow upload' else RuntimeError) + assert tracker.calls == ['start', ('finish', None, None)] + + +def test_model_upload_gets_a_longer_http_timeout_and_a_late_one_is_dropped( + monkeypatch, caplog): + import os + + from test_tracking_contract import CONTEXT + + from kumo.engine._tracking import TrackingSession + from kumo.engine.integrations import mlflow as integration + from kumo.engine.trainer._run_spec import fit_spec + + for name in integration._HTTP_LIMITS: + monkeypatch.delenv(name, raising=False) + tracker = integration.MlflowTracking(log_system_metrics=False) + location = kumo.ModelRegistryLocation('m') + seen = [] + + def log_model(target, client, bundle_dir, manifest, 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: fit_spec(CONTEXT)) + registered, error = session.log_model(None, None, None, location) + assert error is None + during = os.environ.get('MLFLOW_HTTP_REQUEST_TIMEOUT') + 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 + + +def test_a_model_with_text_columns_requires_the_text_extra(): + from kumo_core.encoders.config import EncodersConfig + from kumoapi.encoder import Numerical, TextEmbedding + + from kumo.engine.integrations.mlflow import _requirement + + def requirement(encoders): + return _requirement(EncodersConfig(encoders={'users': encoders})) + + assert requirement({'age': Numerical()}) == ( + f'kumo-engine[s3,gcs]=={kumo.__version__}') + assert requirement({'age': Numerical(), 'bio': TextEmbedding('words') + }) == f'kumo-engine[s3,gcs,text]=={kumo.__version__}' + + +def test_a_tracker_dropped_earlier_says_the_model_was_not_logged(): + import time + + class Tracker(_Tracker): + def start(self, record): + raise RuntimeError('server down') + + def log_model(self, bundle_dir, manifest, snapshot, registry): + raise AssertionError('not called') + + session = _started(Tracker()) + deadline = time.monotonic() + 10 + while session._tracker is not None and time.monotonic() < deadline: + time.sleep(0.01) + with pytest.warns(UserWarning, match=r'Could not log the model: ' + r'RuntimeError: tracking stopped after an earlier ' + r'failure \(RuntimeError: server down\)'): + registered, error = session.log_model(None, None, None) + session.finish(None, error) + assert registered is None and isinstance(error, RuntimeError) + assert str(session.error) == 'server down' + + +@pytest.mark.parametrize('on_error', ['warn', 'raise']) +def test_a_registration_failure_is_tagged_on_the_fit_run_and_can_raise( + table_and_snapshot, quick_plan, tmp_path, monkeypatch, on_error): + from mlflow.exceptions import MlflowException + from mlflow.tracking import MlflowClient + + from kumo.engine.integrations.mlflow import MlflowTracking + + detail = ("CATALOG_DOES_NOT_EXIST: Catalog 'nope' does not exist. " + + 'x' * 300 + ' the actionable end') + + def missing_catalog(*args, **kwargs): + raise MlflowException(detail) + + monkeypatch.setattr(MlflowClient, 'create_model_version', + missing_catalog) + table, snapshot = table_and_snapshot + cfg = kumo.MLOps(str(tmp_path / 'root'), + tracking=MlflowTracking(log_system_metrics=False), + registry=kumo.ModelRegistryLocation('nope.s.m'), + on_error=on_error) + trainer = kumo.Trainer(quick_plan) + cause = f'MlflowException: {detail}' + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter('always') + if on_error == 'warn': + result = trainer.fit(table, snapshot=snapshot, mlops=cfg) + else: + with pytest.raises(RuntimeError) as raised: + trainer.fit(table, snapshot=snapshot, mlops=cfg) + assert str(raised.value).endswith(cause) + assert isinstance(raised.value.__cause__, MlflowException) + result = trainer.result + reported = [str(w.message) for w in caught if 'nope' in str(w.message)] + assert reported == ([f'Could not register the model: {cause}'] + if on_error == 'warn' else []) + assert result.id and result.registered_model is None + + run = _fit_run() + assert run.info.status == 'FAILED' + assert run.data.tags['kumo.error'].startswith('MlflowException: ' + 'CATALOG_DOES_NOT_EXIST') + assert run.data.tags['kumo.bundle_id'] == result.id + + +def test_a_tracking_start_failure_names_its_cause_and_can_raise( + table_and_snapshot, quick_plan, tmp_path, monkeypatch, caplog): + from mlflow.exceptions import MlflowException + from mlflow.tracking import MlflowClient + + from kumo.engine.integrations.mlflow import MlflowTracking + + def no_folder(self, name, *args, **kwargs): + raise MlflowException(f'RESOURCE_DOES_NOT_EXIST: Parent of {name} ' + f'does not exist.') + + table, snapshot = table_and_snapshot + cfg = kumo.MLOps(str(tmp_path / 'root'), + tracking=MlflowTracking('/NoSuchFolder/churn', + log_system_metrics=False), + on_error='raise') + warn = dataclasses.replace(cfg, on_error='warn') + monkeypatch.setattr(MlflowClient, 'get_experiment_by_name', + lambda self, name: None) + monkeypatch.setattr(MlflowClient, 'create_experiment', no_folder) + trainer = kumo.Trainer(quick_plan) + with pytest.raises(RuntimeError, match='RESOURCE_DOES_NOT_EXIST') as raised: + trainer.fit(table, snapshot=snapshot, mlops=cfg) + + assert isinstance(raised.value.__cause__, MlflowException) + assert trainer.result.id + off, = [record for record in caplog.records + if 'off for the rest' in record.getMessage()] + assert 'MlflowException: RESOURCE_DOES_NOT_EXIST' in off.getMessage() + assert not off.exc_info + + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter('always') + trainer.fit(table, snapshot=snapshot, mlops=warn) + reported = [str(w.message) for w in caught + if 'RESOURCE_DOES_NOT_EXIST' in str(w.message)] + assert reported == [ + 'Could not log the model: RuntimeError: tracking stopped after an ' + 'earlier failure (MlflowException: RESOURCE_DOES_NOT_EXIST: Parent ' + 'of /NoSuchFolder/churn does not exist.)' + ] + + +def test_a_logged_model_is_named_after_its_model(): + from kumo.engine.integrations.mlflow import _logged_model_name + + manifest = SimpleNamespace(parents=SimpleNamespace( + training_table=SimpleNamespace( + query='PREDICT users.CHURN FOR EACH users.ID', + task_type='binary_classification'))) + assert _logged_model_name('main.ml.churn', manifest) == 'churn' + assert _logged_model_name('a/b:c%d"e\'f', None) == 'a-b-c-d-e-f' + assert _logged_model_name(None, manifest) == 'users-binary_classification' + assert _logged_model_name(None, None) == 'model' + + +def test_a_later_registration_logs_the_model_in_its_fit_run_else_in_none( + table_and_snapshot, quick_plan, tmp_path): + from kumo.engine.integrations.mlflow import MlflowTracking + + table, snapshot = table_and_snapshot + tracking = MlflowTracking('fits', log_system_metrics=False) + result = kumo.Trainer(quick_plan).fit(table, snapshot=snapshot, mlops=( + kumo.MLOps(str(tmp_path / 'root'), tracking=tracking))) + fits = [mlflow.get_experiment_by_name('fits').experiment_id] + assert mlflow.MlflowClient().search_logged_models(fits) == [] + published = tracking.register(result, kumo.ModelRegistryLocation('m')) + + fit = mlflow.get_run(result.tracking.run_id) + model, = mlflow.MlflowClient().search_logged_models(fits) + assert (published.version, model.source_run_id, model.name) == ( + '1', fit.info.run_id, 'm') + assert {m.key: m.value for m in model.metrics} == registry_metrics(result) + + other = f'sqlite:///{tmp_path / "other.db"}' + mlflow.set_tracking_uri(other) + untracked = MlflowTracking(log_system_metrics=False) + assert untracked.register(result, kumo.ModelRegistryLocation('m')) == ( + published) + client = mlflow.MlflowClient(other) + experiment = client.get_experiment_by_name('m').experiment_id + model, = client.search_logged_models([experiment]) + assert (model.source_run_id, client.search_runs([experiment])) == ( + None, []) diff --git a/packages/kumo-engine/test/test_mlflow_target.py b/packages/kumo-engine/test/test_mlflow_target.py new file mode 100644 index 000000000..51e3e04bb --- /dev/null +++ b/packages/kumo-engine/test/test_mlflow_target.py @@ -0,0 +1,427 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +r"""How `MlflowTracking` reaches its server.""" + +import os +import threading +from types import SimpleNamespace + +import pytest +from packaging.version import Version + +import kumo.engine as kumo +from kumo.engine.connector import ( + DatabricksConnector, + DatabricksIdentity, + S3Identity, +) +from kumo.engine.integrations import _mlflow_target as target_module +from kumo.engine.integrations._mlflow_target import SageMakerMLflow, resolve + +mlflow = pytest.importorskip('mlflow') + +ARN = 'arn:aws:sagemaker:us-west-2:account:mlflow-tracking-server/demo' + + +def _connector(identity): + return DatabricksConnector(catalog='main', schema_name='ml', + identity=identity, warehouse_id='w') + + +def test_no_target_reads_mlflows_own_uris_once(tracking_uri): + target = resolve(None) + assert target.tracking_uri == tracking_uri + assert not target.explicit + assert target.platform == 'mlflow' and target.run_url('1', 'r') is None + assert resolve('databricks').registry_uri == 'databricks-uc' + assert resolve('databricks').platform == 'databricks' + assert resolve('databricks://prod').registry_uri == ( + 'databricks-uc://prod') + server = resolve('http://mlflow:5000') + assert server.registry_uri == 'http://mlflow:5000' + assert server.run_url('1', 'r') == ( + 'http://mlflow:5000/#/experiments/1/runs/r') + assert resolve(ARN).platform == 'sagemaker' + assert resolve(server) is server + + +def test_ambient_identities_set_no_credentials(monkeypatch, tracking_uri): + monkeypatch.setenv('DATABRICKS_TOKEN', 'mine') + monkeypatch.setenv('MLFLOW_TRACKING_URI', tracking_uri) + target = resolve(_connector(DatabricksIdentity())) + before = dict(os.environ) + with target.scope(): + assert mlflow.get_tracking_uri() == 'databricks' + assert {k: v for k, v in os.environ.items() + if k.startswith('DATABRICKS_')} == {'DATABRICKS_TOKEN': 'mine'} + assert mlflow.get_tracking_uri() == tracking_uri + assert dict(os.environ) == before + + assert resolve(_connector(DatabricksIdentity())).assets.catalog == 'main' + for target in (resolve(SageMakerMLflow(ARN, S3Identity())), + resolve(SageMakerMLflow(ARN))): + assert type(target.assets).__name__ == 'S3Connector' + with target.scope(): + assert mlflow.get_tracking_uri() == ARN + assert target.artifact_repository('s3://b/run') is None + assert mlflow.get_tracking_uri() == tracking_uri + assert dict(os.environ) == before + + +def test_a_databricks_service_principal_lends_a_token_for_each_call( + monkeypatch): + client = SimpleNamespace(config=SimpleNamespace( + host='https://dbc.example', authenticate=lambda: { + 'Authorization': 'Bearer short-lived' + })) + monkeypatch.setattr('kumo_databricks.session._client_for', + lambda principal: client) + monkeypatch.setenv('DATABRICKS_CLIENT_SECRET', 'other') + monkeypatch.setenv('DATABRICKS_CONFIG_PROFILE', 'laptop') + monkeypatch.delenv('DATABRICKS_TOKEN', raising=False) + target = resolve( + _connector( + DatabricksIdentity(host='https://dbc.example', client_id='sp', + client_secret='secret'))) + + previous = mlflow.get_tracking_uri() + with target.scope(): + inside = {k: v for k, v in os.environ.items() + if k.startswith('DATABRICKS_')} + assert mlflow.get_tracking_uri() == 'databricks' + assert mlflow.get_tracking_uri() == previous + assert inside == {'DATABRICKS_HOST': 'https://dbc.example', + 'DATABRICKS_TOKEN': 'short-lived'} + assert 'secret' not in inside.values() + assert os.environ['DATABRICKS_CLIENT_SECRET'] == 'other' + assert os.environ['DATABRICKS_CONFIG_PROFILE'] == 'laptop' + assert 'DATABRICKS_TOKEN' not in os.environ + assert (target.tracking_uri, target.registry_uri) == ('databricks', + 'databricks-uc') + assert target.run_url('7', 'r') == ( + 'https://dbc.example/ml/experiments/7/runs/r') + + +def test_a_cluster_run_links_to_the_workspace_url_not_the_regional_host( + monkeypatch): + from mlflow.utils import databricks_utils + + from kumo.engine.integrations._mlflow_target import resolve + + target = resolve('databricks') + monkeypatch.setattr(databricks_utils, 'get_workspace_info_from_dbutils', + lambda: ('https://dbc-1.example', '42')) + assert target.run_url('7', 'r') == ( + 'https://dbc-1.example/ml/experiments/7/runs/r') + + +def test_an_aws_access_key_signs_sagemaker_calls_and_uploads( + monkeypatch, tracking_uri): + boto3 = pytest.importorskip('boto3') + plugin = pytest.importorskip('sagemaker_mlflow.session') + monkeypatch.setenv('MLFLOW_TRACKING_URI', tracking_uri) + monkeypatch.setenv('SAGEMAKER_MLFLOW_ASSUME_ROLE_ARN', 'arn:aws:iam::x') + target = resolve( + SageMakerMLflow(ARN, S3Identity(access_key_id='AKIA', + secret_access_key='s3cret'))) + + with target.scope(): + assert mlflow.get_tracking_uri() == ARN + assert 'SAGEMAKER_MLFLOW_ASSUME_ROLE_ARN' not in os.environ + session = plugin._get_current_session() + assert session.get_credentials().access_key == 'AKIA' + repository = target.artifact_repository('s3://bucket/run/artifacts') + assert mlflow.get_tracking_uri() == tracking_uri + assert os.environ['MLFLOW_TRACKING_URI'] == tracking_uri + assert os.environ['SAGEMAKER_MLFLOW_ASSUME_ROLE_ARN'] == 'arn:aws:iam::x' + assert plugin._get_current_session() is None + assert repository._access_key_id == 'AKIA' + assert target.artifact_repository('dbfs:/x') is None + assert isinstance(session, boto3.Session) + + +def test_your_own_mlflow_calls_reach_the_trackers_server_as_its_identity( + monkeypatch, tracking_uri, tmp_path): + plugin = pytest.importorskip('sagemaker_mlflow.session') + from kumo.engine.integrations.mlflow import MlflowTracking + + monkeypatch.setattr(MlflowTracking, '_require_experiment', lambda _: 0) + tracking = MlflowTracking('e', target=SageMakerMLflow( + ARN, S3Identity(access_key_id='AKIA', secret_access_key='s'))) + with tracking.client() as client: + assert client._tracking_client.tracking_uri == ARN + assert mlflow.get_tracking_uri() == ARN + assert plugin._get_current_session().get_credentials( + ).access_key == 'AKIA' + assert mlflow.get_tracking_uri() == tracking_uri + + other = f'sqlite:///{tmp_path / "other.db"}' + with MlflowTracking('e', target=other).client() as client: + assert mlflow.get_tracking_uri() == other + assert client._tracking_client.tracking_uri == other + assert mlflow.get_tracking_uri() == tracking_uri + + +def test_explicit_calls_take_turns(monkeypatch): + target = target_module.Target('databricks', 'databricks-uc', + swap=lambda: target_module.environ( + {'KUMO_PROBE': 'x'})) + seen = [] + + def call(): + with target.scope(): + seen.append(os.environ.get('KUMO_PROBE')) + + threads = [threading.Thread(target=call) for _ in range(8)] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + assert seen == ['x'] * 8 and 'KUMO_PROBE' not in os.environ + + +def test_ambient_sagemaker_calls_see_its_uri_and_restore_the_global_one( + tracking_uri, monkeypatch): + pytest.importorskip('sagemaker_mlflow') + from kumo.engine.integrations.mlflow import MlflowTracking + + seen = [] + + def lookup(self, name): + seen.append((self.tracking_uri, mlflow.get_tracking_uri())) + return SimpleNamespace(experiment_id='1') + + monkeypatch.setattr(mlflow.MlflowClient, 'get_experiment_by_name', + lookup) + tracking = MlflowTracking(target=SageMakerMLflow(ARN), + experiment_name='e') + kumo.MLOps('/tmp/x', tracking=tracking) + + assert seen == [(ARN, ARN)] + assert mlflow.get_tracking_uri() == tracking_uri + + +NO_LINKS = 'MLFLOW_SUPPRESS_PRINTING_URL_TO_STDOUT' + + +def test_sagemaker_calls_print_no_mlflow_run_links(monkeypatch): + from contextlib import nullcontext + + monkeypatch.delenv(NO_LINKS, raising=False) + seen = [] + for uri in (ARN, 'http://mlflow'): + with target_module.Target(uri, uri, swap=nullcontext).scope(): + seen.append(os.environ.get(NO_LINKS)) + + assert (seen, os.environ.get(NO_LINKS)) == (['true', None], None) + + +def test_pin_mlflow_reads_its_run_link_switch_per_call(): + from mlflow.environment_variables import ( + MLFLOW_SUPPRESS_PRINTING_URL_TO_STDOUT, + ) + + assert MLFLOW_SUPPRESS_PRINTING_URL_TO_STDOUT.name == NO_LINKS + + +@pytest.mark.skipif(Version(mlflow.__version__) < Version('3.17'), + reason='this MLflow still accepts folder stores') +@pytest.mark.parametrize('store', ['file', 'path']) +def test_a_refused_folder_store_says_to_use_sqlite(tmp_path, monkeypatch, + store): + from kumo.engine.integrations.mlflow import MlflowTracking + + monkeypatch.delenv('MLFLOW_ALLOW_FILE_STORE', raising=False) + uri = (tmp_path / 'mlruns').as_uri() if store == 'file' else str( + tmp_path / 'mlruns') + tracking = MlflowTracking(target=uri, experiment_name='e') + with pytest.raises(RuntimeError) as raised: + kumo.MLOps(str(tmp_path / 'root'), tracking=tracking) + assert str(raised.value) == ( + f'MLflow refused the local folder store at {uri}: use its default ' + f'sqlite:///mlflow.db or another tracking backend.') + + +def test_an_unreachable_server_fails_when_mlops_is_built( + tmp_path, monkeypatch): + from kumo.engine.integrations.mlflow import MlflowTracking + + monkeypatch.setenv('MLFLOW_HTTP_REQUEST_MAX_RETRIES', '0') + tracking = MlflowTracking(target='http://127.0.0.1:9', + experiment_name='e') + with pytest.raises(RuntimeError, + match='MLflow at http://127.0.0.1:9 is not usable: '): + kumo.MLOps(str(tmp_path), tracking=tracking) + with pytest.raises(ValueError, match='needs experiment_name'): + kumo.MLOps(str(tmp_path), + tracking=MlflowTracking(target='http://127.0.0.1:9')) + + +def test_a_missing_unity_catalog_schema_fails_when_mlops_is_built( + monkeypatch): + from kumo.engine.integrations import mlflow as integration + + calls = [] + + def http_request(creds, endpoint, method, **kwargs): + calls.append(endpoint) + return SimpleNamespace( + status_code=404, reason='Not Found', json=lambda: { + 'error_code': 'CATALOG_DOES_NOT_EXIST', + 'message': "Catalog 'nope' does not exist." + }) + + monkeypatch.setattr('mlflow.utils.rest_utils.http_request', http_request) + monkeypatch.setattr('mlflow.utils.databricks_utils.' + 'get_databricks_host_creds', lambda uri: None) + with pytest.raises(ValueError, match="cannot register 'nope.ml.churn': " + 'CATALOG_DOES_NOT_EXIST'): + integration.check_registry_schema('nope.ml.churn', 'databricks-uc') + assert calls == ['/api/2.1/unity-catalog/schemas/nope.ml'] + integration.check_registry_schema('churn', 'databricks-uc') + integration.check_registry_schema('a.b.c', 'sqlite:///x.db') + assert len(calls) == 1 + + +def test_pin_set_tracking_uri_keeps_the_module_value(tracking_uri): + utils = pytest.importorskip('mlflow.tracking._tracking_service.utils') + assert utils._tracking_uri == tracking_uri + + +def test_pin_current_experiment_lookup(): + from mlflow.tracking.fluent import _get_experiment_id + + assert callable(_get_experiment_id) + + +def test_pin_sagemaker_signer_takes_the_thread_session(monkeypatch): + boto3 = pytest.importorskip('boto3') + pytest.importorskip('sagemaker_mlflow') + from sagemaker_mlflow.auth import AuthBoto + from sagemaker_mlflow.session import use_session + + session = boto3.Session(aws_access_key_id='AKIA', + aws_secret_access_key='s3cret') + with use_session(session): + auth = AuthBoto('us-west-2', 'sagemaker-mlflow') + assert auth.creds.access_key == 'AKIA' + + +def test_pin_system_metric_monitors(): + from kumo.engine.integrations.mlflow import _metric_monitors + + monitors = _metric_monitors() + assert len(monitors) >= 3 + for monitor in monitors: + monitor.collect_metrics() + assert isinstance(monitor.aggregate_metrics(), dict) + monitor.clear_metrics() + + +def test_pin_databricks_auth_variables(): + from databricks.sdk.config import Config + + names = {attribute.env for attribute in Config.attributes()} + assert {'DATABRICKS_HOST', 'DATABRICKS_TOKEN', 'DATABRICKS_CLIENT_ID', + 'DATABRICKS_CLIENT_SECRET', 'DATABRICKS_CONFIG_PROFILE'} <= names + + +def test_registering_a_registered_bundle_reuses_its_version( + tracking_uri, caplog): + from kumo.engine.integrations.mlflow import _register, _registered_version + + target = resolve(None) + client = target.client() + run = client.create_run('0') + info = SimpleNamespace(model_uri=f'runs:/{run.info.run_id}/model', + run_id=run.info.run_id, model_id=None) + location = kumo.ModelRegistryLocation('churn', alias='champion') + + def register(bundle): + pointers = {'kumo_bundle_id': bundle} + return _register(target, client, info, location, pointers, + _registered_version(client, 'churn', pointers)) + + first = register('a') + with caplog.at_level('INFO'): + again = register('a') + other = register('b') + + assert (first.version, again.version, other.version) == ('1', '1', '2') + assert 'churn version 1 already holds bundle a; reused it' in caplog.text + assert len(client.search_model_versions("name = 'churn'")) == 2 + assert str(client.get_model_version_by_alias('churn', + 'champion').version) == '2' + + +def test_the_smoke_script_compiles(tmp_path): + import py_compile + from pathlib import Path + + for parent in Path(__file__).resolve().parents: + script = parent / 'examples' / 'mlops' / 'smoke.py' + if script.is_file(): + py_compile.compile(str(script), cfile=str(tmp_path / 'out'), + doraise=True) + return + pytest.skip('examples/mlops is not in this checkout') + + +def test_a_service_principal_client_ignores_a_token_lent_to_mlflow( + monkeypatch): + from kumo_databricks import session + + from kumo.engine.integrations._mlflow_target import _principal + + made = [] + monkeypatch.setattr(session, 'WorkspaceClient', + lambda **kwargs: made.append(kwargs)) + monkeypatch.setenv('DATABRICKS_TOKEN', 'short-lived') + session._client_for(_principal(DatabricksIdentity( + host='https://dbc.example', client_id='sp', client_secret='secret'))) + assert made[0]['auth_type'] == 'oauth-m2m' + + +def test_unity_catalog_search_results_with_tags_as_a_method(): + from kumo.engine.integrations.mlflow import _registered_version + + class Search: + def __init__(self, version): + self.version = version + + def tags(self): + raise AssertionError('not a mapping') + + class Page(list): + token = None + + held = {'1': {'kumo_bundle_id': 'other'}, '2': {'kumo_bundle_id': 'b'}} + client = SimpleNamespace( + search_model_versions=lambda *a, **k: Page([Search('1'), + Search('2')]), + get_model_version=lambda name, version: SimpleNamespace( + tags=held[version])) + assert _registered_version(client, 'm', {'kumo_bundle_id': 'b'}) == '2' + assert _registered_version(client, 'm', {'kumo_bundle_id': 'x'}) is None + + +def test_kumo_calls_drop_mlflow_request_header_errors_only(caplog): + import logging + + from kumo.engine.integrations._mlflow_target import resolve + + headers = logging.getLogger('mlflow.tracking.request_header.registry') + message = ('Encountered unexpected error during resolving request ' + 'headers: %s') + headers.addHandler(caplog.handler) + try: + with caplog.at_level('WARNING', logger=headers.name): + with resolve('sqlite:///x.db').scope(): + headers.warning(message, 'inside') + headers.warning('another warning') + headers.warning(message, 'outside') + finally: + headers.removeHandler(caplog.handler) + assert 'inside' not in caplog.text + assert 'another warning' in caplog.text and 'outside' 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..739354b3b --- /dev/null +++ b/packages/kumo-engine/test/test_mlflow_tracking.py @@ -0,0 +1,498 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import json +import os +import re +import socket +import threading +import time +import warnings +from pathlib import Path + +import pytest +from kumo_core.artifacts import PROVENANCE_FILE +from kumo_core.summaries import Dataset +from kumoapi.encoder import Numerical +from kumoapi.model_plan import ModelPlan +from test_tracking_contract import ( + CONTEXT, + INPUTS, + join_tracking_threads, + replay, +) + +import kumo.engine as kumo +from kumo.engine._manifest import TrackingRef +from kumo.engine._tracking import logged_names, run_record +from kumo.engine.trainer._run_spec import fit_spec + +mlflow = pytest.importorskip('mlflow') +autologging_utils = pytest.importorskip('mlflow.utils.autologging_utils') +pytestmark = pytest.mark.usefixtures('tracking_uri') + + +def _runs(client): + return client.search_runs( + [e.experiment_id for e in client.search_experiments()]) + + +def _fit_run(): + run, = [run for run in _runs(mlflow.MlflowClient()) + if 'mlflow.parentRunId' not in run.data.tags] + return run + + +def test_fit_logs_a_run_per_fit_and_per_trial(table_and_snapshot, quick_plan, + tmp_path, caplog): + from kumo.engine.integrations.mlflow import MlflowTracking + + table, snapshot = table_and_snapshot + client = mlflow.MlflowClient() + plan = quick_plan + plan.training_job.num_experiments = 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') + mlops = kumo.MLOps(artifact_root=str(tmp_path / 'root'), + tracking=tracking) + result = kumo.Trainer(plan).fit(table, snapshot=snapshot, + mlops=mlops) + 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 + assert re.fullmatch(r'fit-ents-regression-\d{8}-\d{6}', + parent.info.run_name) + 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.bundle_id', 'kumo.dropped_params', + 'kumo.snapshot_id', 'kumo.task_type', 'kumo.entity', + 'kumo.tune_metric', 'kumo.engine_version', 'kumo.best_trial', + 'kumo.bundle_uri'} + assert parent.data.tags['kumo.bundle_uri'] == str( + tmp_path / 'root' / 'models' / result.id) + assert parent.data.tags['kumo.snapshot_id'] == snapshot.id + assert parent.data.tags['kumo.task_type'] == 'regression' + assert parent.data.tags['kumo.bundle_id'] == result.id + stored = tmp_path / 'root' / 'models' / result.id / PROVENANCE_FILE + assert json.loads(stored.read_text())['tracking'] == { + 'platform': 'mlflow', 'run_id': parent.info.run_id, + 'experiment': 'ents', 'run_name': parent.info.run_name} + assert result.tracking == TrackingRef( + platform='mlflow', run_id=parent.info.run_id, experiment='ents', + run_name=parent.info.run_name) + logged = client.download_artifacts(parent.info.run_id, PROVENANCE_FILE, + str(tmp_path)) + assert json.loads(Path(logged).read_text())['id'] == result.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/') + } + tested = {name: value for name, value in logged_names( + result.metrics).items() if not name.startswith(('train/', 'val/'))} + assert {name for name in metrics + if name.startswith(('test/', 'before_refit/'))} == set(tested) + health = json.loads(mlflow.artifacts.load_text( + f'{parent.info.artifact_uri}/data/health.json')) + assert health['data'][0][:2] == ['ents', 400] + assert {name: metrics[name] for name in tested} == tested + 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] == [ + f'trial-{i}-{parent.info.run_name}' for i in (1, 2)] + for child in children: + assert child.info.status == 'FINISHED' + assert child.data.tags['kumo.status'] == 'succeeded' + assert {'val/best_mae', '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] + selected = min(children, key=lambda run: run.data.metrics['val/best_mae']) + for name in ('train/loss', 'val/mae'): + curve, trial = (sorted((metric.step, metric.value) + for metric in client.get_metric_history( + run.info.run_id, name)) + for run in (parent, selected)) + assert curve == trial and len(curve) == 2 + + +def test_failed_fit_marks_the_run_failed(table_and_snapshot, tmp_path, + 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 kumo.mlops(kumo.MLOps(artifact_root=str(tmp_path / 'root'), + tracking=MlflowTracking())): + with pytest.raises(ValueError, match='stop'): + kumo.Trainer(ModelPlan()).fit(table, snapshot=snapshot) + + 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(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_a_user_run_name_names_every_run(): + from kumo.engine.integrations.mlflow import MlflowTracking + + tracking = MlflowTracking(run_name='nightly', log_system_metrics=False) + replay(tracking) + join_tracking_threads() + + assert _fit_run().info.run_name == 'fit-nightly' + + +SAGEMAKER = ('arn:aws:sagemaker:us-west-2:account:' + 'mlflow-tracking-server/demo') + + +@pytest.mark.parametrize('uri,good,bad', [ + (SAGEMAKER, ['users-churn', 'u--1', 'a' * 57], + ['users_churn', 'main.ml.churn', 'bad/name:x', '-churn', 'a' * 58]), + ('arn:aws:sagemaker:us-east-1:account:mlflow-app/app-1', + ['users-churn'], ['users_churn']), + ('databricks-uc', ['main.ml.users_churn', 'ml.users-churn', 'churn'], + ['main.ml.users churn', 'a.b.c.d', 'main..churn', 'main.ml.a/b']), + ('sqlite:///mlflow.db', ['bad/name:x', 'users_churn', 'a.b.c.d'], []), +]) +def test_registry_model_names_are_checked_when_mlops_is_built( + tmp_path, uri, good, bad): + from kumo.engine.integrations.mlflow import ( + MlflowTracking, + check_model_name, + ) + + for name in good: + check_model_name(name, uri) + tracking = MlflowTracking(log_system_metrics=False, + target=f'sqlite:///{tmp_path}/m.db', + experiment_name='e') + tracking.target.registry_uri = uri + for name in bad: + with pytest.raises(ValueError, match="e.g. '"): + kumo.MLOps(str(tmp_path), tracking=tracking, + registry=kumo.ModelRegistryLocation(name)) + + +def test_building_mlops_creates_the_experiment_or_says_why_it_cannot( + tmp_path, monkeypatch): + from kumo.engine.integrations.mlflow import MlflowTracking + + uri = f'sqlite:///{tmp_path}/m.db' + kumo.MLOps(str(tmp_path), tracking=MlflowTracking( + 'churn', log_system_metrics=False, target=uri)) + assert mlflow.MlflowClient(uri).get_experiment_by_name('churn') + + def refuse(self, name, *args, **kwargs): + raise mlflow.exceptions.MlflowException('folder does not exist') + + monkeypatch.setattr(mlflow.MlflowClient, 'create_experiment', refuse) + with pytest.raises(RuntimeError, match=r"experiment '/Shared/x/churn'.*" + r"folder does not exist.*create its parent folder"): + kumo.MLOps(str(tmp_path), tracking=MlflowTracking( + '/Shared/x/churn', log_system_metrics=False, target=uri)) + + +@pytest.mark.parametrize('uri,command', [ + (SAGEMAKER, 'aws sagemaker create-presigned-mlflow-tracking-server-url ' + '--tracking-server-name demo --region us-west-2'), + ('sqlite:///mlflow.db', None), +]) +def test_a_tracked_fit_on_sagemaker_says_how_to_open_the_ui( + uri, command, monkeypatch): + from kumo.engine.integrations import mlflow as integration + from kumo.engine.integrations.mlflow import MlflowTracking + + monkeypatch.setattr(integration, '_HINTED', set()) + + tracker = MlflowTracking(log_system_metrics=False) + tracker.start(run_record(fit_spec(CONTEXT))) + tracker.target.tracking_uri = uri + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter('always') + tracker.finish(run_record(fit_spec(CONTEXT)), None) + hints = [str(w.message) for w in caught if 'Studio' in str(w.message)] + assert hints == ([] if command is None else [ + "MLflow run 'fit-ents-regression-20261029-061800' in experiment " + "'Default': MLflow's run links do not open on SageMaker: open the " + f'MLflow UI from SageMaker Studio or at the URL `{command}` prints.']) + with warnings.catch_warnings(record=True) as again: + warnings.simplefilter('always') + tracker.finish(run_record(fit_spec(CONTEXT)), None) + assert not [w for w in again if 'Studio' in str(w.message)] + + +def test_user_disabled_autologging_stays_disabled(): + from kumo.engine.integrations.mlflow import MlflowTracking + + tracker = MlflowTracking(log_system_metrics=False) + with autologging_utils.disable_autologging(): + tracker.start(run_record(fit_spec(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(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), summaries=[ + Dataset(name, source, {}, {}, 'source table') + for name, source in INPUTS.items()]) + + run = _fit_run() + 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 _fit_manifest(): + from kumo.engine._manifest import ( + ModelManifest, + ModelParents, + SnapshotManifest, + SnapshotParents, + SourceTableEntry, + TrainingTableManifest, + TrainingTableParents, + ) + + return ModelManifest(id='m', parents=ModelParents( + training_table=TrainingTableManifest( + id='tt-fadf04', created_at='2026-11-01T10:00:00+00:00', + location='s3://root/training_tables/tt-fadf04', + parents=TrainingTableParents(snapshot=SnapshotManifest( + id='g@5ac245', created_at='2026-11-01T09:00:00+00:00', + location='s3://root/snapshots/g/5ac245', + parents=SnapshotParents(source_tables={ + name: SourceTableEntry( + connector='S3', location=f's3://data/{name}', + version='3', version_kind='exact') + for name in ('users', 'orders')})))))) + + +def test_a_fits_inputs_are_logged_in_role_order(): + from kumo.engine._manifest import FIT_ROLES, artifact_summaries + from kumo.engine.integrations.mlflow import MlflowTracking + + replay(MlflowTracking(log_system_metrics=False), + summaries=artifact_summaries(_fit_manifest(), FIT_ROLES)) + + assert [item.dataset.name for item in + _fit_run().inputs.dataset_inputs] == [ + 'training table 2026-11-01-fadf04', + 'graph snapshot 2026-11-01-5ac245', 'orders', 'users'] + + +def test_the_snapshot_input_lists_its_tables_and_rows(tmp_path): + from kumo.engine._data_summaries import table_details + from kumo.engine._manifest import FIT_ROLES, artifact_summaries + from kumo.engine.integrations.mlflow import MlflowTracking + + (tmp_path / 'ingest_stats.json').write_text(json.dumps( + {'stats': {'node_counts': {'users': 300, 'orders': 6000}}})) + (tmp_path / 'diskgraph_config.json').write_text(json.dumps( + {'tables': [{'name': 'users'}, {'name': 'orders'}]})) + replay(MlflowTracking(log_system_metrics=False), + summaries=artifact_summaries(_fit_manifest(), FIT_ROLES, + table_details(tmp_path, 'g@5ac245'))) + + snapshot, = [item.dataset for item in _fit_run().inputs.dataset_inputs + if item.dataset.name.startswith('graph snapshot')] + assert json.loads(snapshot.schema) == {'mlflow_colspec': []} + assert json.loads(snapshot.profile) == { + 'snapshot': 'g@5ac245', 'tables': 2, 'rows': 6300, + 'users rows': 300, 'orders rows': 6000} + + +def test_link_prediction_metric_names_are_logged(): + from kumo.engine._tracking import RunRecord + from kumo.engine.integrations.mlflow import MlflowTracking + + tracker = MlflowTracking(experiment_name='lp', log_system_metrics=False) + tags = {'kumo.kind': 'fit', 'kumo.run_id': 'r'} + tracker.start(RunRecord(name='r', tags=tags)) + tracker.finish( + RunRecord(name='r', tags=tags, metrics={ + 'test_map@10': 0.5, + 'val_precision@10': 0.25 + }), None) + + run = mlflow.get_run(tracker.run_id) + assert run.info.status == 'FINISHED' + assert run.data.metrics == { + 'test_map_at_10': 0.5, + 'val_precision_at_10': 0.25 + } + + +class _Monitor: + def collect_metrics(self): + pass + + def aggregate_metrics(self): + return {'cpu': 1.0} + + def clear_metrics(self): + pass + + +def test_system_metrics_skip_a_failed_log_and_stop_after_repeated_ones( + monkeypatch, caplog): + import logging + + from kumo.engine.integrations import mlflow as integration + + batches = [] + + def log(readings): + batches.append([(values, step) for values, step, _ in readings]) + if len(batches) == 1 or len(batches) >= 3: + raise ConnectionError('blip') + + monkeypatch.setattr(integration, '_metric_monitors', lambda: [_Monitor()]) + monkeypatch.setattr(integration, '_FLUSH_SECONDS', 0.0) + caplog.set_level(logging.INFO, integration.logger.name) + with integration._SystemMetrics(log, interval=0.001) as metrics: + metrics._thread.join(10) + + reading = {'system/cpu': 1.0} + assert batches == [[(reading, 0), (reading, 1)], *[ + [(reading, step)] + for step in range(2, 3 + integration._SYSTEM_METRICS_FAILURES)]] + assert caplog.text.count('Could not log system metrics') == 1 + assert 'ConnectionError: blip' in caplog.text + assert 'stopped after 6 failed logs in a row' in caplog.text + + +def test_system_metrics_sample_every_second_unless_the_environment_says( + monkeypatch): + from kumo.engine.integrations import mlflow as integration + + monkeypatch.setattr(integration, '_metric_monitors', lambda: [_Monitor()]) + batches = [] + with integration._SystemMetrics(batches.append) as metrics: + pass + assert batches == [] + assert metrics._interval == 1.0 + monkeypatch.setenv('MLFLOW_SYSTEM_METRICS_SAMPLING_INTERVAL', '2') + assert integration._SystemMetrics(batches.append)._interval == 2.0 diff --git a/packages/kumo-engine/test/test_mlops.py b/packages/kumo-engine/test/test_mlops.py index 3289c8467..b78762b46 100644 --- a/packages/kumo-engine/test/test_mlops.py +++ b/packages/kumo-engine/test/test_mlops.py @@ -31,6 +31,31 @@ def test_mlops_needs_an_artifact_root_it_can_write_files_to(tmp_path): kumo.MLOps(artifact_root=str(tmp_path), on_error='ignore') +def test_the_tracker_registers_and_its_registry_checks_the_name(tmp_path): + def check_model_location(connector, location): + if '-' in location.model_name: + raise ValueError('bad name') + + tracker = SimpleNamespace(connector=SimpleNamespace( + sink_backend=SimpleNamespace( + check_model_location=check_model_location))) + good = kumo.ModelRegistryLocation('USERS_CHURN', alias='candidate') + kumo.MLOps(str(tmp_path), tracking=tracker, registry=good) + with pytest.raises(ValueError, match='bad name'): + kumo.MLOps(str(tmp_path), tracking=tracker, + registry=kumo.ModelRegistryLocation('users-churn')) + with pytest.raises(ValueError, match='model alias'): + kumo.MLOps(str(tmp_path), tracking=tracker, + registry=kumo.ModelRegistryLocation('USERS_CHURN', + alias='LATEST')) + with pytest.raises(ValueError, match='registry= needs tracking'): + kumo.MLOps(str(tmp_path), registry=good) + with pytest.raises(TypeError, match='ModelRegistryLocation'): + kumo.MLOps(str(tmp_path), tracking=tracker, + registry=kumo.ModelOutputConfig(tracker.connector, + location=good)) + + def test_a_root_is_a_uri_or_a_connector_and_path(tmp_path): uri = kumo.MLOps(artifact_root=str(tmp_path)) assert uri.root.location('models/m') == 'models/m' @@ -45,6 +70,9 @@ def test_a_root_is_a_uri_or_a_connector_and_path(tmp_path): assert pair.root.location('models/m') == 'volume/kumo/models/m' assert pair.wants('model') and not pair.wants('snapshot') assert not kumo.MLOps(artifact_root=str(tmp_path), store=False).kinds + assert uri.snapshot_uri('ab12@cd34') == f'{tmp_path}/snapshots/ab12/cd34' + assert kumo.MLOps('s3://bucket/kumo', store=False).snapshot_uri( + SimpleNamespace(id='ab12@cd34')) == 's3://bucket/kumo/snapshots/ab12/cd34' def test_a_connector_root_outside_its_root_dir_fails_when_built(): diff --git a/packages/kumo-engine/test/test_model_agnostic_mlops.py b/packages/kumo-engine/test/test_model_agnostic_mlops.py new file mode 100644 index 000000000..113ec0015 --- /dev/null +++ b/packages/kumo-engine/test/test_model_agnostic_mlops.py @@ -0,0 +1,434 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import io +import json +import os +import subprocess +import sys +import textwrap +import warnings +from dataclasses import replace +from pathlib import Path +from types import SimpleNamespace + +import numpy as np +import pandas as pd +import pyarrow as pa +import pyarrow.parquet as pq +import pytest +from kumo_core.artifacts import ( + PROVENANCE_FILE, + add_bundle_file, + bundle_id_from_manifest, +) +from kumoapi.encoder import TextEmbedding +from kumoapi.plugin import ModelRegistryLocation, PublishedModel +from kumoapi.typing import Stype + +import kumo.engine as kumo +from kumo.engine import _deps +from kumo.engine._deps import Dependencies, FileAsset, MissingDependencyError +from kumo.engine._manifest import ModelManifest, read_manifest, write_manifest +from kumo.engine._predictor import load_predictor +from kumo.engine._tracking import RunSpec, TrackingSession, run_record +from kumo.engine._version import __version__ +from kumo.engine.tracking import Event +from kumo.engine.trainer._spec import MANIFEST_FILE, PROVENANCE_ROLE +from kumo.engine.trainer.trainer import _store_and_register + + +class ToyPredictor: + def __init__(self, weight): + self.weight = weight + + @classmethod + def load(cls, bundle_dir): + return cls(float((Path(bundle_dir) / 'weight.txt').read_text())) + + def predict(self, rows, snapshot, params): + return pd.DataFrame({'ENTITY': rows['ENTITY'], + 'SCORE': rows['ENTITY'] * self.weight}) + + +def _toy_bundle(root, predictor=f'{__name__}:ToyPredictor'): + root.mkdir(parents=True) + (root / 'weight.txt').write_text('2.0') + (root / MANIFEST_FILE).write_text(json.dumps({'artifacts': []})) + add_bundle_file(root, 'weights', 'weight.txt') + manifest = ModelManifest(id=bundle_id_from_manifest(root), + dependencies=[_deps.predictor(predictor)]) + write_manifest(root / PROVENANCE_FILE, manifest) + add_bundle_file(root, PROVENANCE_ROLE, PROVENANCE_FILE) + return root, manifest + + +class _Tracker: + def __init__(self): + self.calls = [] + + def start(self, record): + self.calls.append(('start', record)) + + def on_trial(self, record, event): + self.calls.append(('on_trial', record)) + + def log_model(self, bundle_dir, manifest, snapshot, registry): + self.calls.append(('log_model', (bundle_dir, manifest.id, snapshot))) + return registry and PublishedModel(name=registry.model_name, + version='2', uri='fake://TOY/2') + + def finish(self, record, error): + self.calls.append(('finish', record)) + + +def test_any_run_is_tracked_from_its_spec(): + tracker = _Tracker() + session = TrackingSession(tracker) + spec = RunSpec(kind='custom', params={'lr': 0.1}) + session.start(lambda: spec) + session.on_event(Event('trial_started', 0.0, trial=1, + params={'attempt': 1})) + session.finish(lambda: replace(spec, metrics={'auc': 0.9}, + parents={'bundle_id': 'b'}), None) + session._worker.join(5) + started, trial, finished = (record for _, record in tracker.calls) + assert started == run_record(spec) + assert trial.name == 'trial-1-attempt-1' + assert trial.tags == {'kumo.kind': 'trial', 'kumo.engine_version': __version__} + assert finished.metrics == {'auc': 0.9} + assert finished.tags == {'kumo.kind': 'custom', 'kumo.bundle_id': 'b', + 'kumo.engine_version': __version__} + + +class _RegistryTracker(_Tracker): + log_model = None + + def __init__(self, connector): + super().__init__() + self.connector = connector + + +@pytest.mark.parametrize('logs', [False, True]) +def test_any_bundle_is_stored_and_registered_by_its_tracker(tmp_path, logs): + bundle, manifest = _toy_bundle(tmp_path / 'bundle') + published = [] + sink = SimpleNamespace( + publish_model=lambda connector, bundle_dir, **kw: published.append( + (bundle_dir, kw)) or PublishedModel(name='TOY', version='1', + uri='fake://TOY/1')) + tracker = (_Tracker() if logs else _RegistryTracker( + SimpleNamespace(sink_backend=sink))) + settings = kumo.MLOps(str(tmp_path / 'root'), tracking=tracker, + registry=ModelRegistryLocation('TOY')) + session = TrackingSession(tracker) + session.start(lambda: RunSpec(kind='custom')) + + registered, error = _store_and_register( + bundle, manifest, settings, session=session, snapshot='snap', + metrics={'auc': 0.9}) + + assert error is None + stored = tmp_path / 'root' / 'models' / manifest.id + assert (stored / 'weight.txt').read_text() == '2.0' + if logs: + assert registered.version == '2' and published == [] + assert ('log_model', (bundle, manifest.id, 'snap')) in tracker.calls + else: + assert registered.version == '1' + assert published == [(str(bundle), { + 'path': None, 'location': ModelRegistryLocation('TOY'), + 'metadata': {'kumo_bundle_id': manifest.id}, + 'metrics': {'auc': 0.9} + })] + session.finish(None, None) + + +def test_a_bundle_is_scored_by_the_predictor_it_names(tmp_path): + pytest.importorskip('mlflow') + from kumo.engine.integrations._mlflow_model import KumoModel + + bundle, _ = _toy_bundle(tmp_path / 'toy') + rows = pd.DataFrame({'ENTITY': [1, 2]}) + assert load_predictor(bundle).predict(rows, None, {})['SCORE'].tolist( + ) == [2.0, 4.0] + model = KumoModel() + model.load_context(SimpleNamespace(artifacts={'bundle': str(bundle)}, + model_config={'columns': []})) + model._snapshots['s'] = 'snapshot' + assert model.predict(None, rows, {'snapshot': 's'})['SCORE'].tolist() == [ + 2.0, 4.0 + ] + + unknown, _ = _toy_bundle(tmp_path / 'unknown', 'nowhere.at_all:Model') + with pytest.raises(ValueError, match="names predictor 'nowhere.at_all:" + "Model', which cannot be imported here"): + load_predictor(unknown) + + +def _vocab(tmp_path): + (tmp_path / 'vocab').mkdir() + (tmp_path / 'vocab' / 'words.txt').write_text('a b c') + return tmp_path / 'vocab' + + +def _volume(filesystem): + return SimpleNamespace(sink_backend=SimpleNamespace( + holds_files=lambda connector, path: path.startswith('/Volumes/'), + filesystem=filesystem)) + + +def test_a_file_asset_is_stored_once_and_resolved_by_its_hash(tmp_path): + asset = FileAsset('toy', 'vocab', _vocab(tmp_path), version='1') + settings = kumo.MLOps(str(tmp_path / 'root')) + recorded = Dependencies(settings=settings) + record = recorded.require(asset) + stored = tmp_path / 'root' / 'assets' / 'toy' / asset.fingerprint() + assert record.location == str(stored) + assert read_manifest(stored / PROVENANCE_FILE, _deps.AssetManifest).name \ + == 'vocab' + + resolved = Dependencies(recorded.records).resolve('toy', 'vocab') + assert (resolved / 'words.txt').read_text() == 'a b c' + + (stored / 'words.txt').write_text('changed') + with pytest.raises(MissingDependencyError, + match=rf"needs toy 'vocab' \(sha256 {record.sha256}\)"): + Dependencies(recorded.records).resolve('toy', 'vocab') + + +def test_a_manifest_without_dependencies_uses_the_trainer_predictor(): + old = ModelManifest.from_dict({'kind': 'model', 'schema_version': 1, + 'id': 'b'}) + assert old.dependencies is None + assert Dependencies.of(old).records == [] + + +WORDS = ['red', 'green', 'blue', 'cyan', 'gold', 'gray', 'pink'] + + +def _vectors(root): + root.mkdir() + rng = np.random.default_rng(0) + pd.DataFrame({'index': pd.Series(WORDS, dtype=object)}).to_parquet( + root / 'glove_tiny_index.parquet') + np.save(root / 'glove_tiny_8d_emb.npy', + np.asfortranarray(rng.normal(size=(len(WORDS), 8)).astype( + np.float32))) + return root + + +def _users(root): + rng = np.random.default_rng(1) + root.mkdir() + pq.write_table( + pa.table({ + 'id': pa.array(range(120), type=pa.int64()), + 'bio': [' '.join(rng.choice(WORDS, 6)) for _ in range(120)], + 'label': [bool(i % 2) if i < 100 else None for i in range(120)], + }), root / 'users.parquet') + return kumo.Table.from_source_table( + kumo.FileConnector(root_dir=str(root))['users'], + primary_key='id').with_column('bio', stype=Stype.text) + + +def test_a_model_serves_from_its_bundle_snapshot_and_assets_alone( + tmp_path, monkeypatch): + monkeypatch.setattr('kumo.engine._scratch._root', tmp_path / 'scratch') + graph = kumo.Graph([_users(tmp_path / 'source')]) + snapshot = graph.snapshot() + query = kumo.PredictiveQuery(graph, + 'PREDICT users.label FOR EACH users.id') + plan = query.suggest_model_plan(snapshot) + plan.training_job.num_experiments = 1 + plan.training_job.disable_explain = True + plan.optimization.max_epochs = 1 + plan.column_processing.encoder_overrides = { + 'users.bio': TextEmbedding('colors') + } + store = kumo.EmbeddingModelStore() + store.add('colors', kumo.GloVeVectors(_vectors(tmp_path / 'vectors'))) + trainer = kumo.Trainer(plan, embedding_models=store, accelerator='cpu') + root = tmp_path / 'root' + result = trainer.fit(query.generate_training_table(snapshot), + snapshot=snapshot, mlops=kumo.MLOps(str(root))) + location = snapshot.export( + kumo.SnapshotOutputConfig( + kumo.FileConnector(root_dir=str(tmp_path / 'exports')), 'snap')) + expected = trainer.predict( + snapshot, binary_classification_threshold=.5).frame.sort_values( + 'ENTITY').reset_index(drop=True) + expected.to_parquet(tmp_path / 'expected.parquet') + dependency, = Dependencies.of(result.manifest).assets('text_encoder') + assert dependency.name == 'colors' and dependency.packing == 'external' + assert Path(dependency.location).parent == root / 'assets' / 'text_encoder' + + serve = tmp_path / 'serve' + (serve / 'home').mkdir(parents=True) + script = textwrap.dedent(f""" + import pandas as pd + import kumo.engine as kumo + snapshot = kumo.GraphSnapshot.load_from_connector( + kumo.FileConnector(root_dir={str(tmp_path / 'exports')!r}), + path={location!r}) + trainer = kumo.Trainer.load({str(root / 'models' / result.id)!r}, + accelerator='cpu') + scored = trainer.predict( + snapshot, binary_classification_threshold=.5).frame.sort_values( + 'ENTITY').reset_index(drop=True) + pd.testing.assert_frame_equal( + scored, pd.read_parquet({str(tmp_path / 'expected.parquet')!r})) + """) + env = {**os.environ, 'HOME': str(serve / 'home'), 'CUDA_VISIBLE_DEVICES': + '', 'HTTPS_PROXY': 'http://127.0.0.1:9', + 'HTTP_PROXY': 'http://127.0.0.1:9'} + for name in ('XDG_CACHE_HOME', 'NO_PROXY', 'no_proxy'): + env.pop(name, None) + done = subprocess.run([sys.executable, '-c', script], cwd=serve, env=env, + capture_output=True, text=True, timeout=600) + assert done.returncode == 0, done.stderr[-3000:] + + (root / 'assets').rename(tmp_path / 'moved') + done = subprocess.run([sys.executable, '-c', script], cwd=serve, env=env, + capture_output=True, text=True, timeout=600) + assert "MissingDependencyError: The model needs text_encoder 'colors'" \ + in done.stderr + + +class _PrefixOnlyS3: + type_name = 's3' + + def __init__(self): + self.objects = {} + + def get_file_info(self, path): + return SimpleNamespace(is_file=path in self.objects) + + def create_dir(self, path, recursive=True): + raise OSError('AWS Error ACCESS_DENIED during HeadBucket operation') + + def open_output_stream(self, path): + objects = self.objects + + class Writer(io.BytesIO): + def close(self): + objects[path] = self.getvalue() + super().close() + + return Writer() + + def delete_file(self, path): + del self.objects[path] + + +def test_an_asset_is_stored_on_s3_with_prefix_scoped_rights( + tmp_path, monkeypatch): + from kumo.engine import connector + + s3 = _PrefixOnlyS3() + monkeypatch.setattr(connector, 'file_store', + lambda connector, path: (s3, f'bucket/kumo/{path}')) + deps = Dependencies(settings=kumo.MLOps(str(tmp_path / 'root'))) + record = deps.require(FileAsset('toy', 'vocab', _vocab(tmp_path))) + assert deps.errors == [] and record.location + assert sorted(s3.objects) == [ + f'bucket/kumo/assets/toy/{record.sha256}/_provenance.json', + f'bucket/kumo/assets/toy/{record.sha256}/words.txt'] + + +@pytest.mark.parametrize('on_error', ['warn', 'raise']) +def test_a_model_that_was_not_stored_is_not_registered( + tmp_path, monkeypatch, on_error): + from kumo.engine._mlops import raise_on_error + from kumo.engine.trainer import trainer + + bundle, manifest = _toy_bundle(tmp_path / 'bundle') + published = [] + sink = SimpleNamespace( + publish_model=lambda *a, **k: published.append(a) or 1 / 0) + settings = kumo.MLOps(str(tmp_path / 'root'), on_error=on_error, + tracking=_RegistryTracker( + SimpleNamespace(sink_backend=sink)), + registry=ModelRegistryLocation('TOY')) + + def refuse(*args): + raise RuntimeError('Could not store model m under the artifact ' + 'root: validate: oauth and pat') + + monkeypatch.setattr(trainer, '_store_model', refuse) + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter('always') + registered, error = _store_and_register( + bundle, manifest, settings, warn=on_error == 'warn') + assert registered is None and published == [] + assert 'oauth and pat' in str(error) + assert any('not registered' in str(w.message) + for w in caught) is (on_error == 'warn') + if on_error == 'raise': + with pytest.raises(RuntimeError, match='oauth and pat'): + raise_on_error(settings, error, 'Trainer.fit') + asset_failure = OSError('ACCESS_DENIED') + registered, error = _store_and_register(bundle, manifest, settings, + warn=False, stored=asset_failure) + assert registered is None and error is asset_failure + + +def test_volume_assets_load_through_the_connector_the_loader_was_given( + tmp_path, monkeypatch): + from pyarrow.fs import LocalFileSystem, SubTreeFileSystem + + from kumo.engine.integrations import _mlflow_model + from kumo.engine.integrations.mlflow import _target_hint + + monkeypatch.setenv('XDG_CACHE_HOME', str(tmp_path / 'cache')) + (tmp_path / 'volume').mkdir() + files = SubTreeFileSystem(str(tmp_path / 'volume'), LocalFileSystem()) + volume = _volume(lambda connector, path: (files, path.lstrip('/'))) + recorded = Dependencies(settings=kumo.MLOps((volume, '/Volumes/c/s/v/a'))) + record = recorded.require(FileAsset('toy', 'vocab', _vocab(tmp_path))) + assert record.location.startswith('/Volumes/c/s/v/a/assets/toy/') + manifest = ModelManifest(id='m', dependencies=[record]) + + with _deps.asset_locations([], _target_hint): + with pytest.raises(MissingDependencyError, + match=r'pass target=DatabricksConnector\(' + r'identity=\.\.\.\) to load_model\(model_uri, ' + r'\.\.\.\), with an identity that reads ' + r'/Volumes/c/s/v/a/assets\.$'): + Dependencies.of(manifest).resolve('toy', 'vocab') + with _deps.asset_locations(_deps.connector_roots(manifest, volume)): + assert (Dependencies.of(manifest).resolve('toy', 'vocab') / + 'words.txt').read_text() == 'a b c' + other_root = kumo.MLOps((volume, '/Volumes/c/s/v/b')) + assert Dependencies.of(manifest, other_root).resolve('toy', 'vocab') + + seen = [] + monkeypatch.setattr( + 'kumo.engine._predictor.load_predictor', lambda bundle: seen.append( + Dependencies.of(manifest).candidates(record)[0])) + monkeypatch.setattr('kumo.engine.trainer._provenance.read_provenance', + lambda bundle: manifest) + _mlflow_model.KumoModel().load_context(SimpleNamespace( + artifacts={'bundle': str(tmp_path)}, + model_config={'columns': [], 'assets_from': volume})) + assert seen == [(volume, record.location)] + + +def test_the_pyfunc_reads_a_volume_snapshot_through_assets_from( + monkeypatch): + from kumo.engine.graph import GraphSnapshot + from kumo.engine.integrations import _mlflow_model + + volume = _volume(lambda connector, path: None) + loads = [] + monkeypatch.setattr(GraphSnapshot, 'load_from_connector', + lambda connector, *, path: loads.append( + (connector, path))) + model = _mlflow_model.KumoModel() + model._snapshots, model._reader = {}, volume + model._snapshot('/Volumes/c/s/v/a/snapshots/g/version') + model._reader = None + model._snapshot('/Volumes/c/s/v/b/snapshots/g/version') + assert loads[0] == (volume, '/Volumes/c/s/v/a/snapshots/g/version') + assert loads[1][0] is not volume and loads[1][1] == 'version' diff --git a/packages/kumo-engine/test/test_optional_imports.py b/packages/kumo-engine/test/test_optional_imports.py index 22d40a95a..02affe547 100644 --- a/packages/kumo-engine/test/test_optional_imports.py +++ b/packages/kumo-engine/test/test_optional_imports.py @@ -27,6 +27,7 @@ 'google', 'azure', 'gcsfs', + 'mlflow', ) PROVIDER_PACKAGES = ( 'kumo_snowflake', @@ -117,6 +118,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_completion.py b/packages/kumo-engine/test/test_prediction_completion.py index 526d87a3e..255ba9a7a 100644 --- a/packages/kumo-engine/test/test_prediction_completion.py +++ b/packages/kumo-engine/test/test_prediction_completion.py @@ -13,6 +13,7 @@ from kumoapi.model_plan import ModelPlan from kumo.engine.pquery.prediction_table import PredictionTable +from kumo.engine.trainer._scoring import ScoringDefaults from kumo.engine.trainer.trainer import Trainer @@ -33,7 +34,8 @@ def prepared_prediction(snapshot, tmp_path, monkeypatch): from kumo_core.types import TaskType trainer = Trainer(ModelPlan()) - trainer._training_result = SimpleNamespace(artifact_paths={}) + trainer._training_result = SimpleNamespace( + artifact_paths={}, scoring_defaults=ScoringDefaults()) trainer._snapshot_id = snapshot.id trainer._training_query = 'PREDICT users.label FOR EACH users.id' table = PredictionTable(trainer._training_query, tmp_path, tmp_path, 8, diff --git a/packages/kumo-engine/test/test_prediction_workers.py b/packages/kumo-engine/test/test_prediction_workers.py index 2b9e7b7c7..d06ee9c92 100644 --- a/packages/kumo-engine/test/test_prediction_workers.py +++ b/packages/kumo-engine/test/test_prediction_workers.py @@ -14,6 +14,7 @@ from kumo_core.types.enums import PredictionArtifactType from kumoapi.model_plan import ModelPlan +from kumo.engine.trainer._scoring import ScoringDefaults from kumo.engine.trainer.trainer import Trainer @@ -80,7 +81,8 @@ def test_bad_class_count_fails_before_prediction_work(count, monkeypatch): def test_valid_class_count_reaches_scorer(task_type, count, tmp_path, monkeypatch): trainer = Trainer(ModelPlan(), accelerator='cpu') - trainer._training_result = SimpleNamespace(artifact_paths={}) + trainer._training_result = SimpleNamespace( + artifact_paths={}, scoring_defaults=ScoringDefaults()) trainer._training_query = 'trained query' monkeypatch.setattr(Trainer, 'scratch_dir', property(lambda self: tmp_path)) monkeypatch.setattr( 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..110cd178b --- /dev/null +++ b/packages/kumo-engine/test/test_scoring_defaults.py @@ -0,0 +1,164 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import hashlib +import shutil +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.plugin import PublishedModel + +import kumo.engine as kumo +from kumo.engine.trainer._scoring import normalized_scoring_defaults + + +def _frame(model, table, snapshot, **kwargs): + result = model.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.result.bundle_dir, destination) + return destination + + +def test_derived_model_thresholds_unless_explicit(fitted): + trainer, table, snapshot = fitted + parent = trainer.result + original = _digest(parent.bundle_dir) + probs = _frame(parent, table, snapshot, + binary_classification_threshold=.5)['True_PROB'] + threshold = float(probs.median()) + child = parent.with_scoring_defaults( + kumo.ScoringDefaults(binary_classification_threshold=threshold)) + + implicit = _frame(child, table, snapshot) + pd.testing.assert_frame_equal( + implicit, + _frame(parent, 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.bundle_dir) == original + assert parent.scoring_defaults == kumo.ScoringDefaults() + assert child.id != parent.id + assert child.manifest.id == child.id + assert child.manifest.parents.training_table is None + derived_from = child.manifest.parents.derived_from + assert derived_from.id == parent.id + assert derived_from.parents.training_table == ( + parent.manifest.parents.training_table) + cleared = child.with_scoring_defaults(kumo.ScoringDefaults()) + assert cleared.id == parent.id + assert cleared.manifest.parents.derived_from.parents.derived_from == ( + kumo.ModelManifest(id=parent.id, reconstructed=True)) + + +def test_equal_defaults_give_equal_bundle_ids(fitted): + trainer, *_ = fitted + children = [ + trainer.result.with_scoring_defaults( + kumo.ScoringDefaults(binary_classification_threshold=value)) + for value in (1, np.float32(1)) + ] + assert children[0].id == children[1].id + assert children[0].scoring_defaults.binary_classification_threshold == 1.0 + + +def test_parent_without_a_manifest_reads_as_a_stub(fitted, tmp_path): + trainer, *_ = fitted + old = _copy(trainer, tmp_path / 'old') + remove_bundle_file(old, 'provenance') + child = kumo.Trainer.load(old).result.with_scoring_defaults( + kumo.ScoringDefaults(binary_classification_threshold=.5)) + assert child.manifest.parents.derived_from == kumo.ModelManifest( + id=trainer.result.id, reconstructed=True) + + +def test_defaults_survive_export_and_connector_load(fitted, tmp_path): + trainer, *_ = fitted + child = trainer.result.with_scoring_defaults( + kumo.ScoringDefaults(binary_classification_threshold=.4)) + registry = kumo.FileConnector(root_dir=str(tmp_path)) + child._trainer.export(kumo.ModelOutputConfig(registry, path='m1')) + loaded = kumo.Trainer.load_from_connector(registry, path='m1') + assert loaded.result.scoring_defaults == child.scoring_defaults + assert loaded.result.id == child.id + + +def test_derived_export_publishes_the_child_bundle_id(fitted): + trainer, *_ = fitted + child = trainer.result.with_scoring_defaults( + kumo.ScoringDefaults(binary_classification_threshold=.4)) + calls = [] + sink = SimpleNamespace(write_table=lambda *a: None, + publish_model=lambda *a, **k: calls.append(k) or + PublishedModel(name='m', version='1', uri='u')) + child._trainer.export( + kumo.ModelOutputConfig(SimpleNamespace(sink_backend=sink), path='m')) + assert calls[0]['metadata']['kumo_bundle_id'] == child.id + + +def test_defaults_are_validated_against_the_task(fitted): + trainer, *_ = fitted + with pytest.raises(ValueError, match='num_classes_to_return'): + trainer.result.with_scoring_defaults( + kumo.ScoringDefaults(num_classes_to_return=2)) + 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.result.scoring_defaults == kumo.ScoringDefaults() + assert old.result.id == trainer.result.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.result.scoring_defaults == kumo.ScoringDefaults() + run = ignored.predict(table, snapshot=snapshot) + assert run.frame['TARGET_PRED'].isna().all() + + +def test_a_fit_stores_its_defaults_in_the_bundle(fitted, tmp_path): + from kumoapi.model_plan import TrainingTableGenerationPlan + + from kumo.engine.trainer.trainer import TrainingResult + + trainer, table, snapshot = fitted + first = trainer.result + defaults = kumo.ScoringDefaults(binary_classification_threshold=0.3) + refit = kumo.Trainer(trainer.model_plan) + result = refit._finish_fit( + TrainingResult({}, dict(first.artifact_paths), + _copy(trainer, tmp_path / 'bundle')), + predictive_query=trainer._training_query, + snapshot_id=snapshot.id, export_config=None, + training_table_generation=TrainingTableGenerationPlan(), + graph_config=snapshot.graph_config, scoring_defaults=defaults) + assert result.scoring_defaults == defaults and result.id != first.id + loaded = kumo.Trainer.load(result.bundle_dir) + assert loaded.result.scoring_defaults == defaults + pd.testing.assert_frame_equal( + _frame(loaded, table, snapshot), + _frame(first, table, snapshot, binary_classification_threshold=0.3)) diff --git a/packages/kumo-engine/test/test_semantic_carry_forward.py b/packages/kumo-engine/test/test_semantic_carry_forward.py index fc5788a9b..81da3228e 100644 --- a/packages/kumo-engine/test/test_semantic_carry_forward.py +++ b/packages/kumo-engine/test/test_semantic_carry_forward.py @@ -24,6 +24,7 @@ PredictionTable, _compile_prediction_table, ) +from kumo.engine.trainer._scoring import ScoringDefaults from kumo.engine.trainer.trainer import Trainer, _stage_model_bundle _SPLIT = SimpleNamespace(stage_counts={'test': 1}) @@ -128,7 +129,6 @@ def test_compile_rejects_invalid_prediction_time(snapshot, tmp_path, value, @pytest.mark.parametrize('query, threshold, message', [ ('PREDICT users.age FOR EACH users.id', .5, 'target/entity'), - ('PREDICT users.label FOR EACH users.id', None, 'must be specified'), ('PREDICT users.label FOR EACH users.id', 1.1, 'between 0 and 1'), ]) def test_prediction_admission_before_materialization(snapshot, tmp_path, query, @@ -136,7 +136,7 @@ def test_prediction_admission_before_materialization(snapshot, tmp_path, query, monkeypatch): from kumoapi.model_plan import ModelPlan trainer = Trainer(ModelPlan()) - trainer._training_result = SimpleNamespace() + trainer._training_result = SimpleNamespace(scoring_defaults=ScoringDefaults()) trainer._snapshot_id = snapshot.id trainer._training_query = 'PREDICT users.label FOR EACH users.id' materialize = Mock(side_effect=AssertionError('must not materialize')) @@ -164,7 +164,8 @@ def test_binary_predictions_preserve_probabilities_and_metadata_contract( _stub_generated_task_dataset_builder(monkeypatch) monkeypatch.setattr(type(trainer), 'scratch_dir', property(lambda _: tmp_path)) - trainer._training_result = SimpleNamespace(artifact_paths={}) + trainer._training_result = SimpleNamespace( + artifact_paths={}, scoring_defaults=ScoringDefaults()) trainer._snapshot_id = snapshot.id trainer._training_query = 'PREDICT users.label FOR EACH users.id' table = PredictionTable(trainer._training_query, tmp_path, tmp_path, 0, @@ -273,7 +274,8 @@ def stats(index_path, table_name, mode): trainer = Trainer(ModelPlan()) monkeypatch.setattr(type(trainer), 'scratch_dir', property(lambda _: tmp_path)) - trainer._training_result = SimpleNamespace(artifact_paths={}) + trainer._training_result = SimpleNamespace( + artifact_paths={}, scoring_defaults=ScoringDefaults()) trainer._snapshot_id = snapshot.id trainer._training_query = query @@ -708,7 +710,8 @@ def run_predict() -> None: trainer = Trainer(ModelPlan()) monkeypatch.setattr(type(trainer), 'scratch_dir', property(lambda _: tmp_path)) - trainer._training_result = SimpleNamespace(artifact_paths={}) + trainer._training_result = SimpleNamespace( + artifact_paths={}, scoring_defaults=ScoringDefaults()) trainer._snapshot_id = snapshot.id trainer._training_query = query trainer.predict(table, snapshot=snapshot, output_config=output, @@ -720,7 +723,8 @@ def run_predict() -> None: same_trainer = Trainer(ModelPlan()) monkeypatch.setattr(type(same_trainer), 'scratch_dir', property(lambda _: tmp_path)) - same_trainer._training_result = SimpleNamespace(artifact_paths={}) + same_trainer._training_result = SimpleNamespace( + artifact_paths={}, scoring_defaults=ScoringDefaults()) same_trainer._snapshot_id = snapshot.id same_trainer._training_query = query same_trainer.predict(table, snapshot=snapshot, output_config=output, 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..5cbcb6feb --- /dev/null +++ b/packages/kumo-engine/test/test_snowflake_tracking.py @@ -0,0 +1,263 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import inspect +import threading +from datetime import datetime, timezone +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 + +from kumo.engine._manifest import TrackingRef +from kumo.engine._tracking import RunRecord, TrackingSession, tracking_ref +from kumo.engine.trainer._run_spec import fit_spec, fitted_spec + +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=FakeSession, **kwargs): + from kumo.engine.integrations.snowflake import SnowflakeTracking + return SnowflakeTracking( + session() if session is FakeSession else session, + experiment_name='ents', **kwargs) + + +def _replay(tracking, result=RESULT, error=None, during=None): + session = TrackingSession(tracking) + session.start(lambda: fit_spec(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( + None if result is None else + lambda: fitted_spec(session.spec, result), + error) + return session + + +def test_native_model_lineage_holds_the_run_open_during_the_fit(fake): + lineage = [] + _replay( + _tracking(database='ml', schema_name='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_and_fails_strict_mode( + fake, busy, table_and_snapshot, quick_plan, tmp_path, monkeypatch): + import kumo.engine as kumo + from kumo.engine.integrations.snowflake import SnowflakeTracking + + monkeypatch.setattr(SnowflakeTracking, 'check', lambda *args: None) + 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, run = len(user.calls), user._run + table, snapshot = table_and_snapshot + trainer = kumo.Trainer(quick_plan) + in_use = 'Snowflake ExperimentTracking is already in use' + + with pytest.raises(RuntimeError, match=in_use): + trainer.fit(table, snapshot=snapshot, mlops=kumo.MLOps( + str(tmp_path / 'root'), tracking=_tracking(session), + store=False, on_error='raise')) + + assert trainer.result is not None + assert user.calls[calls:] == [] + assert (user._run, user.warnings_as_errors) == (run, True) + + +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,run_name,new', [ + ('FIT', 'fit', 'FIT_2'), ('FIT_MYXRUN', 'my_run', 'FIT_MY_RUN_2')]) +def test_an_existing_run_is_never_resumed(fake, taken, run_name, new, caplog): + session = FakeSession() + user = fake(session) + user.set_experiment('ents') + user.start_run(taken) + if taken == 'FIT': + user.end_run() + else: + user._run = None + + tracking = _tracking(session, run_name=run_name) + _replay(tracking) + + assert tracking.run_id == tracking.logged_run_name == new + assert tracking_ref(tracking) == TrackingRef( + platform='snowflake', run_id=new, experiment='ENTS', run_name=new) + assert set(user.runs) == {taken, new} + assert user.runs[taken]['params'] == {} + assert user.runs[new]['status'] == 'FINISHED' + assert 'recording this run 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)(**{**vars(RESULT), 'metrics': { + **RESULT.metrics, 'r2': float('nan')}}) + _replay(_tracking(tags={'note': 'a$$b', 'level': 2}), result) + + run, = fake._instance.runs.values() + assert sorted({name for name, _, _ in run['metrics']}) == [ + 'TRIAL_1/val/best_mae', 'TRIAL_2_ATTEMPT_2/val/best_mae', + 'duration_seconds', 'test/mae' + ] + assert run['params']['kumo.dropped_metrics'] == '2' + assert run['params']['kumo.dropped_params'] == '1' + assert 'note' not in run['params'] + assert run['params']['level'] == '2' + + +def test_the_users_session_is_not_altered(fake): + session = FakeSession() + + _replay(_tracking(session)) + + assert session.statements == [] + + +def _connected(): + from kumo.engine.connector.base import Connector + + session = FakeSession(schema='RUNS', schemas=[('RUNS', )]) + connector = mock.Mock(spec=Connector, session=lambda: session) + return _tracking(connector), connector, session + + +def test_a_connector_opens_the_session_and_the_schema_is_checked_first( + fake): + import kumo.engine as kumo + + tracking, connector, session = _connected() + tracking.check() + assert tracking.session is session and tracking.connector is connector + assert session.statements == ["SHOW SCHEMAS LIKE 'RUNS' IN DATABASE DB"] + + tracking.session = FakeSession(schema='RUNS') + with pytest.raises(ValueError, match='schema DB.RUNS does not exist'): + kumo.MLOps('/tmp/root', tracking=tracking) + with pytest.raises(ValueError, match='or a Snowpark session'): + _tracking(None) + + +def test_mlops_loads_a_registered_model_through_the_tracker( + fake, monkeypatch, tmp_path): + import kumo.engine as kumo + + loaded = [] + original = kumo.Trainer.load_from_connector.__func__ + monkeypatch.setattr( + kumo.Trainer, 'load_from_connector', + classmethod(lambda cls, source, **kwargs: loaded.append( + (source, kwargs)) if not isinstance(source, kumo.MLOps) else + original(cls, source, **kwargs))) + tracking, connector, _ = _connected() + cfg = kumo.MLOps(str(tmp_path), tracking=tracking) + location = kumo.ModelRegistryLocation('CHURN', alias='champion') + kumo.Trainer.load_from_connector(cfg, location=location) + assert loaded == [(connector, { + 'location': location, 'embedding_models': None, + 'accelerator': 'auto' + })] + with pytest.raises(ValueError, match='SnowflakeTracking\\(connector'): + _tracking().load_model(location) + + +@pytest.mark.parametrize('run_name, record, name', [ + (None, RunRecord(tags={'kumo.kind': 'fit'}, model_label='kumo-demo-churn', + started=datetime(2026, 10, 29, 6, 18, 5, + tzinfo=timezone.utc)), + 'FIT_KUMO_DEMO_CHURN_20261029_061805'), + ('nightly', RunRecord(tags={'kumo.kind': 'fit'}), 'FIT_NIGHTLY'), +]) +def test_runs_are_named_by_kind_model_and_date_or_the_users_name( + fake, run_name, record, name): + tracking = _tracking(run_name=run_name) + tracking.start(record) + tracking.finish(record, None) + + assert list(fake._instance.runs) == [name] diff --git a/packages/kumo-engine/test/test_summary_tables.py b/packages/kumo-engine/test/test_summary_tables.py new file mode 100644 index 000000000..fdff90149 --- /dev/null +++ b/packages/kumo-engine/test/test_summary_tables.py @@ -0,0 +1,23 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import json + +from kumo_core.summaries import Table + +from kumo.engine._summary_tables import Routed, route, table_json + + +def test_tables_round_measurements_and_keep_definitions_and_counts(): + table = Table('holdout/by_threshold', ('threshold', 'tp', 'precision'), ( + (0.07 + 1e-17 * 3, 12, 0.48774387193112), + (1 / 3, 7, 27.610799999999998), + )) + + assert json.loads(table_json(table))['data'] == [ + [0.07, 12, 0.487744], [0.3333333333, 7, 27.6108]] + + +def test_a_summary_type_no_tracker_knows_is_skipped(): + assert route([object()]) == Routed() 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..0e53d39fa --- /dev/null +++ b/packages/kumo-engine/test/test_tracking_contract.py @@ -0,0 +1,516 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import hashlib +import json +import threading +import time +from dataclasses import dataclass, field, replace +from datetime import datetime, timezone +from pathlib import Path +from types import SimpleNamespace +from typing import Any, Callable, Dict, List, Mapping, Optional, Tuple + +import pytest +from kumo_core.artifacts import PROVENANCE_FILE +from kumo_core.summaries import Dataset, Scalar, Table + +from kumo.engine._tracking import ( + DROPPED_METRICS, + DROPPED_PARAMS, + TrackingSession, + run_record, +) +from kumo.engine._version import __version__ +from kumo.engine.tracking import Event, RunRecord +from kumo.engine.trainer._run_spec import FitContext, fit_spec, fitted_spec + +FIT = 'fit' +USER_TAGS = {'team': 'ents', 'kumo.kind': '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'}, +} +SNAPSHOT = f"{'s' * 64}@{'v' * 16}" +CONTEXT = FitContext( + query=QUERY, snapshot_id=SNAPSHOT, task_type='regression', + started=datetime(2026, 10, 29, 6, 18, tzinfo=timezone.utc), + 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, 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) + ] + + +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'), +] +SELECTED = {'train_loss': ((1, 0.6), (2, 0.4)), 'val_mae': ((1, 0.5), (2, 0.3))} +BUNDLE = 'b' * 64 +RESULT = SimpleNamespace(metrics={'val_mae': 0.3, 'test_mae': 0.32}, + id=BUNDLE, bundle_dir=Path('no-bundle'), + registered_model=None, registered_alias=None) + + +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 log_metrics(self, metrics, step): + self.calls.append(('log_metrics', (metrics, step))) + getattr(self.tracker, 'log_metrics', lambda *_: None)(metrics, step) + + 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, + context=CONTEXT, series=None, + summaries=()) -> Dict[str, RunRecord]: + recorder = _Recorder(tracker) + session = TrackingSession(recorder) + session.start(lambda: fit_spec(context)) + for event in events: + session.on_event(event) + session.finish( + None if result is None else lambda: fitted_spec( + session.spec, result, summaries=summaries), error) + records = { + record.name: record + for hook, record in recorder.calls if hook == 'on_trial' + } + spec = replace(fit_spec(context), tags=USER_TAGS) + if result is not None: + spec = fitted_spec(spec, result) + spec = replace(spec, steps=series if series is not None else { + **({} if error else SELECTED)}) + if series is None and not error: + spec = replace(spec, parents={**spec.parents, 'best_trial': 2}) + if 'error' in events[-1].params: + spec = replace(spec, error=str(events[-1].params['error'])) + records[FIT] = run_record(spec) + 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 + run_name: Callable[[str], Optional[str]] = str + expected: Callable[[Dict[str, RunRecord]], Dict[str, RunRecord]] = dict + + +def gaps(expected: Dict[str, RunRecord], observed: Dict[str, Observed], + adapter: Adapter) -> List[str]: + kept_params = adapter.kept_params + expected = adapter.expected(expected) + 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', adapter.run_name(seen.name), + adapter.run_name(record.name)) + differ('params', dict(seen.params), dict(kept_params(record.params))) + differ('metrics', {k: seen.metrics.get(k) for k in record.metrics}, + dict(record.metrics)) + 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, + FIT if record.tags['kumo.kind'] == 'trial' else None) + if adapter.artifacts: + differ('artifacts', dict(seen.artifacts), dict(record.artifacts)) + return found + + +def _trial(run): + trial = run.info.run_name.split('-')[1] + return f"trial-{trial}-attempt-{run.data.params['attempt']}" + + +@pytest.fixture +def mlflow_adapter(tracking_uri): + import mlflow + + from kumo.engine.integrations.mlflow import MlflowTracking, _mlflow_params + + def observe(): + client = mlflow.MlflowClient() + runs = client.search_runs( + [e.experiment_id for e in client.search_experiments()]) + kumo_ids = { + r.info.run_id: + FIT if r.data.tags['kumo.kind'] == 'fit' else _trial(r) + 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 if run.data.tags['kumo.kind'] != 'trial' + else _trial(run)), 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) + if not a.is_dir + }) + for run in runs + } + + return 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) + + +@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_name, run in fake._instance.runs.items(): + names: Dict[Optional[str], Dict[str, Dict[str, Any]]] = {} + + def bucket(key): + prefix, _, name = key.partition('/') + if not prefix.startswith('TRIAL_'): + prefix, name = '', key + return names.setdefault(prefix or None, { + 'params': {}, 'tags': {}, 'metrics': {}, 'steps': {} + }), name + + for key, value in run['params'].items(): + runs, name = bucket(key) + 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']: + if key == 'duration_seconds': + continue + runs, name = bucket(key) + runs['metrics'][name] = value + runs['steps'][name] = (*runs['steps'].get(name, ()), + (step, value)) + for prefix, runs in names.items(): + observed[prefix or FIT] = Observed( + name=prefix or run_name, + artifacts=run['artifacts'] if prefix is None else {}, + **runs) + return observed + + def expected(records): + kept = {} + for key, record in records.items(): + if record.tags['kumo.kind'] == 'trial': + trial, attempt = key.split('-')[1::2] + key = f'TRIAL_{trial}' + ( + f'_ATTEMPT_{attempt}' if attempt != '1' else '') + record = replace(record, steps={}, metrics={ + name: value for name, value in record.metrics.items() + if name == 'val/best_mae'}, tags={ + name: value for name, value in record.tags.items() + if name in ('kumo.status', 'kumo.error')}) + kept[key] = record + return kept + + return Adapter( + make=lambda: SnowflakeTracking(_snowflake_experiment_fake.FakeSession(), + experiment_name='ents', tags=USER_TAGS), + observe=observe, step_metrics=True, artifacts=True, + run_name=lambda name: None, expected=expected) + + +@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[FIT] + assert fit.name == 'fit-ents-regression-20261029-061800' + assert (fit.model_label, fit.started) == ('ents-regression', + CONTEXT.started) + assert fit.tags == {'team': 'ents', 'kumo.kind': 'fit', + 'kumo.bundle_id': BUNDLE, 'kumo.snapshot_id': SNAPSHOT, + 'kumo.task_type': 'regression', + 'kumo.entity': 'ents', 'kumo.tune_metric': 'mae', + 'kumo.engine_version': __version__, + 'kumo.best_trial': '2'} + 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(replace(fit_spec(CONTEXT), tags={ + 'kumo.bundle_id': 'x', 'training_job.tune_metric': 'y', 'a': 1})) + assert spoofed.tags.keys() - fit.tags.keys() == {'a'} + assert fit.metrics == {'test/mae': 0.32} + assert fit.steps == {'train/loss': ((1, 0.6), (2, 0.4)), + 'val/mae': ((1, 0.5), (2, 0.3))} + registered = SimpleNamespace(**{ + **vars(RESULT), 'registered_alias': 'candidate', + 'registered_model': SimpleNamespace(name='ents', version='3')}) + tags = run_record(fitted_spec(fit_spec(CONTEXT), registered)).tags + assert tags.items() >= {'kumo.model_name': 'ents', + 'kumo.model_version': '3'}.items() + assert 'kumo.model_alias' not in tags + assert fit.artifacts['query.pql'] == QUERY + assert json.loads(fit.artifacts['model_plan.json']) == CONTEXT.model_plan + + oom = records['trial-2-attempt-1'] + assert oom == RunRecord( + name='trial-2-attempt-1', model_label='ents-regression', + params={'optimization.base_lr': '0.001', 'attempt': '1'}, tags={ + 'kumo.kind': 'trial', + 'kumo.engine_version': __version__, + 'kumo.snapshot_id': SNAPSHOT, + 'kumo.task_type': 'regression', + 'kumo.status': 'failed', + 'kumo.error': 'OutOfMemoryError: CUDA out of memory', + }) + retry = records['trial-2-attempt-2'] + assert retry.params['attempt'] == '2' + assert retry.metrics == {'val/best_mae': 0.3} + 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_FILE).write_text('{"bundle": "b"}') + result = SimpleNamespace(**{**vars(RESULT), 'bundle_dir': tmp_path}) + expected = replay(adapter.make(), result=result) + + assert expected[FIT].artifacts[PROVENANCE_FILE] == '{"bundle": "b"}' + assert gaps(expected, adapter.observe(), adapter) == [] + + +def test_adapter_records_a_single_trial_fit_epoch_by_epoch_without_trials( + adapter): + plan = {**CONTEXT.model_plan, + 'training_job': {'tune_metric': 'mae', 'num_experiments': 1}} + events = [EVENTS[0], *_epochs(1, (0.5, 0.4), (0.3, 0.35)), EVENTS[3], + EVENTS[-1]] + expected = replay(adapter.make(), events, + context=replace(CONTEXT, model_plan=plan), series={ + 'train_loss': ((1, 0.5), (2, 0.3)), + 'val_mae': ((1, 0.4), (2, 0.35))}) + + assert list(expected) == [FIT] + assert gaps(expected, adapter.observe(), adapter) == [] + + +SUMMARIES = ( + Scalar('test/brier', 0.25), + Table('graph/col_stats', ('table', 'rows'), (('ents', 400),)), + Dataset('ents', INPUTS['ents'], {'ID': {'stype': 'id', 'dtype': 'int64'}}, + {'rows': 400}, 'source table'), +) + + +def test_adapter_logs_each_summary_type(adapter, mlflow_adapter): + replay(adapter.make(), summaries=SUMMARIES) + + fit = adapter.observe()[FIT] + assert fit.steps['test/brier'] == ((0, 0.25),) + tables = {'col_stats.json': {'columns': ['table', 'rows'], + 'data': [['ents', 400]]}} + if adapter is not mlflow_adapter: + assert {name: json.loads(fit.artifacts[f'graph/{name}']) + for name in tables} == tables + return + import mlflow + + client = mlflow.MlflowClient() + run, = (r for r in client.search_runs( + [e.experiment_id for e in client.search_experiments()]) + if r.data.tags['kumo.kind'] == 'fit') + assert {name: json.loads(mlflow.artifacts.load_text( + f'{run.info.artifact_uri}/{section}/{name}')) + for section, name in (('graph', 'col_stats.json'),)} == tables + assert json.loads(run.data.tags['mlflow.loggedArtifacts']) == [ + {'path': 'graph/col_stats.json', 'type': 'table'}] + dataset, = run.inputs.dataset_inputs + assert (dataset.dataset.name, dataset.dataset.digest, + dataset.tags[0].value) == ('ents', hashlib.sha256(json.dumps( + ['ents', INPUTS['ents'], SUMMARIES[-1].schema], + sort_keys=True).encode()).hexdigest()[:8], 'source table') + assert json.loads(dataset.dataset.schema) == { + 'mlflow_colspec': [{'name': 'ID', 'type': 'long'}]} + assert json.loads(dataset.dataset.profile) == {'rows': 400} + + +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[FIT].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: fit_spec(CONTEXT)) + session.finish(lambda: fitted_spec(session.spec, RESULT), None) + + assert recorder.calls == [] + + +def test_a_hung_tracker_does_not_stall_the_fit(monkeypatch): + 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) + + session = TrackingSession(Hung()) + began = time.monotonic() + try: + session.start(lambda: fit_spec(CONTEXT)) + session.finish(None, None) + elapsed = time.monotonic() - began + finally: + release.set() + join_tracking_threads() + + assert elapsed < 1.5 + assert isinstance(session.error, TimeoutError) diff --git a/packages/kumo-engine/test/test_trainer_registry.py b/packages/kumo-engine/test/test_trainer_registry.py index 7ea9ba950..1d47e25a1 100644 --- a/packages/kumo-engine/test/test_trainer_registry.py +++ b/packages/kumo-engine/test/test_trainer_registry.py @@ -29,6 +29,7 @@ SPEC_FILE, write_training_spec, ) +from kumo.engine.trainer.output_config import ModelOutputConfig from kumo.engine.trainer.trainer import Trainer, TrainingResult @@ -230,7 +231,7 @@ def test_finish_fit_forwards_publication_details_and_returns_assigned_version( tmp_path, path, location): bundle = _fake_published_bundle(tmp_path) trainer = Trainer(ModelPlan()) - original = TrainingResult({'val_auroc': .9}, {}, bundle) + original = TrainingResult({'val_auroc': .9, 'test_auroc': .8}, {}, bundle) published = PublishedModel(name='CHURN_GNN', version='assigned-v7', uri='provider://models/CHURN_GNN/assigned-v7') calls = [] @@ -238,7 +239,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') } @@ -257,7 +262,9 @@ def publish_model(self, connector, bundle_dir, **kwargs): assert calls[0][1] == str(bundle) assert calls[0][2]['path'] == path assert calls[0][2]['location'] == location - assert calls[0][2]['metrics'] == {'val_auroc': .9} + assert calls[0][2]['metrics'] == {'test/auroc': .8} + assert {key: value for key, value in calls[0][2]['metadata'].items() + if key.startswith('kumo_')} == {'kumo_bundle_id': result.bundle_id} assert result.published_model is published assert result.metrics == original.metrics assert result.bundle_dir == bundle @@ -572,3 +579,182 @@ def fake_predict(**kwargs) -> None: binary_classification_threshold=.5) assert result.frame['TARGET_PRED'].tolist() == [True] + + +def _export(sink, bundle, location): + return Trainer(ModelPlan())._finish_fit( + TrainingResult({}, {}, bundle), predictive_query='PREDICT COUNT(*)', + snapshot_id='snap-123', training_table_generation=_generation_plan(), + export_config=SimpleNamespace( + connector=SimpleNamespace(sink_backend=sink), path=None, + location=location)) + + +def _specced_bundle(tmp_path: Path) -> Path: + 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()) + return bundle + + +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 = _specced_bundle(tmp_path) + 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.result.id == Trainer.load(bundle).result.id + + +def test_training_result_exports_like_its_trainer(tmp_path: Path) -> None: + bundle = _specced_bundle(tmp_path) + sink = _AliasingSink() + published = TrainingResult({}, {}, bundle).export( + ModelOutputConfig(_FakeConnector(sink), + location=ModelRegistryLocation('CHURN_GNN', + alias='prod'))) + assert published.version == 'V1' and sink.aliases == {'prod': 'V1'} + assert list(sink.published) == ['CHURN_GNN@latest'] + + +def test_alias_needs_an_aliasing_sink() -> None: + connector = _FakeConnector(_FakeSink()) + with pytest.raises(NotImplementedError, match='alias'): + Trainer.load_from_connector( + connector, location=ModelRegistryLocation('m', alias='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 == {} + + +class _RegistryTracker: + def __init__(self, sink): + self.connector = _FakeConnector(sink) + + def start(self, record): + pass + + def on_trial(self, record, event): + pass + + def finish(self, record, error): + pass + + +@pytest.mark.parametrize('sink,registered', [ + (_AliasingSink(), 'V1'), + (SimpleNamespace(write_table=lambda *a: None, + publish_model=lambda *a, **k: 1 / 0), None), +]) +def test_mlops_registry_registers_after_the_fit(tmp_path: Path, monkeypatch, + sink, registered) -> None: + import warnings + + import kumo.engine as kumo + + bundle = _specced_bundle(tmp_path) + monkeypatch.setattr( + Trainer, '_fit', lambda self, *a, **k: TrainingResult( + {}, {}, bundle, bundle_id='b1')) + tracker = _RegistryTracker(sink) + registry = ModelRegistryLocation('CHURN_GNN', alias='prod') + trainer = Trainer(ModelPlan()) + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter('always') + result = trainer.fit( + SimpleNamespace(query='PREDICT COUNT(*)'), + snapshot=SimpleNamespace(id='s'), + mlops=kumo.MLOps(str(tmp_path / 'root'), tracking=tracker, + registry=registry, store=False)) + failures = [w for w in caught if 'Could not register' in str(w.message)] + assert getattr(result.registered_model, 'version', None) == registered + assert result.published_model is None + if registered: + assert trainer.result is result and not failures + assert sink.aliases == {'prod': 'V1'} + else: + assert len(failures) == 1 + assert 'ZeroDivisionError' in str(failures[0].message) + with pytest.raises(RuntimeError, match='ZeroDivisionError') as raised: + trainer.fit( + SimpleNamespace(query='PREDICT COUNT(*)'), + snapshot=SimpleNamespace(id='s'), + mlops=kumo.MLOps(str(tmp_path / 'root'), tracking=tracker, + registry=registry, store=False, + on_error='raise')) + assert isinstance(raised.value.__cause__, ZeroDivisionError) + + +def test_mlops_registry_alias_is_checked_before_the_fit(tmp_path: Path, + monkeypatch) -> None: + import kumo.engine as kumo + + monkeypatch.setattr(Trainer, '_fit', lambda *a, **k: 1 / 0) + with pytest.raises(ValueError, match='model alias'): + Trainer(ModelPlan()).fit( + SimpleNamespace(query='q'), snapshot=SimpleNamespace(id='s'), + mlops=kumo.MLOps( + str(tmp_path), tracking=_RegistryTracker(_AliasingSink()), + registry=ModelRegistryLocation('CHURN_GNN', alias='LATEST'), + store=False)) diff --git a/packages/kumo-engine/test/test_trainer_warnings.py b/packages/kumo-engine/test/test_trainer_warnings.py index af488eacb..b95b7ab75 100644 --- a/packages/kumo-engine/test/test_trainer_warnings.py +++ b/packages/kumo-engine/test/test_trainer_warnings.py @@ -8,16 +8,21 @@ from kumoapi.model_plan import ModelPlan from kumo.engine import Trainer -from kumo.engine.trainer.trainer import _without_torch_scatter_hint +from kumo.engine.trainer.trainer import _without_unactionable_warnings HINT = ("The usage of `scatter(reduce='min')` can be accelerated via the " "'torch-scatter' package, but it was not found") -def test_torch_scatter_hint_is_hidden_and_other_warnings_are_not() -> None: - @_without_torch_scatter_hint +def test_unactionable_warnings_are_hidden_and_other_warnings_are_not() -> None: + @_without_unactionable_warnings def run() -> str: warnings.warn(HINT, UserWarning) + warnings.warn('This process (pid=1) is multi-threaded, use of fork() ' + 'may lead to deadlocks in the child.', + DeprecationWarning) + warnings.warn('Your `IterableDataset` has `__len__` defined.', + UserWarning) warnings.warn('an unrelated warning', UserWarning) return 'done' diff --git a/packages/kumo-engine/test/test_warm_start.py b/packages/kumo-engine/test/test_warm_start.py index 86d274770..6d37621f2 100644 --- a/packages/kumo-engine/test/test_warm_start.py +++ b/packages/kumo-engine/test/test_warm_start.py @@ -21,7 +21,7 @@ from kumo_core.train.trainer import get_best_ckpt_path from kumo_core.types import TaskType from kumoapi.model_plan import ModelPlan -from kumoapi.plugin import ModelFetchingSink +from kumoapi.plugin import ModelFetchingSink, PublishedModel from kumo.engine import FileConnector, Graph, PredictiveQuery, Table, Trainer @@ -111,8 +111,9 @@ def test_bundle_warm_start_fit_reload_and_predict(tmp_path, monkeypatch, refit, shutil.rmtree(initial.scratch_dir) if refit == 'full': sink = Mock(spec=ModelFetchingSink) - sink.fetch_model.side_effect = lambda connector, target, **kwargs: shutil.copytree( - durable, target) + sink.fetch_model.side_effect = lambda connector, target, **kwargs: ( + shutil.copytree(durable, target) and PublishedModel( + name='m', version='1', uri=str(durable))) source = Trainer.load_from_connector(SimpleNamespace(sink_backend=sink), path='retained-model') else: diff --git a/packages/kumo-snowflake/pyproject.toml b/packages/kumo-snowflake/pyproject.toml index 960687071..2f7591943 100644 --- a/packages/kumo-snowflake/pyproject.toml +++ b/packages/kumo-snowflake/pyproject.toml @@ -27,13 +27,14 @@ 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 +# experiment tracking, >=2.0; <3 because kumo-engine reads its 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/registry.py b/packages/kumo-snowflake/src/kumo_snowflake/registry.py index fd220653e..32d23e589 100644 --- a/packages/kumo-snowflake/src/kumo_snowflake/registry.py +++ b/packages/kumo-snowflake/src/kumo_snowflake/registry.py @@ -33,11 +33,15 @@ """ import json +import logging +import math +import re import shutil from contextlib import contextmanager from pathlib import Path from typing import ( TYPE_CHECKING, + Any, ContextManager, Dict, Iterator, @@ -53,9 +57,13 @@ resolved, ) from kumoapi.plugin import ( + BUNDLE_ID_KEY, + ENGINE_MARKER, ModelRegistryLocation, PublishedModel, SourceConnector, + check_model_alias, + registry_lineage, ) if TYPE_CHECKING: @@ -74,14 +82,19 @@ #: 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'}) +_IDENTIFIER = r'(?:[A-Za-z_][A-Za-z0-9_$]*|"(?:[^"]|"")+")' +_MODEL_NAME = re.compile(rf'(?:{_IDENTIFIER}\.){{0,2}}{_IDENTIFIER}') + +_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 #: a reader can look at. _EXPORT_DIR = '_registry-export' +_METRICS_LINE = 'Holdout: ' + def _require_registry() -> type["Registry"]: r"""Returns the ``Registry`` class, or explains how to install it.""" @@ -113,7 +126,8 @@ class _BundleModel(custom_model.CustomModel): Deliberately not runnable. Scoring a Kumo model needs a materialized DiskGraph task over the graph being scored, built where the data is, so this refuses rather than answering. The version exists to carry, - version and expose the bundle; :func:`kumoai_studio.import_model` runs it. + version and expose the bundle; load it with + ``Trainer.load_from_connector``. """ # Unquoted, and evaluated here rather than deferred: the registry # compares this annotation against `pandas.DataFrame` itself, so a @@ -121,14 +135,32 @@ class _BundleModel(custom_model.CustomModel): @custom_model.inference_api def predict(self, df: pd.DataFrame) -> pd.DataFrame: raise NotImplementedError( - "This model version carries a Kumo model bundle; it does not " - "score in Snowflake. Import it with kumoai_studio.import_model() and " - "predict against a graph.") + "This model version holds a Kumo bundle and does not score in " + "Snowflake; load it with Trainer.load_from_connector.") return _BundleModel( custom_model.ModelContext(artifacts={'bundle': bundle_dir})) +_TASKS = { + 'binary_classification': 'TABULAR_BINARY_CLASSIFICATION', + 'multiclass_classification': 'TABULAR_MULTI_CLASSIFICATION', + 'regression': 'TABULAR_REGRESSION', +} + + +def _task(bundle: Path) -> Dict[str, Any]: + r"""``log_model``'s ``task=`` for the bundle, if Snowflake has one.""" + try: + from snowflake.ml.model import type_hints + + schema = json.loads((bundle / 'prediction-schema.json').read_text()) + name = _TASKS.get(schema.get('task_type')) + return {} if name is None else {'task': getattr(type_hints.Task, name)} + except Exception: + return {} + + def _carrier_signature() -> Dict[str, "model_signature.ModelSignature"]: r"""Declares the carrier's signature instead of inferring one. @@ -154,6 +186,14 @@ def _carrier_signature() -> Dict[str, "model_signature.ModelSignature"]: } +def _model_args(bundle: Path) -> Dict[str, Any]: + r"""``log_model`` arguments for the bundle's carrier model.""" + return { + 'model': _bundle_model(str(bundle)), + 'signatures': _carrier_signature() + } + + def _locate_bundle(root: Path) -> Path: r"""Finds the published bundle inside a downloaded model tree. @@ -209,6 +249,64 @@ 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 _finite(metrics: Mapping[str, float]) -> Dict[str, float]: + return {name: float(value) for name, value in metrics.items() + if math.isfinite(value)} + + +def _comment(user: Optional[str], + metrics: Mapping[str, float]) -> Optional[str]: + r"""The user's comment, then a line of the model's holdout metrics.""" + line = ' · '.join(f'{name} {value:.3g}' + for name, value in _finite(metrics).items()) + return '\n'.join(text for text in (user, line and _METRICS_LINE + line) + if text) or None + + +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) + + +def _next_version(registry: Any, model_name: str) -> str: + try: + names = [str(version.version_name) + for version in registry.get_model(model_name).versions()] + except ValueError: + names = [] + numbers = [int(name[1:]) for name in names if re.fullmatch(r'V\d+', name)] + return f'V{max(numbers, default=0) + 1}' + + class SnowflakeModelRegistryMixin: r"""The model half of :class:`SnowflakeSinkBackend`. @@ -229,6 +327,17 @@ def _session( ) -> ContextManager[SnowflakeSession]: ... + def check_model_location(self, connector: SourceConnector, + location: ModelRegistryLocation) -> None: + r"""Raises ``ValueError`` unless the model name is an identifier.""" + del connector + if not _MODEL_NAME.fullmatch(location.model_name): + raise ValueError( + f'Snowflake cannot register a model named ' + f'{location.model_name!r}: start with a letter or underscore ' + f'and use letters, digits, underscores and $, or double-quote ' + f"the name, e.g. 'USERS_CHURN'.") + @contextmanager def _registry(self, connector: SourceConnector) -> Iterator["Registry"]: r"""A registry statement is a session statement, so the caller keeps @@ -264,13 +373,10 @@ def publish_model( path: Refused. A registry addresses a model by name. location: The :class:`~kumoapi.plugin.ModelRegistryLocation` naming - 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. - metrics: Returned with the publication result. Training metrics - remain in the bundle's metrics.json. + the model, and optionally the version; by default the next + ``V``. + metadata: Returned with the result; its bundle id is recorded. + metrics: The model's holdout metrics, recorded on the version. project_id: Ignored; a local runtime has no projects. Returns: @@ -288,15 +394,27 @@ 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 = {**_finite(metrics), **_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'}, - signatures=_carrier_signature(), - ) + args = _model_args(bundle) + name = location.version or _next_version(registry, + location.model_name) + for attempt in range(3): + try: + version = registry.log_model( + model_name=location.model_name, + version_name=name, + comment=_comment(location.comment, metrics), + metrics=record, + **args, + **_task(bundle), + ) + break + except Exception: + taken, name = name, location.version or _next_version( + registry, location.model_name) + if name == taken or attempt == 2: + raise return _published(connector, version, metadata=metadata, metrics=metrics) @@ -308,7 +426,7 @@ 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. @@ -319,9 +437,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 +465,13 @@ 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 a version or alias (else the + registry's default version). 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 and the recorded bundle id. """ del project_id _refuse_path(path, 'fetching') @@ -366,8 +485,17 @@ 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)) + if location.alias is not None: + if resolved(location.alias) not in _alias_holders(model): + raise ValueError( + f'No version of {location.model_name} holds alias ' + f'{location.alias}.') + version = model.version(location.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 +513,48 @@ 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.""" + 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/tests/test_registry.py b/packages/kumo-snowflake/tests/test_registry.py index 9d07e6c28..ac856e203 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,12 +81,21 @@ 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): +def _bundle(tmp_path): bundle = tmp_path / 'bundle' bundle.mkdir() (bundle / 'training-spec.json').write_text('{}') + return bundle + + +@pytest.mark.parametrize('metadata', [{}, { + 'predictive_query': 'x' * 300, + 'kumo_bundle_id': 'b1', + 'kumo_engine': '0' +}]) +def test_publish_records_pointers_and_holdout_metrics_in_catalog( + tmp_path, monkeypatch, metadata): + bundle = _bundle(tmp_path) calls = [] version = SimpleNamespace(model_name='CHURN', version_name='V1', fully_qualified_model_name='DB.SCHEMA.CHURN') @@ -89,26 +104,79 @@ def log_model(*args, **kwargs): calls.append(kwargs) return version + def get_model(_): + raise ValueError('Unable to find model CHURN') + @contextmanager def registry(_): - yield SimpleNamespace(log_model=log_model) + yield SimpleNamespace(log_model=log_model, get_model=get_model) sink = SnowflakeModelRegistryMixin() monkeypatch.setattr(sink, '_registry', registry) + monkeypatch.setattr(registry_module, '_model_args', lambda *_: {}) + metrics = {'auroc': 0.9213, 'holdout_loss': 0.5, 'brier': float('nan')} + published = sink.publish_model( + SimpleNamespace(), str(bundle), + location=ModelRegistryLocation('CHURN', comment='Weekly retrain.'), + metadata=metadata, metrics=metrics) + assert calls[0]['metrics'] == { + 'auroc': 0.9213, 'holdout_loss': 0.5, + **({'kumo_bundle_id': 'b1'} if metadata else {}), 'kumo_engine': '1' + } + assert calls[0]['version_name'] == 'V1' + assert calls[0]['comment'] == ('Weekly retrain.\n' + 'Holdout: auroc 0.921 · holdout_loss 0.5') + assert published.metadata == metadata + assert published.metrics == metrics + + +def test_publish_names_the_next_version_and_retries_a_taken_one( + tmp_path, monkeypatch): + bundle = _bundle(tmp_path) + names = ['V1', 'SLIMY_DONKEY_1', 'V3'] + tried = [] + + def log_model(*args, version_name, **kwargs): + tried.append(version_name) + if version_name == 'V4': + names.append('V4') + raise ValueError('Model CHURN version V4 already existed.') + return SimpleNamespace(model_name='CHURN', version_name=version_name, + fully_qualified_model_name='DB.SCHEMA.CHURN') + + model = SimpleNamespace(versions=lambda: [ + SimpleNamespace(version_name=name) for name in names]) + sink = _sink(monkeypatch, log_model=log_model, get_model=lambda _: model) monkeypatch.setattr(registry_module, '_bundle_model', lambda _: object()) 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'} - assert published.metadata == metadata - assert published.metrics == {'auc': 0.92} + location=ModelRegistryLocation('CHURN')) + assert tried == ['V4', 'V5'] + assert published.version == 'V5' + with pytest.raises(ValueError, match='already existed'): + sink.publish_model(SimpleNamespace(), str(bundle), + location=ModelRegistryLocation('CHURN', + version='V4')) + def fail_after_creating(*args, version_name, **kwargs): + tried.append(version_name) + names.append(version_name) + raise RuntimeError('upload failed') -def test_listing_filters_engine_marker_and_returns_empty_metrics(monkeypatch): + tried.clear() + sink = _sink(monkeypatch, log_model=fail_after_creating, + get_model=lambda _: model) + with pytest.raises(RuntimeError, match='upload failed'): + sink.publish_model(SimpleNamespace(), str(bundle), + location=ModelRegistryLocation('CHURN')) + assert len(tried) == 3 + + +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 +202,158 @@ 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 _sink(monkeypatch, **registry_methods): + @contextmanager + def registry(_): + yield SimpleNamespace(**registry_methods) + + 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 = _sink(monkeypatch, get_model=lambda _: 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): + model = _Model({'V1': ['DEFAULT'], 'V2': ['PROD']}) + sink = _sink(monkeypatch, get_model=lambda _: model) + 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 = _sink(monkeypatch, get_model=lambda _: model) + with pytest.raises(RuntimeError, match='lost connection'): + sink.set_model_alias(SimpleNamespace(), + ModelRegistryLocation('CHURN', version='V2'), + 'prod') + assert model.aliases == {'V1': ['PROD'], 'V2': []} + + +@pytest.mark.parametrize('task, expected', [ + ('binary_classification', 'TABULAR_BINARY_CLASSIFICATION'), + ('link_prediction', None), +]) +def test_publish_records_the_task_type_monitors_need(tmp_path, monkeypatch, + task, expected): + import sys + + bundle = _bundle(tmp_path) + (bundle / 'prediction-schema.json').write_text( + json.dumps({'task_type': task})) + tasks = SimpleNamespace( + TABULAR_BINARY_CLASSIFICATION='TABULAR_BINARY_CLASSIFICATION') + monkeypatch.setitem(sys.modules, 'snowflake.ml.model', + SimpleNamespace(type_hints=SimpleNamespace(Task=tasks))) + calls = [] + + @contextmanager + def registry(_): + yield SimpleNamespace(log_model=lambda *a, **kw: calls.append(kw) or + SimpleNamespace( + model_name='CHURN', version_name='V1', + fully_qualified_model_name='DB.S.CHURN')) + + sink = SnowflakeModelRegistryMixin() + monkeypatch.setattr(sink, '_registry', registry) + monkeypatch.setattr(registry_module, '_model_args', lambda *_: {}) + sink.publish_model(SimpleNamespace(), str(bundle), + location=ModelRegistryLocation('CHURN', version='V1')) + assert calls[0].get('task') == expected + + +@pytest.mark.parametrize('name', ['USERS_CHURN', 'users_churn$2', '_m', + 'ML.USERS_CHURN', 'DB.ML.USERS_CHURN', + '"users-churn"', '"a ""b"""']) +def test_snowflake_identifiers_are_accepted_as_model_names(name): + SnowflakeModelRegistryMixin().check_model_location( + SimpleNamespace(), ModelRegistryLocation(name)) + + +@pytest.mark.parametrize('name', ['users-churn', '1churn', 'users churn', + 'A.B.C.D', 'ML.', '""']) +def test_other_model_names_are_refused_with_the_rule(name): + with pytest.raises(ValueError, match="letter or underscore.*USERS_CHURN"): + SnowflakeModelRegistryMixin().check_model_location( + SimpleNamespace(), ModelRegistryLocation(name)) diff --git a/setup.py b/setup.py index a8446e82b..f80ad52d9 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.6", + # Full mlflow, for its local SQLite store. + "mlflow>=3.13,<4", ] dev_requires = test_requires + [ diff --git a/uv.lock b/uv.lock index 50911245a..331bcd44c 100644 --- a/uv.lock +++ b/uv.lock @@ -141,7 +141,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" } @@ -305,7 +305,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 = [ @@ -377,7 +377,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 = [ @@ -573,12 +573,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 = [ @@ -606,7 +615,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 = [ @@ -658,8 +667,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 = [ @@ -680,7 +689,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 = [ @@ -961,7 +970,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 = [ @@ -982,7 +991,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 = [ @@ -1087,7 +1096,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 = [ @@ -1278,7 +1287,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]] @@ -1286,7 +1295,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 = [ @@ -1342,7 +1351,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" }, @@ -1391,34 +1400,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]] @@ -1437,43 +1446,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]] @@ -1595,6 +1604,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" @@ -1718,6 +1741,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" @@ -1927,6 +1980,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" @@ -2152,6 +2229,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" @@ -2327,6 +2440,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" @@ -2377,6 +2499,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" @@ -2474,6 +2605,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" @@ -2752,8 +2892,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" }, @@ -2765,11 +2905,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" }, @@ -2789,7 +2929,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 = [ @@ -2806,6 +2946,7 @@ dev = [ { name = "flake8" }, { name = "isort" }, { name = "kubernetes-stubs" }, + { name = "mlflow" }, { name = "moto" }, { name = "packaging" }, { name = "pick" }, @@ -2884,6 +3025,7 @@ rest = [ test = [ { name = "category-encoders" }, { name = "coverage" }, + { name = "mlflow" }, { name = "moto" }, { name = "pytest" }, { name = "pytest-asyncio" }, @@ -2907,7 +3049,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')" }, ] @@ -2989,6 +3131,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" }, @@ -3247,11 +3391,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" }, @@ -3384,9 +3528,19 @@ databricks = [ gcs = [ { name = "kumo-gcs" }, ] +mlflow = [ + { name = "mlflow-skinny" }, + { name = "nvidia-ml-py" }, +] s3 = [ { name = "kumo-s3" }, ] +sagemaker = [ + { name = "kumo-s3" }, + { name = "mlflow-skinny" }, + { name = "nvidia-ml-py" }, + { name = "sagemaker-mlflow" }, +] snowflake = [ { name = "kumo-snowflake", extra = ["sdk"] }, ] @@ -3397,11 +3551,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" }, @@ -3421,20 +3575,24 @@ requires-dist = [ { name = "kumo-bigquery", marker = "extra == 'bigquery'", editable = "packages/kumo-bigquery" }, { name = "kumo-core", editable = "packages/kumo-core" }, { name = "kumo-databricks", marker = "extra == 'databricks'", editable = "packages/kumo-databricks" }, + { name = "kumo-engine", extras = ["mlflow", "s3"], marker = "extra == 'sagemaker'" }, { name = "kumo-gcs", marker = "extra == 'gcs'", editable = "packages/kumo-gcs" }, { 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-skinny", marker = "extra == 'mlflow'", specifier = ">=3.5,<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" }, { name = "pytest", marker = "extra == 'test'" }, { name = "pytest-mock", marker = "extra == 'test'" }, + { name = "sagemaker-mlflow", marker = "extra == 'sagemaker'", specifier = ">=0.5.0" }, { name = "torch", marker = "extra == 'test'", specifier = ">=2.9,<2.13" }, { name = "typing-extensions", specifier = ">=4.5.0" }, ] -provides-extras = ["cpu", "test", "snowflake", "databricks", "bigquery", "s3", "text", "adls", "gcs"] +provides-extras = ["cpu", "test", "snowflake", "databricks", "bigquery", "s3", "text", "adls", "gcs", "mlflow", "sagemaker"] [[package]] name = "kumo-gcs" @@ -3580,7 +3738,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"] @@ -3815,8 +3973,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" }, @@ -3903,6 +4061,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" @@ -4130,7 +4371,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 = [ @@ -4328,7 +4569,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" }, @@ -4471,7 +4712,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" }, @@ -4509,7 +4750,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" }, @@ -4621,9 +4862,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" }, @@ -4661,7 +4902,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" }, @@ -4679,6 +4920,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" @@ -5132,7 +5382,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 = [ @@ -5299,6 +5549,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" @@ -5812,7 +6074,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 = [ @@ -5845,7 +6107,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" }, @@ -5862,7 +6124,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 = [ @@ -5965,6 +6227,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" @@ -5996,11 +6267,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" } @@ -6017,11 +6288,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" }, @@ -6040,6 +6311,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" @@ -6248,6 +6541,19 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/03/19/90d7d4ed51932c022d53f1d02d564b62d10e272692a1f9b76425c1ad2a02/s3transfer-0.16.1-py3-none-any.whl", hash = "sha256:61bcd00ccb83b21a0fe7e91a553fff9729d46c83b4e0106e7c314a733891f7c2", size = 86825, upload-time = "2026-04-22T20:36:04.992Z" }, ] +[[package]] +name = "sagemaker-mlflow" +version = "0.5.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "boto3" }, + { name = "mlflow-skinny" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/d5/e6/c147424bb1b1d16963037940949533dc893b511cf5919ecbc82388dfcf0d/sagemaker_mlflow-0.5.0.tar.gz", hash = "sha256:6786b9aa738bf0ec551dc99bd44069ac39ad202fa234f99b2fc914556f1af2f2", size = 22371, upload-time = "2026-06-27T00:06:19.285Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/69/50/4d0a43a18e0596f2fd1a506a4467577c45d2697b245cb5f61e8704209c29/sagemaker_mlflow-0.5.0-py3-none-any.whl", hash = "sha256:f55dc39c6ed345d4359c7d5cddd99fb59c852609c463a397d18e8bbaba672dac", size = 31159, upload-time = "2026-06-27T00:06:18.253Z" }, +] + [[package]] name = "scikit-learn" version = "1.9.1" @@ -6256,8 +6562,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" } @@ -6320,7 +6626,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 = [ @@ -6414,7 +6720,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 = [ @@ -6525,6 +6831,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" @@ -6534,6 +6857,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" @@ -6631,8 +6963,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" }, @@ -6689,7 +7021,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" }, @@ -6843,7 +7175,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" } @@ -6910,7 +7242,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 = [ @@ -7116,14 +7448,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" }, @@ -7200,14 +7532,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" }, @@ -7250,21 +7582,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 = [ @@ -7299,21 +7631,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" }, @@ -7364,11 +7696,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 = [ @@ -7380,7 +7712,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 = [ @@ -7448,7 +7780,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 = [ @@ -7465,7 +7797,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 = [ @@ -7533,7 +7865,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 = [ @@ -7577,6 +7909,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" @@ -7621,6 +7962,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"