diff --git a/.github/workflows/testing.yml b/.github/workflows/testing.yml index 179decec6..1c59eb78b 100644 --- a/.github/workflows/testing.yml +++ b/.github/workflows/testing.yml @@ -720,6 +720,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..0d4b1cfcf --- /dev/null +++ b/examples/mlops/smoke.py @@ -0,0 +1,133 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Manual live check, not run in CI: fit a small model on synthetic data on +this machine, record it in Databricks, SageMaker or Snowflake with an +explicit identity, register it, and load it back. A few minutes on a CPU. +Install kumo-engine[mlflow,databricks], kumo-engine[mlflow,sagemaker] or +kumo-engine[snowflake]. + + python smoke.py databricks --host https:// --client-id \ + --catalog main --schema ml --warehouse --experiment /Users//kumo-smoke + python smoke.py sagemaker --arn --role-arn + python smoke.py snowflake --database ML --schema RUNS --warehouse \ + --account --user --private-key-file key.p8 + +Secrets come from DATABRICKS_CLIENT_SECRET (Databricks) or the key file +(Snowflake); --role-arn may be left out to use the ambient AWS identity, +with no MLFLOW_TRACKING_URI needed. +""" + +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 b649e302b..018755408 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,14 @@ from kumoapi.typing import Dtype __all__ = [ + 'BUNDLE_ID_KEY', + 'ENGINE_MARKER', + 'REGISTRY_KEY_PREFIX', 'ArtifactLocation', 'ArtifactSink', 'FileStoreSink', 'IngestResult', + 'ModelAliasingSink', 'ModelFetchingSink', 'ModelListingSink', 'ModelPublishingSink', @@ -56,8 +62,20 @@ 'SourceBackend', 'SourceConnector', 'SourceCredentials', + 'check_model_alias', + 'registry_lineage', ] +#: Prefix of the keys an engine export records in registry metadata. +REGISTRY_KEY_PREFIX = 'kumo_' +#: Registry metadata key marking a version as a Kumo engine bundle. +ENGINE_MARKER = 'kumo_engine' +#: Registry metadata key holding the version's bundle id; the bundle's own +#: provenance.json holds its lineage. +BUNDLE_ID_KEY = 'kumo_bundle_id' +_ALIAS = re.compile(r'[A-Za-z_][A-Za-z0-9_]{0,254}') +_RESERVED_ALIAS = re.compile(r'DEFAULT|FIRST|LAST|LATEST|V[0-9]+', re.I) + @dataclass(frozen=True) class SourceCredentials: @@ -149,13 +167,16 @@ class ModelRegistryLocation: Args: model_name: Registry-local model name. version: Exact version to write or read. ``None`` lets the registry - assign the next one on export, and selects the latest on import. - An assigned version comes back on the export result. + assign the next one on export; on import it selects Snowflake's + default or Databricks' latest Kumo version. comment: Optional description recorded on the version. + alias: Set on the new version after export; on import, the version + it points at. Exclusive with :obj:`version`. """ model_name: str version: Optional[str] = None comment: Optional[str] = None + alias: Optional[str] = None def __post_init__(self) -> None: if (not isinstance(self.model_name, str) or not self.model_name @@ -166,6 +187,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) @@ -705,7 +730,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: @@ -715,6 +741,44 @@ def fetch_model( ... +@runtime_checkable +class ModelAliasingSink(ArtifactSink, Protocol): + r"""A registry sink that sets aliases and fetches by + :attr:`ModelRegistryLocation.alias`.""" + def set_model_alias( + self, + connector: SourceConnector, + location: ModelRegistryLocation, + alias: str, + ) -> None: + r"""Points :obj:`alias` at ``location.version``, moving it if held.""" + ... + + +def check_model_alias(alias: str) -> None: + r"""Raises :class:`ValueError` unless every model registry can set + :obj:`alias`: a letter or underscore, then letters, digits or underscores + (at most 255), and not ``DEFAULT``, ``FIRST``, ``LAST``, ``LATEST`` or + ``v`` in any case.""" + if (not isinstance(alias, str) or not _ALIAS.fullmatch(alias) + or _RESERVED_ALIAS.fullmatch(alias)): + raise ValueError( + f'{alias!r} cannot be set as a model alias: use letters, digits ' + f'and underscores, starting with a letter or underscore, and ' + f'avoid the reserved DEFAULT, FIRST, LAST, LATEST and v ' + f'(for DEFAULT use ALTER MODEL ... SET DEFAULT_VERSION).') + + +def registry_lineage(recorded: Mapping[str, Any]) -> Dict[str, str]: + r"""The ``kumo_*`` entries of a version's registry metadata, such as + :data:`BUNDLE_ID_KEY`, as strings and without :data:`ENGINE_MARKER`.""" + return { + key: str(value) + for key, value in recorded.items() + if key.startswith(REGISTRY_KEY_PREFIX) and key != ENGINE_MARKER + } + + @runtime_checkable class ModelListingSink(ArtifactSink, Protocol): r"""A sink that can enumerate the models published through it. diff --git a/packages/kumo-api/test/test_plugin_registry.py b/packages/kumo-api/test/test_plugin_registry.py new file mode 100644 index 000000000..52ca7a69f --- /dev/null +++ b/packages/kumo-api/test/test_plugin_registry.py @@ -0,0 +1,33 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +import pytest + +from kumoapi.plugin import check_model_alias, registry_lineage + + +@pytest.mark.parametrize('alias', [ + '', ' ', 'default', 'LAST', 'Latest', 'v3', 'my-alias', '2nd', 'a' * 256, + None +]) +def test_check_model_alias_refuses_what_a_registry_reserves(alias): + with pytest.raises(ValueError): + check_model_alias(alias) + + +def test_check_model_alias_accepts_a_plain_name(): + check_model_alias('champion') + check_model_alias('_v2_prod') + + +def test_registry_lineage_keeps_kumo_keys_as_strings(): + recorded = { + 'kumo_engine': '1', + 'kumo_bundle_id': 'b', + 'auc': 0.9, + 'kumo_seconds': 3 + } + assert registry_lineage(recorded) == { + 'kumo_bundle_id': 'b', + 'kumo_seconds': '3' + } diff --git a/packages/kumo-core/src/kumo_core/artifact_export.py b/packages/kumo-core/src/kumo_core/artifact_export.py index 311b972bc..e438e3e79 100644 --- a/packages/kumo-core/src/kumo_core/artifact_export.py +++ b/packages/kumo-core/src/kumo_core/artifact_export.py @@ -304,10 +304,16 @@ 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() + # In Arrow: a pandas round trip turns an all-null bool column + # into Arrow's null type. + 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.array([value] * table.num_rows, + type=pa.scalar(value).type) + 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 10b33ee11..d9f57fa41 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 + #: Every engine artifact's manifest file; underscore-prefixed so Parquet #: readers skip it. PROVENANCE_FILE = '_provenance.json' @@ -382,6 +386,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' @@ -479,6 +503,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..f5a23de69 100644 --- a/packages/kumo-core/src/kumo_core/automl/optuna_searcher.py +++ b/packages/kumo-core/src/kumo_core/automl/optuna_searcher.py @@ -66,6 +66,7 @@ def suggest(self) -> PlanT: for _ in range(self.fallback_sampler.max_attempts): self.current_trial = self.study.ask() cfg = self.fallback_sampler.sample_instance() + self.sampled_params = self.fallback_sampler.sampled_params if cfg not in [prev_cfg for prev_cfg, _ in self.state]: return cfg return cfg diff --git a/packages/kumo-core/src/kumo_core/automl/runtime.py b/packages/kumo-core/src/kumo_core/automl/runtime.py index 209680edf..880246456 100644 --- a/packages/kumo-core/src/kumo_core/automl/runtime.py +++ b/packages/kumo-core/src/kumo_core/automl/runtime.py @@ -196,9 +196,10 @@ def resolve_warm_start( encoders; they do not retain hosted AutoML directory layouts. """ if artifacts is not None: - source_schema = load_json(artifacts["prediction_schema"], dict) + from kumo_core.artifacts import bundle_task_type + validate_warm_start_task_type(task.task_type, - TaskType(source_schema["task_type"])) + bundle_task_type(artifacts)) checkpoint = artifacts["selected_checkpoint"] plan_path = artifacts["resolved_model_plan"] encoder_dir = posixpath.dirname(artifacts["encoders"]) diff --git a/packages/kumo-core/src/kumo_core/automl/searcher.py b/packages/kumo-core/src/kumo_core/automl/searcher.py index f8810e7b6..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..abdbd7803 100644 --- a/packages/kumo-core/src/kumo_core/callbacks/train_progress_logger.py +++ b/packages/kumo-core/src/kumo_core/callbacks/train_progress_logger.py @@ -5,7 +5,7 @@ import logging import os.path as osp from datetime import datetime, timezone -from typing import Any +from typing import Any, Callable import kumoapi.jobs as jobs_api from kumo_core.io.json import load_json, save_json @@ -140,17 +140,22 @@ class AutoTrainingProgressLogger(MonitoringCallback): # pragma: no cover # Constant to account for overhead not captured in trial estimation: OVERHEAD_CONSTANT = 1.30 - def __init__(self, log_dir: str, num_experiments: int): + def __init__(self, log_dir: str, num_experiments: int, + on_event: Callable[..., None] | None = None): super().__init__() self.path = osp.join(log_dir, self.FILE_NAME) + self.on_event = on_event + self.trial: int | None = None self.progress = load_json(self.path, AutoTrainerProgress, raise_on_missing=False) if self.progress is None: self.progress = AutoTrainerProgress(total_trials=num_experiments) - def set_current_trainer(self, trainer: Any | None): + def set_current_trainer(self, trainer: Any | None, + trial: int | None = None): self.trainer = trainer + self.trial = trial def on_train_start(self, pl_trainer, pl_module): if self.progress.start_time is None: @@ -184,6 +189,20 @@ def on_train_epoch_end(self, pl_trainer, pl_module): save_json(self.progress, self.path) + if self.on_event is not None and self.trial is not None: + try: + progress = self.trainer.training_progress_logger.progress + epoch = progress.num_epochs_completed + metrics = progress.metrics[epoch - 1] + self.on_event('epoch_ended', trial=self.trial, epoch=epoch, + metrics={ + **metrics.train_metrics, + **metrics.validation_metrics + }) + except Exception: + logger.warning('Could not report epoch_ended', + exc_info=True) + def on_validation_end(self, pl_trainer, pl_module): if pl_trainer.state.fn == TrainerFn.VALIDATING: save_json(self.progress, self.path) diff --git a/packages/kumo-core/src/kumo_core/diskgraph/lightning_data.py b/packages/kumo-core/src/kumo_core/diskgraph/lightning_data.py index 3c8e30cec..aa28e8e73 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 @@ -235,6 +236,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: + # Sandboxes such as SageMaker hosting refuse pidfd_open. + if error.errno not in (errno.ENOSYS, errno.EPERM): + os.close(sentinel_fd) + raise except BaseException: try: os.close(sentinel_fd) @@ -261,9 +267,8 @@ 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: + # Without a pidfd the live sentinel above guards the pid. 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 effb6b69e..14f18de0d 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/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..0d5585d73 100644 --- a/packages/kumo-core/src/kumo_core/train/formatting.py +++ b/packages/kumo-core/src/kumo_core/train/formatting.py @@ -31,6 +31,15 @@ class LinkPredictionTable(Protocol): _ANCHOR_TIME_COL = "ANCHOR_TIMESTAMP" +_CLASS_TASKS = frozenset({ + TaskType.MULTICLASS_CLASSIFICATION, + TaskType.MULTILABEL_CLASSIFICATION, +}) +_LONG_FORMAT_TASKS = _CLASS_TASKS | { + TaskType.MULTILABEL_RANKING, + TaskType.TEMPORAL_LINK_PREDICTION, + TaskType.STATIC_LINK_PREDICTION, +} def _get_col_names( @@ -67,6 +76,23 @@ def _get_col_names( return result +def prediction_output_columns(task_type: TaskType) -> list[str]: + r"""The prediction columns a scored row carries after its entity column, + for a target named ``TARGET`` and a time column named ``TIMESTAMP``.""" + if task_type in _CLASS_TASKS: + return ['PREDICTED', 'CLASS', 'SCORE'] + if task_type.is_ranking: + return ['CLASS', 'SCORE'] + if task_type == TaskType.FORECASTING: + return ['TIMESTAMP', 'TARGET_PRED'] + if task_type == TaskType.BINARY_CLASSIFICATION: + return [ + 'TARGET_PRED', + *_get_col_names(pd.Index([]), task_type, is_wide_format=True) + ] + return ['TARGET_PRED'] + + def _format_prediction_sanity_checks( target_encoder: TargetEncoder, task_type: TaskType, @@ -284,13 +310,7 @@ def format_prediction( 'metadata. Rename the time column before forecasting.') format: Literal['wide', 'long'] - if task_type in { - TaskType.MULTICLASS_CLASSIFICATION, - TaskType.MULTILABEL_CLASSIFICATION, - TaskType.MULTILABEL_RANKING, - TaskType.TEMPORAL_LINK_PREDICTION, - TaskType.STATIC_LINK_PREDICTION, - }: + if task_type in _LONG_FORMAT_TASKS: format = 'long' else: assert top_k is None, f"`top_k` is not supported for {task_type}." diff --git a/packages/kumo-core/src/kumo_core/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..e74b28ff2 100644 --- a/packages/kumo-core/tests/test_artifact_export.py +++ b/packages/kumo-core/tests/test_artifact_export.py @@ -138,6 +138,21 @@ 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())) + + @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_prediction_output_columns.py b/packages/kumo-core/tests/test_prediction_output_columns.py new file mode 100644 index 000000000..eee36d837 --- /dev/null +++ b/packages/kumo-core/tests/test_prediction_output_columns.py @@ -0,0 +1,23 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +from __future__ import annotations + +import pytest +from kumo_core.train.formatting import prediction_output_columns +from kumo_core.types import TaskType + + +@pytest.mark.parametrize('task_type, columns', [ + (TaskType.BINARY_CLASSIFICATION, + ['TARGET_PRED', 'False_PROB', 'True_PROB']), + (TaskType.REGRESSION, ['TARGET_PRED']), + (TaskType.FORECASTING, ['TIMESTAMP', 'TARGET_PRED']), + (TaskType.MULTICLASS_CLASSIFICATION, ['PREDICTED', 'CLASS', 'SCORE']), + (TaskType.MULTILABEL_CLASSIFICATION, ['PREDICTED', 'CLASS', 'SCORE']), + (TaskType.MULTILABEL_RANKING, ['CLASS', 'SCORE']), + (TaskType.TEMPORAL_LINK_PREDICTION, ['CLASS', 'SCORE']), + (TaskType.STATIC_LINK_PREDICTION, ['CLASS', 'SCORE']), +]) +def test_prediction_output_columns(task_type, columns) -> None: + assert prediction_output_columns(task_type) == columns diff --git a/packages/kumo-core/tests/test_searched_params.py b/packages/kumo-core/tests/test_searched_params.py new file mode 100644 index 000000000..c59c8390a --- /dev/null +++ b/packages/kumo-core/tests/test_searched_params.py @@ -0,0 +1,38 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import json + +from kumo_core.automl.optuna_searcher import TPESearcher +from kumo_core.model_plan import ModelPlan + + +def test_searcher_records_only_varied_fields_as_json(): + space = ModelPlan() + space.training_job.tune_metric = 'mae' + space.optimization.base_lr = [0.1, 0.01] + space.optimization.batch_size = [512] + space.optimization.optimizer = ['Adam', 'AdamW'] + space.neighbor_sampling.num_neighbors = [[16, 8], [8, 4]] + space.model_architecture.channels = [64, 128] + searcher = TPESearcher(space, seed=0) + + def chosen(cfg): + return { + 'optimization.base_lr': cfg.optimization.base_lr[0], + 'optimization.optimizer': cfg.optimization.optimizer[0], + 'model_architecture.channels': cfg.model_architecture.channels[0], + 'neighbor_sampling.num_neighbors': + str(cfg.neighbor_sampling.num_neighbors[0]), + } + + cfg = searcher.suggest() + assert searcher.sampled_params == chosen(cfg) + json.dumps(searcher.sampled_params) + + searcher.max_attempts = 0 + cfg = searcher.suggest() + assert searcher.sampled_params == chosen(cfg) + assert (searcher.sampled_params + is searcher.fallback_sampler.sampled_params) diff --git a/packages/kumo-core/tests/test_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 d2d88de09..d51e45474 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__) @@ -41,6 +45,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 at ' + f'most 255 characters with no space, "/", "." or control ' + f"character, e.g. 'users_churn'.") return f'{config.catalog}.{schema}.{name}' @@ -56,6 +65,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], @@ -99,6 +119,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.') @@ -133,7 +167,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) @@ -200,7 +242,7 @@ def _bundle_metadata(bundle: Path) -> tuple[dict[str, str], dict[str, float]]: def _download(client: WorkspaceClient, remote: str, destination: Path) -> None: for entry in client.files.list_directory_contents(remote): # A server-provided path must stay an immediate child of this directory. - path = entry.path or '' + path = (entry.path or '').rstrip('/') prefix = remote + '/' child = path[len(prefix):] if path.startswith(prefix) else '' if not child or '/' in child or child in ('.', '..') or '\\' in child: @@ -219,13 +261,18 @@ def _download(client: WorkspaceClient, remote: str, destination: Path) -> None: def fetch_bundle(client: WorkspaceClient, name: str, local_dir: Path, *, - version: Optional[str] = None) -> PublishedModel: + version: Optional[str] = None, + alias: Optional[str] = None) -> PublishedModel: from kumo_core.artifacts import verify_model_bundle - if version is None: + if alias is not None: + aliased = cast(dict[str, Any], + client.api_client.do('GET', _alias_path(name, alias))) + model_version = _get_version(client, name, aliased['version']) + elif version is None: versions = [ v for v in _versions(client, name) if v.get('status') == 'READY' - and _tags(v).get(_ENGINE_MARKER) == '1' + and _tags(v).get(ENGINE_MARKER) == '1' ] if not versions: raise ValueError('No ready Kumo model version in this registry.') @@ -233,7 +280,7 @@ def fetch_bundle(client: WorkspaceClient, name: str, local_dir: Path, *, else: model_version = _get_version(client, name, version) if model_version.get('status') != 'READY' or _tags(model_version).get( - _ENGINE_MARKER) != '1': + ENGINE_MARKER) != '1': raise ValueError('This version is not a ready Kumo model bundle.') destination = local_dir destination.parent.mkdir(parents=True, exist_ok=True) @@ -251,9 +298,12 @@ def fetch_bundle(client: WorkspaceClient, name: str, local_dir: Path, *, 'The registered model must contain one Kumo bundle.') bundle = bundles[0] verify_model_bundle(bundle) - published = _published(name, model_version) metadata, metrics = _bundle_metadata(bundle) - published = replace(published, metadata=metadata, metrics=metrics) + published = _published( + name, model_version, metadata={ + **metadata, + **registry_lineage(_tags(model_version)) + }, metrics=metrics) backup = Path(tmp) / 'previous' if destination.exists(): os.replace(destination, backup) @@ -277,13 +327,21 @@ 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: + """Raise ``ValueError`` unless ``location`` names a model this + connector's catalog and schema can hold.""" + _address(connector, None, location) + def publish_model( self, connector: SourceConnector, @@ -321,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=getattr(location, 'alias', None)) + + def set_model_alias(self, connector: SourceConnector, + location: ModelRegistryLocation, alias: str) -> None: + """Point an alias at a version; UC moves it from any other.""" + name = _address(connector, None, location) + number = _version_number(location) + if number is None: + raise ValueError('Setting an alias needs a ModelRegistryLocation ' + 'naming a version.') + check_model_alias(alias) + self._client(connector).api_client.do('PUT', _alias_path(name, alias), + body=dict(version_num=number)) def list_models(self, connector: SourceConnector, *, project_id: Optional[str] = None) -> list[PublishedModel]: diff --git a/packages/kumo-databricks/src/kumo_databricks/session.py b/packages/kumo-databricks/src/kumo_databricks/session.py index b274cd234..ceb299d0b 100644 --- a/packages/kumo-databricks/src/kumo_databricks/session.py +++ b/packages/kumo-databricks/src/kumo_databricks/session.py @@ -91,10 +91,13 @@ 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: a DATABRICKS_TOKEN another thread lends MLflow must not + # make this client's configuration ambiguous. 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..e0172b877 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,11 @@ def list_directory(self, path): if key.startswith(path + '/'): relative = key[len(path) + 1:] child = relative.split('/')[0] - children[child] = SimpleNamespace(path=path + '/' + child, - is_directory='/' in relative) + directory = '/' in relative + # The Files API lists directories with a trailing slash. + children[child] = SimpleNamespace( + path=f"{path}/{child}{'/' if directory else ''}", + is_directory=directory) return children.values() def download(self, path): @@ -485,3 +507,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 a22ce00c6..1a7ab98c8 100644 --- a/packages/kumo-engine/docs/mlops.md +++ b/packages/kumo-engine/docs/mlops.md @@ -1,127 +1,305 @@ # 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 with `pip install "kumo-engine[mlflow]"` (add `s3`, `gcs`, `databricks`, `sagemaker` or `snowflake` as needed; Snowflake needs only `kumo-engine[snowflake]`). GPU, CUDA and Databricks Runtime notes are under [Install](#install). + ## Quickstart -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 # MLflow's own tracking URI, else ./mlruns -cfg = kumo.MLOps("artifacts") # or s3://..., gs://..., /Volumes/..., (connector, "@DB.SCHEMA.STAGE/kumo") +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) -# Later, on any machine that reaches the root: -snapshot = kumo.GraphSnapshot.load_from_connector(cfg, id=result.snapshot_id) -trainer = kumo.Trainer.load_from_connector(cfg, path=f"models/{result.id}") + 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 + +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, 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 -| artifact | id | stored at (MLOps on) | +- 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 | `snapshot.id` | `/snapshots///` | -| training table | `table.id` | `/training_tables//` | -| model | `result.id` (the bundle id) | `/models//` | +| `snapshot.id` | a graph snapshot | `/snapshots///`; `result.snapshot_id` | +| `table.id` | a training table | `/training_tables//` | +| `result.id` | a model (its bundle id) | `/models//`; the `kumo.bundle_id` run tag; `kumo_bundle_id` on the registered version | +| `result.registered_model.version` | the registered version | the registry | +| `result.tracking` | the tracker's own run of the fit: `platform`, `run_id`, `url` | the tracker's UI | -Each carries a `_provenance.json` manifest that embeds its parents' manifests (see [Lineage](#lineage)). +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: +`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") +cfg = kumo.MLOps("s3://bucket/kumo", tracking=..., registry=..., on_error="raise") with kumo.mlops(cfg): snapshot = graph.snapshot() trainer.fit(table, snapshot=snapshot, mlops=cfg) ``` -- **`artifact_root`** is 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, such as `(DatabricksConnector(...), "/Volumes////kumo")` or `(SnowflakeConnector(...), "@DB.SCHEMA.STAGE/kumo")`. A connector whose sink cannot hold files there raises when `kumo.MLOps` is built. -- **`store`** names the artifact kinds stored when they are produced: `snapshot`, `training_table` and `model` (default: all). +- **`artifact_root`** (required): a URI (`s3://`, `gs://`, `file://`, or a local path, including `/Volumes/...` on a Databricks cluster) or a `(connector, path)` pair for a store a URI cannot reach: a volume through a `DatabricksConnector` (`"/Volumes////kumo"`) or a named stage through a `SnowflakeConnector` (`"@DB.SCHEMA.STAGE/kumo"`) once the installed sink holds files there. +- **`tracking`**: `MlflowTracking` or `SnowflakeTracking` (below), or your own [`Tracking`](#custom-trackers). +- **`registry`**: a `ModelRegistryLocation(name, alias=...)`; every fit registers its model there, in the tracker's registry. +- **`store`**: the artifact kinds stored when produced: `snapshot`, `training_table` and `model` (default: all). +- **`on_error`**: `"warn"` (default) warns about a failed store, tracking call or registration and returns; `"raise"`, for scheduled jobs, raises once the call has finished, with the failure as the `__cause__`. After `fit`, `trainer.result` still holds the model. A model whose bundle or assets could not be stored under the root is never registered. - **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, and plain exports (`export(config)`, `export_config=`) never write to the artifact root. +- **Off by default.** Without a block or `mlops=`, every call behaves as it always has; plain exports (`export(config)`, `export_config=`) never write to the artifact root. +- `kumo.MLOps(...)` checks the tracker and registry when it is built (one experiment lookup, one registry listing, the model name's rules), so wrong credentials or a bad name fail before any work. -## Exporting and loading artifacts +### Reading artifacts back -Every artifact has one way out, `.export(OutputConfig(...))`, and one way back, `.load_from_connector(...)`. Pass the `kumo.MLOps` in place of the connector to read from the artifact root, with paths relative to it: +Pass the `kumo.MLOps` where a connector goes, with paths relative to the artifact root: -| Artifact | Export | Load | -|---|---|---| -| Snapshot | `snapshot.export(SnapshotOutputConfig(connector, path))`, returns the location | `GraphSnapshot.load_from_connector(connector, path=location)`, or `(cfg, id=snapshot_id)` | -| Training table | `table.export(TrainingTableOutputConfig(connector, name))` | `TrainingTable.load_from_connector(connector, path=name, pq=pq, snapshot=snapshot)`, or `(cfg, path=f"training_tables/{id}", ...)` | -| Model | `result.export(ModelOutputConfig(connector, path=... or location=...))` | `Trainer.load_from_connector(connector, path=... or location=...)`, or `(cfg, path=f"models/{id}")` | +```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}", pq=pquery, snapshot=snapshot) +``` -`fit(export_config=...)` is shorthand for fitting, then exporting. Loading needs a file destination (S3, GCS, local files, a volume or stage). Exporting again to a destination that already holds a complete copy of the same artifact reuses it (logged at info) instead of uploading again. +`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()`). -## Reference +## Train and register -### Lineage +`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. -Every artifact carries a `_provenance.json` manifest, metadata only. The leading underscore keeps Spark and Hive Parquet readers from treating it as data. +- `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 every platform, registering a bundle the registry already holds under that name returns its version instead of adding one. +- If registration fails, `fit` warns with the full error and returns with `registered_model=None` (with `on_error="raise"`, it raises), and the fit run ends failed with `kumo.error`. Registering a bundle the registry already holds under the name returns that version (logged at info) and moves the alias to it. -A graph snapshot's manifest, next to its index, records each source table's connector, location and version, plus the engine build that wrote it. `graph.snapshot(record_sources=False)` keeps only connector and table names, plus a SHA-256 of each full location so a source moved elsewhere is still noticed, and keeps Custom SQL only as a hash. +**Model names** follow the registry's rules, checked when `kumo.MLOps` is built: -#### Training tables +| 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` | -A generated or caller-supplied training table carries `_provenance.json` in its staging directory: -- the query, the effective generation plan and the task type; -- the split seed; -- the snapshot it was generated from (that snapshot's manifest, embedded); -- for `TrainingTable.from_source_table`, the source table's connector, location and version. +An alias is a letter or underscore followed by letters, digits or underscores; `DEFAULT`, `FIRST`, `LAST`, `LATEST` and `v` are reserved. -`TrainingTable.id` is deterministic: the same snapshot, query, plan, split seed and engine build give the same id. A caller-supplied source without a version gets a random id. An export to files (`TrainingTable.export`) writes the manifest next to the Parquet parts. `TrainingTable.load_from_connector(connector, path=..., pq=..., snapshot=...)` materializes an export again. +### Promote a model -With MLOps on, each table is stored once under `/training_tables//`. +Register every fit under a candidate alias and move the serving alias yourself once the new model is better, comparing `result.metrics` with the champion's: -#### Models +```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) +``` -Every fitted bundle carries `_provenance.json` (`ModelManifest`): -- its `id`, the bundle id (a content hash that leaves the manifest file out); -- when it was fitted and with which fit mode; -- its parents: the training table's manifest, embedded with that table's snapshot, and the warm-start model's manifest, if any. +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: -An embedded model keeps only the ids of its own model parents, so a long chain of warm-start retrains does not grow each manifest. `fit` returns a new `TrainingResult` per call (`.id`, `.manifest`); `Trainer.result` is the latest fit's result or the loaded bundle's (`trainer.result.id`, `trainer.result.manifest`). +```sql +ALTER MODEL ML.MODELS.CHURN_GNN VERSION champion UNSET ALIAS; -- skip on the first promotion +ALTER MODEL ML.MODELS.CHURN_GNN VERSION V3 UNSET ALIAS; -- V3 holds candidate +ALTER MODEL ML.MODELS.CHURN_GNN VERSION V3 SET ALIAS = champion; +``` -With MLOps on, each bundle is stored under `/models//`. A complete copy that is already there is kept; a partial one is replaced. +Jobs that load `alias="champion"` pick up the new version on their next run. -The bundle also carries the published model's holdout predictions (`holdout.parquet`, with its truncation flag), so `result.holdout_df()` works on a model loaded with `Trainer.load` or `Trainer.load_from_connector` (a path or a registry), not only on the result of `fit`. Bundles fitted before this have no copy, and `holdout_df()` raises `FileNotFoundError` for them. +### Trackers -### Moving snapshots between machines +`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 configuration points (`MLFLOW_TRACKING_URI`, `mlflow.set_tracking_uri`; automatic on Databricks). `target` names another server with its credentials: a `DatabricksConnector`, a `SageMakerMLflow(arn)`, or a tracking URI. With a `target`, `experiment_name` is required to record runs. +- `log_model=True` logs the model on the fit run without registering it; `MLOps(registry=)` implies it. + +`SnowflakeTracking(source, *, experiment_name, database=None, schema_name=None, run_name=None, tags=None, native_model_lineage=False)`: +- `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. -To build an index on one machine and train or score on another, export it to a file destination and load it there: +### MLflow pyfunc + +A model registered through `MlflowTracking` is an MLflow pyfunc: it serves with `mlflow.pyfunc.load_model(uri)`, `spark_udf` and Model Serving. ```python -from kumo.engine import GraphSnapshot, S3Connector, SnapshotOutputConfig +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 +``` + +- Input rows have an `ENTITY` column and optionally a `TIMESTAMP`; scores come back in input order, one per scored row. Binary classification returns `True_PROB`; `TARGET_PRED` needs `params={"binary_classification_threshold": ...}` or a stored threshold, and is null without one. +- Each replica loads a snapshot on its first `predict` with that URI and keeps it. `log_model(trainer, snapshot=..., serving_snapshot=...)` logs a model whose replicas load that snapshot at startup and use it when no `snapshot` param is given; `serving_snapshot` takes a URI, or a `GraphSnapshot` or snapshot id stored under the block's artifact root. +- Serving needs an image with kumo-engine installed: `mlflow models serve -m --env-manager local`. Its pip requirement is `kumo-engine[s3,gcs]` of the same version (plus `[text]` for text columns). + +## 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. -location = graph.snapshot().export(SnapshotOutputConfig(S3Connector(), "s3://bucket/snapshots")) -# On the training machine: -snapshot = GraphSnapshot.load_from_connector(S3Connector(), path=location) +```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") ``` -- **Layout:** `export` writes `///`, the two halves of `snapshot.id`, and returns that location. It holds `kumo_snapshot.json` and one directory per export generation (`index/` and `_provenance.json`). -- **Transfer:** files move in parallel with a CRC32 per range. Each export uploads into a new generation and commits by writing `kumo_snapshot.json` last, naming that generation. Until then the previous export stays loadable, and concurrent exports never mix files: the last commit wins whole. Superseded generations, and abandoned ones over a day old, are removed best-effort. A load either completes or leaves nothing; if a newer export commits while it downloads, it switches to that one. -- **Reuse:** loading again reuses the local copy only if it holds this export's bytes (the same range checksums) and passes DiskGraph's integrity check; otherwise it is downloaded again. Exporting again skips the upload when the destination already holds a complete export (its `kumo_snapshot.json` and `_DONE`) with the local index's bytes, and logs that it reused it. -- **Staleness:** `snapshot.source_changes(graph)` reports, per table of `graph`, whether its source is `unchanged`, `changed` or `unknown` since the snapshot was built. `load_from_connector(..., graph=graph)` warns when the snapshot was taken of another graph, naming the schema changes (tables, columns, types, keys, edges), or when any of its tables changed. A graph id also covers each table's connector and location, so a graph that differs only there also warns, saying the schema is the same. -- **Manifest:** the snapshot's `_provenance.json` travels with the export. -- **Trust:** load only from paths whose writers you trust. Checksums detect corruption, not tampering. -- **Destinations:** any destination whose sink holds files at the path works: S3, GCS, local files, or a volume or stage once those sinks are installed. A warehouse connector raises. +`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`. +2. Give the job's identity `sagemaker-mlflow:*` on the server ARN 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: open the UI from SageMaker Studio, or at the URL `aws sagemaker create-presigned-mlflow-tracking-server-url --tracking-server-name ` prints. A tracked fit warns with this command, and `result.tracking.url` is `None`. + +## Reference -With MLOps on (`graph.snapshot(mlops=...)` or a `kumo.mlops` block), each snapshot is exported under `/snapshots/`, unless a complete export of it is already there. +### Install -### Scoring without the graph +- A bare `kumo-engine[mlflow]` with `--find-links ` can resolve to an unrelated `kumo-engine` on PyPI; install the engine wheel by path: `pip install --find-links "/[mlflow]"`. +- **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`). -A scoring job does not need the training job's `Graph` code. It loads a snapshot and the model from the artifact root, then scores: +### 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. + +| Artifact | Export | Load | +|---|---|---| +| Snapshot | `snapshot.export(SnapshotOutputConfig(connector, path))`, returns the location | `GraphSnapshot.load_from_connector(connector, path=location)`, `(cfg, id=...)` | +| Training table | `table.export(TrainingTableOutputConfig(connector, name))` | `TrainingTable.load_from_connector(connector, path=name, pq=pq, snapshot=snapshot)` | +| Model | `result.export(ModelOutputConfig(connector, path=... or location=...))`, or `fit(export_config=...)` | `Trainer.load_from_connector(connector, path=... or location=...)` | + +Loading needs a file destination; a table export is for consumers outside kumo-engine. Exporting again to a destination that holds a complete copy of the same artifact reuses it (logged at info). `ModelOutputConfig(connector, location=...)` publishes to a connector's registry outside a block. + +### Lineage and manifests + +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, 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/`. A stored asset that cannot be found raises `MissingDependencyError`. +- **Text encoders** are the first assets: word-vector models (`GloVeVectors`) are stored; other `TokenEmbeddingModel`s are recorded by fingerprint and must be registered wherever the model loads: ```python -snapshot = GraphSnapshot.load_from_connector(cfg, id=snapshot_id) -trainer = Trainer.load_from_connector(cfg, path=f"models/{bundle_id}") -result = trainer.predict(snapshot, binary_classification_threshold=0.5) +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 ``` -- `predict(snapshot)` reuses the trained query and table-generation plan, and scores its population at `prediction_time`. -- `entities=` scores given rows instead. -- It warns when the snapshot was taken of a different graph than the one the model was trained on, naming the schema changes, for example `Schema changes: events edges.` A graph that differs only in table connectors or locations (a dev graph and a prod graph of the same schema) warns that the schema is the same. The model records its training graph from this release on; for an older model the warning cannot say what differs. -- `predict(prediction_table, snapshot=...)` scores a prediction table generated beforehand, as before. -- `fit` results score too: `result.predict(snapshot, ...)`. +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 + +- **Layout:** `export` writes `///`, holding `kumo_snapshot.json` and one directory per export generation (`index/` and `_provenance.json`). +- **Transfer:** parallel, with a CRC32 per range. Each export commits by writing `kumo_snapshot.json` last, so the previous export stays loadable and concurrent exports never mix: the last commit wins. Superseded generations are removed best-effort. A load completes or leaves nothing, and follows a newer export that commits meanwhile. +- **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. `RunRecord.inputs` maps each source table to its `connector`, `location`, `version` and `version_kind`; a tracker's `tags` are added except `kumo.*` and param names. + +### Progress events + +`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()))]`. + +### 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: `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 a random suffix; an active run elsewhere in the process is never resumed. +- The fit's source tables are MLflow run inputs: Databricks tables as Delta tables at the snapshot's version (so Unity Catalog links them to the registered model), file sources by URI. diff --git a/packages/kumo-engine/kumo/engine/__init__.py b/packages/kumo-engine/kumo/engine/__init__.py index db8c0743e..ca5a6b00f 100644 --- a/packages/kumo-engine/kumo/engine/__init__.py +++ b/packages/kumo-engine/kumo/engine/__init__.py @@ -21,7 +21,9 @@ 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 @@ -57,6 +59,7 @@ ModelOutputConfig, PredictionOutputConfig, PredictionResult, + ScoringDefaults, SnapshotOutputConfig, Trainer, TrainingResult, @@ -68,6 +71,8 @@ "configure", "EmbeddingModelStore", "GloVeVectors", + "TextEmbedding", + "WriteMode", "TokenEmbeddingModel", "ModelRegistryLocation", "Connector", @@ -101,6 +106,7 @@ "Trainer", "TrainingResult", "PredictionResult", + "ScoringDefaults", "ModelOutputConfig", "PredictionOutputConfig", "TrainingTableOutputConfig", diff --git a/packages/kumo-engine/kumo/engine/_deps.py b/packages/kumo-engine/kumo/engine/_deps.py new file mode 100644 index 000000000..f8d6b5713 --- /dev/null +++ b/packages/kumo-engine/kumo/engine/_deps.py @@ -0,0 +1,304 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +r"""What a model needs to predict besides its bundle and a snapshot: its +dependencies, recorded at fit, and the assets among them, resolved at +predict time.""" + +import hashlib +import importlib +import os +import shutil +import tempfile +from pathlib import Path +from typing import ( + Any, + 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, +) + +#: A place assets are stored: a URI or path, or a ``(connector, path)``. +Location = Union[str, Tuple[Any, str]] + + +class Asset(Protocol): + r"""A file or directory a model needs at predict time, addressed by its + content. ``files`` maps names to local files; the class's + ``load(directory, dependency)`` restores one from a directory holding + them, and must give the same ``fingerprint()``.""" + + 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 (``require``) and resolves a fitted + model's assets (``resolve``). With MLOps ``settings`` whose export + includes models, each required asset is stored once under the + artifact root; ``locations`` are asset roots looked in first.""" + def __init__(self, records: Sequence[Dependency] = (), + settings: Any = None, + locations: Sequence[Location] = ()) -> None: + self.records: List[Dependency] = list(records) + self.errors: List[BaseException] = [] + self._settings = settings + self._locations = list(locations) + + @classmethod + def of(cls, manifest: Optional[ModelManifest], settings: Any = None, + locations: Sequence[Location] = ()) -> 'Dependencies': + return cls(getattr(manifest, 'dependencies', None) or (), settings, + locations) + + 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 against its + sha256: from the given locations, then where it was stored, then + the current artifact root.""" + 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 + 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'}. Store it by fitting " + f"with MLOps on (it is kept under /assets/), " + f"predict inside a kumo.mlops block whose artifact root holds " + f"it, or register it where the model looks it up.") + + 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: + 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}'))) + return found + + +def asset_cache() -> Path: + base = os.environ.get('XDG_CACHE_HOME') or Path.home() / '.cache' + return Path(base) / 'kumo' / 'assets' + + +def _store(asset: Asset, sha256: str, settings: Any) -> Optional[str]: + r"""Writes ``asset`` once to ``/assets///``, + its ``_provenance.json`` last; returns where.""" + 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) + with filesystem.open_output_stream(marker) as writer: + writer.write(manifest_bytes(_manifest(asset, sha256))) + return settings.root.uri(relative) + 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) -> AssetManifest: + return AssetManifest(id=sha256, 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) + target = asset_cache() / 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 + loader = imported(record.entry_point or entry_point(FileAsset)) + return bool(loader.load(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..ffeb42ff7 100644 --- a/packages/kumo-engine/kumo/engine/_embeddings.py +++ b/packages/kumo-engine/kumo/engine/_embeddings.py @@ -4,18 +4,26 @@ 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. + # and a trained model encodes columns as it was trained to. Given the + # columns whose encoders the model plan overrides, fit stays silent on + # those alone. 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 +56,90 @@ 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 = getattr(model, 'path', None) + + 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 word vectors of a fitted model's text encoders in + ``deps``; returns whether it has 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 stored word vectors of the model's text encoders, + resolved where ``deps`` looks for assets. Unless a place was given to + look, ``store`` alone serves, as before; a missing model it does not + have either is an error.""" + 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..98e8c2394 --- /dev/null +++ b/packages/kumo-engine/kumo/engine/_events.py @@ -0,0 +1,82 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +r"""Progress events reported by `Trainer.fit`.""" + +import hashlib +import logging +import math +import time +from dataclasses import asdict, dataclass, field +from typing import ( + Any, + Callable, + Dict, + Literal, + Mapping, + Optional, + Sequence, + cast, +) + +from kumo_core.instrumentation import NoOpTrainingInstrumentation + +logger = logging.getLogger(__name__) + +EventName = Literal['fit_started', 'trial_started', 'epoch_ended', + 'trial_ended', 'fit_ended'] + + +@dataclass(frozen=True) +class Event: + r"""One progress event. ``trial`` and ``epoch`` count from 1.""" + + name: EventName + time: float + trial: Optional[int] = None + epoch: Optional[int] = None + metrics: Mapping[str, float] = field(default_factory=dict) + params: Mapping[str, Any] = field(default_factory=dict) + + def to_dict(self) -> Dict[str, Any]: + return asdict(self) + + +def query_sha256(query: str) -> str: + return hashlib.sha256(query.encode()).hexdigest() + + +def _finite_floats(metrics: Mapping[str, Any]) -> Dict[str, float]: + result = {} + for key, value in metrics.items(): + try: + number = float(value) + except (TypeError, ValueError): + continue + if math.isfinite(number): + result[key] = number + return result + + +class EventEmitter(NoOpTrainingInstrumentation): + def __init__(self, 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 fit', 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 a6e8a871d..574e84836 100644 --- a/packages/kumo-engine/kumo/engine/_manifest.py +++ b/packages/kumo-engine/kumo/engine/_manifest.py @@ -233,6 +233,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) @@ -242,16 +243,63 @@ class FitTiming(Record): seconds: float +@dataclass(frozen=True, kw_only=True) +class TrackingRef(Record): + r"""A run in an experiment tracker: its ``platform`` (``mlflow``, + ``databricks``, ``sagemaker``, ``snowflake``), the platform's own + ``run_id`` and, where the platform has one, the run's ``url``.""" + + platform: str + run_id: str + url: Optional[str] = omit_none() + + +@dataclass(frozen=True, kw_only=True) +class Dependency(Record): + r"""One thing a model needs to predict: its ``predictor`` entry point, + a ``package`` requirement, or an ``asset`` (content-addressed by + ``sha256``, loaded by ``entry_point``, stored at ``location``).""" + + 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() + + @dataclass(frozen=True, kw_only=True) class ModelManifest(Manifest): """How a model bundle was fitted: from which training table (embedded, with its snapshot) and warm-start model, when, and with which fit - mode. Its ``id`` is the bundle id.""" + mode. Its ``id`` is the bundle id; ``dependencies`` is what predicting + needs besides the bundle and a snapshot (none listed: the built-in + `Trainer` predictor).""" kind: str = 'model' 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, stored once under + ``/assets///``; 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 b4848c15e..c95843f95 100644 --- a/packages/kumo-engine/kumo/engine/_mlops.py +++ b/packages/kumo-engine/kumo/engine/_mlops.py @@ -5,8 +5,21 @@ import contextlib import contextvars +import re +import warnings from dataclasses import dataclass, field -from typing import Any, FrozenSet, Iterable, Iterator, Optional, Tuple, Union +from typing import ( + Any, + FrozenSet, + Iterable, + Iterator, + Literal, + Optional, + Tuple, + Union, +) + +from kumoapi.plugin import ModelRegistryLocation, check_model_alias from kumo.engine._manifest import KINDS @@ -29,24 +42,39 @@ def uri(self, relative: str) -> str: @dataclass(frozen=True) class MLOps: - r"""Where artifacts go, for :func:`mlops` blocks and ``mlops=`` - arguments. + r"""Where artifacts go and what records them, for :func:`mlops` blocks + and ``mlops=`` arguments. Args: artifact_root: an ``s3://``, ``gs://``, ``file://`` or local path (``/Volumes/...`` on a Databricks cluster), or a ``(connector, path)`` pair for a store a URI cannot reach, such as a volume through a ``DatabricksConnector`` or a Snowflake stage. + tracking: the experiment tracker, a + ``kumo.engine.tracking.Tracking``. Its ``check(registry)``, if + any, runs here, so a server or registry that is unreachable or + refuses these credentials fails before any work. + registry: a ``ModelRegistryLocation`` each fit registers its model + at, in the tracker's registry: MLflow's (Unity Catalog, SageMaker + MLflow) for ``MlflowTracking``, the Snowflake Model Registry for + ``SnowflakeTracking``. store: the artifact kinds stored under the root when produced (``snapshot``, ``training_table``, ``model``); ``True`` for all, ``False`` for none. + on_error: what a failure to track a run or register a model does: + ``'warn'`` warns and the call returns; ``'raise'`` raises once + the call has finished and stored its result, with the failure as + the cause. In ``load_from_connector`` calls, an ``MLOps`` stands in for the connector, with paths relative to its artifact root. """ 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) kinds: FrozenSet[str] = field(init=False, repr=False, compare=False) @@ -58,6 +86,9 @@ def __post_init__(self) -> None: if not self.artifact_root: raise ValueError('MLOps needs an artifact_root.') + if self.on_error not in ('warn', 'raise'): + raise ValueError(f"on_error must be 'warn' or 'raise', not " + f"{self.on_error!r}.") if isinstance(self.artifact_root, str): if self.artifact_root.startswith('@'): raise ValueError( @@ -79,9 +110,45 @@ def __post_init__(self) -> None: if unknown: raise ValueError(f'Unknown store kinds {unknown}; choose from ' f'{list(KINDS)}.') + if self.registry is not None: + self._check_registry() + check = getattr(self.tracking, 'check', None) + if callable(check): + check(self.registry) object.__setattr__(self, 'root', root) object.__setattr__(self, 'kinds', kinds) + def _check_registry(self) -> 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 callable(getattr(self.tracking, 'log_model', None)): + if registry.version is not None: + raise ValueError('MLflow assigns registered model versions; ' + 'leave ModelRegistryLocation.version unset.') + return + connector = getattr(self.tracking, 'connector', None) + if connector is None: + raise ValueError( + 'registry= registers through the tracker: pass ' + '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 under this + artifact root, e.g. for the pyfunc's ``snapshot`` param.""" + return self.root.uri(snapshot_path(snapshot)) + def wants(self, kind: str) -> bool: r"""Whether artifacts of ``kind`` are stored under the root.""" return kind in self.kinds @@ -116,7 +183,8 @@ def mlops(settings: MLOps) -> Iterator[MLOps]: r"""Turns MLOps on for the calls in the block (in this thread or asyncio task): each snapshot, training table and model is stored under ``settings.artifact_root``. An inner block replaces an outer one; - an ``mlops=`` argument replaces both for its call.""" + an ``mlops=`` argument replaces both for its call. Runs are tracked and + models registered as ``settings`` says.""" if not isinstance(settings, MLOps): raise TypeError('kumo.mlops() takes a kumo.MLOps') token = _CURRENT.set(settings) @@ -132,6 +200,39 @@ def resolve(settings: Optional[MLOps] = None) -> Optional[MLOps]: return settings if settings is not None else _CURRENT.get() +def raise_on_error(settings: Optional[MLOps], error: Optional[BaseException], + call: str, stacklevel: int = 2) -> None: + r"""Reports ``error``, a storing, tracking or registration failure of a + finished ``call``: raises when ``settings.on_error`` is ``'raise'``, + else warns unless `warn_failure` already warned about it.""" + if error is None or settings is None: + return + message = (f'{call} finished, but MLOps storing, tracking or ' + f'registration 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: + r"""``Type: message``, not shortened.""" + return f'{type(error).__name__}: {error}' + + +def warn_failure(message: str, error: BaseException, stacklevel: int = 2, + also: Optional[BaseException] = None) -> None: + r"""Warns ``message`` (a :class:`UserWarning`) once per ``error``; it + also reports ``also``, an earlier error that caused ``error``.""" + 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]: r"""The connector and path a ``load_from_connector(source, path=...)`` reads: an `MLOps` ``source`` reads ``path`` under its artifact root.""" @@ -144,3 +245,23 @@ def snapshot_path(snapshot: Any) -> str: r"""Where a snapshot, or a snapshot id, is stored under a root.""" 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 to serve against, as a location: a URI or stage path as + given; a `GraphSnapshot` or snapshot id as stored under the artifact + root of ``settings`` (else of the enclosing block).""" + 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 given by object or id is read from the artifact ' + 'root: call this in a kumo.mlops block, or pass the location ' + 'snapshot.export(...) returned.') + return settings.root.uri(snapshot_path(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..b65ebb906 --- /dev/null +++ b/packages/kumo-engine/kumo/engine/_predictor.py @@ -0,0 +1,65 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +r"""What scores a model bundle: the ``predictor`` dependency its manifest +names.""" + +from pathlib import Path +from typing import Any, Mapping, Protocol, Tuple, Union + +import pandas as pd + +from kumo.engine._deps import imported + +#: The predictor of bundles whose manifest names none. +TRAINER_PREDICTOR = 'kumo.engine._predictor:TrainerPredictor' + + +class Predictor(Protocol): + r"""A loaded bundle: scores ``rows`` (an ``ENTITY`` frame, or a + `PredictionTable` for files) against ``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 with + ``.load(bundle_dir)``.""" + from kumo.engine.trainer._provenance import read_provenance + + manifest = read_provenance(Path(bundle_dir)) + entry = next((record.name + for record in getattr(manifest, 'dependencies', None) 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.integrations._mlflow_model import predict + from kumo.engine.pquery.prediction_table import PredictionTable + + if isinstance(rows, PredictionTable): + return self.trainer.predict(snapshot, prediction_table=rows, + **params).files + return predict(self.trainer, snapshot, rows, params) diff --git a/packages/kumo-engine/kumo/engine/_scratch.py b/packages/kumo-engine/kumo/engine/_scratch.py index 7381e5cb5..8a73ce5fa 100644 --- a/packages/kumo-engine/kumo/engine/_scratch.py +++ b/packages/kumo-engine/kumo/engine/_scratch.py @@ -35,11 +35,15 @@ import hashlib import uuid +from contextlib import contextmanager +from contextvars import ContextVar from pathlib import Path -from typing import Union +from typing import Iterator, Optional, Union _DEFAULT_ROOT = Path.home() / ".cache" / "kumo-engine" _root = _DEFAULT_ROOT +_scoped_root: ContextVar[Optional[Path]] = ContextVar('kumo_scratch_root', + default=None) def set_scratch_dir(path: Union[str, Path]) -> None: @@ -65,7 +69,18 @@ def set_scratch_dir(path: Union[str, Path]) -> None: def get_scratch_dir() -> Path: - return _root + return _scoped_root.get() or _root + + +@contextmanager +def scoped_scratch_dir(path: Path) -> Iterator[None]: + r"""Nests this thread's (context's) new scratch directories under + ``path`` instead of the configured root.""" + token = _scoped_root.set(path) + try: + yield + finally: + _scoped_root.reset(token) def _spool_dir() -> str: diff --git a/packages/kumo-engine/kumo/engine/_snapshot_io.py b/packages/kumo-engine/kumo/engine/_snapshot_io.py index 7f83970c2..4c037397e 100644 --- a/packages/kumo-engine/kumo/engine/_snapshot_io.py +++ b/packages/kumo-engine/kumo/engine/_snapshot_io.py @@ -194,9 +194,7 @@ def export_location(snapshot: Any, path: str) -> str: def store(snapshot: Any, settings: Any) -> None: r"""Export ``snapshot`` under ``/snapshots/`` unless a complete - export of it is there; failures warn.""" - import warnings - + export of it is there.""" try: connector = settings.root.connector path = settings.root.location('snapshots') @@ -210,9 +208,8 @@ def store(snapshot: Any, settings: Any) -> None: and marker.get('snapshot_id') == snapshot.id): export(snapshot, connector, path) except Exception as error: - warnings.warn( - f"Could not store snapshot {snapshot.id} under the artifact " - f"root: {error}", UserWarning, stacklevel=3) + raise RuntimeError(f"Could not store snapshot {snapshot.id} under the " + f"artifact root: {error}") from error def export(snapshot: Any, connector: Any, path: str) -> str: diff --git a/packages/kumo-engine/kumo/engine/_tracking.py b/packages/kumo-engine/kumo/engine/_tracking.py new file mode 100644 index 000000000..159c9a83a --- /dev/null +++ b/packages/kumo-engine/kumo/engine/_tracking.py @@ -0,0 +1,478 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +r"""Platform-neutral run records and the `Tracking` protocol.""" + +import dataclasses +import json +import logging +import os +import queue +import threading +import time +import warnings +from dataclasses import dataclass, field +from pathlib import Path +from typing import ( + TYPE_CHECKING, + Any, + Callable, + Dict, + Iterator, + List, + Mapping, + Optional, + Protocol, + Sequence, + Tuple, +) + +from kumo_core.instrumentation import format_error + +from kumo.engine._events import Event, _finite_floats +from kumo.engine._mlops import full_error, warn_failure + +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' +# Tracker calls run on one background thread so a slow tracking server +# never stalls training; these bound how long the fit waits for it. +QUEUE_SIZE = 1000 +FINISH_SECONDS = 30.0 +MODEL_SECONDS = 3600.0 + + +@dataclass(frozen=True) +class RunRecord: + r"""One tracked run; ``steps`` are per-epoch ``(epoch, value)`` + histories, ``artifacts`` map file names to text and ``inputs`` map + source table names to their ``connector``, ``location``, ``version`` and + ``version_kind``.""" + + name: Optional[str] = None + params: Mapping[str, str] = field(default_factory=dict) + metrics: Mapping[str, float] = field(default_factory=dict) + steps: Mapping[str, Tuple[Tuple[int, float], + ...]] = field(default_factory=dict) + tags: Mapping[str, str] = field(default_factory=dict) + artifacts: Mapping[str, str] = field(default_factory=dict) + inputs: Mapping[str, Mapping[str, Optional[str]]] = field( + default_factory=dict) + + +@dataclass(frozen=True) +class RunSpec: + r"""What a run records, whatever trained or scored it: ``params`` are + flat, ``parents`` become ``kumo.`` pointer tags and ``tags`` are + the user's own.""" + + 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) + inputs: Mapping[str, Mapping[str, Optional[str]]] = field( + default_factory=dict) + parents: Mapping[str, Optional[str]] = field(default_factory=dict) + artifacts: Mapping[str, str] = field(default_factory=dict) + status: Optional[str] = None + error: Optional[str] = None + + +INPUT_FIELDS = ('connector', 'location', 'version', 'version_kind') + + +def fit_inputs(manifest: Any) -> Dict[str, Dict[str, Optional[str]]]: + r"""Source tables of a training table or model manifest, read from its + embedded parents; stub parents add none.""" + try: + table = manifest + if getattr(manifest, 'kind', None) == 'model': + table = getattr(_parents(manifest), 'training_table', None) + parents = _parents(table) + snapshot = _parents(getattr(parents, 'snapshot', None)) + tables = dict(getattr(snapshot, 'source_tables', None) or {}) + source = getattr(parents, 'source_table', None) + if source is not None: + tables['training_table'] = source + return { + str(name): {key: _text(getattr(entry, key)) + for key in INPUT_FIELDS} + for name, entry in tables.items() + } + except Exception as error: + logger.warning('Could not read the fit\'s source tables: %s', + format_error(error), exc_info=_debug()) + return {} + + +def _parents(manifest: Any) -> Any: + if manifest is None or getattr(manifest, 'reconstructed', False): + return None + return getattr(manifest, 'parents', None) + + +def _text(value: Any) -> Optional[str]: + return None if value is None else str(value) + + +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=_finite_floats(spec.metrics), + steps={name: tuple(history) + for name, history in spec.steps.items()}, + tags={ + **_user_tags(spec.tags, params), + 'kumo.kind': spec.kind, + **{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), + inputs=dict(spec.inputs), + ) + + +def _user_tags(tags: Optional[Mapping[str, Any]], + params: Mapping[str, str]) -> Dict[str, str]: + dropped = sorted(key for key in tags or {} + if key.startswith('kumo.') or key in params) + if dropped: + warnings.warn( + f'Tracking tags {dropped} are reserved and not recorded.', + UserWarning, stacklevel=3) + return { + key: str(value) + for key, value in (tags or {}).items() if key not in dropped + } + + +def trial_spec(events: Sequence[Event]) -> RunSpec: + r"""One trial attempt of a fit, from its events so far.""" + started = next(event for event in events if event.name == 'trial_started') + trial, attempt = started.trial, int(started.params['attempt']) + 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({ + name if name == 'seconds' else f'best_{name}': value + for name, value in event.metrics.items() + }) + return RunSpec( + kind='trial', + name=f'trial-{trial}-attempt-{attempt}', + params=dict(started.params), + metrics=metrics, + steps=steps, + **ended, + ) + + +class Tracking(Protocol): + r"""Records one `Trainer.fit` in an experiment tracker. + + Calls come from one background thread. `start` gets the fit run's + record without metrics, `on_trial` a trial attempt's record so far with + the event that updated it, and `finish` the fit run's final record + (`None` if tracking stopped early, which should end the run as failed). + Optional members: ``prepare()``, called first on the fit's own thread; + ``tags`` (extra fit-run tags); ``log_model(bundle_dir, manifest, + snapshot, registry)``, called after a fit that succeeded and before + `finish`, which the fit waits for while the tracking thread is alive, up + to an hour (an error, a stopped thread or the time limit ends the run as + failed), and which returns the version it registered at ``registry``, a + ``ModelRegistryLocation``; ``connector``, a connector whose model + registry ``registry`` registers in when there is no ``log_model``; and + ``platform``, ``run_id`` and ``run_url``, the tracker's own run once + `start` returned, recorded as the model manifest's ``tracking``. + """ + 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 RegistrationError(Exception): + r"""Raised by a tracker's ``log_model`` from the error that registering + the logged model raised.""" + + +class TrackingSession: + def __init__(self, tracking: Optional[Tracking], + warn: bool = True) -> None: + self._tracker = tracking if _rank() == 0 else None + self._warn = warn + self.spec: Optional[RunSpec] = None + self._trials: Dict[int, List[Event]] = {} + 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._call(lambda tracker: getattr(tracker, 'prepare', _none)()) + self._worker = threading.Thread(target=self._work, + name='kumo-tracking', daemon=True) + self._worker.start() + + def begin(tracker: Tracking) -> None: + tracker.start(run_record(dataclasses.replace(value, tags={ + **value.tags, **(getattr(tracker, 'tags', None) 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: its platform, run id + and URL.""" + from kumo.engine._manifest import TrackingRef + + if self.spec is None or not self._started.wait(FINISH_SECONDS): + return None + tracker = self._tracker + run_id = getattr(tracker, 'run_id', None) + if tracker is None or not run_id: + return None + return TrackingRef( + platform=str(getattr(tracker, 'platform', type(tracker).__name__)), + run_id=str(run_id), url=getattr(tracker, 'run_url', None)) + + 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) + record = run_record(trial_spec(events)) + self._submit(lambda: self._call( + lambda tracker: tracker.on_trial(record, event))) + + def log_model( + self, + bundle_dir: Path, + manifest: Optional['ModelManifest'], + snapshot: Any, + registry: Optional['ModelRegistryLocation'] = None, + ) -> Tuple[Optional['PublishedModel'], Optional[BaseException]]: + r"""Runs the tracker's ``log_model`` on the tracking thread and waits + for it; returns the registered version and the error it raised.""" + if self._worker is None or (self._tracker is not None and not callable( + getattr(self._tracker, 'log_model', 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( + getattr(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 with ``spec()``, its final spec (`None`: the started + one), 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() + 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(): + logger.warning( + 'Experiment tracking did not finish within %ss; it goes on ' + 'in the background and the fit returns', FINISH_SECONDS) + + def _submit(self, job: Callable[[], Any]) -> bool: + if self._behind: + return False + try: + self._jobs.put_nowait(job) + return True + except queue.Full: + self._behind = True + logger.warning('Experiment tracking is falling behind; the rest ' + 'of this fit is not tracked') + return False + + def _work(self) -> None: + while True: + job = self._jobs.get() + if job is None: + return + try: + job() + except Exception as error: + logger.warning('Experiment tracking failed: %s', + full_error(error), exc_info=_debug()) + + def _call(self, call: Callable[[Tracking], 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 _debug() -> bool: + r"""Tracebacks of tracking failures only at DEBUG: the failure itself is + reported once, by the call that tracked it.""" + return logger.isEnabledFor(logging.DEBUG) + + +def _none(*args: Any) -> None: + return None + + +def _abandon(tracker: Tracking) -> None: + try: + tracker.finish(None, None) + except Exception: + logger.debug('Could not end the run of %r', tracker, exc_info=True) + + +def _rank() -> int: + for name in ('RANK', 'LOCAL_RANK'): + if os.environ.get(name, '').strip() not in ('', '0'): + return 1 + return 0 + + +def _flatten(value: Any, prefix: str = '') -> Iterator[Tuple[str, Any]]: + if isinstance(value, Mapping): + for key, item in value.items(): + yield from _flatten(item, f'{prefix}{key}.') + else: + yield prefix[:-1], value + + +def _param_value(value: Any) -> str: + if isinstance(value, str): + return value + return json.dumps(value, default=str) diff --git a/packages/kumo-engine/kumo/engine/connector/snowflake.py b/packages/kumo-engine/kumo/engine/connector/snowflake.py index 0e75a9947..bdfc97387 100644 --- a/packages/kumo-engine/kumo/engine/connector/snowflake.py +++ b/packages/kumo-engine/kumo/engine/connector/snowflake.py @@ -11,7 +11,7 @@ import os from functools import cached_property from pathlib import Path -from typing import Dict, Optional +from typing import Any, Dict, Optional from kumoapi.data_source import ( DataSourceType, @@ -100,6 +100,12 @@ def sink_backend(self) -> ArtifactSink: "snowflake").SnowflakeSinkBackend return SnowflakeSinkBackend(identity=self.identity) + def session(self) -> Any: + r"""The Snowpark session this connector reads through: its identity, + pointed at its warehouse, database and schema, kept alive.""" + with self.source_backend._session(self) as session: + return session + def source_table_for(self, table_name: str) -> SnowflakeSourceTable: return SnowflakeSourceTable( snowflake_connector_id=self.id, diff --git a/packages/kumo-engine/kumo/engine/graph.py b/packages/kumo-engine/kumo/engine/graph.py index ea5c66bed..982ee731b 100644 --- a/packages/kumo-engine/kumo/engine/graph.py +++ b/packages/kumo-engine/kumo/engine/graph.py @@ -720,9 +720,13 @@ def snapshot(self, *, record_sources: bool = True, snapshot = self._snapshot_no_lock(graph_key, persist_dir, record_sources) if settings is not None and settings.wants('snapshot'): + from kumo.engine._mlops import raise_on_error from kumo.engine._snapshot_io import store - store(snapshot, settings) + try: + store(snapshot, settings) + except Exception as error: + raise_on_error(settings, error, 'Graph.snapshot') return snapshot def _snapshot_no_lock(self, graph_key: str, persist_dir: 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..75d063c23 --- /dev/null +++ b/packages/kumo-engine/kumo/engine/integrations/_mlflow_model.py @@ -0,0 +1,207 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +import logging +import tempfile +from pathlib import Path +from typing import Any, Dict, Mapping, Optional, Tuple + +import pandas as pd + +from kumo.engine.connector._extras import import_provider +from kumo.engine.connector.base import Connector +from kumo.engine.graph import GraphSnapshot + +_pyfunc = import_provider('mlflow.pyfunc', 'mlflow') +_schema = import_provider('mlflow.types.schema', 'mlflow') +_signature = import_provider('mlflow.models.signature', 'mlflow') + +logger = logging.getLogger(__name__) +_ANCHOR = 'ANCHOR_TIMESTAMP' +_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._predictor import load_predictor + + self._predictor = load_predictor(context.artifacts['bundle']) + self._columns = list(context.model_config['columns']) + self._snapshots: Dict[str, GraphSnapshot] = {} + self._default = context.model_config.get('snapshot') + if self._default: + self._snapshot(self._default) + + def predict(self, context: Any, model_input: pd.DataFrame, + params: Optional[Dict[str, Any]] = None) -> pd.DataFrame: + params = params or {} + if len(model_input) == 0 and not params.get('population'): + return pd.DataFrame(columns=self._columns) + uri = params.get('snapshot') or self._default + if not uri: + raise ValueError("Pass params={'snapshot': }.") + return self._predictor.predict(model_input, self._snapshot(uri), + params) + + def _snapshot(self, uri: str) -> GraphSnapshot: + if uri not in self._snapshots: + connector, path = snapshot_location(uri) + self._snapshots[uri] = GraphSnapshot.load_from_connector( + connector, path=path) + return self._snapshots[uri] + + +def predict(trainer: Any, snapshot: GraphSnapshot, rows: pd.DataFrame, + params: Mapping[str, Any]) -> pd.DataFrame: + from kumo.engine._scratch import scoped_scratch_dir + from kumo.engine.connector.file import FileConnector + + time = params.get('prediction_time') + time = None if time is None else pd.Timestamp(time) + options = dict( + binary_classification_threshold=params.get( + 'binary_classification_threshold'), + num_classes_to_return=params.get('num_classes_to_return'), + ) + _require_inputs(trainer, snapshot, rows, time) + trainer.scratch_dir.mkdir(parents=True, exist_ok=True) + with tempfile.TemporaryDirectory(dir=trainer.scratch_dir) as tmp, \ + scoped_scratch_dir(Path(tmp) / 'scratch'): + if len(rows) == 0: + return trainer.predict( + snapshot, prediction_time=None + if time is None else time.to_pydatetime(), **options).frame + rows = rows.rename(columns={'ENTITY': 'ENTITY1'}) + if time is not None and 'TIMESTAMP' not in rows: + rows = rows.assign(TIMESTAMP=time) + (Path(tmp) / 'rows').mkdir() + rows.to_parquet(Path(tmp) / 'rows' / 'rows.parquet', index=False) + entities = FileConnector(root_dir=str(Path(tmp) / 'rows'))[ + 'rows.parquet'] + timed = 'TIMESTAMP' in rows and _temporal(trainer, snapshot) + scored = trainer._predict( + trainer._prediction_table(snapshot, entities, None), + snapshot=snapshot, extra_fields=[_ANCHOR] if timed else [], + **options).frame + return _in_input_order(rows, scored) + + +def _in_input_order(rows: pd.DataFrame, scored: pd.DataFrame) -> pd.DataFrame: + r"""``scored`` in the order of ``rows``, which scoring does not keep; + rows scoring skipped stay out, and repeated rows pair up in order.""" + keys = ['ENTITY'] + left = pd.DataFrame({'ENTITY': rows['ENTITY1'].to_numpy()}) + right = scored[['ENTITY']].copy() + if _ANCHOR in scored: + keys.append(_ANCHOR) + left[_ANCHOR] = _seconds(rows['TIMESTAMP']) + right[_ANCHOR] = _seconds(scored[_ANCHOR]) + try: + left = left.astype({'ENTITY': right['ENTITY'].dtype}) + left['_n'] = left.groupby(keys, sort=False, dropna=False).cumcount() + right['_n'] = right.groupby(keys, sort=False, dropna=False).cumcount() + right['_row'] = range(len(right)) + order = left.merge(right, on=[*keys, '_n'])['_row'] + except (TypeError, ValueError): + order = None + scored = scored.drop(columns=_ANCHOR, errors='ignore') + if order is None or len(order) != len(scored): + logger.warning('Could not match the scores to the input rows; they ' + 'are returned in scoring order: join them on ENTITY') + return scored + return scored.iloc[order.to_numpy()].reset_index(drop=True) + + +def _temporal(trainer: Any, snapshot: GraphSnapshot) -> bool: + from kumo_core.pquery import query_is_temporal + + from kumo.engine.pquery import PredictiveQuery + + return query_is_temporal( + PredictiveQuery._from_snapshot(snapshot, + trainer._training_query)._parsed, + snapshot._native_graph) + + +def _seconds(values: pd.Series) -> pd.Series: + stamps = pd.to_datetime(values) + if stamps.dt.tz is not None: + stamps = stamps.dt.tz_convert(None) + return stamps.astype('datetime64[s]') + + +def _require_inputs(trainer: Any, snapshot: GraphSnapshot, rows: pd.DataFrame, + time: Optional[pd.Timestamp]) -> None: + from kumo_core.pquery import query_is_temporal + + from kumo.engine.pquery import PredictiveQuery + + if len(rows) == 0 or time is not None or 'TIMESTAMP' in rows: + return + parsed = PredictiveQuery._from_snapshot(snapshot, + trainer._training_query)._parsed + if query_is_temporal(parsed, snapshot._native_graph): + raise ValueError("temporal query: pass params={'prediction_time': " + "...} or a TIMESTAMP column") + + +def snapshot_location(uri: str) -> Tuple[Connector, str]: + from kumo.engine.connector import connector_for_root + + parent, slash, name = uri.rstrip('/').rpartition('/') + return connector_for_root(parent or slash or '.'), name + + +def signature(snapshot: GraphSnapshot, query: str) -> Tuple[Any, pd.DataFrame]: + from kumo_core.train.formatting import prediction_output_columns + from kumo_core.types import TaskType + + from kumo.engine.pquery import PredictiveQuery + + bound = PredictiveQuery._from_snapshot(snapshot, query) + parsed = bound._parsed + entity = snapshot._native_graph.table(parsed.entity_table) + entity_type = _entity_type(str(entity.column(entity.primary_key).dtype)) + task_type = TaskType(parsed.task_type) + dtypes = { + **_DTYPES, 'TARGET_PRED': + 'boolean' if task_type == TaskType.BINARY_CLASSIFICATION else 'double' + } + inputs = _schema.Schema([ + _schema.ColSpec(entity_type, 'ENTITY'), + _schema.ColSpec('datetime', 'TIMESTAMP', required=False), + ]) + outputs = _schema.Schema([ + _schema.ColSpec(entity_type, 'ENTITY'), + *(_schema.ColSpec(dtypes.get(name) or _schema.AnyType(), name) + for name in prediction_output_columns(task_type)), + ]) + params = _schema.ParamSchema([ + _schema.ParamSpec('snapshot', 'string', None), + _schema.ParamSpec('prediction_time', 'string', None), + _schema.ParamSpec('binary_classification_threshold', 'double', None), + _schema.ParamSpec('num_classes_to_return', 'long', None), + _schema.ParamSpec('population', 'boolean', False), + ]) + example = pd.DataFrame({ + 'ENTITY': { + 'long': [0, 1], + 'double': [0.0, 1.0], + }.get(entity_type, ['entity-0', 'entity-1']) + }) + return _signature.ModelSignature(inputs, outputs, params), example + + +def _entity_type(dtype: str) -> str: + if dtype.startswith('int') or dtype == 'byte': + return 'long' + if dtype.startswith('float'): + return 'double' + return 'string' diff --git a/packages/kumo-engine/kumo/engine/integrations/_mlflow_model_code.py b/packages/kumo-engine/kumo/engine/integrations/_mlflow_model_code.py new file mode 100644 index 000000000..03cedae08 --- /dev/null +++ b/packages/kumo-engine/kumo/engine/integrations/_mlflow_model_code.py @@ -0,0 +1,7 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +from kumo.engine.connector._extras import import_provider +from kumo.engine.integrations._mlflow_model import KumoModel + +import_provider('mlflow.models', 'mlflow').set_model(KumoModel()) diff --git a/packages/kumo-engine/kumo/engine/integrations/_mlflow_target.py b/packages/kumo-engine/kumo/engine/integrations/_mlflow_target.py new file mode 100644 index 000000000..62ecac2aa --- /dev/null +++ b/packages/kumo-engine/kumo/engine/integrations/_mlflow_target.py @@ -0,0 +1,269 @@ +# 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. + +MLflow clients take a tracking and a registry URI but no credentials: the +Databricks store reads them from ``DATABRICKS_*`` variables, a profile or the +notebook (``mlflow.utils.databricks_utils.get_databricks_host_creds``), and +the SageMaker plugin signs with the process-wide tracking URI and a boto3 +session (``sagemaker_mlflow.auth_provider.AuthProvider.get_auth``). With +explicit credentials, `Target.scope` sets exactly that state around each +call, under one lock, and restores it. Ambient credentials set no +credentials; a Databricks or SageMaker target still sets the process-wide +tracking URI. +""" + +import os +import re +import threading +from contextlib import ExitStack, 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() + + +@dataclass(frozen=True) +class SageMakerMLflow: + r"""A SageMaker managed MLflow tracking server or app, for + ``MlflowTracking(target=...)``. + + Args: + tracking_server_arn: the server's (or app's) ARN. + identity: a ``kumo.engine.connector.S3Identity`` with an access key + or a role to assume; `None` or ``S3Identity()`` uses the ambient + AWS credential chain (an instance or execution role, + ``aws configure``). + """ + + 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, for explicit credentials, the + state MLflow reads them from.""" + + def __init__(self, tracking_uri: str, registry_uri: Optional[str], + swap: Optional[Callable[[], ContextManager[Any]]] = None, + aws_session: Any = None, host: Optional[str] = None) -> None: + self.tracking_uri = tracking_uri + 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; None where no link opens it + (SageMaker, local stores).""" + 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.""" + if self._swap is None: + yield + return + 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"""The repository for artifacts at ``uri``, or `None` to go through + the client: MLflow's S3 store signs with the default AWS chain.""" + 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)) + return Target(target, registry_for(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'a SageMakerMLflow or a tracking URI, not ' + f'{type(target).__name__}.') + + +def _workspace_host() -> Optional[str]: + try: + return import_provider('databricks.sdk.config', + 'databricks').Config().host + except Exception: + return None + + +def registry_for(tracking_uri: str) -> str: + r"""The registry a tracking URI implies: Unity Catalog on Databricks, + the tracking server itself elsewhere.""" + scheme, colon, rest = tracking_uri.partition(':') + return f'databricks-uc{colon}{rest}' if scheme == 'databricks' else ( + tracking_uri) + + +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): + # Unity Catalog registration downloads the logged model through the + # process-wide tracking URI, so it must point at the workspace. + return Target('databricks', 'databricks-uc', + swap=lambda: global_tracking_uri('databricks')) + client = import_provider('kumo_databricks.session', + 'databricks')._client_for(principal) + return Target('databricks', 'databricks-uc', + swap=lambda: _databricks_token(client), + host=client.config.host) + + +@contextmanager +def _databricks_token(client: Any) -> Iterator[None]: + r"""``DATABRICKS_HOST`` and a short-lived ``DATABRICKS_TOKEN`` minted + from the identity's own client, with every other Databricks auth + variable unset, and the process-wide tracking URI set to + ``databricks``; the client secret never enters the environment.""" + 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'): + # Unity Catalog registration downloads the logged model through the + # process-wide tracking URI, so it must point at the workspace too. + yield + + +def _sagemaker(target: SageMakerMLflow) -> Target: + from kumoapi.data_source import S3AmbientAuthentication + + arn = target.tracking_server_arn + principal = _principal(target.identity) + if principal is None or isinstance(principal.config.authentication, + S3AmbientAuthentication): + return Target(arn, arn, swap=lambda: global_tracking_uri(arn)) + session = import_provider('kumo_s3.backend', + 's3')._session_for(principal) + return Target(arn, arn, swap=lambda: _sagemaker_session(arn, session), + aws_session=session) + + +@contextmanager +def _sagemaker_session(arn: str, session: Any) -> Iterator[None]: + r"""The plugin's signer and router read the process-wide tracking URI + and take a boto3 session from a per-thread context variable; + ``SAGEMAKER_MLFLOW_ASSUME_ROLE_ARN`` would make it assume a role.""" + plugin = import_provider('sagemaker_mlflow.session', 'sagemaker') + with ExitStack() as stack: + stack.enter_context(environ({'SAGEMAKER_MLFLOW_ASSUME_ROLE_ARN': None})) + stack.enter_context(global_tracking_uri(arn)) + stack.enter_context(plugin.use_session(session)) + yield + + +@contextmanager +def global_tracking_uri(uri: str) -> Iterator[None]: + r"""``mlflow.set_tracking_uri(uri)`` for the block. MLflow's own + ``_use_tracking_uri`` restores by unsetting ``MLFLOW_TRACKING_URI``, + so this restores the module value and the variable itself.""" + 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/mlflow.py b/packages/kumo-engine/kumo/engine/integrations/mlflow.py new file mode 100644 index 000000000..774eb1bfb --- /dev/null +++ b/packages/kumo-engine/kumo/engine/integrations/mlflow.py @@ -0,0 +1,988 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +r"""MLflow tracking and models for `Trainer.fit` +(``pip install 'kumo-engine[mlflow]'``).""" + +import hashlib +import json +import logging +import os +import posixpath +import re +import tempfile +import threading +import time +import warnings +from contextlib import ExitStack, contextmanager +from pathlib import Path +from typing import ( + TYPE_CHECKING, + Any, + Callable, + Dict, + FrozenSet, + Iterator, + List, + Literal, + Mapping, + Optional, + Set, + Tuple, +) +from urllib.parse import quote + +from kumo_core.instrumentation import format_error +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._tracking import DROPPED_PARAMS, RegistrationError, RunRecord +from kumo.engine._version import __version__ +from kumo.engine.connector._extras import import_provider +from kumo.engine.integrations._mlflow_target import ( + SAGEMAKER_ARN, + SageMakerMLflow, + 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 +# MLflow's portable name charset (its Windows rule, which omits ':'). +_PARAM_KEY = re.compile(r'[/\w.\- ]+') +# MLflow's defaults (7 retries, 120 s timeout) can hold a call for minutes. +_HTTP_LIMITS = { + 'MLFLOW_HTTP_REQUEST_MAX_RETRIES': '2', + 'MLFLOW_HTTP_REQUEST_TIMEOUT': '30', +} +# Model upload only, and only where the limits above apply. +_UPLOAD_HTTP = {'MLFLOW_HTTP_REQUEST_TIMEOUT': '600'} +_PATH_CONNECTORS = ('LOCAL_FILE', 'S3', 'GCS') +_SAGEMAKER_MODEL = re.compile(r'[a-zA-Z0-9](-*[a-zA-Z0-9]){0,56}') +_UC_FORBIDDEN = re.compile(r'[ /\x00-\x1f\x7f]') + + +class MlflowTracking: + r"""Records `Trainer.fit` in MLflow: one run per fit (under the active + run, if any, on the same server) with a child run per trial attempt. + Source tables are logged as run inputs: Databricks tables as Delta + tables (Unity Catalog then links them to a model registered from the + run), file and object-store sources by URI. Params MLflow rejects are + counted in ``kumo.dropped_params``. Concurrent fits in one process are + not supported. + + Args: + experiment_name: experiment for the fit run, created if missing. + `None` uses MLflow's current experiment; required to record + runs at a ``target``. + run_name: name of the fit run. `None` names it after its UTC start + time, as ``fit-20261027T061838Z``. + log_system_metrics: sample CPU, memory, disk, network and (with + ``nvidia-ml-py``) GPU usage on the fit run. + tags: extra tags for the fit run. + log_model: log the fitted model on the fit run (see `log_model`) + before the fit returns. ``kumo.MLOps(registry= + ModelRegistryLocation(...))`` also registers it, which implies + this. + target: where to record, with the credentials that reach it: a + ``DatabricksConnector`` (Databricks MLflow and Unity Catalog, + as its identity), a `SageMakerMLflow`, or another MLflow + server's tracking URI. `None` uses MLflow's own tracking and + registry URIs, read once here. + + With explicit credentials (a Databricks service principal, an AWS + access key or role), each MLflow call sets the state MLflow reads them + from for its duration (Databricks: ``DATABRICKS_HOST`` and a short-lived + ``DATABRICKS_TOKEN``; SageMaker: the process-wide tracking URI), so do + not call MLflow or Databricks yourself on other threads during a fit. + Ambient credentials (a notebook, job, profile or instance role) set no + credentials, but a Databricks or SageMaker target still sets the + process-wide tracking URI for each call: Unity Catalog downloads and the + SageMaker plugin route with it. + + ``run_id`` is the MLflow id of the latest fit run, once it has started, + and ``run_url`` its link in the MLflow UI, where the server has one. + """ + + 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.run_id: Optional[str] = None + self._reset() + + def _reset(self) -> None: + self._client: Any = None + self._parent: Optional[str] = None + self._experiment_id: Optional[str] = None + self._children: Dict[Tuple[int, int], str] = {} + self._artifacts: Set[str] = set() + self._model_id: Optional[str] = None + self._http_owned: FrozenSet[str] = frozenset() + self._guard = ExitStack() + + @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`` for this tracker's server, with its + credentials and tracking URI set for the block, so your own MLflow + calls (the client's, or ``mlflow.search_runs``) reach the server + kumo-engine records to, as the same identity.""" + 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 this tracker's registry + accepts ``name``.""" + check_model_name(name, self.target.registry_uri or '') + + def check(self, registry: Optional[ModelRegistryLocation] = None) -> None: + r"""Raises before any work unless the server is reachable with these + credentials and, given ``registry``, can register its model: one + experiment lookup, one registry listing and, on Unity Catalog, the + model's schema.""" + self._require_experiment() + if registry is not None: + self.check_model_name(registry.model_name) + uri = self.target.tracking_uri + owned = frozenset(name for name in _HTTP_LIMITS + if name not in os.environ) + missing = False + try: + with _limited_http(owned), 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 RuntimeError(f'MLflow at {uri} is not usable with these ' + f'credentials: {full_error(error)}') from error + if missing: + self._create_experiment(client, owned) + + def _create_experiment(self, client: Any, owned: FrozenSet[str]) -> None: + with _limited_http(owned), 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)}. Its ' + f'parent folder must exist: on Databricks, create it ' + f'with `databricks workspace mkdirs ` or in the ' + f'workspace UI.') 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 _require_experiment(self) -> None: + if self._named and self.experiment_name is None: + raise ValueError('MlflowTracking(target=...) records runs in an ' + 'experiment: pass experiment_name.') + + def load_model(self, location: ModelRegistryLocation, + **kwargs: Any) -> 'Trainer': + r"""The `Trainer` of the version of ``location`` its alias or + version names, from this tracker's registry.""" + if (location.alias is None) == (location.version is None): + raise ValueError('Loading from an MLflow registry takes a ' + 'ModelRegistryLocation with an alias or a ' + 'version.') + reference = (f'@{location.alias}' + if location.alias else f'/{location.version}') + return load_model(f'models:/{location.model_name}{reference}', + target=self.target, **kwargs) + + def register(self, result: Any, + location: ModelRegistryLocation) -> PublishedModel: + r"""Registers a fitted or loaded model (a `TrainingResult`) at + ``location``: the version that already holds its bundle, else a new + version of the model logged in a new run of this experiment.""" + from kumo.engine._manifest import stub + from kumo.engine.graph import GraphSnapshot + from kumo.engine.trainer._spec import read_training_spec + from kumo.engine.trainer.trainer import registry_pointers + + self._require_experiment() + client = self.target.client() + pointers = registry_pointers(result.id) + with _limited_http(frozenset(name for name in _HTTP_LIMITS + if name not in os.environ)): + with self.target.scope(): + held = _registered_version(client, location.model_name, + pointers) + info = None + if held is None: + spec = read_training_spec(result.bundle_dir) + snapshot = GraphSnapshot(spec.snapshot_id, Path(), + spec.graph_config or {}) + with self.target.scope(): + experiment_id = self._default_experiment or '0' + if self.experiment_name is not None: + experiment = client.get_experiment_by_name( + self.experiment_name) + experiment_id = (experiment.experiment_id + if experiment else + client.create_experiment( + self.experiment_name)) + run_id = client.create_run( + experiment_id, + tags={'kumo.bundle_id': str(result.id)}).info.run_id + info = _log_model(self.target, client, result.bundle_dir, + result.manifest or stub('model', result.id), + snapshot, name='model', run_id=run_id, + experiment_id=experiment_id, + serving_snapshot=None) + with self.target.scope(): + client.set_terminated(run_id, 'FINISHED') + return _register(self.target, client, info, location, pointers) + + def start(self, record: RunRecord) -> None: + self._require_experiment() + self._reset() + self.run_id = None + self._http_owned = frozenset(name for name in _HTTP_LIMITS + if name not in os.environ) + self._guard.enter_context(_limited_http(self._http_owned)) + try: + autologging = import_provider('mlflow.utils.autologging_utils', + 'mlflow') + # Not reentrant: leaving it clears the flag. + if not getattr(autologging, '_AUTOLOGGING_GLOBALLY_DISABLED', + False): + self._guard.enter_context(autologging.disable_autologging()) + except Exception 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_id = self._default_experiment + if self.experiment_name is not None: + experiment = client.get_experiment_by_name( + self.experiment_name) + experiment_id = (experiment.experiment_id if experiment else + client.create_experiment( + self.experiment_name)) + parent = {} if self._outer is None else { + 'mlflow.parentRunId': self._outer + } + run = client.create_run( + experiment_id or '0', + run_name=record.name or self.run_name + or time.strftime('fit-%Y%m%dT%H%M%SZ', time.gmtime()), + tags=context.resolve_tags({**record.tags, **parent})) + self._parent = self.run_id = run.info.run_id + self._experiment_id = run.info.experiment_id + if self.log_system_metrics and record.tags.get('kumo.kind') == 'fit': + self._guard.enter_context( + _SystemMetrics(self._log_system_metrics)) + self._log_params(self._parent, record.params) + self._log_artifacts(record.artifacts) + with self.target.scope(): + _log_inputs(self._mlflow, client, self._parent, record.inputs) + + def on_trial(self, record: RunRecord, event: Event) -> None: + key = (event.trial or 0, int(event.params['attempt'])) + if event.name == 'trial_started': + with self.target.scope(): + run = self._client.create_run( + self._experiment_id, run_name=record.name, + tags={**record.tags, 'mlflow.parentRunId': self._parent}) + self._children[key] = run.info.run_id + self._log_params(run.info.run_id, record.params) + elif event.name == 'epoch_ended': + self._log_metrics(self._children.get(key), event.metrics, + event.epoch or 0) + elif event.name == 'trial_ended': + run_id = self._children.pop(key) + self._log_metrics(run_id, record.metrics) + self._log_tags(run_id, record.tags) + succeeded = event.params.get('status') == 'succeeded' + 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='model', 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 + try: + return _register(self.target, self._client, info, registry, + registry_pointers(_id(manifest))) + except Exception as error: + raise RegistrationError(full_error(error)) from error + + def finish(self, record: Optional[RunRecord], + error: Optional[BaseException]) -> None: + status = 'FAILED' + fit_run = (record is not None and self._parent is not None + and record.tags.get('kumo.kind') == 'fit') + try: + if record is not None and self._parent is not None: + self._log_metrics(self._parent, record.metrics, + model_id=self._model_id) + self._log_tags(self._parent, record.tags) + self._log_artifacts(record.artifacts) + if error is None: + status = 'FINISHED' + finally: + 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) + self._reset() + hint = _sagemaker_ui_hint(self.target.tracking_uri) + if fit_run and hint: + warnings.warn(hint, UserWarning, stacklevel=2) + + def _log_system_metrics(self, metrics: Mapping[str, float], + step: int) -> None: + self._log_metrics(self._parent, metrics, step) + + def _log_artifacts(self, artifacts: Mapping[str, str]) -> None: + for name, text in artifacts.items(): + if name in self._artifacts: + continue + with self.target.scope(): + repository = self.target.artifact_repository( + self._client.get_run(self._parent).info.artifact_uri + ) if self.target.explicit else None + if repository is None: + self._client.log_text(self._parent, text, name) + else: + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / name + path.write_text(text) + repository.log_artifact(str(path)) + self._artifacts.add(name) + + 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) + with self.target.scope(): + self._client.log_batch( + run_id, metrics=[ + entities.Metric(metric_key(key), float(value), timestamp, + step, model_id=model_id) + for key, value in metrics.items() + ]) + + +class _SystemMetrics: + r"""Samples MLflow's CPU, disk, network and GPU monitors every + ``interval`` seconds on its own thread and hands each aggregate to + ``log(metrics, step)``, which logs it through the tracker's + credentials. (MLflow's own ``SystemMetricsMonitor`` logs with + process-wide ones.)""" + + def __init__(self, log: Callable[[Mapping[str, float], int], None], + interval: float = 10.0) -> None: + self._log = log + self._interval = interval + self._monitors = _metric_monitors() + 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(self._interval + 5) + + def _run(self) -> None: + step = 0 + while not 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() + try: + self._log({f'system/{k}': v + for k, v in metrics.items()}, step) + except Exception as error: + logger.debug('Could not log system metrics: %s', + format_error(error)) + return + step += 1 + + +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 the MLflow registry at + ``registry_uri`` accepts ``name``: SageMaker managed MLflow takes letters, + digits and hyphens; Unity Catalog takes ``catalog.schema.model``, each + part at most 255 characters without spaces, ``/``, ``.`` or control + characters.""" + 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 letters, digits and hyphens, starting and ending with ' + f'a letter or digit, at most 57 letters and digits ' + f'(^[a-zA-Z0-9](-*[a-zA-Z0-9]){{0,56}}$), 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 (on a cluster, a shorter name takes ' + f'its current catalog and schema), each part at most 255 ' + f'characters with no space, "/" or control character, e.g. ' + f"'main.ml.users_churn'.") + + +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 metric_key(name: str) -> str: + r"""The MLflow name of a Kumo metric: MLflow rejects ``@``, so + ``map@10`` is logged as ``map_at_10``.""" + return name.replace('@', '_at_') + + +def log_model( + trainer: 'Trainer', + *, + snapshot: 'GraphSnapshot', + name: str = 'model', + run_id: Optional[str] = None, + serving_snapshot: Any = None, + target: Any = None, + experiment_name: Optional[str] = None, +) -> Any: + r"""Log a trainer's bundle as an MLflow pyfunc model in ``run_id`` (or + the active run, else a new one in ``experiment_name``) with its source + tables as run inputs; returns its `ModelInfo`. ``target`` says where, + with what credentials, as in `MlflowTracking`; `None` uses MLflow's own + tracking URI. Scores with ``predict(entities, params={'snapshot': + uri})``; with ``serving_snapshot`` (an exported snapshot URI, or a + `GraphSnapshot` or snapshot id stored under the artifact root of the + enclosing `kumo.mlops` block), each serving replica downloads that + snapshot once at startup and scores against it when no ``snapshot`` + param is given.""" + from kumo.engine._manifest import stub + from kumo.engine._mlops import snapshot_uri + from kumo.engine._tracking import fit_inputs + + 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 and tracking._named and experiment_name is None: + tracking._require_experiment() + 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: + experiment = (client.get_experiment_by_name( + tracking.experiment_name) if tracking.experiment_name else + None) + run_id = client.create_run( + experiment.experiment_id if experiment else + tracking._default_experiment or '0').info.run_id + experiment_id = client.get_run(run_id).info.experiment_id + _log_inputs(mlflow, client, run_id, fit_inputs(result.manifest)) + info = _log_model(target, client, result.bundle_dir, result.manifest or + (None if result.id is None else stub('model', result.id)), + snapshot, name=name, run_id=run_id, + experiment_id=experiment_id, + serving_snapshot=snapshot_uri(serving_snapshot)) + with target.scope(): + if result.metrics: + client.log_batch(run_id, metrics=[ + mlflow.entities.Metric(metric_key(key), float(value), + int(time.time() * 1000), 0, + model_id=info.model_id) + for key, value in result.metrics.items() + ]) + if created: + 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"""Restore the `Trainer` of a model `log_model` logged, from the + server ``target`` reaches, as in `MlflowTracking` (`None`: MLflow's own + tracking URI).""" + from kumo.engine._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'] + trainer = Trainer.load(local / bundle, embedding_models, + accelerator=accelerator) + warn_if_other_bundle(model.metadata, trainer) + return trainer + + +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 _log_model(target: Any, client: Any, bundle_dir: Path, + manifest: Optional['ModelManifest'], + snapshot: Optional['GraphSnapshot'], *, name: str, run_id: str, + experiment_id: Optional[str], + serving_snapshot: Optional[str]) -> Any: + r"""What ``mlflow.pyfunc.log_model`` does, through ``client``: it only + logs to the process-wide tracking URI. A bundle without a training spec + is logged without a signature.""" + 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, + listed_file, + read_training_spec, + ) + from kumo.engine.trainer.trainer import 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(): + model = client.create_logged_model( + experiment_id or client.get_run(run_id).info.experiment_id, + name=name, source_run_id=run_id, + params=dict(client.get_run(run_id).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, + }, + artifacts={ + 'bundle': str(bundle_dir), + **({'provenance': str(provenance)} + if provenance.is_file() else {}), + }, + signature=signature, + input_example=example, + metadata=pointers, + pip_requirements=[ + _requirement(listed_file(bundle_dir, 'encoders')) + ], + ) + 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') + 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(encoders: Optional[Path]) -> str: + from kumo_core.encoders.config import EncodersConfig + from kumo_core.io.json import load_json + from kumoapi.encoder import TextEmbedding + + extras = ['s3', 'gcs'] + path = None if encoders is None else str(encoders) + config = (load_json(path, EncodersConfig) + if path and Path(path).is_file() else None) + if config is not None and any( + isinstance(encoder, TextEmbedding) + for table in config.encoders.values() + for encoder in table.values()): + extras.append('text') + return f"kumo-engine[{','.join(extras)}]=={__version__}" + + +def _log_inputs(mlflow: Any, client: Any, run_id: str, + inputs: Mapping[str, Mapping[str, Optional[str]]]) -> None: + try: + meta = import_provider('mlflow.data.meta_dataset', 'mlflow') + entities = mlflow.entities + datasets = [] + for name, entry in inputs.items(): + source = _dataset_source(entry) + if source is None: + continue + digest = hashlib.sha256( + json.dumps(dict(entry), sort_keys=True).encode()).hexdigest() + dataset = meta.MetaDataset(source, name=name, digest=digest[:8]) + datasets.append( + entities.DatasetInput( + dataset._to_mlflow_entity(), + tags=[entities.InputTag('mlflow.data.context', 'train')])) + if datasets: + client.log_inputs(run_id, datasets=datasets) + except Exception as error: + logger.warning('Could not log the fit\'s source tables to MLflow: %s', + format_error(error), exc_info=True) + + +def _register(target: Any, client: Any, info: Any, + location: ModelRegistryLocation, + metadata: Mapping[str, str]) -> PublishedModel: + name = location.model_name + with target.scope(): + version = _registered_version(client, name, metadata) + 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]: + r"""The version of ``name`` that already holds this bundle, if any.""" + 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's search results carry 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 :class:`ValueError` unless the catalog and + schema of a ``catalog.schema.model`` name exist and these credentials + can use them (one metadata call, no warehouse).""" + 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)}. The catalog and schema " + f"must exist, and the registering principal needs USE CATALOG, " + f"USE SCHEMA and CREATE MODEL on them.") + response.raise_for_status() + + +def _no_serving_check(record: logging.LogRecord) -> bool: + # MLflow's input-example check predicts without a snapshot; no switch. + return not record.getMessage().startswith( + 'Failed to validate serving input example') + + +def _dataset_source(entry: Mapping[str, Optional[str]]) -> Any: + connector = (entry.get('connector') or '').upper() + location = entry.get('location') + if not location or location.startswith('sql:'): + return None + if connector == 'DATABRICKS': + delta = import_provider('mlflow.data.delta_dataset_source', 'mlflow') + return delta.DeltaDatasetSource( + delta_table_name=_unquoted(location), + delta_table_version=_delta_version(entry)) + if connector in _PATH_CONNECTORS: + registry = import_provider('mlflow.data.dataset_source_registry', + 'mlflow') + with warnings.catch_warnings(): + warnings.simplefilter('ignore') + return registry.resolve_dataset_source(location) + return None + + +def _unquoted(name: str) -> str: + parts = re.findall(r'`((?:[^`]|``)*)`', name) + return '.'.join(part.replace('``', '`') for part in parts) or name + + +def _delta_version(entry: Mapping[str, Optional[str]]) -> Optional[int]: + version = entry.get('version') or '' + head = version.split('_', 1)[0] + if entry.get('version_kind') == 'exact' and head.isdigit(): + return int(head) + return None + + +def _mlflow_params(params: Mapping[str, str]) -> Dict[str, str]: + return { + key: value + for key, value in params.items() + if len(key) <= _MAX_PARAM_KEY and _PARAM_KEY.fullmatch(key) is not None + and posixpath.normpath(key) == key and not key.startswith('/') + and len(value) <= _MAX_PARAM_VALUE + } + + +def _quietly(function: Callable[..., Any], *args: Any) -> None: + try: + function(*args) + except Exception: + logger.debug('MLflow cleanup call failed', exc_info=True) + + +@contextmanager +def _limited_http(names: FrozenSet[str]) -> Iterator[None]: + os.environ.update({name: _HTTP_LIMITS[name] for name in names}) + try: + yield + finally: + for name in names: + os.environ.pop(name, None) diff --git a/packages/kumo-engine/kumo/engine/integrations/snowflake.py b/packages/kumo-engine/kumo/engine/integrations/snowflake.py new file mode 100644 index 000000000..78b3e23b6 --- /dev/null +++ b/packages/kumo-engine/kumo/engine/integrations/snowflake.py @@ -0,0 +1,299 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +r"""Snowflake Experiment Tracking for `Trainer.fit` +(``pip install 'kumo-engine[snowflake]'``).""" + +import logging +import re +import tempfile +import uuid +from contextlib import ExitStack +from pathlib import Path +from typing import Any, Dict, List, Mapping, Optional, Set, Tuple + +from kumo_core.instrumentation import format_error +from kumo_core.sources.sql import resolve_snowflake_identifier + +from kumo.engine._events import Event +from kumo.engine._tracking import DROPPED_METRICS, DROPPED_PARAMS, RunRecord +from kumo.engine.connector._extras import import_provider + +logger = logging.getLogger(__name__) + +_MAX_METRICS = 200 +_MAX_PARAMS = 1000 +# Metric name budgets: per-epoch metrics give way first, then trial summaries. +_EPOCH_METRICS = _MAX_METRICS - 60 +_SUMMARY_METRICS = _MAX_METRICS - 30 +_FINISH_PARAMS = 20 + + +class SnowflakeTracking: + r"""Records `Trainer.fit` in Snowflake Experiment Tracking as one run; + trial attempts are logged in it under ``trial--attempt-/`` names and + tags are logged as params. Snowflake has no dataset inputs, so a fit's + source tables are not recorded as inputs. The Snowpark session's + parameters are not changed. + + Args: + source: a ``SnowflakeConnector``, whose identity opens the session + and whose Snowflake Model Registry ``kumo.MLOps(registry=...)`` + registers models in; or a Snowpark session, which tracks only. + It needs CREATE EXPERIMENT on the schema, or MODIFY on an + existing experiment. + experiment_name: experiment for the fit run, created if missing. + database: `None` uses the session's current database. + schema_name: `None` uses the session's current schema, or ``PUBLIC`` + when `database` is given or the session has none. + run_name: name of the fit run. `None` lets Snowflake pick one; a + taken name gets a random suffix. + tags: extra tags for the fit run. + native_model_lineage: keep the run open as a context for the whole + fit, so models registered meanwhile are linked to it. + + ``run_id`` is the latest fit run's name (Snowflake runs are identified + by name), once it has started. + """ + + platform = 'snowflake' + run_url = None + + 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.run_id: Optional[str] = None + self._reset() + + def _reset(self) -> None: + self._tracking: Any = None + self._run: Any = None + self._run_context = ExitStack() + self._trial_tags: Dict[Tuple[int, int], Set[str]] = {} + self._metrics: Set[str] = set() + self._params = 0 + self._start_tags: Set[str] = set() + self._artifacts: Set[str] = set() + self._dropped_metrics: Set[str] = set() + self._dropped_params = 0 + + def load_model(self, location: Any, **kwargs: Any) -> Any: + r"""The `Trainer` of the Snowflake Model Registry version + ``location`` names, through this tracker's connector.""" + 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 check(self, registry: Any = None) -> None: + r"""Raises before any work unless the session can see the + experiment's schema (one metadata query, no warehouse).""" + database = self.database or self.session.get_current_database() + schema = self.schema or (None if self.database else + self.session.get_current_schema()) + if database is None: + return + name = resolve_snowflake_identifier(schema or 'PUBLIC') + 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'{type(error).__name__}: {error}') from error + if not rows: + raise ValueError( + f'Snowflake Experiment Tracking needs schema {database}.' + f'{schema or "PUBLIC"}, which does not exist or this role ' + f'cannot see.') + + def start(self, record: RunRecord) -> None: + self._reset() + self.run_id = None + cls = self._module.ExperimentTracking + tracking = getattr(cls, '_instance', None) + if tracking is None: + tracking = cls(self.session) + elif (not hasattr(tracking, '_run') or tracking._run is not None + or getattr(tracking, '_session', None) is not self.session): + logger.warning('Snowflake ExperimentTracking is in use in this ' + 'process; this fit is not tracked in Snowflake') + return + self._tracking = tracking + database = self.database or self.session.get_current_database() + schema = self.schema or (None if self.database else + self.session.get_current_schema()) + experiment = tracking.set_experiment(self.experiment_name, + database_name=database, + schema_name=schema or 'PUBLIC') + self._run = tracking.start_run( + self._new_run_name(tracking, experiment.name)) + self.run_id = str(self._run.name) + if self.native_model_lineage: + self._run_context.enter_context(self._run) + self._start_tags = set(record.tags) + self._log_params({**record.params, **record.tags}) + self._log_artifacts(record.artifacts) + + def on_trial(self, record: RunRecord, event: Event) -> None: + if self._run is None: + return + prefix = f'{record.name}/' + key = (event.trial or 0, int(event.params['attempt'])) + if event.name == 'trial_started': + self._trial_tags[key] = set(record.tags) + self._log_params({prefix + k: v + for k, v in {**record.params, + **record.tags}.items()}) + elif event.name == 'epoch_ended': + self._log_metrics({prefix + k: v + for k, v in event.metrics.items()}, + _EPOCH_METRICS, step=event.epoch or 0) + elif event.name == 'trial_ended': + self._log_metrics({prefix + k: v + for k, v in record.metrics.items()}, + _SUMMARY_METRICS) + logged = self._trial_tags.pop(key, set()) + self._log_params({prefix + k: v + for k, v in record.tags.items() + if k not in logged}) + + def finish(self, record: Optional[RunRecord], + error: Optional[BaseException]) -> None: + status = 'FAILED' + try: + if self._run is not None and record is not None: + self._log_metrics(record.metrics, _MAX_METRICS) + self._log_params({k: v + for k, v in record.tags.items() + if k not in self._start_tags}, _MAX_PARAMS) + self._log_artifacts(record.artifacts) + self._log_dropped() + if error is None: + status = 'FINISHED' + finally: + try: + if self._run is not None: + self._tracking.end_run(status=status) + except Exception 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 _new_run_name(self, tracking: Any, + experiment: Any) -> Optional[str]: + if self.run_name is None: + return None + # start_run(name) resumes, or raises on, any run matching name as a + # SQL LIKE pattern and cannot create an existing run, so pick a name + # that is neither. + existing = [ + str(row['name']) + for row in tracking._sql_client.show_runs_in_experiment( + experiment_name=experiment) + ] + short = uuid.uuid4().hex + for name in (self.run_name, _suffixed(self.run_name, short[:8]), + _suffixed(self.run_name, short)): + if not _taken(name, existing): + break + else: + name = None + if name != self.run_name: + logger.warning( + 'Snowflake run %s already exists in experiment %s; recording ' + 'this fit as %s', self.run_name, self.experiment_name, + name or 'a generated run name') + return name + + def _log_artifacts(self, artifacts: Mapping[str, str]) -> None: + with tempfile.TemporaryDirectory() as directory: + for name, text in artifacts.items(): + if name in self._artifacts: + continue + path = Path(directory) / name + path.write_text(text, encoding='utf-8') + self._tracking.log_artifact(str(path)) + self._artifacts.add(name) + + def _log_dropped(self) -> None: + dropped = { + DROPPED_METRICS: len(self._dropped_metrics), + DROPPED_PARAMS: self._dropped_params, + } + dropped = {name: count for name, count in dropped.items() if count} + if dropped: + self._tracking.log_params(dropped) + + def _log_params(self, params: Mapping[str, str], + limit: int = _MAX_PARAMS - _FINISH_PARAMS) -> None: + kept = {} + for name, value in params.items(): + if '$$' in name or '$$' in value or self._params >= limit: + self._dropped_params += 1 + else: + kept[name] = value + self._params += 1 + if kept: + self._tracking.log_params(kept) + + def _log_metrics(self, metrics: Mapping[str, float], limit: int, + step: int = 0) -> None: + kept = {} + for name, value in metrics.items(): + if '$$' in name: + self._dropped_metrics.add(name) + continue + if name not in self._metrics: + if len(self._metrics) >= limit: + self._dropped_metrics.add(name) + continue + self._metrics.add(name) + kept[name] = value + if kept: + self._tracking.log_metrics(kept, step=step) + + +def _suffixed(name: str, suffix: str) -> str: + if len(name) > 1 and name[0] == name[-1] == '"': + return f'{name[:-1]}_{suffix}"' + return f'{name}_{suffix}' + + +def _taken(name: str, existing: List[str]) -> bool: + resolved = resolve_snowflake_identifier(name) + regex = ''.join('.*' if c == '%' else '.' if c == '_' else re.escape(c) + for c in name) + return any( + other == resolved + or re.fullmatch(regex, other, re.IGNORECASE | re.DOTALL) is not None + for other in existing) diff --git a/packages/kumo-engine/kumo/engine/pquery/training_table.py b/packages/kumo-engine/kumo/engine/pquery/training_table.py index 417295022..54b892c59 100644 --- a/packages/kumo-engine/kumo/engine/pquery/training_table.py +++ b/packages/kumo-engine/kumo/engine/pquery/training_table.py @@ -148,9 +148,9 @@ def _store(self, settings: Any) -> None: self.export( TrainingTableOutputConfig(settings.root.connector, path)) except Exception as error: - warnings.warn( + raise RuntimeError( f"Could not store training table {self.id} under the " - f"artifact root: {error}", UserWarning, stacklevel=3) + f"artifact root: {error}") from error @classmethod def from_source_table( @@ -312,7 +312,13 @@ def _finish(table: TrainingTable, snapshot: GraphSnapshot, request: Any, _record_manifest(table, snapshot, request, seed, source_table=source_table) settings = resolve(mlops) if settings is not None and settings.wants('training_table'): - table._store(settings) + from kumo.engine._mlops import raise_on_error + + try: + table._store(settings) + except Exception as error: + raise_on_error(settings, error, 'Training table generation', + stacklevel=3) return table diff --git a/packages/kumo-engine/kumo/engine/tracking.py b/packages/kumo-engine/kumo/engine/tracking.py new file mode 100644 index 000000000..af95f228f --- /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-tracking adapter for ``kumo.MLOps(tracking=...)``.""" + +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 67bb84d56..72c83129e 100644 --- a/packages/kumo-engine/kumo/engine/trainer/_provenance.py +++ b/packages/kumo-engine/kumo/engine/trainer/_provenance.py @@ -6,12 +6,14 @@ from dataclasses import 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, @@ -30,6 +32,8 @@ def write_provenance( warm_start: Optional[ModelManifest], started_at: datetime, fit_mode: Optional[str], + tracking: Optional[TrackingRef] = None, + dependencies: Optional[List[Dependency]] = None, ) -> None: from kumo_core.artifacts import PROVENANCE_FILE, add_bundle_file @@ -50,6 +54,8 @@ def write_provenance( seconds=round((finished_at - started_at).total_seconds(), 3), ), fit_mode=fit_mode, + tracking=tracking, + dependencies=dependencies, **stamp(), )) add_bundle_file(bundle_dir, PROVENANCE_ROLE, PROVENANCE_FILE) @@ -67,8 +73,9 @@ def cut(model: Optional[ModelManifest]) -> Optional[ModelManifest]: return None if model is None else ModelManifest(id=model.id, reconstructed=True) - return replace(copy, - 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 parent_model(bundle_dir: Path, bundle_id: Optional[str] @@ -84,3 +91,27 @@ def parent_model(bundle_dir: Path, bundle_id: Optional[str] 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 ``bundle_dir``, a copy of the bundle at + ``parent_dir`` with other scoring defaults.""" + from kumo_core.artifacts import ( + PROVENANCE_FILE, + add_bundle_file, + bundle_id_from_manifest, + ) + + parent = parent_model(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..1a916aad4 --- /dev/null +++ b/packages/kumo-engine/kumo/engine/trainer/_run_spec.py @@ -0,0 +1,54 @@ +# 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, field, replace +from pathlib import Path +from typing import Any, Mapping, Optional + +from kumo_core.artifacts import PROVENANCE_FILE + +from kumo.engine._tracking import RunSpec, _flatten + + +@dataclass(frozen=True) +class FitContext: + r"""What a fit trains, known once its model plan is resolved.""" + + query: str + model_plan: Mapping[str, Any] + inputs: Mapping[str, Mapping[str, Optional[str]]] = field( + default_factory=dict) + + +def fit_spec(context: FitContext, *, + tags: Optional[Mapping[str, Any]] = None) -> RunSpec: + return RunSpec( + kind='fit', + params=dict(_flatten(context.model_plan)), + tags=dict(tags or {}), + inputs=dict(context.inputs), + 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, + seconds: Optional[float] = None) -> RunSpec: + r"""``spec`` once the fit returned ``result``, a `TrainingResult`, + after ``seconds``.""" + 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={ + **result.metrics, **({} if seconds is None else { + 'seconds': seconds + }) + }, parents={'bundle_id': result.id}, artifacts=artifacts) 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 6c7d4a660..830937a37 100644 --- a/packages/kumo-engine/kumo/engine/trainer/_spec.py +++ b/packages/kumo-engine/kumo/engine/trainer/_spec.py @@ -26,6 +26,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' @@ -146,6 +147,7 @@ def read_bundle_artifact_paths( return { entry['role']: str(bundle_dir / entry['relative_path']) for entry in artifacts - if entry['role'] not in (SPEC_ROLE, PROVENANCE_ROLE, HOLDOUT_ROLE, + if entry['role'] not in (SPEC_ROLE, PROVENANCE_ROLE, + SCORING_DEFAULTS_ROLE, HOLDOUT_ROLE, HOLDOUT_META_ROLE) }, bundle_id diff --git a/packages/kumo-engine/kumo/engine/trainer/output_config.py b/packages/kumo-engine/kumo/engine/trainer/output_config.py index ddc9a9ef3..41c0c7ba5 100644 --- a/packages/kumo-engine/kumo/engine/trainer/output_config.py +++ b/packages/kumo-engine/kumo/engine/trainer/output_config.py @@ -8,7 +8,11 @@ from kumo_core.sources.object_store_sessions import ObjectStoreSessions from kumoapi.jobs import MetadataField, PredictionArtifactType, WriteMode -from kumoapi.plugin import ModelRegistryLocation, SourceConnector +from kumoapi.plugin import ( + ModelRegistryLocation, + SourceConnector, + check_model_alias, +) from kumoapi.prediction_outputs import resolve_prediction_output_paths from kumo.engine.connector import is_file_destination @@ -102,3 +106,17 @@ 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`` holds an alias and + a model name its registry accepts.""" + 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 e50cca8b0..04219f639 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,12 @@ Any, Callable, Dict, + List, Literal, Mapping, Optional, + Sequence, + Tuple, TypeVar, Union, cast, @@ -28,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 @@ -41,6 +46,8 @@ from kumo_core.types import ExplainMethod, PredictionArtifactType from kumoapi.model_plan import ModelPlan, TrainingTableGenerationPlan from kumoapi.plugin import ( + BUNDLE_ID_KEY, + ModelAliasingSink, ModelFetchingSink, ModelRegistryLocation, PublishedModel, @@ -49,18 +56,33 @@ from kumoapi.train import TrainingStage from typing_extensions import Self +from kumo.engine import _deps as dependency +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._manifest import ModelManifest +from kumo.engine._events import Event, EventEmitter, query_sha256 +from kumo.engine._manifest import Dependency, ModelManifest, TrackingRef, stub +from kumo.engine._mlops import full_error, warn_failure +from kumo.engine._predictor import TRAINER_PREDICTOR from kumo.engine._scratch import new_run_id, trainer_scratch_dir +from kumo.engine._tracking import RunSpec, TrackingSession, fit_inputs from kumo.engine.connector.base import Connector from kumo.engine.connector.source_table import SourceTable from kumo.engine.graph import GraphSnapshot, report_other_graph from kumo.engine.pquery.prediction_table import PredictionTable from kumo.engine.pquery.training_table import TrainingTable +from kumo.engine.trainer._run_spec import FitContext, fit_spec, fitted_spec +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, HOLDOUT_ROLE, @@ -73,6 +95,7 @@ from kumo.engine.trainer.output_config import ( ModelOutputConfig, PredictionOutputConfig, + check_model_location, ) logger = logging.getLogger(__name__) @@ -84,6 +107,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 @@ -101,6 +125,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) @@ -123,10 +149,58 @@ 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: its ``platform``, + ``run_id`` and ``url``; None when the fit was not 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: + """Publish this model's bundle to a ``ModelOutputConfig``, or + register it at a `kumo.MLOps` ``registry`` through its tracker; a + registry that already holds the bundle returns that version.""" + from kumo.engine._mlops import MLOps + + if isinstance(config, MLOps): + return _register_result(self, config) + _require_publisher(config) + return _publish(self.bundle_dir, self.bundle_id, self.metrics, config) + + def with_scoring_defaults(self, + defaults: ScoringDefaults) -> 'TrainingResult': + """A copy of this model whose ``predict`` applies ``defaults`` when a + call omits them; a new bundle, whose manifest names this one in + ``parents.derived_from``. To register a model with its defaults, + pass ``fit(scoring_defaults=...)`` instead.""" + from kumo_core.artifacts import bundle_task_type + + from kumo.engine.trainer._provenance import record_derivation + + 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: @@ -276,30 +350,73 @@ def _listed_holdout(bundle_dir: Path) -> tuple[Optional[Path], Optional[bool]]: return path, truncated -def _store_model(result: TrainingResult, settings: Any) -> None: +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]]: + r"""Store a fitted bundle under the artifact root, log it in the + session's run and register it, as MLOps ``settings`` say; returns the + registered version and the first failure. A model whose bundle or + assets (``stored``, the asset failure) could not be stored is not + registered.""" + 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 else _mlflow_registry(settings)) + error = stored or error + registry = None if stored else _connector_registry(settings) + if registry is None: + return registered, error + errors: List[BaseException] = [] + try: + registered = _publish(bundle_dir, None if manifest is None else + manifest.id, metrics, registry, errors, 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 (errors or [None])[0] + + +def _store_model(bundle_dir: Path, manifest: Optional[ModelManifest], + settings: Any) -> None: r"""Copy the bundle to ``/models//`` unless a complete - copy is there; failures warn.""" + copy is there.""" from kumo_core.object_store_sink import holds_bundle, upload_bundle from pyarrow.fs import FileSelector from kumo.engine.connector import file_store - if result.bundle_id is None: + if manifest is None: return try: filesystem, root = file_store( settings.root.connector, - settings.root.location(f'models/{result.bundle_id}')) - if holds_bundle(filesystem, root, str(result.bundle_dir)): + settings.root.location(f'models/{manifest.id}')) + if holds_bundle(filesystem, root, str(bundle_dir)): return if filesystem.get_file_info( FileSelector(root, allow_not_found=True)): filesystem.delete_dir_contents(root, missing_dir_ok=True) - upload_bundle(filesystem, root, str(result.bundle_dir)) + upload_bundle(filesystem, root, str(bundle_dir)) except Exception as error: - warnings.warn( - f"Could not store model {result.bundle_id} under the artifact " - f"root: {error}", UserWarning, stacklevel=4) + raise RuntimeError(f"Could not store model {manifest.id} under the " + f"artifact root: {error}") from error def _published_holdout(trainer: Any) -> tuple[Optional[Path], Optional[bool]]: @@ -330,6 +447,93 @@ def _to_core_model_plan(model_plan: ModelPlan) -> CoreModelPlan: return internal_model_plan_from_public(model_plan) +def _link_or_copy(source: str, destination: str) -> None: + try: + os.link(source, destination) + except OSError: + shutil.copy2(source, destination) + + +def _require_publisher(export_config: ModelOutputConfig) -> None: + if not callable( + getattr(export_config.connector.sink_backend, 'publish_model', + None)): + raise NotImplementedError( + f'{type(export_config.connector).__name__} cannot publish ' + f'models') + check_model_location(export_config) + + +def _register_result(result: TrainingResult, settings: Any) -> PublishedModel: + if settings.registry is None: + raise ValueError('export(cfg) registers at cfg.registry, which is ' + 'not set.') + register = getattr(settings.tracking, 'register', None) + if callable(register): + return register(result, settings.registry) + registry = cast(ModelOutputConfig, _connector_registry(settings)) + _require_publisher(registry) + return _publish(result.bundle_dir, result.bundle_id, result.metrics, + registry) + + +def _mlflow_registry(settings: Any) -> Optional[ModelRegistryLocation]: + if settings is None or not _logs_models(settings.tracking): + return None + return settings.registry + + +def _connector_registry(settings: Any) -> Optional[ModelOutputConfig]: + if (settings is None or settings.registry is None + or _logs_models(settings.tracking)): + return None + return ModelOutputConfig(settings.tracking.connector, + location=settings.registry) + + +def _logs_models(tracking: Any) -> bool: + return callable(getattr(tracking, 'log_model', None)) + + +def _publish(bundle_dir: Path, bundle_id: Optional[str], + metrics: Mapping[str, float], export_config: ModelOutputConfig, + errors: Optional[List[BaseException]] = None, + warn: bool = True) -> PublishedModel: + 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 location and 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=location, + metadata=metadata, + metrics=metrics, + ) + if location and alias: + try: + _set_alias(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) + if errors is not None: + errors.append(error) + return published + + def _stage_model_bundle(artifacts: Mapping[str, str], dest: Path) -> Path: from kumo_core.artifacts import stage_model_bundle return stage_model_bundle(artifacts, dest) @@ -419,6 +623,18 @@ def result(self) -> Optional[TrainingResult]: """The last fit's result, or the loaded bundle's.""" return self._training_result + 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: + """Publish the fitted or loaded model bundle, as + :meth:`TrainingResult.export`.""" + return self._fitted('export').export(config) + @_without_torch_scatter_hint def fit( self, @@ -426,8 +642,10 @@ def fit( *, 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`. @@ -448,6 +666,9 @@ 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: prediction arguments the bundle applies when + ``predict`` omits them, such as a classification threshold; + stored, registered and served with the model. 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; @@ -457,16 +678,91 @@ 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 under the artifact root and + the fit is recorded by its ``tracking``. + callbacks: called with each progress + `kumo.engine.tracking.Event`; a callback that raises is logged + and not called again. Each call returns a new `TrainingResult`; `Trainer.result` holds the latest. """ + if export_config is not None: + _require_publisher(export_config) + 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()}' + 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() + try: + started = { + 'snapshot_id': snapshot.id, + 'query_sha256': query_sha256(train_table.query), + } + except Exception: + started = {} + events.emit('fit_started', params=started) + try: + result = self._fit(train_table, snapshot=snapshot, + export_config=export_config, + scoring_defaults=scoring_defaults, + warm_start=warm_start, run_id=run_id, + events=events, tracking=session, + deps=deps) + except BaseException as error: + events.emit('fit_ended', params={ + 'status': 'failed', + 'error': format_error(error), + }) + session.finish(None, error) + raise + seconds = time.monotonic() - fit_start + events.emit('fit_ended', metrics={ + **result.metrics, 'seconds': seconds + }, params={'status': 'succeeded'}) + registered, error = _store_and_register( + result.bundle_dir, result.manifest or ( + None if result.bundle_id is None else stub( + 'model', result.bundle_id)), settings, session=session, + snapshot=snapshot, metrics=result.metrics, 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, seconds), + error) + raise_on_error(settings, session.error or error, + 'Trainer.fit (its model is in trainer.result)') + return result + + def _fit( + self, + train_table: TrainingTable, + *, + snapshot: GraphSnapshot, + export_config: Optional[ModelOutputConfig], + scoring_defaults: Optional[ScoringDefaults], + warm_start: Optional[Self], + run_id: str, + events: EventEmitter, + tracking: TrackingSession, + deps: Dependencies, + ) -> TrainingResult: from kumo_core.automl.runtime import ( inherit_warm_start_model_plan, validate_warm_start_task_type, ) - from kumo_core.io.json import load_json + from kumo_core.io.json import load_json, to_json_dict from kumo_core.training.orchestration import ( CallableTrainingOptions, ResolvedTrainingInput, @@ -481,20 +777,10 @@ def fit( ) from kumo_core.types import TaskType - from kumo.engine.trainer._runtime import local_training_bindings - - if export_config is not None and not callable( - getattr(export_config.connector.sink_backend, 'publish_model', - None)): - raise NotImplementedError( - f'{type(export_config.connector).__name__} cannot publish ' - f'models') - from kumo.engine._mlops import resolve from kumo.engine.trainer._provenance import parent_model + from kumo.engine.trainer._runtime import local_training_bindings - settings = resolve(mlops) started_at = datetime.now(timezone.utc) - run_id = f'{self._id}-fit-{new_run_id()}' run_dir = self.scratch_dir / 'runs' / run_id model_plan = _to_core_model_plan(self._model_plan) @@ -528,12 +814,15 @@ 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()) + from kumo_core.artifacts import bundle_task_type + validate_warm_start_task_type( TaskType(workflow.query.task_type), - TaskType(source_schema['task_type']), + bundle_task_type(warm_start_artifacts), ) from kumo_core.pquery import ( query_is_temporal, @@ -574,18 +863,20 @@ def fit( 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=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=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) @@ -605,6 +896,11 @@ def fit( and not materialized.task.task_type.is_link_pred): raise ValueError('The test split is empty: shorten the query ' 'window or add data.') + tracking.start(lambda: fit_spec(FitContext( + query=train_table.query, + model_plan=to_json_dict(materialized.model_plan), + inputs=fit_inputs(train_table.manifest), + ))) training_input = ResolvedTrainingInput( task=materialized.task, model_plan=materialized.model_plan, @@ -617,7 +913,7 @@ def fit( 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(), ) @@ -663,14 +959,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=dict( train_table=train_table, warm_start=None if warm_start is None else parent_model( warm_start._training_result.bundle_dir, warm_start_bundle_id), - started_at=started_at, fit_mode=fit_mode), - settings=settings, + started_at=started_at, fit_mode=fit_mode, + tracking=tracking.pointer()), + deps=deps, ) def _finish_fit( @@ -684,7 +982,8 @@ def _finish_fit( graph_config: Optional[Dict[str, Any]] = None, model_plan: Optional[ModelPlan] = None, provenance: Optional[Mapping[str, Any]] = None, - settings: Optional[Any] = 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 @@ -700,6 +999,11 @@ def _finish_fit( training_table_generation=training_table_generation, graph_config=graph_config, ) + if scoring_defaults is not None: + write_scoring_defaults(training_result.bundle_dir, + scoring_defaults) + training_result = replace(training_result, + scoring_defaults=scoring_defaults) training_result = _bundle_holdout(training_result) try: training_result = replace( @@ -710,31 +1014,21 @@ def _finish_fit( UserWarning, stacklevel=3) if provenance is not None and training_result.bundle_id is not None: try: - write_provenance(training_result.bundle_dir, - bundle_id=training_result.bundle_id, - **provenance) + write_provenance( + training_result.bundle_dir, + bundle_id=training_result.bundle_id, + dependencies=self._dependencies(training_result, deps + or Dependencies()), + **provenance) except Exception as error: warnings.warn( f"Publishing the model bundle without _provenance.json: " f"{error}", UserWarning, stacklevel=3) if export_config is not None: - spec = read_training_spec(training_result.bundle_dir) - published = export_config.connector.sink_backend.publish_model( - export_config.connector, - str(training_result.bundle_dir), - path=export_config.path, - location=export_config.location, - metadata={ - name: str(getattr(spec, name)) - for name in ('spec_version', 'predictive_query', - 'snapshot_id') - }, - metrics=training_result.metrics, - ) - training_result = replace(training_result, - published_model=published) - if settings is not None and settings.wants('model'): - _store_model(training_result, settings) + training_result = replace( + training_result, published_model=_publish( + training_result.bundle_dir, training_result.bundle_id, + training_result.metrics, export_config)) training_result = replace(training_result, _trainer=self) self._training_result = training_result @@ -745,6 +1039,18 @@ def _finish_fit( return training_result + 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_torch_scatter_hint def predict( self, @@ -783,9 +1089,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, @@ -852,11 +1161,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: - if self._training_result is None: - raise RuntimeError( - "Trainer.predict() requires a prior Trainer.fit() call on " - "this same Trainer instance.") + training_result = self._fitted('predict') + if binary_classification_threshold is None: + binary_classification_threshold = ( + 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, @@ -874,8 +1184,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( @@ -890,6 +1202,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, @@ -916,7 +1230,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) @@ -980,12 +1294,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, @@ -1062,12 +1379,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`` from the registry its + tracker registers in. """ 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: + load = getattr(model_registry.tracking, 'load_model', None) + 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, @@ -1077,10 +1402,18 @@ 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) + trainer = cls._load_bundle(bundle_dir, embedding_models, + accelerator=accelerator) + warn_if_other_bundle(getattr(published, 'metadata', None), trainer) + return trainer @classmethod def _load_bundle( @@ -1104,9 +1437,49 @@ 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 + from kumo.engine.trainer._provenance import read_provenance + + 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 _set_alias(connector: Connector, location: ModelRegistryLocation, + alias: str) -> None: + sink = connector.sink_backend + if not isinstance(sink, ModelAliasingSink): + raise NotImplementedError( + f'{type(connector).__name__} cannot set model aliases') + sink.set_model_alias(cast(SourceConnector, connector), location, alias) + + +def warn_if_other_bundle(metadata: Optional[Mapping[str, Any]], + trainer: Trainer) -> None: + recorded = (metadata or {}).get(BUNDLE_ID_KEY) + 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 fb5741d0b..c55d4f8f5 100644 --- a/packages/kumo-engine/pyproject.toml +++ b/packages/kumo-engine/pyproject.toml @@ -71,6 +71,18 @@ text=[ 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..27b1cc778 --- /dev/null +++ b/packages/kumo-engine/test/_snowflake_experiment_fake.py @@ -0,0 +1,192 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""A recording stand-in for `snowflake.ml.experiment.ExperimentTracking` that +keeps its observable rules (singleton, one active run, LIKE-matched resume).""" + +import os +import re +import sys +import types +from typing import Any, Dict, List, Optional, Tuple + +from kumo_core.sources.sql import resolve_snowflake_identifier as _identifier + +STATUSES = ('FAILED', 'FINISHED') + + +def _like(pattern: str, name: str) -> bool: + regex = ''.join('.*' if c == '%' else '.' if c == '_' else re.escape(c) + for c in pattern) + return re.fullmatch(regex, name, re.IGNORECASE | re.DOTALL) is not None + + +class FakeRun: + def __init__(self, tracking: 'FakeExperimentTracking', experiment: str, + name: str) -> None: + self._tracking = tracking + self.experiment_name = experiment + self.name = name + + def __enter__(self) -> 'FakeRun': + self._tracking.calls.append(('run.__enter__', self.name)) + self._tracking.lineage_stack.append((self.experiment_name, self.name)) + return self + + def __exit__(self, exc_type: Any, exc: Any, tb: Any) -> None: + self._tracking.calls.append(('run.__exit__', self.name)) + self._tracking.lineage_stack.pop() + if self._tracking._run is self: + self._tracking.end_run( + status='FAILED' if exc_type is not None else None) + + +class FakeExperimentTracking: + _instance: Optional['FakeExperimentTracking'] = None + + def __new__(cls, *args: Any, **kwargs: Any) -> 'FakeExperimentTracking': + if cls._instance is None: + cls._instance = super().__new__(cls) + return cls._instance + + def __init__(self, session: Any, *, database_name: Optional[str] = None, + schema_name: Optional[str] = None, + capture_source_info: bool = True, model_registry: Any = None, + warnings_as_errors: bool = False) -> None: + self.capture_source_info = capture_source_info + self.warnings_as_errors = warnings_as_errors + if hasattr(self, '_initialized'): + self.calls.append(('__init__', database_name, schema_name)) + return + self._session = session + self.database = _identifier(database_name + or session.get_current_database()) + self.schema = _identifier(schema_name or session.get_current_schema() + or 'PUBLIC') + self.calls: List[Tuple[Any, ...]] = [('__init__', database_name, + schema_name)] + self.fail_on: set = set() + self.experiment: Optional[str] = None + self._run: Optional[FakeRun] = None + self.runs: Dict[str, Dict[str, Any]] = {} + self.lineage_stack: List[Tuple[str, str]] = [] + self._sql_client = types.SimpleNamespace( + show_runs_in_experiment=self._show_runs) + self._initialized = True + + def _show_runs(self, *, experiment_name: str, + like: Optional[str] = None) -> List[Dict[str, Any]]: + self._record('show_runs', experiment_name, like) + return self._matching(experiment_name, like) + + def _matching(self, experiment: str, + like: Optional[str]) -> List[Dict[str, Any]]: + return [{'name': name} for name, run in self.runs.items() + if run['experiment'] == experiment + and (like is None or _like(like, name))] + + def _record(self, *call: Any) -> None: + self.calls.append(call) + if call[0] in self.fail_on: + raise RuntimeError(f'fake Snowflake error in {call[0]}') + + def set_experiment(self, experiment_name: str, + database_name: Optional[str] = None, + schema_name: Optional[str] = None) -> Any: + self._record('set_experiment', experiment_name, database_name, + schema_name) + if database_name is not None and schema_name is None: + schema_name = 'PUBLIC' + location = (_identifier(database_name or self.database), + _identifier(schema_name or self.schema)) + name = _identifier(experiment_name) + if (name, *location) != (self.experiment, self.database, self.schema): + self._run = None + self.experiment = name + self.database, self.schema = location + return types.SimpleNamespace(name=name) + + def start_run(self, run_name: Optional[str] = None) -> FakeRun: + self._record('start_run', run_name) + if self._run is not None: + raise RuntimeError('A run is already active. Please end the ' + 'current run before starting a new one.') + if self.experiment is None: + self.experiment = 'DEFAULT' + matches = self._matching(self.experiment, + run_name) if run_name else [] + if matches: + name = matches[0]['name'] + if self.runs[name]['status'] != 'RUNNING': + raise RuntimeError(f'Run {run_name} exists but cannot be ' + f'resumed as it is no longer running.') + else: + name = _identifier(run_name or f'RUN_{len(self.runs) + 1}') + self.runs[name] = {'experiment': self.experiment, + 'status': 'RUNNING', 'params': {}, + 'metrics': [], 'artifacts': {}} + self._run = FakeRun(self, self.experiment, name) + return self._run + + def end_run(self, run_name: Optional[str] = None, + status: Optional[str] = None) -> None: + self._record('end_run', run_name, status) + if self.experiment is None: + raise RuntimeError('No experiment set.') + if self._run is None: + raise RuntimeError('No run is active.') + if status is not None and status not in STATUSES: + raise ValueError(f"Unsupported status value: '{status}'.") + self.runs[self._run.name]['status'] = status or 'FINISHED' + self._run = None + + def _active(self) -> Dict[str, Any]: + if self._run is None: + self.start_run() + return self.runs[self._run.name] + + def log_params(self, params: Dict[str, Any]) -> None: + self._record('log_params', dict(params)) + self._active()['params'].update( + {key: str(value) for key, value in params.items()}) + + def log_metrics(self, metrics: Dict[str, float], step: int = 0) -> None: + self._record('log_metrics', dict(metrics), step) + self._active()['metrics'].extend( + (key, float(value), step) for key, value in metrics.items()) + + def log_artifact(self, local_path: str, + artifact_path: Optional[str] = None) -> None: + self._record('log_artifact', local_path, artifact_path) + with open(local_path, encoding='utf-8') as file: + self._active()['artifacts'][os.path.basename(local_path)] = ( + file.read()) + + +class FakeSession: + def __init__(self, database: Optional[str] = 'DB', + schema: Optional[str] = 'PUBLIC') -> None: + self.database, self.schema = database, schema + self.statements: List[str] = [] + + def sql(self, query: str) -> Any: + self.statements.append(query) + return types.SimpleNamespace(collect=lambda: []) + + def get_current_database(self) -> Optional[str]: + return self.database + + def get_current_schema(self) -> Optional[str]: + return self.schema + + +def install(monkeypatch: Any) -> type: + """Makes `import snowflake.ml.experiment` return a fresh fake.""" + FakeExperimentTracking._instance = None + module = types.ModuleType('snowflake.ml.experiment') + module.ExperimentTracking = FakeExperimentTracking + for name in ('snowflake', 'snowflake.ml'): + if name not in sys.modules: + monkeypatch.setitem(sys.modules, name, types.ModuleType(name)) + monkeypatch.setitem(sys.modules, 'snowflake.ml.experiment', module) + return FakeExperimentTracking diff --git a/packages/kumo-engine/test/test_bundle_provenance.py b/packages/kumo-engine/test/test_bundle_provenance.py index 854e6e0dc..f07b93a10 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 embedded_model from kumo.engine.trainer._spec import MANIFEST_FILE -from kumo.engine.trainer.trainer import TrainingResult +from kumo.engine.trainer.trainer import TrainingResult, _store_and_register def test_fit_records_a_self_contained_manifest(fitted): @@ -103,13 +103,14 @@ def _finish_stub_fit(bundle, trainer, warm_start=None, settings=None): (bundle / MANIFEST_FILE).write_text(json.dumps({'artifacts': []})) add_bundle_file(bundle, 'trainer_config', 'model/config.json') table = SimpleNamespace(tt_id='tt-raw', id='tt-1', manifest=None) - return trainer._finish_fit( + result = trainer._finish_fit( TrainingResult({}, {}, bundle), predictive_query='PREDICT x', snapshot_id='snap', export_config=None, training_table_generation=TrainingTableGenerationPlan(), provenance=dict(train_table=table, warm_start=warm_start, - started_at=datetime.now(timezone.utc), fit_mode=None), - settings=settings) + started_at=datetime.now(timezone.utc), fit_mode=None)) + _store_and_register(result.bundle_dir, result.manifest, settings) + return result def test_missing_parents_are_stubs_and_retrain_chains_stay_flat(tmp_path): 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_fit_events.py b/packages/kumo-engine/test/test_fit_events.py new file mode 100644 index 000000000..08a22729c --- /dev/null +++ b/packages/kumo-engine/test/test_fit_events.py @@ -0,0 +1,153 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import json + +import numpy as np +import pandas as pd +import pytest +from kumoapi.model_plan import ModelPlan + +import kumo.engine as kumo +from kumo.engine._scratch import get_scratch_dir +from kumo.engine.tracking import Event + + +@pytest.fixture +def table_and_snapshot(tmp_path): + previous_scratch = get_scratch_dir() + kumo.configure(scratch_dir=tmp_path / 'scratch') + try: + rng = np.random.default_rng(0) + (tmp_path / 'data').mkdir() + pd.DataFrame({ + 'ID': np.arange(400), + 'Y': rng.random(400), + 'F1': rng.normal(size=400), + }).to_parquet(tmp_path / 'data' / 'ents.parquet') + connector = kumo.FileConnector(root_dir=str(tmp_path / 'data')) + ents = kumo.Table.from_source_table(connector['ents.parquet'], + primary_key='ID').infer_metadata() + graph = kumo.Graph(tables={'ents': ents}, edges=[]) + snapshot = graph.snapshot() + query = kumo.PredictiveQuery(graph=graph, + query='PREDICT ents.Y FOR EACH ents.ID') + yield query.generate_training_table(snapshot), snapshot + finally: + kumo.configure(scratch_dir=previous_scratch) + + +def test_fit_reports_progress_events(table_and_snapshot, monkeypatch, + caplog): + from kumo_core.automl.optuna_searcher import TPESearcher + + update = TPESearcher.update + + def fail_second_trial(self, *args, **kwargs): + if self._num_updates: + raise RuntimeError('searcher down') + self._num_updates = 1 + return update(self, *args, **kwargs) + + monkeypatch.setattr(TPESearcher, '_num_updates', 0, raising=False) + monkeypatch.setattr(TPESearcher, 'update', fail_second_trial) + table, snapshot = table_and_snapshot + plan = ModelPlan() + plan.training_job.num_experiments = 2 + plan.training_job.disable_explain = True + plan.optimization.max_epochs = 2 + plan.optimization.max_steps_per_epoch = 2 + events = [] + broken_calls = [] + + def broken(event): + broken_calls.append(event) + raise RuntimeError('callback bug') + + result = kumo.Trainer(plan).fit(table, snapshot=snapshot, + callbacks=[broken, events.append]) + + assert len(broken_calls) == 1 + assert 'callback bug' in caplog.text + 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[1].params['attempt'] == 1 + assert events[4].params == {'status': 'succeeded', 'attempt': 1} + assert {'val_mae', 'seconds'} <= set(events[4].metrics) + assert events[8].params == { + 'status': 'failed', + 'attempt': 1, + 'error': 'RuntimeError: searcher down' + } + assert events[-1].metrics == {**result.metrics, + 'seconds': events[-1].metrics['seconds']} + assert events[-1].params == {'status': 'succeeded'} + assert [json.loads(json.dumps(event.to_dict()))['name'] + for event in events] == names + + +def test_failed_fit_reports_fit_ended(table_and_snapshot, monkeypatch): + import kumo.engine.trainer.trainer as trainer_module + + def fail(*args, **kwargs): + raise ValueError('stop') + + monkeypatch.setattr(trainer_module, 'build_generated_training_task_dataset', + fail) + table, snapshot = table_and_snapshot + events = [] + + with pytest.raises(ValueError, match='stop'): + kumo.Trainer(ModelPlan()).fit(table, snapshot=snapshot, + callbacks=[events.append]) + + assert [event.name for event in events] == ['fit_started', 'fit_ended'] + assert events[-1].params == { + 'status': 'failed', + 'error': 'ValueError: stop' + } + + +def test_event_metrics_keep_finite_numbers(): + from kumo.engine._events import EventEmitter + + received = [] + EventEmitter([received.append]).emit( + 'epoch_ended', metrics={'val_loss': 'n/a', 'val_auc': float('nan'), + 'val_mae': 0.5}) + + assert received[0].metrics == {'val_mae': 0.5} + + +def test_epoch_event_without_recorded_metrics_does_not_raise(tmp_path): + from datetime import datetime, timezone + from types import SimpleNamespace + + from kumo_core.callbacks.train_progress_logger import ( + AutoTrainingProgressLogger, + ) + + received = [] + logger = AutoTrainingProgressLogger( + str(tmp_path), 1, on_event=lambda *a, **k: received.append(a)) + progress = SimpleNamespace(num_epochs_completed=1, metrics=[], + estimated_train_time=0.0) + logger.set_current_trainer( + SimpleNamespace(training_progress_logger=SimpleNamespace( + progress=progress)), trial=1) + logger.progress.start_time = datetime.now(timezone.utc) + + logger.on_train_epoch_end(None, None) + + assert received == [] diff --git a/packages/kumo-engine/test/test_mlflow_model.py b/packages/kumo-engine/test/test_mlflow_model.py new file mode 100644 index 000000000..51eb51fb7 --- /dev/null +++ b/packages/kumo-engine/test/test_mlflow_model.py @@ -0,0 +1,522 @@ +# 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.model_plan import ModelPlan +from kumoapi.plugin import PublishedModel +from test_mlflow_tracking import ( # noqa: F401 + _runs, + table_and_snapshot, + tracking_uri, +) + +import kumo.engine as kumo +from kumo.engine._scratch import get_scratch_dir + +mlflow = pytest.importorskip('mlflow') + + +def _plan(): + plan = ModelPlan() + plan.training_job.num_experiments = 1 + plan.training_job.disable_explain = True + plan.optimization.max_epochs = 2 + plan.optimization.max_steps_per_epoch = 2 + return plan + + +def _sorted(frame): + return frame.sort_values('ENTITY').reset_index(drop=True) + + +def test_fit_logs_registers_and_serves_the_model(table_and_snapshot, + tmp_path, caplog, + monkeypatch): + from kumo.engine.integrations.mlflow import MlflowTracking, load_model + + table, snapshot = table_and_snapshot + plan = _plan() + cfg = kumo.MLOps(str(tmp_path / 'root'), + tracking=MlflowTracking(log_system_metrics=False), + registry=kumo.ModelRegistryLocation('m', + alias='champion')) + trainer = kumo.Trainer(plan) + result = trainer.fit(table, snapshot=snapshot, mlops=cfg) + + assert 'MLflow tracking failed' not in caplog.text + parent, = [ + run for run in _runs(mlflow.MlflowClient()) + if 'mlflow.parentRunId' not in run.data.tags + ] + model, = mlflow.search_logged_models(output_format='list') + assert model.source_run_id == parent.info.run_id + metrics = {metric.key: metric.value for metric in model.metrics} + assert metrics == {**result.metrics, 'seconds': metrics['seconds']} + version = mlflow.MlflowClient().get_model_version_by_alias('m', 'champion') + assert version.source == model.model_uri + assert {k: v for k, v in parent.data.tags.items() + if k.startswith('kumo.')} == { + 'kumo.kind': 'fit', + 'kumo.bundle_id': result.id, + } + assert result.registered_model == PublishedModel( + name='m', version=str(version.version), + uri=f'models:/m/{version.version}', + metadata={'kumo_bundle_id': result.id}) + 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 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_the_query_population_and_timed_rows( + fitted, tracking_uri, tmp_path, monkeypatch): + from kumo.engine.integrations.mlflow import log_model + + trainer, table, snapshot = fitted + location = snapshot.export( + kumo.SnapshotOutputConfig( + kumo.FileConnector(root_dir=str(tmp_path / 'exports')), + 'snapshots')) + + info = log_model(trainer, snapshot=snapshot, + serving_snapshot=str(tmp_path / 'exports' / location)) + + assert mlflow.get_run(info.run_id).info.status == 'FINISHED' + requirements = mlflow.pyfunc.get_model_dependencies(info.model_uri) + assert f'kumo-engine[gcs,s3]=={kumo.__version__}' in open( + requirements).read().splitlines() + assert {metric.key + for metric in mlflow.get_logged_model(info.model_id).metrics + } == set(trainer.result.metrics) + assert [c['name'] for c in info.signature.outputs.to_dict() + ] == ['ENTITY', 'TARGET_PRED', 'False_PROB', 'True_PROB'] + pyfunc = mlflow.pyfunc.load_model(info.model_uri) + params = { + 'snapshot': str(tmp_path / 'exports' / location), + 'prediction_time': str(pd.Timestamp(table.timestamps[0])), + 'binary_classification_threshold': 0.5, + } + expected = _sorted( + trainer.predict(table, snapshot=snapshot, + binary_classification_threshold=0.5).frame) + empty = pd.DataFrame({'ENTITY': pd.Series([], dtype='int64')}) + tables = get_scratch_dir() / 'prediction_tables' + before = set(tables.glob('*')) + predict = kumo.Trainer.predict + + def predict_beside_another_call(self, *args, **kwargs): + (tables / 'other-call').mkdir(parents=True, exist_ok=True) + return predict(self, *args, **kwargs) + + monkeypatch.setattr(kumo.Trainer, 'predict', predict_beside_another_call) + + assert pyfunc.predict(empty, params=params).empty + served = pyfunc.unwrap_python_model()._snapshots + assert list(served) == [params['snapshot']] + pd.testing.assert_frame_equal( + _sorted( + pyfunc.predict(empty, params={ + **params, 'snapshot': None, 'population': True + })), expected) + for rows, population in ((empty, True), + (pd.DataFrame({'ENTITY': range(50)}), False)): + pd.testing.assert_frame_equal( + _sorted( + pyfunc.predict(rows, + params={**params, 'population': population})), + expected) + scratch = pyfunc.unwrap_python_model()._predictor.trainer.scratch_dir + assert list(scratch.iterdir()) == [] + assert set(tables.glob('*')) == before | {tables / 'other-call'} + + +def test_pyfunc_names_a_missing_prediction_time_and_needs_no_threshold( + fitted): + from kumo.engine.integrations._mlflow_model import predict + + trainer, table, snapshot = fitted + rows = pd.DataFrame({'ENTITY': range(5)}) + with pytest.raises(ValueError) as error: + predict(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(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.integrations._mlflow_model import predict + + 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(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.integrations._mlflow_model 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(tracking_uri, 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 + from kumo.engine._tracking import TrackingSession + from kumo.engine.trainer._run_spec import FitContext, fit_spec + + monkeypatch.setattr(_tracking, 'FINISH_SECONDS', 0.01) + + class Tracker: + def __init__(self, fail): + self.fail, self.calls = fail, [] + + def start(self, record): + self.calls.append('start') + + def on_trial(self, record, event): + pass + + def log_model(self, bundle_dir, manifest, snapshot, registry): + time.sleep(0.2) + self.calls.append('log_model') + if self.fail: + raise RuntimeError('upload failed') + + def finish(self, record, error): + self.calls.append(('finish', error)) + + result = SimpleNamespace(metrics={}, bundle_dir='missing', bundle_id='b') + for fail in (False, True): + tracker = Tracker(fail) + session = TrackingSession(tracker) + session.start(lambda: fit_spec(FitContext('q', {}))) + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter('always') + _, error = session.log_model(result.bundle_dir, None, None) + session.finish(None, error) + session._worker.join(5) + assert tracker.calls[:2] == ['start', 'log_model'] + assert tracker.calls[2] == ('finish', error) + assert (error is not None) is fail + assert any('upload failed' in str(w.message) for w in caught) is fail + + +@pytest.mark.parametrize('stall', ['slow upload', 'dead thread']) +def test_model_logging_wait_is_bounded(monkeypatch, stall): + import threading + + from kumo.engine import _tracking + from kumo.engine._tracking import TrackingSession + from kumo.engine.trainer._run_spec import FitContext, fit_spec + + monkeypatch.setattr(_tracking, 'MODEL_SECONDS', 0.5) + release = threading.Event() + + class Tracker: + calls = [] + + def start(self, record): + pass + + def on_trial(self, record, event): + pass + + def log_model(self, bundle_dir, manifest, snapshot, registry): + release.wait(30) + + def finish(self, record, error): + self.calls.append((record, error)) + + tracker = Tracker() + session = TrackingSession(tracker) + session.start(lambda: fit_spec(FitContext('q', {}))) + if stall == 'dead thread': + session._submit(lambda: (_ for _ in ()).throw(SystemExit())) + session._worker.join(5) + with pytest.warns(UserWarning, match='Could not log the model'): + _, error = session.log_model(None, None, None) + release.set() + session.finish(None, error) + + assert type(error) is (TimeoutError + if stall == 'slow upload' else RuntimeError) + assert tracker.calls == [(None, None)] + + +def test_model_upload_gets_a_longer_http_timeout_and_a_late_one_is_dropped( + tracking_uri, monkeypatch, caplog): + import os + + from test_tracking_contract import CONTEXT + + from kumo.engine._tracking import TrackingSession + from kumo.engine.integrations import mlflow as integration + 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(tmp_path): + from kumo_core.encoders.config import EncodersConfig + from kumo_core.io.json import save_json + from kumoapi.encoder import Numerical, TextEmbedding + + from kumo.engine.integrations.mlflow import _requirement + + def requirement(encoders): + path = tmp_path / 'encoders.json' + save_json(EncodersConfig(encoders={'users': encoders}), str(path)) + return _requirement(path) + + 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 + + from kumo.engine._tracking import TrackingSession + from kumo.engine.trainer._run_spec import FitContext, fit_spec + + class Tracker: + def start(self, record): + raise RuntimeError('server down') + + def log_model(self, bundle_dir, manifest, snapshot, registry): + raise AssertionError('not called') + + def finish(self, record, error): + pass + + session = TrackingSession(Tracker()) + session.start(lambda: fit_spec(FitContext('q', {}))) + 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, 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(_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, = [run for run in _runs(mlflow.MlflowClient()) + if 'mlflow.parentRunId' not in run.data.tags] + 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, 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(_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.)' + ] 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..4b9b360a8 --- /dev/null +++ b/packages/kumo-engine/test/test_mlflow_target.py @@ -0,0 +1,363 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +r"""How `MlflowTracking` reaches its server: explicit credentials set exactly +the state MLflow reads, only around each call; ambient ones set nothing. The +last tests pin the MLflow and plugin internals this relies on.""" + +import os +import threading +from types import SimpleNamespace + +import pytest + +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' + + +@pytest.fixture +def tracking_uri(tmp_path): + previous = mlflow.get_tracking_uri() + uri = f'sqlite:///{tmp_path / "mlflow.db"}' + mlflow.set_tracking_uri(uri) + yield uri + mlflow.set_tracking_uri(previous) + + +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(): + # A profile or notebook identity: Unity Catalog registration + # downloads through the process-wide URI, so it points at Databricks. + 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 + + for target in (resolve(SageMakerMLflow(ARN, S3Identity())), + resolve(SageMakerMLflow(ARN))): + 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_')} + # Unity Catalog registration downloads through the global URI. + 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_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') # registers the arn store + from kumo.engine.integrations.mlflow import MlflowTracking + + seen = [] + + def lookup(self, name): + seen.append((self.tracking_uri, mlflow.get_tracking_uri())) + + 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 + + +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 with these credentials'): + kumo.MLOps(str(tmp_path), tracking=tracking) + with pytest.raises(ValueError, match='pass 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 + + +# Pins: MLflow and plugin internals the code above relies on. + + +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 + + 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') + + first = _register(target, client, info, location, + {'kumo_bundle_id': 'a'}) + with caplog.at_level('INFO'): + again = _register(target, client, info, location, + {'kumo_bundle_id': 'a'}) + other = _register(target, client, info, location, + {'kumo_bundle_id': '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 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..709b9885f --- /dev/null +++ b/packages/kumo-engine/test/test_mlflow_tracking.py @@ -0,0 +1,425 @@ +# 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 numpy as np +import pandas as pd +import pytest +from kumo_core.artifacts import PROVENANCE_FILE +from kumoapi.encoder import Numerical +from kumoapi.model_plan import ModelPlan +from test_tracking_contract import CONTEXT, join_tracking_threads, replay + +import kumo.engine as kumo +from kumo.engine._scratch import get_scratch_dir +from kumo.engine._tracking import run_record +from kumo.engine.trainer._run_spec import fit_spec + +mlflow = pytest.importorskip('mlflow') +autologging_utils = pytest.importorskip('mlflow.utils.autologging_utils') + +QUERY = 'PREDICT ents.Y FOR EACH ents.ID' + + +@pytest.fixture +def tracking_uri(tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) # the default artifact root is ./mlruns + previous = mlflow.get_tracking_uri() + uri = f'sqlite:///{tmp_path / "mlflow.db"}' + mlflow.set_tracking_uri(uri) + try: + yield uri + finally: + while mlflow.active_run() is not None: + mlflow.end_run() + mlflow.set_tracking_uri(previous) + + +@pytest.fixture +def table_and_snapshot(tmp_path, tracking_uri): + previous_scratch = get_scratch_dir() + kumo.configure(scratch_dir=tmp_path / 'scratch') + try: + rng = np.random.default_rng(0) + (tmp_path / 'data').mkdir() + pd.DataFrame({ + 'ID': np.arange(400), + 'Y': rng.random(400), + 'F1': rng.normal(size=400), + 'amount ($)': rng.normal(size=400), + }).to_parquet(tmp_path / 'data' / 'ents.parquet') + connector = kumo.FileConnector(root_dir=str(tmp_path / 'data')) + ents = kumo.Table.from_source_table(connector['ents.parquet'], + primary_key='ID').infer_metadata() + graph = kumo.Graph(tables={'ents': ents}, edges=[]) + snapshot = graph.snapshot() + query = kumo.PredictiveQuery(graph=graph, query=QUERY) + yield query.generate_training_table(snapshot), snapshot + finally: + kumo.configure(scratch_dir=previous_scratch) + + +def _runs(client): + return client.search_runs( + [e.experiment_id for e in client.search_experiments()]) + + +def test_fit_logs_a_run_per_fit_and_per_trial(table_and_snapshot, tmp_path, + caplog): + from kumo.engine.integrations.mlflow import MlflowTracking + + table, snapshot = table_and_snapshot + client = mlflow.MlflowClient() + plan = ModelPlan() + plan.training_job.num_experiments = 2 + plan.training_job.disable_explain = True + plan.optimization.max_epochs = 2 + plan.optimization.max_steps_per_epoch = 2 + plan.optimization.base_lr = [0.01, 0.001] + plan.column_processing.encoder_overrides = { + 'ents.amount ($)': Numerical('standard') + } + mlflow.pytorch.autolog() + try: + with mlflow.start_run() as outer: + tracking = MlflowTracking(experiment_name='ents') + 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-\d{8}T\d{6}Z', 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'} + 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} + assert result.tracking.run_id == parent.info.run_id + 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/') + } + assert metrics == {**result.metrics, 'seconds': metrics['seconds']} + assert 'off for the rest of this run' not in caplog.text + ents, = parent.inputs.dataset_inputs + assert (ents.dataset.name, ents.dataset.source_type) == ('ents', 'local') + assert json.loads(ents.dataset.source) == { + 'uri': str((tmp_path / 'data' / 'ents.parquet').resolve()) + } + + assert [run.info.run_name + for run in children] == ['trial-1-attempt-1', 'trial-2-attempt-1'] + for child in children: + assert child.info.status == 'FINISHED' + assert child.data.tags['kumo.status'] == 'succeeded' + assert {'best_val_mae', 'seconds', + 'val_mae'} <= set(child.data.metrics) + history = client.get_metric_history(child.info.run_id, 'val_mae') + assert sorted(metric.step for metric in history) == [1, 2] + + +def test_failed_fit_marks_the_run_failed(table_and_snapshot, 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( + tracking_uri, monkeypatch, caplog): + from mlflow.tracking import MlflowClient + + from kumo.engine.integrations.mlflow import MlflowTracking + + original = MlflowClient.log_batch + + def down(self, *args, metrics=(), **kwargs): + if metrics: + raise RuntimeError('tracking server down') + return original(self, *args, **kwargs) + + monkeypatch.setattr(MlflowClient, 'log_batch', down) + replay(MlflowTracking(log_system_metrics=False)) + + assert caplog.text.count('off for the rest of this run') == 1 + assert mlflow.active_run() is None + statuses = {run.info.status for run in _runs(mlflow.MlflowClient())} + assert statuses == {'FAILED'} + + +def test_a_user_run_name_wins_over_the_default(tracking_uri): + from kumo.engine.integrations.mlflow import MlflowTracking + + replay(MlflowTracking(run_name='nightly', log_system_metrics=False)) + join_tracking_threads() + + names = {run.info.run_name for run in _runs(mlflow.MlflowClient()) + if 'mlflow.parentRunId' not in run.data.tags} + assert names == {'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.*parent folder must exist"): + 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( + tracking_uri, uri, command): + from kumo.engine.integrations.mlflow import MlflowTracking + + 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's run links do not open on SageMaker: open the MLflow UI " + f'from SageMaker Studio or at the URL `{command}` prints.']) + + +def test_user_disabled_autologging_stays_disabled(tracking_uri): + from kumo.engine.integrations.mlflow import MlflowTracking + + tracker = MlflowTracking(log_system_metrics=False) + with autologging_utils.disable_autologging(): + tracker.start(run_record(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( + tracking_uri, monkeypatch): + from mlflow.data.delta_dataset_source import DeltaDatasetSource + + from kumo.engine.integrations.mlflow import MlflowTracking + + monkeypatch.setattr(DeltaDatasetSource, '_lookup_table_id', + lambda self, name: 'table-id') + replay(MlflowTracking(log_system_metrics=False)) + + run, = [run for run in _runs(mlflow.MlflowClient()) + if run.data.tags['kumo.kind'] == 'fit'] + inputs = { + item.dataset.name: (item.dataset.source_type, + json.loads(item.dataset.source)) + for item in run.inputs.dataset_inputs + } + assert inputs == { + 'ents': ('s3', {'uri': 's3://bucket/ents.parquet'}), + 'orders': ('delta_table', { + 'delta_table_name': 'main.sales.orders', + 'delta_table_version': 7, + 'is_databricks_uc_table': True, + 'delta_table_id': 'table-id', + }), + } + orders = next(item for item in run.inputs.dataset_inputs + if item.dataset.name == 'orders') + assert mlflow.data.get_source(orders).delta_table_id == 'table-id' + + +def test_fit_inputs_come_from_the_embedded_parents(): + from kumo.engine._manifest import ( + ModelManifest, + ModelParents, + SnapshotManifest, + SnapshotParents, + SourceTableEntry, + TrainingTableManifest, + TrainingTableParents, + stub, + ) + from kumo.engine._tracking import fit_inputs + + entry = {'connector': 'LOCAL_FILE', 'location': '/data/ents.parquet', + 'version': '1790000000', 'version_kind': 'timestamp'} + snapshot = SnapshotManifest(id='s', parents=SnapshotParents( + source_tables={'ents': SourceTableEntry(**entry, sql='SELECT 1')})) + custom = SourceTableEntry(**{**entry, 'version': None}) + table = TrainingTableManifest(id='t', parents=TrainingTableParents( + snapshot=snapshot, source_table=custom)) + model = ModelManifest(id='m', parents=ModelParents( + training_table=table, warm_start=ModelManifest( + id='w', parents=ModelParents(training_table=TrainingTableManifest( + id='u', parents=TrainingTableParents(snapshot=SnapshotManifest( + id='v', parents=SnapshotParents(source_tables={ + 'old': SourceTableEntry(**entry)})))))))) + + expected = {'ents': entry, 'training_table': {**entry, 'version': None}} + assert fit_inputs(table) == expected + assert fit_inputs(model) == expected + assert fit_inputs(TrainingTableManifest(id='t', parents=TrainingTableParents( + snapshot=stub('snapshot', 's'), source_table=custom))) == { + 'training_table': {**entry, 'version': None}} + assert fit_inputs(stub('model', 'm')) == {} + assert fit_inputs(None) == {} + + +def test_link_prediction_metric_names_are_logged(tracking_uri): + from kumo.engine._tracking import RunRecord + from kumo.engine.integrations.mlflow import MlflowTracking + + tracker = MlflowTracking(experiment_name='lp') + tags = {'kumo.kind': 'fit', 'kumo.run_id': 'r'} + 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 + } diff --git a/packages/kumo-engine/test/test_mlops.py b/packages/kumo-engine/test/test_mlops.py index e289fc5ee..728ea19fb 100644 --- a/packages/kumo-engine/test/test_mlops.py +++ b/packages/kumo-engine/test/test_mlops.py @@ -31,6 +31,33 @@ def test_mlops_needs_an_artifact_root_it_can_write_files_to(tmp_path): kumo.MLOps(artifact_root='@DB.SCHEMA.STAGE/kumo') with pytest.raises(ValueError, match='Unknown store kinds'): kumo.MLOps(artifact_root=str(tmp_path), store={'models'}) + with pytest.raises(ValueError, match='on_error'): + 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='registers through the tracker'): + 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): @@ -47,6 +74,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').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..4bfa258ad --- /dev/null +++ b/packages/kumo-engine/test/test_model_agnostic_mlops.py @@ -0,0 +1,366 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The MLOps seams take any bundle, run and predictor, not only the +Trainer's.""" + +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.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}, + inputs={'users': {'connector': 'S3'}}) + 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'} + assert finished.metrics == {'auc': 0.9} + assert finished.tags == {'kumo.kind': 'custom', 'kumo.bundle_id': 'b'} + + +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 test_a_file_asset_is_stored_once_and_resolved_by_its_hash(tmp_path): + (tmp_path / 'vocab').mkdir() + (tmp_path / 'vocab' / 'words.txt').write_text('a b c') + asset = FileAsset('toy', 'vocab', tmp_path / 'vocab', 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: + r"""An S3 filesystem whose credentials cover only the root's prefix: + HeadBucket and parent marker objects are refused.""" + 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 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}')) + (tmp_path / 'vocab').mkdir() + (tmp_path / 'vocab' / 'words.txt').write_text('a b c') + deps = Dependencies(settings=kumo.MLOps(str(tmp_path / 'root'))) + record = deps.require(FileAsset('toy', 'vocab', tmp_path / 'vocab')) + 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 diff --git a/packages/kumo-engine/test/test_optional_imports.py b/packages/kumo-engine/test/test_optional_imports.py index c70368fed..c18d0310d 100644 --- a/packages/kumo-engine/test/test_optional_imports.py +++ b/packages/kumo-engine/test/test_optional_imports.py @@ -26,6 +26,7 @@ 's3fs', 'google', 'gcsfs', + 'mlflow', ) PROVIDER_PACKAGES = ( 'kumo_snowflake', @@ -114,6 +115,22 @@ def find_spec(self, name, path=None, target=None): # and only the stand-in `self` is missing. Fine: nothing # broke at import time. pass + + from kumo.engine.integrations.mlflow import MlflowTracking + try: + MlflowTracking() + except ImportError as e: + assert "kumo-engine[mlflow]" in str(e), e + else: + raise AssertionError("MlflowTracking() needs mlflow") + + from kumo.engine.integrations.snowflake import SnowflakeTracking + try: + SnowflakeTracking(None, experiment_name="ents") + except ImportError as e: + assert "kumo-engine[snowflake]" in str(e), e + else: + raise AssertionError("SnowflakeTracking() needs snowflake-ml-python") print("ok") ''') diff --git a/packages/kumo-engine/test/test_prediction_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..3ba5d2537 --- /dev/null +++ b/packages/kumo-engine/test/test_scoring_defaults.py @@ -0,0 +1,163 @@ +# 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(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_snapshot_export.py b/packages/kumo-engine/test/test_snapshot_export.py index 4452bc06d..01cb268d7 100644 --- a/packages/kumo-engine/test/test_snapshot_export.py +++ b/packages/kumo-engine/test/test_snapshot_export.py @@ -787,3 +787,17 @@ def test_a_mlops_block_stores_each_snapshot_once(tmp_path, scratch, graph, graph.snapshot(mlops=kumo.MLOps(artifact_root=str(tmp_path / 'off'), store={'model'})) assert not (tmp_path / 'off').exists() + + +@pytest.mark.parametrize('on_error', ['warn', 'raise']) +def test_a_failed_store_follows_on_error(tmp_path, scratch, graph, + monkeypatch, on_error): + monkeypatch.setattr(snapshot_io, 'export', + Mock(side_effect=OSError('ACCESS_DENIED'))) + cfg = kumo.MLOps(str(tmp_path / 'root'), on_error=on_error) + if on_error == 'raise': + with pytest.raises(RuntimeError, match='Could not store snapshot'): + graph.snapshot(mlops=cfg) + else: + with pytest.warns(UserWarning, match='ACCESS_DENIED'): + graph.snapshot(mlops=cfg) 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..7cbbacec9 --- /dev/null +++ b/packages/kumo-engine/test/test_snowflake_tracking.py @@ -0,0 +1,283 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import inspect +import re +import threading +import time +from unittest import mock + +import _snowflake_experiment_fake +import pytest +from _snowflake_experiment_fake import FakeExperimentTracking, FakeSession +from test_tracking_contract import ( + CONTEXT, + EVENTS, + RESULT, + join_tracking_threads, +) + +from kumo.engine._tracking import TrackingSession +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, EVENTS[-1].metrics['seconds']), + error) + + +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(fake, busy, caplog): + session = FakeSession() + user = fake(session, warnings_as_errors=True) + user.set_experiment('mine') + if busy == 'active run': + user.start_run('mine_1') + else: + session = FakeSession() + calls = len(user.calls) + + _replay(_tracking(session)) + + assert user.calls[calls:] == [] + assert user.warnings_as_errors is True + assert caplog.text.count('not tracked in Snowflake') == 1 + + +def test_location_defaults_to_the_session_not_the_last_fit(fake): + session = FakeSession(database='DEV', schema='WORK') + _replay(_tracking(session, database='ml')) + _replay(_tracking(session)) + + assert [ + call[1:] for call in fake._instance.calls if call[0] == 'set_experiment' + ] == [('ents', 'ml', 'PUBLIC'), ('ents', 'DEV', 'WORK')] + assert not any(call[0].startswith('run.') for call in fake._instance.calls) + + +@pytest.mark.parametrize('taken', ['FIT', 'MYXRUN']) +def test_an_existing_run_is_never_resumed(fake, taken, caplog): + session = FakeSession() + user = fake(session) + user.set_experiment('ents') + user.start_run(taken) + if taken == 'FIT': + user.end_run() + else: + user._run = None # still running, in another process + run_name = 'fit' if taken == 'FIT' else 'my_run' + + tracking = _tracking(session, run_name=run_name) + _replay(tracking) + + new = tracking.run_id + assert re.fullmatch(f'{run_name}_[0-9a-f]{{8}}'.upper(), new) + assert set(user.runs) == {taken, new} + assert user.runs[taken]['params'] == {} + assert user.runs[new]['status'] == 'FINISHED' + assert 'recording this fit as' in caplog.text + + +def test_metric_and_param_limits(fake, monkeypatch): + import kumo.engine.integrations.snowflake as snowflake + + monkeypatch.setattr(snowflake, '_EPOCH_METRICS', 2) + monkeypatch.setattr(snowflake, '_SUMMARY_METRICS', 3) + monkeypatch.setattr(snowflake, '_MAX_METRICS', 5) + result = type(RESULT)(metrics={**RESULT.metrics, 'r2': float('nan')}, + id=RESULT.id, bundle_dir=RESULT.bundle_dir) + _replay(_tracking(tags={'note': 'a$$b', 'level': 2}), result) + + run, = fake._instance.runs.values() + assert sorted({name for name, _, _ in run['metrics']}) == [ + 'test_mae', 'trial-1-attempt-1/best_val_mae', + 'trial-1-attempt-1/train_loss', 'trial-1-attempt-1/val_mae', 'val_mae' + ] + assert run['params']['kumo.dropped_metrics'] == '6' + assert run['params']['kumo.dropped_params'] == '1' + assert 'note' not in run['params'] + assert run['params']['level'] == '2' + + +def test_quoted_run_names_are_taken_only_by_their_exact_name(): + from kumo.engine.integrations.snowflake import _taken + + assert _taken('"MyRun"', ['MyRun']) + assert not _taken('"MyRun"', ['MYRUN']) + assert _taken('my_run', ['MYXRUN']) + + +def test_the_users_session_is_not_altered(fake): + session = FakeSession() + + _replay(_tracking(session)) + + assert session.statements == [] + + +def test_a_hung_snowflake_call_does_not_stall_the_fit(fake, monkeypatch, + caplog): + import kumo.engine._tracking as tracking + + monkeypatch.setattr(tracking, 'FINISH_SECONDS', 0.5) + release = threading.Event() + monkeypatch.setattr(FakeExperimentTracking, 'log_metrics', + lambda *args, **kwargs: release.wait(30)) + began = time.monotonic() + try: + _replay(_tracking()) + elapsed = time.monotonic() - began + finally: + release.set() + join_tracking_threads() + + assert elapsed < 1.5 + assert 'did not finish within' in caplog.text + run, = fake._instance.runs.values() + assert run['status'] == 'FINISHED' + + +class _SchemaSession: + def __init__(self, schemas): + self.schemas = schemas + self.statements = [] + + def get_current_database(self): + return 'DB' + + def get_current_schema(self): + return 'RUNS' + + def sql(self, statement): + self.statements.append(statement) + return mock.Mock(collect=lambda: self.schemas) + + +def test_a_connector_opens_the_session_and_the_schema_is_checked_first( + fake): + import kumo.engine as kumo + from kumo.engine.connector.base import Connector + + session = _SchemaSession([('RUNS', )]) + connector = mock.Mock(spec=Connector, session=lambda: session) + tracking = _tracking(connector) + tracking.check() + assert tracking.session is session and tracking.connector is connector + assert session.statements == ["SHOW SCHEMAS LIKE 'RUNS' IN DATABASE DB"] + + tracking.session = _SchemaSession([]) + with pytest.raises(ValueError, match='needs schema DB.RUNS'): + 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 + from kumo.engine.connector.base import Connector + + 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))) + session = _SchemaSession([('RUNS', )]) + connector = mock.Mock(spec=Connector, session=lambda: session) + tracking = _tracking(connector) + 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) 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..0d7bedbf1 --- /dev/null +++ b/packages/kumo-engine/test/test_tracking_contract.py @@ -0,0 +1,409 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Every tracking adapter records the same run records as a recorded fit.""" + +import json +import threading +import time +from dataclasses import dataclass, field, replace +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.engine._tracking import ( + DROPPED_METRICS, + DROPPED_PARAMS, + TrackingSession, + run_record, +) +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'}, +} +CONTEXT = FitContext( + query=QUERY, inputs=INPUTS, + model_plan={ + 'optimization': {'base_lr': [0.01, 0.001], 'max_epochs': 2}, + 'training_job': {'tune_metric': 'mae'}, + 'column_processing': { + 'encoder_overrides': {'ents.amount ($)': 'numerical'} + }, + }) + + +def _event(name, trial=None, epoch=None, metrics=None, **params): + return Event(name=name, time=0.0, trial=trial, epoch=epoch, + metrics=metrics or {}, params=params) + + +def _epochs(trial, *values): + return [ + _event('epoch_ended', trial, epoch, {'train_loss': loss, + 'val_mae': mae}) + for epoch, (loss, mae) in enumerate(values, 1) + ] + + +# Trial 2's first attempt fails; epoch events carry no attempt, as in core. +EVENTS = [ + _event('trial_started', 1, attempt=1, **{'optimization.base_lr': 0.01}), + *_epochs(1, (0.5, 0.4), (0.3, 0.35)), + _event('trial_ended', 1, metrics={'val_mae': 0.35, 'seconds': 3.0}, + status='succeeded', attempt=1), + _event('trial_started', 2, attempt=1, **{'optimization.base_lr': 0.001}), + _event('trial_ended', 2, status='failed', attempt=1, + error='OutOfMemoryError: CUDA out of memory'), + _event('trial_started', 2, attempt=2, **{'optimization.base_lr': 0.001, + 'optimization.batch_size': 256}), + *_epochs(2, (0.6, 0.5), (0.4, 0.3)), + _event('trial_ended', 2, metrics={'val_mae': 0.3, 'seconds': 4.0}, + status='succeeded', attempt=2), + _event('fit_ended', metrics={'val_mae': 0.3, 'test_mae': 0.32, + 'seconds': 9.0}, status='succeeded'), +] +BUNDLE = 'b' * 64 +RESULT = SimpleNamespace(metrics={'val_mae': 0.3, 'test_mae': 0.32}, + id=BUNDLE, bundle_dir=Path('no-bundle')) + + +class _Recorder: + def __init__(self, tracker=None): + self.calls: List[Tuple[str, Any]] = [] + self.tracker = tracker + self.tags = getattr(tracker, 'tags', None) + + def prepare(self): + getattr(self.tracker, 'prepare', lambda: None)() + + def start(self, record): + self.calls.append(('start', record)) + if self.tracker is not None: + self.tracker.start(record) + + def on_trial(self, record, event): + self.calls.append(('on_trial', record)) + if self.tracker is not None: + self.tracker.on_trial(record, event) + + def finish(self, record, error): + self.calls.append(('finish', record)) + if self.tracker is not None: + self.tracker.finish(record, error) + + +def join_tracking_threads(): + for thread in threading.enumerate(): + if thread.name == 'kumo-tracking': + thread.join(30) + + +def replay(tracker=None, events=EVENTS, result=RESULT, + error=None) -> Dict[str, RunRecord]: + r"""Feeds the recorded fit to `tracker`; returns the records expected + of it.""" + recorder = _Recorder(tracker) + session = TrackingSession(recorder) + session.start(lambda: 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, events[-1].metrics.get('seconds')), error) + records = { + record.name: record + for hook, record in recorder.calls if hook == 'on_trial' + } + spec = fit_spec(CONTEXT, tags=USER_TAGS) + if result is not None: + spec = fitted_spec(spec, result, events[-1].metrics.get('seconds')) + 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 + + +def gaps(expected: Dict[str, RunRecord], observed: Dict[str, Observed], + adapter: Adapter) -> List[str]: + kept_params = adapter.kept_params + found = [] + if set(observed) != set(expected): + return [f'runs {sorted(observed)} != {sorted(expected)}'] + dropped = sum(len(record.params) - len(kept_params(record.params)) + for record in expected.values()) + counted = sum(int(seen.tags.get(DROPPED_PARAMS, 0)) + for seen in observed.values()) + if counted != dropped: + found.append(f'{DROPPED_PARAMS}: {counted} != {dropped}') + for run_id, record in expected.items(): + seen = observed[run_id] + + def differ(what, have, want): + if have != want: + found.append(f'{run_id} {what}: {have!r} != {want!r}') + + if record.name is not None: + differ('name', seen.name, record.name) + differ('params', dict(seen.params), dict(kept_params(record.params))) + differ('metrics', {k: seen.metrics.get(k) for k in record.metrics}, + dict(record.metrics)) + names = set(record.metrics) + if adapter.step_metrics: + names |= set(record.steps) + differ('steps', {k: seen.steps.get(k) for k in record.steps}, + dict(record.steps)) + differ('metric names', sorted(seen.metrics), sorted(names)) + differ('tags', {k: seen.tags.get(k) for k in record.tags}, + dict(record.tags)) + differ( + 'kumo tags', + sorted(k for k in seen.tags if k.startswith('kumo.') + and k not in (DROPPED_PARAMS, DROPPED_METRICS)), + sorted(k for k in record.tags if k.startswith('kumo.'))) + if adapter.nested_runs: + differ('parent', seen.parent, + FIT if record.tags['kumo.kind'] == 'trial' else None) + if adapter.artifacts: + differ('artifacts', dict(seen.artifacts), dict(record.artifacts)) + return found + + +@pytest.fixture +def mlflow_adapter(tmp_path, monkeypatch): + mlflow = pytest.importorskip('mlflow') + from kumo.engine.integrations.mlflow import MlflowTracking, _mlflow_params + + monkeypatch.chdir(tmp_path) + previous = mlflow.get_tracking_uri() + mlflow.set_tracking_uri(f'sqlite:///{tmp_path / "mlflow.db"}') + + def observe(): + client = mlflow.MlflowClient() + runs = client.search_runs( + [e.experiment_id for e in client.search_experiments()]) + kumo_ids = { + r.info.run_id: + FIT if r.data.tags['kumo.kind'] == 'fit' else r.info.run_name + for r in runs + } + + def history(run_id, key): + return tuple((m.step, m.value) for m in sorted( + client.get_metric_history(run_id, key), + key=lambda m: (m.step, m.timestamp))) + + return { + kumo_ids[run.info.run_id]: Observed( + name=run.info.run_name, params=run.data.params, + metrics=run.data.metrics, tags=run.data.tags, + steps={k: history(run.info.run_id, k) + for k in run.data.metrics}, + parent=kumo_ids.get(run.data.tags.get('mlflow.parentRunId')), + artifacts={ + a.path: mlflow.artifacts.load_text( + f'{run.info.artifact_uri}/{a.path}') + for a in client.list_artifacts(run.info.run_id) + }) + for run in runs + } + + try: + yield Adapter( + make=lambda: MlflowTracking(log_system_metrics=False, + tags=USER_TAGS), observe=observe, + kept_params=_mlflow_params, nested_runs=True, step_metrics=True, + artifacts=True) + finally: + while mlflow.active_run() is not None: + mlflow.end_run() + mlflow.set_tracking_uri(previous) + + +@pytest.fixture +def snowflake_adapter(monkeypatch): + import _snowflake_experiment_fake + + fake = _snowflake_experiment_fake.install(monkeypatch) + from kumo.engine.integrations.snowflake import SnowflakeTracking + + def observe(): + observed = {} + for run in fake._instance.runs.values(): + names: Dict[Optional[str], Dict[str, Dict[str, Any]]] = {} + + def bucket(key): + prefix, _, name = key.rpartition('/') + return names.setdefault(prefix or None, { + 'params': {}, 'tags': {}, 'metrics': {}, 'steps': {} + }), name + + for key, value in run['params'].items(): + runs, name = bucket(key) + # Snowflake has no tags: they are params of the same name. + if name.startswith('kumo.') or name in USER_TAGS: + runs['tags'][name] = value + else: + runs['params'][name] = value + for key, value, step in run['metrics']: + runs, name = bucket(key) + runs['metrics'][name] = value + runs['steps'][name] = (*runs['steps'].get(name, ()), + (step, value)) + for prefix, runs in names.items(): + observed[prefix or FIT] = Observed( + name=prefix, + artifacts=run['artifacts'] if prefix is None else {}, + **runs) + return observed + + return Adapter( + make=lambda: SnowflakeTracking(_snowflake_experiment_fake.FakeSession(), + experiment_name='ents', tags=USER_TAGS), + observe=observe, step_metrics=True, artifacts=True) + + +@pytest.fixture(params=['mlflow', 'snowflake']) +def adapter(request) -> Adapter: + return request.getfixturevalue(f'{request.param}_adapter') + + +def test_records_of_the_recorded_fit(): + records = replay() + + fit = records[FIT] + assert fit.tags == {'team': 'ents', 'kumo.kind': 'fit', + 'kumo.bundle_id': BUNDLE} + 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(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 == {'val_mae': 0.3, 'test_mae': 0.32, 'seconds': 9.0} + assert fit.artifacts['query.pql'] == QUERY + assert json.loads(fit.artifacts['model_plan.json']) == CONTEXT.model_plan + assert fit.inputs == INPUTS + + oom = records['trial-2-attempt-1'] + assert oom == RunRecord( + name='trial-2-attempt-1', + params={'optimization.base_lr': '0.001', 'attempt': '1'}, tags={ + 'kumo.kind': 'trial', + 'kumo.status': 'failed', + 'kumo.error': 'OutOfMemoryError: CUDA out of memory', + }) + retry = records['trial-2-attempt-2'] + assert retry.params['attempt'] == '2' + assert retry.metrics == {'best_val_mae': 0.3, 'seconds': 4.0} + assert retry.steps == {'train_loss': ((1, 0.6), (2, 0.4)), + 'val_mae': ((1, 0.5), (2, 0.3))} + + +def test_adapter_records_the_same_runs(adapter, tmp_path): + (tmp_path / PROVENANCE_FILE).write_text('{"bundle": "b"}') + result = SimpleNamespace(metrics=RESULT.metrics, id=BUNDLE, + 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_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, caplog): + import kumo.engine._tracking as tracking + + monkeypatch.setattr(tracking, 'FINISH_SECONDS', 0.5) + release = threading.Event() + + class Hung(_Recorder): + def start(self, record): + super().start(record) + release.wait(30) + + hung = Hung() + began = time.monotonic() + try: + replay(hung) + elapsed = time.monotonic() - began + finally: + release.set() + join_tracking_threads() + + assert elapsed < 1.5 + assert 'did not finish within' in caplog.text diff --git a/packages/kumo-engine/test/test_trainer_registry.py b/packages/kumo-engine/test/test_trainer_registry.py index 7ea9ba950..fc1fd6185 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 @@ -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') } @@ -258,6 +263,7 @@ def publish_model(self, connector, bundle_dir, **kwargs): assert calls[0][2]['path'] == path assert calls[0][2]['location'] == location assert calls[0][2]['metrics'] == {'val_auroc': .9} + assert calls[0][2]['metadata']['kumo_bundle_id'] == result.bundle_id assert result.published_model is published assert result.metrics == original.metrics assert result.bundle_dir == bundle @@ -572,3 +578,196 @@ def fake_predict(**kwargs) -> None: binary_classification_threshold=.5) assert result.frame['TARGET_PRED'].tolist() == [True] + + +def _export(sink, bundle, location=None, artifact_paths=None): + return Trainer(ModelPlan())._finish_fit( + TrainingResult({}, artifact_paths or {}, + bundle), predictive_query='PREDICT COUNT(*)', + snapshot_id='snap-123', training_table_generation=_generation_plan(), + export_config=SimpleNamespace( + connector=SimpleNamespace(sink_backend=sink), + path=None if location else 'm', location=location)) + + +def _recording_sink(calls): + return SimpleNamespace(publish_model=lambda *a, **k: calls.append(k) or + PublishedModel(name='m', version='1', uri='u')) + + +def test_publish_records_only_the_bundle_id_pointer(tmp_path: Path) -> None: + calls = [] + result = _export(_recording_sink(calls), _fake_published_bundle(tmp_path)) + assert {k: v + for k, v in calls[0]['metadata'].items() + if k.startswith('kumo_')} == {'kumo_bundle_id': result.bundle_id} + + +class _AliasingSink(_FakeSink): + def __init__(self, recorded: Optional[str] = None, + fail: bool = False) -> None: + super().__init__() + self.recorded = recorded + self.fail = fail + self.aliases: dict = {} + + def set_model_alias(self, connector, location, alias) -> None: + if self.fail: + raise RuntimeError('registry unavailable') + self.aliases[alias] = location.version + + def fetch_model(self, connector, local_dir, **kwargs) -> PublishedModel: + shutil.copytree(self.published['CHURN_GNN@latest'], local_dir, + dirs_exist_ok=True) + return PublishedModel(name='CHURN_GNN', version='V1', uri='fake://', + metadata={'kumo_bundle_id': self.recorded}) + + +@pytest.mark.parametrize('sink,warning', [ + (_AliasingSink(), None), + (_AliasingSink(fail=True), 'registry unavailable'), + (_FakeSink(), 'cannot set model aliases'), +]) +def test_export_with_alias_publishes_then_sets_it_best_effort( + tmp_path: Path, sink, warning) -> None: + import warnings + + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter('always') + result = _export(sink, _fake_published_bundle(tmp_path), + ModelRegistryLocation('CHURN_GNN', alias='prod')) + assert result.published_model.version == 'V1' + assert list(sink.published) == ['CHURN_GNN@latest'] + messages = [str(w.message) for w in caught if 'alias' in str(w.message)] + if warning is None: + assert sink.aliases == {'prod': 'V1'} and messages == [] + else: + assert len(messages) == 1 and warning in messages[0] + + +def test_load_by_alias_warns_on_recorded_bundle_mismatch( + tmp_path: Path) -> None: + bundle = _fake_published_bundle(tmp_path) + write_training_spec(bundle, predictive_query='PREDICT COUNT(*)', + model_plan=ModelPlan(), snapshot_id='snap-123', + training_table_generation=_generation_plan()) + sink = _AliasingSink('other') + sink.publish_model(None, str(bundle), + location=ModelRegistryLocation('CHURN_GNN')) + with pytest.warns(UserWarning, match='registry records bundle other'): + loaded = Trainer.load_from_connector( + _FakeConnector(sink), + location=ModelRegistryLocation('CHURN_GNN', alias='prod')) + assert loaded.result.id == Trainer.load(bundle).result.id + + +def test_training_result_exports_like_its_trainer(tmp_path: Path) -> None: + bundle = _fake_published_bundle(tmp_path) + write_training_spec(bundle, predictive_query='PREDICT COUNT(*)', + model_plan=ModelPlan(), snapshot_id='snap-123', + training_table_generation=_generation_plan()) + sink = _AliasingSink() + 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(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 = _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()) + 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-snowflake/pyproject.toml b/packages/kumo-snowflake/pyproject.toml index 960687071..47547c8b9 100644 --- a/packages/kumo-snowflake/pyproject.toml +++ b/packages/kumo-snowflake/pyproject.toml @@ -27,13 +27,15 @@ dependencies = [ [project.optional-dependencies] # What the SDK plugin half needs beyond the platform reader: Snowpark for the -# session it opens, the model registry for what a sink publishes, and pyyaml -# for the manifest that carries. The platform installs this distribution +# session it opens, the model registry for what a sink publishes (and, from +# 2.0, the experiment tracking kumo-engine records fits in; capped below 3 +# because that adapter reads ExperimentTracking internals), and pyyaml for the +# manifest that carries. The platform installs this distribution # without this extra and imports `SnowflakeSourceReader` alone, so none of it # is loaded there -- `__init__` exports the plugin lazily to keep that true. sdk = [ "snowflake-snowpark-python", - "snowflake-ml-python", + "snowflake-ml-python>=2.0,<3", "pyyaml", ] test = [ diff --git a/packages/kumo-snowflake/src/kumo_snowflake/registry.py b/packages/kumo-snowflake/src/kumo_snowflake/registry.py index fd220653e..c4d88b580 100644 --- a/packages/kumo-snowflake/src/kumo_snowflake/registry.py +++ b/packages/kumo-snowflake/src/kumo_snowflake/registry.py @@ -33,11 +33,14 @@ """ import json +import logging +import re import shutil from contextlib import contextmanager from pathlib import Path from typing import ( TYPE_CHECKING, + Any, ContextManager, Dict, Iterator, @@ -53,9 +56,13 @@ resolved, ) from kumoapi.plugin import ( + BUNDLE_ID_KEY, + ENGINE_MARKER, ModelRegistryLocation, PublishedModel, SourceConnector, + check_model_alias, + registry_lineage, ) if TYPE_CHECKING: @@ -74,8 +81,11 @@ #: 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 @@ -113,7 +123,11 @@ 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(SnowflakeConnector(...), + location=ModelRegistryLocation(...))`` and predict against a snapshot + (kumo-engine), or with :func:`kumoai_studio.import_model` (hosted + service). """ # Unquoted, and evaluated here rather than deferred: the registry # compares this annotation against `pandas.DataFrame` itself, so a @@ -122,13 +136,36 @@ class _BundleModel(custom_model.CustomModel): 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.") + "score in Snowflake. Load it with Trainer.load_from_connector(" + "SnowflakeConnector(...), location=ModelRegistryLocation(...)) " + "and predict against a snapshot (kumo-engine), or " + "kumoai_studio.import_model() (hosted service).") return _BundleModel( custom_model.ModelContext(artifacts={'bundle': bundle_dir})) +#: Snowflake ML task types by Kumo task type; Model Monitor needs one. +_TASKS = { + 'binary_classification': 'TABULAR_BINARY_CLASSIFICATION', + 'multiclass_classification': 'TABULAR_MULTI_CLASSIFICATION', + 'regression': 'TABULAR_REGRESSION', +} + + +def _task(bundle: Path) -> Dict[str, Any]: + r"""``task=`` for ``log_model`` from the bundle's prediction schema; + empty for other tasks, or when it cannot be read.""" + 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. @@ -209,6 +246,40 @@ def _refuse_path(path: Optional[str], action: str) -> None: f'a path.') +def _version_record(metadata: Mapping[str, str]) -> Dict[str, str]: + record = {ENGINE_MARKER: '1'} + if metadata.get(BUNDLE_ID_KEY): + record[BUNDLE_ID_KEY] = str(metadata[BUNDLE_ID_KEY]) + return record + + +def _recorded_lineage(version: "ModelVersion") -> Dict[str, str]: + try: + return registry_lineage(dict(version.show_metrics() or {})) + except Exception as error: + _LOG.warning('Could not read registry metadata: %s', error) + return {} + + +def _alias_holders(model: Any) -> Dict[str, str]: + holders = {} + for row in model.show_versions().to_dict('records'): + for alias in json.loads(row.get('aliases') or '[]'): + holders[alias] = str(row['name']) + return holders + + +def _restore_alias(model: Any, version: str, alias: str) -> None: + try: + model.version(version).set_alias(alias) + except Exception as error: + _LOG.warning('Alias %s is now unset: restoring it on version %s ' + 'failed: %s', alias, version, error) + else: + _LOG.warning('Could not move alias %s; it stays on version %s.', + alias, version) + + class SnowflakeModelRegistryMixin: r"""The model half of :class:`SnowflakeSinkBackend`. @@ -229,6 +300,18 @@ def _session( ) -> ContextManager[SnowflakeSession]: ... + def check_model_location(self, connector: SourceConnector, + location: ModelRegistryLocation) -> None: + r"""Raises :class:`ValueError` unless ``location.model_name`` is a + Snowflake 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 @@ -267,8 +350,8 @@ def publish_model( the model, and optionally the version. A version of ``None`` lets the registry assign one -- the usual case, and the reason this reports back what it created. - metadata: Returned with the publication result. Training metadata - remains in the bundle and is not recorded in registry tags. + metadata: Returned with the publication result; its + ``kumo_bundle_id`` is also recorded in the version's metrics. metrics: Returned with the publication result. Training metrics remain in the bundle's metrics.json. project_id: Ignored; a local runtime has no projects. @@ -288,14 +371,16 @@ def publish_model( f"{bundle} does not look like a published model bundle: it " f"has no {_SPEC_FILE}, which an export writes into every one.") + record = _version_record(metadata) with self._registry(connector) as registry: version = registry.log_model( _bundle_model(str(bundle)), model_name=location.model_name, version_name=location.version, comment=location.comment, - metrics={_ENGINE_MARKER: '1'}, + metrics=record, signatures=_carrier_signature(), + **_task(bundle), ) return _published(connector, version, metadata=metadata, metrics=metrics) @@ -308,10 +393,11 @@ def list_models( ) -> List[PublishedModel]: r"""Lists Kumo model versions visible in the connector's registry. - A Kumo export records only ``kumo_engine=1`` in registry metadata. + A Kumo export records ``kumo_engine=1`` in registry metadata. That marker makes discovery both cheap and precise: versions logged by another tool are omitted, and no bundle is downloaded merely to decide whether it belongs in this list. + ``metadata`` carries the recorded ``kumo_bundle_id``. """ del project_id found: List[PublishedModel] = [] @@ -319,9 +405,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 +433,15 @@ def fetch_model( though the caller always hands this a fresh, empty ``mkdtemp``. path: Refused, as in :meth:`publish_model`. - location: The model, and optionally the version. ``None`` takes the - registry's default version, which is what a caller who names - none means. + location: The model, and optionally the version or an alias. With + neither, takes the registry's default version, which is what a + caller who names none means. project_id: Ignored. Returns: The name, the version actually read -- resolved, when the caller - named none -- where it lives, and metrics read from metrics.json. + named none -- where it lives, metrics read from metrics.json, and + the version's recorded ``kumo_bundle_id`` as ``metadata``. """ del project_id _refuse_path(path, 'fetching') @@ -366,8 +455,18 @@ def fetch_model( # Resolved before anything on disk is touched: a name the registry # does not know must not cost the caller the bundle already there. model = registry.get_model(location.model_name) - version = (model.default if location.version is None else - model.version(location.version)) + alias = getattr(location, 'alias', None) + if alias is not None: + if resolved(alias) not in _alias_holders(model): + raise ValueError( + f'No version of {location.model_name} holds alias ' + f'{alias}.') + version = model.version(alias) + elif location.version is None: + version = model.default + else: + version = model.version(location.version) + recorded = _recorded_lineage(version) destination.parent.mkdir(parents=True, exist_ok=True) staged = destination.parent / f'{destination.name}{_EXPORT_DIR}' @@ -385,9 +484,52 @@ def fetch_model( isinstance(key, str) and isinstance(value, (int, float)) for key, value in metrics.items())): raise ValueError('Invalid training metrics in model bundle.') - published = _published(connector, version, metrics=metrics) + published = _published(connector, version, metadata=recorded, + metrics=metrics) if destination.exists(): shutil.rmtree(destination) shutil.move(str(bundle), str(destination)) shutil.rmtree(staged, ignore_errors=True) return published + + def set_model_alias( + self, + connector: SourceConnector, + location: ModelRegistryLocation, + alias: str, + ) -> None: + r"""Points :obj:`alias` at ``location.version``, moving it if held. + + Snowflake allows one alias per version. A failed move puts the alias + back on the version that held it. + """ + if location.version is None: + raise ValueError('Setting an alias needs a ModelRegistryLocation ' + 'naming an exact version.') + check_model_alias(alias) + wanted = resolved(alias) + with self._registry(connector) as registry: + model = registry.get_model(location.model_name) + target = model.version(location.version) + name = resolved(str(target.version_name)) + holders = { + held: version + for held, version in _alias_holders(model).items() + if held not in _SYSTEM_ALIASES + } + if holders.get(wanted) == name: + return + held = sorted(a for a, v in holders.items() if v == name) + if held: + raise ValueError( + f'Version {name} of {location.model_name} already holds ' + f'alias {held[0]}; unset it first.') + previous = holders.get(wanted) + if previous is not None: + model.version(previous).unset_alias(alias) + try: + target.set_alias(alias) + except Exception: + if previous is not None: + _restore_alias(model, previous, alias) + raise diff --git a/packages/kumo-snowflake/tests/test_registry.py b/packages/kumo-snowflake/tests/test_registry.py index 9d07e6c28..7c6e80dbc 100644 --- a/packages/kumo-snowflake/tests/test_registry.py +++ b/packages/kumo-snowflake/tests/test_registry.py @@ -8,6 +8,7 @@ from pathlib import Path from types import SimpleNamespace +import pandas as pd import pytest from kumo_snowflake import registry as registry_module from kumo_snowflake.registry import SnowflakeModelRegistryMixin @@ -30,7 +31,11 @@ def export(directory): version = SimpleNamespace( model_name='CHURN', version_name='V1', fully_qualified_model_name='DB.SCHEMA.CHURN', export=export, - show_metrics=lambda: pytest.fail('read registry metrics')) + show_metrics=lambda: { + 'kumo_engine': '1', + 'kumo_bundle_id': 'b1', + 'auc': 0.5 + }) model = SimpleNamespace(default=version, version=lambda _: version) @contextmanager @@ -60,6 +65,7 @@ def test_fetch_returns_bundle_metrics(fetch, requested_version, metrics): payload = json.dumps(metrics) if metrics is not None else None published = run(payload, requested_version) assert published.metrics == (metrics or {}) + assert published.metadata == {'kumo_bundle_id': 'b1'} assert published.version == 'V1' assert (destination / 'weights.bin').read_bytes() == b'model weights' assert not (destination / 'previous').exists() @@ -75,9 +81,13 @@ def test_invalid_metrics_preserve_previous_destination(fetch, payload): assert (destination / 'previous').read_bytes() == b'previous bundle' -@pytest.mark.parametrize('metadata', [{}, {'predictive_query': 'x' * 300}]) -def test_publish_records_only_engine_marker_in_catalog(tmp_path, monkeypatch, - metadata): +@pytest.mark.parametrize('metadata', [{}, { + 'predictive_query': 'x' * 300, + 'kumo_bundle_id': 'b1', + 'kumo_engine': '0' +}]) +def test_publish_records_only_pointers_in_catalog( + tmp_path, monkeypatch, metadata): bundle = tmp_path / 'bundle' bundle.mkdir() (bundle / 'training-spec.json').write_text('{}') @@ -99,16 +109,20 @@ def registry(_): monkeypatch.setattr(registry_module, '_carrier_signature', lambda: {}) published = sink.publish_model(SimpleNamespace(), str(bundle), location=ModelRegistryLocation('CHURN'), - metadata=metadata, metrics={'auc': 0.92}) - assert calls[0]['metrics'] == {'kumo_engine': '1'} + metadata=metadata, + metrics={'auc': 0.92, 'loss': float('nan')}) + assert calls[0]['metrics'] == { + **({'kumo_bundle_id': 'b1'} if metadata else {}), 'kumo_engine': '1' + } assert published.metadata == metadata - assert published.metrics == {'auc': 0.92} + assert published.metrics['auc'] == 0.92 -def test_listing_filters_engine_marker_and_returns_empty_metrics(monkeypatch): +def test_listing_filters_engine_marker_and_returns_lineage(monkeypatch): versions = [] for name, recorded in [('V1', { 'kumo_engine': '1', + 'kumo_bundle_id': 'b1', 'auc': 0.92 }), ('V2', { 'other': '1', @@ -134,4 +148,163 @@ def registry(_): found = sink.list_models(SimpleNamespace()) assert [item.version for item in found] == ['V1', 'V4'] assert [item.metrics for item in found] == [{}, {}] - assert all(item.metadata == {} for item in found) + assert [item.metadata for item in found] == [{'kumo_bundle_id': 'b1'}, {}] + + +class _Model: + def __init__(self, aliases): + self.aliases = aliases + + def show_versions(self): + return pd.DataFrame([ + dict(name=name, aliases=json.dumps(held)) + for name, held in self.aliases.items() + ]) + + def version(self, version_or_alias): + name = next((name for name, held in self.aliases.items() + if version_or_alias.upper() in held), version_or_alias) + held = self.aliases[name] + + def export(directory): + bundle = Path(directory) / 'bundle' + bundle.mkdir() + (bundle / 'training-spec.json').write_text(name) + (bundle / 'publication-manifest.json').write_text('{}') + + return SimpleNamespace( + version_name=name, model_name='CHURN', + fully_qualified_model_name='DB.SCHEMA.CHURN', export=export, + set_alias=lambda alias: held.append(alias.upper()), + unset_alias=lambda alias: held.remove(alias.upper()), + show_metrics=lambda: { + 'kumo_engine': '1', + 'kumo_bundle_id': f'id-{name}' + }) + + +def _aliasing_sink(monkeypatch, model): + @contextmanager + def registry(_): + yield SimpleNamespace(get_model=lambda _: model) + + sink = SnowflakeModelRegistryMixin() + monkeypatch.setattr(sink, '_registry', registry) + return sink + + +def test_set_alias_moves_it_and_refuses_what_snowflake_cannot(monkeypatch): + from kumo_snowflake.sink import SnowflakeSinkBackend + from kumoapi.plugin import ModelAliasingSink + + assert isinstance(SnowflakeSinkBackend(), ModelAliasingSink) + model = _Model({'V1': ['FIRST'], 'V2': ['LAST', 'DEFAULT'], 'V3': []}) + sink = _aliasing_sink(monkeypatch, model) + + def set_alias(version, alias): + sink.set_model_alias(SimpleNamespace(), + ModelRegistryLocation('CHURN', version=version), + alias) + + set_alias('V2', 'prod') + set_alias('V3', 'prod') + set_alias('V3', 'prod') + assert model.aliases == { + 'V1': ['FIRST'], + 'V2': ['LAST', 'DEFAULT'], + 'V3': ['PROD'] + } + for version, alias, match in [('V3', 'staging', 'already holds'), + ('V1', 'default', 'reserved'), + (None, 'prod', 'exact version')]: + with pytest.raises(ValueError, match=match): + set_alias(version, alias) + assert model.aliases['V3'] == ['PROD'] + + +def test_fetch_by_alias_reads_the_aliased_version(monkeypatch, tmp_path): + sink = _aliasing_sink(monkeypatch, _Model({ + 'V1': ['DEFAULT'], + 'V2': ['PROD'] + })) + published = sink.fetch_model( + SimpleNamespace(), str(tmp_path / 'out'), + location=ModelRegistryLocation('CHURN', alias='prod')) + assert published.version == 'V2' + assert published.metadata == {'kumo_bundle_id': 'id-V2'} + assert (tmp_path / 'out' / 'training-spec.json').read_text() == 'V2' + with pytest.raises(ValueError, match='holds alias'): + sink.fetch_model(SimpleNamespace(), str(tmp_path / 'other'), + location=ModelRegistryLocation('CHURN', alias='V1')) + + +def test_failed_alias_move_puts_it_back(monkeypatch): + model = _Model({'V1': ['PROD'], 'V2': []}) + + def fail(alias): + raise RuntimeError('lost connection') + + def failing_version(name): + version = _Model.version(model, name) + if name == 'V2': + version.set_alias = fail + return version + + model.version = failing_version + sink = _aliasing_sink(monkeypatch, model) + with pytest.raises(RuntimeError, match='lost connection'): + sink.set_model_alias(SimpleNamespace(), + ModelRegistryLocation('CHURN', version='V2'), + 'prod') + assert model.aliases == {'V1': ['PROD'], 'V2': []} + + +@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 = tmp_path / 'bundle' + bundle.mkdir() + (bundle / 'training-spec.json').write_text('{}') + (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, '_bundle_model', lambda _: object()) + monkeypatch.setattr(registry_module, '_carrier_signature', lambda: {}) + sink.publish_model(SimpleNamespace(), str(bundle), + location=ModelRegistryLocation('CHURN')) + 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..8b1f65e0e 100644 --- a/setup.py +++ b/setup.py @@ -293,6 +293,9 @@ def parse_project_version(project_file_path: Path) -> str: "respx==0.22.0", "pytorch-frame==0.3.0", "pytorch-lightning==2.6.6", + # A superset of kumo-engine's `mlflow` extra (mlflow-skinny): the tests use + # MLflow's local SQLite store. Its wheel is tested with --no-deps. + "mlflow>=3.13,<4", ] dev_requires = test_requires + [ diff --git a/uv.lock b/uv.lock index fae1d0a85..321bd266a 100644 --- a/uv.lock +++ b/uv.lock @@ -140,7 +140,7 @@ dependencies = [ { name = "frozenlist" }, { name = "multidict" }, { name = "propcache" }, - { name = "typing-extensions", marker = "python_full_version < '3.13' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "typing-extensions", marker = "python_full_version < '3.13'" }, { name = "yarl" }, ] sdist = { url = "https://files.pythonhosted.org/packages/6c/4c/bdccd81e9ee225b69c60e7766c9a5b05364f118f4d383713b89a682d772d/aiohttp-3.14.5.tar.gz", hash = "sha256:5558a7f5a05af9ecf744af91e5baefc436f93c9333e656c27ec253f9a6bbe178", size = 8078052, upload-time = "2026-10-11T01:05:12.408Z" } @@ -304,7 +304,7 @@ version = "1.4.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "frozenlist" }, - { name = "typing-extensions", marker = "python_full_version < '3.13' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "typing-extensions", marker = "python_full_version < '3.13'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/61/62/06741b579156360248d1ec624842ad0edf697050bbaf7c3e46394e106ad1/aiosignal-1.4.0.tar.gz", hash = "sha256:f47eecd9468083c2029cc99945502cb7708b082c232f9aca65da147157b251c7", size = 25007, upload-time = "2025-07-03T22:54:43.528Z" } wheels = [ @@ -376,7 +376,7 @@ version = "4.14.2" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "idna" }, - { name = "typing-extensions", marker = "python_full_version < '3.13' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "typing-extensions", marker = "python_full_version < '3.13'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/61/cc/a381afa6efea9f496eff839d4a6a1aed3bfafc7b3ab4b0d1b243a12573dd/anyio-4.14.2.tar.gz", hash = "sha256:cfa139f3ed1a23ee8f88a145ddb5ac7605b8bbfd8592baacd7ce3d8bb4313c7f", size = 260176, upload-time = "2026-07-12T20:29:07.082Z" } wheels = [ @@ -572,12 +572,21 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/88/c6/92fcd42f1ba33e1184263f25bfabf3d27c383410470f169e4b8163bf9c17/beautifulsoup4-4.15.0-py3-none-any.whl", hash = "sha256:d6f88de62e1d4e38ecb1077eb9724cd0eff29d2a08ca16a401e9b9e93f117cf9", size = 109924, upload-time = "2026-06-07T16:44:21.566Z" }, ] +[[package]] +name = "blinker" +version = "1.9.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/21/28/9b3f50ce0e048515135495f198351908d99540d69bfdc8c1d15b73dc55ce/blinker-1.9.0.tar.gz", hash = "sha256:b4ce2265a7abece45e7cc896e98dbebe6cead56bcf805a3d23136d145f5445bf", size = 22460, upload-time = "2024-11-08T17:25:47.436Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/10/cb/f2ad4230dc2eb1a74edf38f1a38b9b52277f75bef262d8908e60d957e13c/blinker-1.9.0-py3-none-any.whl", hash = "sha256:ba0efaa9080b619ff2f3459d1d500c57bddea4a6b424b60a91141db6fd2f08bc", size = 8458, upload-time = "2024-11-08T17:25:46.184Z" }, +] + [[package]] name = "blockbuster" version = "1.5.28" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "forbiddenfruit", marker = "implementation_name == 'cpython' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "forbiddenfruit", marker = "implementation_name == 'cpython'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/28/14/de0062a8e3758fe1880729c1828196e6e951b8cd3299dfdd1e68aeae6e63/blockbuster-1.5.28.tar.gz", hash = "sha256:5af2ed9de6c874339a76cc831d1d11db79800a974573580e52ad827a70803482", size = 98235, upload-time = "2026-09-18T10:51:11.263Z" } wheels = [ @@ -605,7 +614,7 @@ source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "botocore-stubs" }, { name = "types-s3transfer" }, - { name = "typing-extensions", marker = "python_full_version < '3.12' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "typing-extensions", marker = "python_full_version < '3.12'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/7b/94/9dd7426ce339c15ff6bae15b609977aad2f8f778244f9a06e99db821c974/boto3_stubs-1.42.90.tar.gz", hash = "sha256:d0ea4bc1352ba5ef8fae13cfb2fef511f372113ede0388b50cdd19811d39ce01", size = 102688, upload-time = "2026-04-16T21:25:36.455Z" } wheels = [ @@ -657,8 +666,8 @@ dependencies = [ { name = "numpy" }, { name = "pandas" }, { name = "scikit-learn" }, - { name = "scipy", version = "1.17.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, - { name = "scipy", version = "1.18.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "scipy", version = "1.17.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" }, + { name = "scipy", version = "1.18.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/6d/28/3325463e27238bacefce34902ce2e5883ea3bd4044ce93e744ef499a6d72/category_encoders-2.11.1.tar.gz", hash = "sha256:f21b99185800c19b562b2928a0fe67dbf5658d048cd70451bfe8811c4123fdb0", size = 76600, upload-time = "2026-09-08T15:29:09.823Z" } wheels = [ @@ -679,7 +688,7 @@ name = "cffi" version = "2.1.1" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "pycparser", marker = "implementation_name != 'PyPy' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "pycparser", marker = "implementation_name != 'PyPy'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/9e/ef/008a1939e372c06329a3fce4279c02f328488f3526744906eeec3da7ad5f/cffi-2.1.1.tar.gz", hash = "sha256:dd31f52ea1086513bb9df30f8fcee9b8918323ae067a3d5b78bc826a000712be", size = 530807, upload-time = "2026-08-03T21:21:18.939Z" } wheels = [ @@ -960,7 +969,7 @@ name = "colorlog" version = "6.12.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "colorama", marker = "sys_platform == 'win32' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "colorama", marker = "sys_platform == 'win32'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/8c/55/ba79756cb90c8d69d599d57785398ac87bba7b19c80e87f4e8a562197c93/colorlog-6.12.0.tar.gz", hash = "sha256:2a7924c1dadf18b22a0eb8b06d1c7b01d5341707ec1641eb6fcc4fde0c3e8e5f", size = 18151, upload-time = "2026-07-23T13:40:40.71Z" } wheels = [ @@ -981,7 +990,7 @@ resolution-markers = [ "python_full_version < '3.12' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu130' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests'", ] dependencies = [ - { name = "numpy", marker = "python_full_version < '3.12' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "numpy" }, ] sdist = { url = "https://files.pythonhosted.org/packages/58/01/1253e6698a07380cd31a736d248a3f2a50a7c88779a1813da27503cadc2a/contourpy-1.3.3.tar.gz", hash = "sha256:083e12155b210502d0bca491432bb04d56dc3432f95a979b429f2848c3dbe880", size = 13466174, upload-time = "2025-07-26T12:03:12.549Z" } wheels = [ @@ -1086,7 +1095,7 @@ resolution-markers = [ "python_full_version == '3.12.*' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu130' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests'", ] dependencies = [ - { name = "numpy", marker = "python_full_version >= '3.12' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "numpy" }, ] sdist = { url = "https://files.pythonhosted.org/packages/83/5a/a55177dd22553a277388e8a1b3220e92de91bacb28356cdc73caa240121d/contourpy-1.4.0.tar.gz", hash = "sha256:20156f5a1ac4f8ce02656e39a61e82164a3d359796dc8026f75b062783d500e1", size = 13323726, upload-time = "2026-09-11T19:05:05.808Z" } wheels = [ @@ -1277,7 +1286,7 @@ wheels = [ [package.optional-dependencies] toml = [ - { name = "tomli", marker = "python_full_version <= '3.11' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "tomli", marker = "python_full_version <= '3.11'" }, ] [[package]] @@ -1285,7 +1294,7 @@ name = "cryptography" version = "50.0.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "cffi", marker = "platform_python_implementation != 'PyPy' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "cffi", marker = "platform_python_implementation != 'PyPy'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/de/41/6cbdcf9142d00fe82836fbb51e503e58088575cf7a0fe1dbff6695bf0840/cryptography-50.0.0.tar.gz", hash = "sha256:eeac2acb5a20ed25e0ad6d1df9891a520b78b404266b6d11778f25d5d691a6c9", size = 880201, upload-time = "2026-07-31T14:25:10.11Z" } wheels = [ @@ -1341,7 +1350,7 @@ name = "cuda-bindings" version = "13.4.2" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "cuda-pathfinder", marker = "extra == 'extra-4-kumo-cu130' or extra == 'extra-4-kumo-cu132' or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra != 'extra-4-kumo-cpu' and extra != 'group-4-kumo-public-tests')" }, + { name = "cuda-pathfinder", marker = "extra == 'extra-4-kumo-cu130' or extra == 'extra-4-kumo-cu132' or (extra != 'extra-4-kumo-cpu' and extra != 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests')" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/19/6f/e00ffcebcad6405a2326a52612eff98240a67c637dfed516e61b881e2db3/cuda_bindings-13.4.2-cp311-cp311-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:fc0a18f65b26459cd3c49c0250a44b419cdf790cffcb096ff4d46a1973f6f4b7", size = 6484069, upload-time = "2026-09-17T22:05:52.594Z" }, @@ -1390,34 +1399,34 @@ wheels = [ [package.optional-dependencies] cudart = [ - { name = "nvidia-cuda-runtime", version = "13.0.96", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform == 'linux' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'linux' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-cuda-runtime", version = "13.0.96", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform == 'linux' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, ] cufft = [ - { name = "nvidia-cufft", version = "12.0.0.61", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform == 'linux' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'linux' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-cufft", version = "12.0.0.61", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform == 'linux' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, ] cufile = [ - { name = "nvidia-cufile", version = "1.15.1.6", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform == 'linux' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'linux' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-cufile", version = "1.15.1.6", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform == 'linux' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, ] cupti = [ - { name = "nvidia-cuda-cupti", version = "13.0.85", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform == 'linux' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'linux' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-cuda-cupti", version = "13.0.85", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform == 'linux' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, ] curand = [ - { name = "nvidia-curand", version = "10.4.0.35", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform == 'linux' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'linux' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-curand", version = "10.4.0.35", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform == 'linux' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, ] cusolver = [ - { name = "nvidia-cusolver", version = "12.0.4.66", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform == 'linux' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'linux' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-cusolver", version = "12.0.4.66", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform == 'linux' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, ] cusparse = [ - { name = "nvidia-cusparse", version = "12.6.3.3", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform == 'linux' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'linux' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-cusparse", version = "12.6.3.3", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform == 'linux' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, ] nvjitlink = [ - { name = "nvidia-nvjitlink", version = "13.0.88", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform == 'linux' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'linux' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-nvjitlink", version = "13.0.88", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform == 'linux' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, ] nvrtc = [ - { name = "nvidia-cuda-nvrtc", version = "13.0.88", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform == 'linux' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'linux' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-cuda-nvrtc", version = "13.0.88", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform == 'linux' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, ] nvtx = [ - { name = "nvidia-nvtx", version = "13.0.85", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform == 'linux' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'linux' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-nvtx", version = "13.0.85", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform == 'linux' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu132' and extra != 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, ] [[package]] @@ -1436,43 +1445,43 @@ wheels = [ [package.optional-dependencies] cublas = [ - { name = "nvidia-cublas", version = "13.4.0.1", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'AMD64' and sys_platform == 'win32' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, - { name = "nvidia-cuda-nvrtc", version = "13.2.78", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'AMD64' and sys_platform == 'win32' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-cublas", version = "13.4.0.1", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'AMD64' and sys_platform == 'win32' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-cuda-nvrtc", version = "13.2.78", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'AMD64' and sys_platform == 'win32' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, ] cudart = [ - { name = "nvidia-cuda-runtime", version = "13.2.75", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'AMD64' and sys_platform == 'win32' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-cuda-runtime", version = "13.2.75", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'AMD64' and sys_platform == 'win32' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, ] cufft = [ - { name = "nvidia-cufft", version = "12.2.0.46", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'AMD64' and sys_platform == 'win32' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, - { name = "nvidia-nvjitlink", version = "13.4.92", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'AMD64' and sys_platform == 'win32' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-cufft", version = "12.2.0.46", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'AMD64' and sys_platform == 'win32' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-nvjitlink", version = "13.4.92", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'AMD64' and sys_platform == 'win32' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, ] cufile = [ - { name = "nvidia-cufile", version = "1.17.1.22", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'aarch64' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (platform_machine == 'aarch64' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'aarch64' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (platform_machine == 'aarch64' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'aarch64' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (platform_machine == 'x86_64' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (platform_machine == 'x86_64' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (platform_machine == 'x86_64' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform != 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'linux' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'linux' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-cufile", version = "1.17.1.22", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, ] cupti = [ - { name = "nvidia-cuda-cupti", version = "13.2.75", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'AMD64' and sys_platform == 'win32' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-cuda-cupti", version = "13.2.75", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'AMD64' and sys_platform == 'win32' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, ] curand = [ - { name = "nvidia-curand", version = "10.4.2.55", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'AMD64' and sys_platform == 'win32' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-curand", version = "10.4.2.55", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'AMD64' and sys_platform == 'win32' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, ] cusolver = [ - { name = "nvidia-cublas", version = "13.4.0.1", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'AMD64' and sys_platform == 'win32' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, - { name = "nvidia-cusolver", version = "12.2.0.1", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'AMD64' and sys_platform == 'win32' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, - { name = "nvidia-cusparse", version = "12.7.10.1", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'AMD64' and sys_platform == 'win32' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, - { name = "nvidia-nvjitlink", version = "13.4.92", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'AMD64' and sys_platform == 'win32' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-cublas", version = "13.4.0.1", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'AMD64' and sys_platform == 'win32' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-cusolver", version = "12.2.0.1", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'AMD64' and sys_platform == 'win32' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-cusparse", version = "12.7.10.1", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'AMD64' and sys_platform == 'win32' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-nvjitlink", version = "13.4.92", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'AMD64' and sys_platform == 'win32' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, ] cusparse = [ - { name = "nvidia-cusparse", version = "12.7.10.1", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'AMD64' and sys_platform == 'win32' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, - { name = "nvidia-nvjitlink", version = "13.4.92", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'AMD64' and sys_platform == 'win32' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-cusparse", version = "12.7.10.1", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'AMD64' and sys_platform == 'win32' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-nvjitlink", version = "13.4.92", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'AMD64' and sys_platform == 'win32' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, ] nvjitlink = [ - { name = "nvidia-nvjitlink", version = "13.4.92", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'AMD64' and sys_platform == 'win32' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-nvjitlink", version = "13.4.92", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'AMD64' and sys_platform == 'win32' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, ] nvrtc = [ - { name = "nvidia-cuda-nvrtc", version = "13.2.78", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'AMD64' and sys_platform == 'win32' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-cuda-nvrtc", version = "13.2.78", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'AMD64' and sys_platform == 'win32' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, ] nvtx = [ - { name = "nvidia-nvtx", version = "13.2.75", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'AMD64' and sys_platform == 'win32' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'linux' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (sys_platform == 'win32' and extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, + { name = "nvidia-nvtx", version = "13.2.75", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-4-kumo-cu132') or (platform_machine == 'AMD64' and sys_platform == 'win32' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, ] [[package]] @@ -1594,6 +1603,20 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/02/08/9c41fb51ab5b43eb21674aff13df270e8ba6c4b29c8624e328dc7a9482af/distlib-0.4.3-py2.py3-none-any.whl", hash = "sha256:4b0ce306c966eb73bc3a7b6abad017c556dadd92c44701562cd528ac7fde4d5b", size = 470628, upload-time = "2026-06-12T08:04:50.506Z" }, ] +[[package]] +name = "docker" +version = "7.2.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "pywin32", marker = "sys_platform == 'win32'" }, + { name = "requests" }, + { name = "urllib3" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/88/7f/731ff914b0255d3d065f45fd4e626d4b8c95dbcbaada049f337a6ac16410/docker-7.2.0.tar.gz", hash = "sha256:cebb93773d334f778e023a7ee352a8d6e13ab1bd3b863a4d4a59dec897df43ac", size = 118731, upload-time = "2026-07-09T14:53:46.39Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/75/23/529140fe1aab80fc6992f93a706deec709140a6397439139a054e1515c45/docker-7.2.0-py3-none-any.whl", hash = "sha256:a3f45fdeb9165e2d25d9a1d02ddf3bc70fb572cf5ebbf9b58558c22caf29b71f", size = 148775, upload-time = "2026-07-09T14:53:45.224Z" }, +] + [[package]] name = "docutils" version = "0.21.2" @@ -1717,6 +1740,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" @@ -1926,6 +1979,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" @@ -2151,6 +2228,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" @@ -2326,6 +2439,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" @@ -2376,6 +2498,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" @@ -2473,6 +2604,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" @@ -2751,8 +2891,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" }, @@ -2764,11 +2904,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" }, @@ -2788,7 +2928,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 = [ @@ -2805,6 +2945,7 @@ dev = [ { name = "flake8" }, { name = "isort" }, { name = "kubernetes-stubs" }, + { name = "mlflow" }, { name = "moto" }, { name = "packaging" }, { name = "pick" }, @@ -2883,6 +3024,7 @@ rest = [ test = [ { name = "category-encoders" }, { name = "coverage" }, + { name = "mlflow" }, { name = "moto" }, { name = "pytest" }, { name = "pytest-asyncio" }, @@ -2906,7 +3048,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')" }, ] @@ -2988,6 +3130,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" }, @@ -3212,11 +3356,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" }, @@ -3346,9 +3490,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"] }, ] @@ -3359,11 +3513,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" }, @@ -3382,20 +3536,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", "gcs"] +provides-extras = ["cpu", "test", "snowflake", "databricks", "bigquery", "s3", "text", "gcs", "mlflow", "sagemaker"] [[package]] name = "kumo-gcs" @@ -3541,7 +3699,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"] @@ -3776,8 +3934,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" }, @@ -3864,6 +4022,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" @@ -4091,7 +4332,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 = [ @@ -4289,7 +4530,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" }, @@ -4432,7 +4673,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" }, @@ -4470,7 +4711,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" }, @@ -4582,9 +4823,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" }, @@ -4622,7 +4863,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" }, @@ -4640,6 +4881,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" @@ -5093,7 +5343,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 = [ @@ -5260,6 +5510,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" @@ -5773,7 +6035,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 = [ @@ -5806,7 +6068,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" }, @@ -5823,7 +6085,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 = [ @@ -5926,6 +6188,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" @@ -5957,11 +6228,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" } @@ -5978,11 +6249,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" }, @@ -6001,6 +6272,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" @@ -6209,6 +6502,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" @@ -6217,8 +6523,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" } @@ -6281,7 +6587,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 = [ @@ -6375,7 +6681,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 = [ @@ -6486,6 +6792,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" @@ -6495,6 +6818,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" @@ -6592,8 +6924,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" }, @@ -6650,7 +6982,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" }, @@ -6804,7 +7136,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" } @@ -6871,7 +7203,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 = [ @@ -7077,14 +7409,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" }, @@ -7161,14 +7493,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" }, @@ -7211,21 +7543,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 = [ @@ -7260,21 +7592,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" }, @@ -7325,11 +7657,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 = [ @@ -7341,7 +7673,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 = [ @@ -7409,7 +7741,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 = [ @@ -7426,7 +7758,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 = [ @@ -7494,7 +7826,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 = [ @@ -7538,6 +7870,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" @@ -7582,6 +7923,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"