From f535341b63a20382164d4111da1ed794cc4b3f3d Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Wed, 28 Oct 2026 16:02:54 +0000 Subject: [PATCH 01/16] Record fits in experiment tracking and publish models to registries Co-Authored-By: Claude Opus 5.5 --- examples/mlops/smoke.py | 138 +++ packages/kumo-api/kumoapi/plugin.py | 70 +- .../kumo-api/test/test_plugin_registry.py | 33 + packages/kumo-core/src/kumo_core/artifacts.py | 32 + .../src/kumo_core/automl/auto_trainer.py | 37 +- .../src/kumo_core/automl/optuna_searcher.py | 1 + .../kumo-core/src/kumo_core/automl/runtime.py | 5 +- .../src/kumo_core/automl/searcher.py | 14 +- .../callbacks/train_progress_logger.py | 25 +- .../src/kumo_core/diskgraph/lightning_data.py | 11 +- .../src/kumo_core/instrumentation.py | 10 + .../src/kumo_core/train/formatting.py | 34 +- .../src/kumo_core/training/runtime.py | 4 +- packages/kumo-core/tests/test_artifacts.py | 8 + .../tests/test_prediction_output_columns.py | 23 + .../kumo-core/tests/test_searched_params.py | 38 + .../kumo-core/tests/test_worker_handles.py | 46 + .../src/kumo_databricks/registry.py | 103 ++- .../kumo-databricks/tests/test_registry.py | 96 +- .../kumo-databricks/tests/test_uc_registry.py | 2 +- packages/kumo-engine/docs/mlops.md | 356 ++++++-- packages/kumo-engine/kumo/engine/__init__.py | 2 + packages/kumo-engine/kumo/engine/_events.py | 86 ++ packages/kumo-engine/kumo/engine/_manifest.py | 9 + packages/kumo-engine/kumo/engine/_mlops.py | 93 +- packages/kumo-engine/kumo/engine/_scratch.py | 19 +- packages/kumo-engine/kumo/engine/_tracking.py | 482 ++++++++++ .../kumo/engine/connector/snowflake.py | 8 +- .../kumo/engine/integrations/__init__.py | 3 + .../kumo/engine/integrations/_mlflow_model.py | 217 +++++ .../engine/integrations/_mlflow_model_code.py | 7 + .../engine/integrations/_mlflow_target.py | 230 +++++ .../kumo/engine/integrations/mlflow.py | 859 ++++++++++++++++++ .../kumo/engine/integrations/snowflake.py | 281 ++++++ packages/kumo-engine/kumo/engine/tracking.py | 13 + .../kumo/engine/trainer/__init__.py | 2 + .../kumo/engine/trainer/_provenance.py | 32 +- .../kumo/engine/trainer/_runtime.py | 6 +- .../kumo/engine/trainer/_scoring.py | 81 ++ .../kumo-engine/kumo/engine/trainer/_spec.py | 4 +- .../kumo/engine/trainer/output_config.py | 20 +- .../kumo/engine/trainer/trainer.py | 332 ++++++- packages/kumo-engine/pyproject.toml | 11 + .../test/_snowflake_experiment_fake.py | 192 ++++ packages/kumo-engine/test/test_fit_events.py | 154 ++++ .../kumo-engine/test/test_mlflow_model.py | 516 +++++++++++ .../kumo-engine/test/test_mlflow_target.py | 290 ++++++ .../kumo-engine/test/test_mlflow_tracking.py | 404 ++++++++ packages/kumo-engine/test/test_mlops.py | 24 + .../kumo-engine/test/test_optional_imports.py | 17 + .../test/test_prediction_completion.py | 4 +- .../test/test_prediction_workers.py | 4 +- .../kumo-engine/test/test_scoring_defaults.py | 139 +++ .../test/test_semantic_carry_forward.py | 15 +- .../test/test_snowflake_connector.py | 15 + .../test/test_snowflake_tracking.py | 250 +++++ .../test/test_tracking_contract.py | 400 ++++++++ .../kumo-engine/test/test_trainer_registry.py | 187 +++- packages/kumo-snowflake/pyproject.toml | 8 +- .../src/kumo_snowflake/registry.py | 178 +++- .../kumo-snowflake/tests/test_registry.py | 191 +++- setup.py | 3 + uv.lock | 639 ++++++++++--- 63 files changed, 7186 insertions(+), 327 deletions(-) create mode 100644 examples/mlops/smoke.py create mode 100644 packages/kumo-api/test/test_plugin_registry.py create mode 100644 packages/kumo-core/tests/test_prediction_output_columns.py create mode 100644 packages/kumo-core/tests/test_searched_params.py create mode 100644 packages/kumo-core/tests/test_worker_handles.py create mode 100644 packages/kumo-engine/kumo/engine/_events.py create mode 100644 packages/kumo-engine/kumo/engine/_tracking.py create mode 100644 packages/kumo-engine/kumo/engine/integrations/__init__.py create mode 100644 packages/kumo-engine/kumo/engine/integrations/_mlflow_model.py create mode 100644 packages/kumo-engine/kumo/engine/integrations/_mlflow_model_code.py create mode 100644 packages/kumo-engine/kumo/engine/integrations/_mlflow_target.py create mode 100644 packages/kumo-engine/kumo/engine/integrations/mlflow.py create mode 100644 packages/kumo-engine/kumo/engine/integrations/snowflake.py create mode 100644 packages/kumo-engine/kumo/engine/tracking.py create mode 100644 packages/kumo-engine/kumo/engine/trainer/_scoring.py create mode 100644 packages/kumo-engine/test/_snowflake_experiment_fake.py create mode 100644 packages/kumo-engine/test/test_fit_events.py create mode 100644 packages/kumo-engine/test/test_mlflow_model.py create mode 100644 packages/kumo-engine/test/test_mlflow_target.py create mode 100644 packages/kumo-engine/test/test_mlflow_tracking.py create mode 100644 packages/kumo-engine/test/test_scoring_defaults.py create mode 100644 packages/kumo-engine/test/test_snowflake_tracking.py create mode 100644 packages/kumo-engine/test/test_tracking_contract.py diff --git a/examples/mlops/smoke.py b/examples/mlops/smoke.py new file mode 100644 index 000000000..47e152564 --- /dev/null +++ b/examples/mlops/smoke.py @@ -0,0 +1,138 @@ +# 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=connector, + experiment_name='KUMO_SMOKE'), None + + +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')) + cfg = kumo.MLOps( + args.artifact_root or str(work / 'root'), tracking=tracker, + registry=None if model_name is None else ModelRegistryLocation( + model_name, alias='candidate'), 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 {tracker.run_id}, bundle {result.id}, ' + f'registered {result.registered_model}') + if model_name is not None: + from kumo.engine.integrations.mlflow import load_model + + loaded = load_model(f'models:/{model_name}@candidate', + tracking=tracker) + 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/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/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/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_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_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/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 338c2eb73..0a79573eb 100644 --- a/packages/kumo-engine/docs/mlops.md +++ b/packages/kumo-engine/docs/mlops.md @@ -1,69 +1,32 @@ # MLOps with Kumo Engine -## Quickstart +## Turning MLOps on -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: +`kumo.MLOps` says where artifacts go and what records them. Use it as a block, which notebooks and scheduled scripts alike can do, or pass it as `mlops=` on a single call: ```python import kumo.engine as kumo -cfg = kumo.MLOps("artifacts") # or s3://..., gs://..., /Volumes/..., (connector, "@DB.SCHEMA.STAGE/kumo") -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}") -``` - -## How it fits together - -| artifact | id | stored at (MLOps on) | -|---|---|---| -| snapshot | `snapshot.id` | `/snapshots///` | -| training table | `table.id` | `/training_tables//` | -| model | `result.id` (the bundle id) | `/models//` | - -Each carries a `_provenance.json` manifest that embeds its parents' manifests (see [Lineage](#lineage)). - -## Turning MLOps on - -`kumo.MLOps` says where artifacts go. Use it as a block, or pass it as `mlops=` on a single call: - -```python -cfg = kumo.MLOps("s3://bucket/kumo") +cfg = kumo.MLOps(artifact_root="s3://bucket/kumo") # or (connector, path), e.g. a volume or stage 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`** is required. It is 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. +- **Volumes and stages:** `(DatabricksConnector(...), "/Volumes////kumo")` and `(SnowflakeConnector(...), "@DB.SCHEMA.STAGE/kumo")` work as roots once the installed kumo-databricks or kumo-snowflake sink holds files at those paths; with an older plugin, building `kumo.MLOps` raises. +- **`export`** names the artifact kinds stored under the root when they are produced: `snapshot`, `training_table` and `model` (default: all). +- **`tracking`** and **`registry`** record fits and register models. +- **`on_error`** says what a tracking or registration failure in `fit` does. `"warn"` (default): the call warns, naming the error, and returns as usual. `"raise"`, for scheduled jobs: the call finishes and stores its result first, then raises a `RuntimeError` whose `__cause__` is the failure; after `fit`, `trainer.result` holds the fitted model. - **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=`, `output_config=`) never write to the artifact root. A path export that finds a complete copy of the same artifact at its destination reuses it instead of uploading again (logged at info): a snapshot with the same bytes, a training table or prediction directory with `_SUCCESS` and the same `_KUMO_EXPORT` identity, or a model bundle with the same files and `_provenance.json`. -## Exporting and loading artifacts - -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: - -| 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, pquery=pquery, 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}")` | - -`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. - -## Reference - -### Lineage +## Lineage Every artifact carries a `_provenance.json` manifest, metadata only. The leading underscore keeps Spark and Hive Parquet readers from treating it as data. 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. -#### Training tables +### Training tables A generated or caller-supplied training table carries `_provenance.json` in its staging directory: - the query, the effective generation plan and the task type; @@ -71,11 +34,11 @@ A generated or caller-supplied training table carries `_provenance.json` in its - the snapshot it was generated from (that snapshot's manifest, embedded); - for `TrainingTable.from_source_table`, the source table's connector, location and version. -`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=..., pquery=..., snapshot=...)` materializes an export again. +`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. With MLOps on, each table is stored once under `/training_tables//`. -#### Models +### Models Every fitted bundle carries `_provenance.json` (`ModelManifest`): - its `id`, the bundle id (a content hash that leaves the manifest file out); @@ -88,6 +51,75 @@ With MLOps on, each bundle is stored under `/models//`. A compl 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. +### Scoring defaults + +A model can carry the prediction arguments `predict` uses when a call omits them: + +```python +from kumo.engine import ModelOutputConfig, ScoringDefaults + +tuned = result.with_scoring_defaults( + ScoringDefaults(binary_classification_threshold=0.37)) +tuned.predict(snapshot) # thresholds at 0.37 +tuned.scoring_defaults # ScoringDefaults(binary_classification_threshold=0.37, ...) +``` + +`with_scoring_defaults` returns a `TrainingResult` for a copy of the bundle with its own id. The copy's manifest embeds the original's in `parents.derived_from`. The threshold applies to binary and multilabel classification, `num_classes_to_return` to multiclass, multilabel and link prediction. Explicit `predict` arguments and a query's `TOP k` win; `ScoringDefaults()` removes the defaults. `Trainer.export(config)` publishes a fitted or loaded bundle like `fit(export_config=...)`. + +### Registering models + +With `registry=ModelOutputConfig(connector, location=...)`, every fit in the block registers its model in the connector's registry (Snowflake Model Registry, or Unity Catalog through a `DatabricksConnector`) after it is stored and tracked: + +```python +from kumo.engine import ModelOutputConfig, ModelRegistryLocation + +cfg = kumo.MLOps("s3://bucket/kumo", registry=ModelOutputConfig( + registry, location=ModelRegistryLocation("CHURN_GNN", alias="challenger"))) +with kumo.mlops(cfg): + result = trainer.fit(training_table, snapshot=snapshot) +result.registered_model # PublishedModel(name='CHURN_GNN', version='V3', uri=...) + +# Scoring jobs load whichever version an alias points at. +trainer = kumo.Trainer.load_from_connector(registry, location=ModelRegistryLocation( + "CHURN_GNN", alias="challenger")) +``` + +- Registration is independent of `export`. If it fails, `fit` warns with the error and returns the result with `registered_model=None` (with `on_error="raise"`, it raises instead), and a tracked fit run ends failed with `kumo.error`; an invalid alias or a connector that cannot publish models raises before training, and an invalid alias or model name raises when `kumo.MLOps` is built. Snowflake model names are Snowflake identifiers (a letter or underscore, then letters, digits, underscores or `$`, e.g. `USERS_CHURN`; or a double-quoted name); Unity Catalog model names are at most 255 characters with no space, `/`, `.` or control character, optionally qualified with the connector's own catalog and schema. +- `result.export(ModelOutputConfig(connector, location=...))` on a `TrainingResult`, or `trainer.export(...)` for its latest result, publishes the model with or without a block and also accepts an `alias`; `Trainer.load_from_connector` loads it back. `fit(export_config=...)` is shorthand for fitting then exporting, and its version is `result.published_model`. +- An alias is a letter or underscore followed by letters, digits or underscores; `DEFAULT`, `FIRST`, `LAST`, `LATEST` and `v` are reserved. The alias is set on the new version after publishing, moving it from any version that held it; if that fails, the publish still stands, with a warning. Snowflake allows one alias per version; if moving an alias fails, it is put back on the version that held it. Snowflake's `DEFAULT`, `FIRST` and `LAST` can be loaded but not set. File, S3 and GCS connectors have no aliases. +- Each registered version records only pointers: the `kumo_engine` marker and `kumo_bundle_id`, as version metrics in the Snowflake Model Registry and as version tags in Unity Catalog through a `DatabricksConnector`. (A version registered through MLflow, below, gets a `kumo_bundle_id` version tag only.) The bundle's own `_provenance.json` holds its lineage. `list_models` and `fetch_model` return `kumo_bundle_id` in `PublishedModel.metadata`; `load_from_connector` warns if the fetched bundle differs from it. Through a `DatabricksConnector`, publishing a bundle that a ready version of the model already holds (same `kumo_bundle_id`) returns that version, logged at info, instead of adding another; an alias then moves to it. + +### Promote a Kumo model + +To register every fit as a candidate and promote only a better one, register under a candidate alias (or none) and move the serving alias yourself with the registry's own API. Compare the new model's test metrics (`result.metrics`, also on its tracked fit run) with the current champion's (`load_from_connector(...).result.metrics`): + +```python +champion = kumo.Trainer.load_from_connector( + registry, location=ModelRegistryLocation("CHURN_GNN", alias="champion")) +better = result.metrics["test_auroc"] > champion.result.metrics["test_auroc"] +version = result.registered_model.version +``` + +Unity Catalog (registered through MLflow or a `DatabricksConnector`) and SageMaker managed MLflow: + +```python +from mlflow import MlflowClient + +if better: + MlflowClient().set_registered_model_alias("main.ml.churn_gnn", "champion", version) +``` + +`set_registered_model_alias` moves the alias from the version that held it. For models registered through MLflow, the champion's metrics are also on its run: `client.get_run(client.get_model_version_by_alias(name, "champion").run_id).data.metrics`. + +Snowflake Model Registry, where a version holds at most one alias, so the alias is unset first: + +```sql +ALTER MODEL DB.SCHEMA.CHURN_GNN VERSION champion UNSET ALIAS; +ALTER MODEL DB.SCHEMA.CHURN_GNN VERSION V3 SET ALIAS = champion; +``` + +Scoring jobs that load `alias="champion"` (or `models:/@champion`) pick up the new version on their next run. + ### Moving snapshots between machines To build an index on one machine and train or score on another, export it to a file destination and load it there: @@ -112,11 +144,11 @@ With MLOps on (`graph.snapshot(mlops=...)` or a `kumo.mlops` block), each snapsh ### Scoring without the graph -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: +A scoring job does not need the training job's `Graph` code. It loads the snapshot the training job exported and the model, then scores: ```python -snapshot = GraphSnapshot.load_from_connector(cfg, id=snapshot_id) -trainer = Trainer.load_from_connector(cfg, path=f"models/{bundle_id}") +snapshot = GraphSnapshot.load_from_connector(store, path=location) +trainer = Trainer.load_from_connector(store, path="models/") result = trainer.predict(snapshot, binary_classification_threshold=0.5) ``` @@ -125,3 +157,221 @@ result = trainer.predict(snapshot, binary_classification_threshold=0.5) - 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, ...)`. + +## Progress events + +`fit(callbacks=[...])` calls each callback with every progress +`kumo.engine.tracking.Event`: `fit_started`, then per search trial +`trial_started`, `epoch_ended` and `trial_ended`, then `fit_ended`. Ended events +carry `params["status"]` and, on failure, a short `params["error"]`. A callback +that raises is logged and not called again; it never fails the fit. For JSON +lines on stdout: + +```python +import json + +trainer.fit(training_table, snapshot=snapshot, + callbacks=[lambda e: print(json.dumps(e.to_dict()))]) +``` + +## Experiment tracking + +`kumo.MLOps(tracking=...)` records each fit in an experiment tracker: one run +per fit, with each hyperparameter-search trial attempt as a child run (MLflow) +or inside the fit run (Snowflake). Tracking is on only +inside a `kumo.mlops(cfg)` block or for a `fit` given `mlops=cfg`. + +```python +from kumo.engine.integrations.mlflow import MlflowTracking # kumo-engine[mlflow] + +cfg = kumo.MLOps(artifact_root="s3://bucket/kumo", + tracking=MlflowTracking(experiment_name="users-spend")) +with kumo.mlops(cfg): + result = trainer.fit(training_table, snapshot=snapshot) +``` + +- A tracking or registration failure does not fail the fit: it is logged, and reported once as a `UserWarning` naming the full error, so both notebooks and job logs show it. With `on_error="raise"` the call raises once instead, after storing its result, with the error as the cause. Only global rank 0 tracks. +- Tracker calls run on a background thread, so a slow or unreachable tracking server does not hold up training. At the end of the fit, `fit` waits at most 30 s for tracking to finish; the rest continues in the background. +- During the fit, MLflow HTTP calls use 2 retries and a 30 s timeout unless `MLFLOW_HTTP_REQUEST_MAX_RETRIES` / `MLFLOW_HTTP_REQUEST_TIMEOUT` are set. +- `MlflowTracking()` records where MLflow's own configuration points when it is built (`MLFLOW_TRACKING_URI`, or `mlflow.set_tracking_uri`): on Databricks this is automatic, and with nothing configured MLflow writes a local store in the working directory. `MlflowTracking(target=...)` names the server and the credentials instead; see [Running from a laptop](#running-from-a-laptop). +- `kumo.MLOps(...)` checks the tracker when it is built: one experiment lookup, plus one registry listing and (Unity Catalog) the model's catalog and schema when it registers. An unreachable server, refused credentials or a missing catalog or schema raise there, before any work. +- The `View run ... at` links MLflow prints do not open on SageMaker; after a tracked fit kumo-engine warns how to open the UI: from SageMaker Studio, or at the URL `aws sagemaker create-presigned-mlflow-tracking-server-url --tracking-server-name ` prints. +- MLflow rejects `@` in metric names, so link-prediction metrics such as `test_map@10` are logged as `test_map_at_10`. +- Per-epoch metrics use step = epoch. After the fit, `tracking.run_id` is the fit run's MLflow run id. +- The fit run is named `fit-` (for example `fit-20261027T061838Z`) unless `run_name` is given; trial runs are `trial--attempt-`. +- The fit's source tables are logged as MLflow run inputs, read from the training table's embedded snapshot manifest (plus a custom training table's source). Databricks tables are logged as Delta tables at the snapshot's Delta version, so Unity Catalog links each table to a model registered from the run; file and object-store sources by URI. Other warehouse tables and Custom SQL sources are not logged as inputs; they stay in the model's `_provenance.json`. + +Snowflake records the fit in Snowflake Experiment Tracking (Snowsight: AI & ML +» Experiments), using a Snowpark session: + +```python +from kumo.engine.integrations.snowflake import SnowflakeTracking # kumo-engine[snowflake] + +cfg = kumo.MLOps(artifact_root="s3://bucket/kumo", + tracking=SnowflakeTracking(session, experiment_name="users_spend")) +with kumo.mlops(cfg): + result = trainer.fit(training_table, snapshot=snapshot) +``` + +- Snowflake has no child runs or tags, so trial attempts are logged in the fit run as `trial-1-attempt-1/` params and metrics, and tags are params of the same name. +- `native_model_lineage=True` links a model the fit publishes through Snowflake lineage. +- Names beyond Snowflake's 200 metrics and 1000 params per run are counted in `kumo.dropped_metrics` / `kumo.dropped_params`. +- If a Snowflake run is already active in the process, the fit is not tracked there; an existing run is never resumed (a taken `run_name` gets a suffix). +- Snowflake Experiment Tracking has no dataset inputs, so source tables are not recorded there. +- The Snowpark session's parameters are left as they are. After the fit, `tracking.run_id` is the fit run's name. + +Runs carry pointers, not lineage: + +| tag | on | value | +|---|---|---| +| `kumo.kind` | every run | `fit` or `trial` | +| `kumo.run_id` | every run | the kumo run id (trials: `-trial--attempt-`) | +| `kumo.bundle_id` | a finished fit run | the model's bundle id (`models//` under the artifact root) | +| `kumo.parent_run`, `kumo.status`, `kumo.error` | trial runs; `kumo.error` also on a failed fit (including a failed model logging or registration) | outcome | + +The model's `_provenance.json` records the tracked fit as `tracking.run_id` +(the `kumo.run_id` tag) and is logged on the fit run with `query.pql` and +`model_plan.json`. + +Which id is which: + +| id | what it names | where to read it | +|---|---|---| +| kumo run id | one `fit` call | the `kumo.run_id` tag; `tracking.run_id` in the model's `_provenance.json` | +| bundle id | the fitted model's content | `result.id`; the `kumo.bundle_id` tag; `kumo_bundle_id` in the registry; `models//` under the artifact root | +| tracker run id | the run in the tracker | the tracker's `run_id` attribute (`tracking.run_id` on the `MlflowTracking` object: the MLflow run id; Snowflake: the run name) | +| registry version | the registered model version | `result.registered_model.version` | + +To find a bundle's MLflow fit run, search for its tag: ``mlflow.search_runs(filter_string="tags.`kumo.bundle_id` = '' AND tags.`kumo.kind` = 'fit'", search_all_experiments=True)``. + +Other systems plug in through the `kumo.engine.tracking.Tracking` protocol, +which receives ready-built `RunRecord`s so every tracker records the same +names: `start(record)` with the fit run, `on_trial(record, event)` with a trial +attempt's run after each of its events, and `finish(record, error)` with the +fit run's final record (`None` when tracking stopped early). A tracker that +raises is not called again for that fit and its run is ended as failed. +Optional members: `prepare()`, `tags`, `check(registry)` (called when `kumo.MLOps` is built) and `log_model(result, snapshot, registry)` (see MLflow models below). Params are the resolved model plan +(fit) or searched hyperparameters and `attempt` (trial). `RunRecord.inputs` +maps each source table to its `connector`, `location`, `version` and +`version_kind`. A tracker's `tags=` are added as given, except `kumo.*` names +and param names. + +### MLflow models + +With `registry=ModelRegistryLocation(name, alias=...)` and an `MlflowTracking` tracker, every fit in the block logs its model on the fit run as an MLflow pyfunc model and registers it as a new version of `name` (Unity Catalog with `MLFLOW_REGISTRY_URI=databricks-uc`, SageMaker managed MLflow, or any MLflow registry). One alias is set on the new version. A bundle the registry already holds under `name` (same `kumo_bundle_id` version tag) is not registered again: its version is returned, the alias moves to it, and an INFO log says it was reused. A bare `ModelRegistryLocation` without such a tracker, a fixed `version`, an invalid alias, or a name the registry (`mlflow.get_registry_uri()`) rejects raises when `kumo.MLOps` is built: + +- SageMaker managed MLflow (a tracking server or MLflow app ARN): letters, digits and hyphens, starting and ending with a letter or digit, at most 57 letters and digits (`^[a-zA-Z0-9](-*[a-zA-Z0-9]){0,56}$`), so no `_` or `.`: `users-churn`. +- Unity Catalog (`databricks-uc`): `catalog.schema.model`, each part at most 255 characters with no space, `/` or control character: `main.ml.users_spend`. On a cluster a shorter name takes the current catalog and schema. +- Other MLflow registries: no extra rule. + +The rules conflict, so no one name fits every registry: SageMaker forbids the `.` Unity Catalog needs between its parts and the `_` Snowflake needs instead of `-`. The examples use `users-churn` for SageMaker and other MLflow registries, `..users_churn` for Unity Catalog and `USERS_CHURN` for Snowflake. + +```python +import mlflow +import pandas as pd +from kumo.engine import ModelRegistryLocation +from kumo.engine.integrations.mlflow import MlflowTracking, load_model + +cfg = kumo.MLOps("s3://bucket/kumo", + tracking=MlflowTracking(experiment_name="/Shared/users_spend"), + registry=ModelRegistryLocation("main.ml.users_spend", alias="champion")) +with kumo.mlops(cfg): + result = trainer.fit(training_table, snapshot=snapshot) +result.registered_model # PublishedModel(name='main.ml.users_spend', version='4', uri='models:/main.ml.users_spend/4') + +trainer = load_model("models:/main.ml.users_spend@champion") # the Trainer back +model = mlflow.pyfunc.load_model("models:/main.ml.users_spend@champion") +scores = model.predict(pd.DataFrame({"ENTITY": [1, 2, 3]}), params={ + "snapshot": "s3://bucket/kumo/snapshots//", + "prediction_time": "2024-06-01", # temporal queries + "binary_classification_threshold": 0.5, # binary classification only +}) +``` + +- `fit` waits for the upload and registration before it returns, for up to an hour, with a 600 s MLflow HTTP timeout during the upload (unless `MLFLOW_HTTP_REQUEST_TIMEOUT` is set). If they fail, the tracking thread stops or the hour passes, `fit` warns once (`Could not log the model: ...` or `Could not register the model: ...`, with the full error) and still returns its result with `registered_model=None` (with `on_error="raise"`, it raises instead), the run ends `FAILED` with the error in `kumo.error`, and a model that finishes uploading later is not registered. +- `MlflowTracking(log_model=True)` logs the model without registering it. `log_model(trainer, snapshot=...)` logs a fitted or loaded model in any run, with the bundle's source tables as run inputs. +- The model's metadata, the registered version's tags and `result.registered_model.metadata` hold `kumo_bundle_id`. The model's `provenance` artifact is the bundle's `_provenance.json`. Its pip requirement is `kumo-engine[s3,gcs]` of the same version, plus `[text]` when the model encodes a text column. A model trained with a custom embedding model also needs that model's package in the serving environment. +- The pyfunc model loads a snapshot on its first `predict` with that URI and keeps it for the life of the process: each serving replica, and each Spark executor of `mlflow.pyfunc.spark_udf`, downloads the snapshot once and keeps every snapshot it has loaded (no eviction), so its disk must hold them. With `serving_snapshot=` (on `log_model` or `MlflowTracking`), each replica downloads that snapshot when it starts, and `predict` uses it when no `snapshot` param is given. `params={"population": True}` scores every entity and returns the whole frame in memory; for large populations use `Trainer.predict` with a `PredictionOutputConfig`. Each `predict` call works in its own scratch directory and removes it before returning. +- Input rows have an `ENTITY` column and, optionally, a `TIMESTAMP` column. Scores come back in the input's row order, one per input row (repeated rows included); rows whose entity is not in the snapshot are left out. A temporal query scores entities at a time: pass `params={"prediction_time": ...}` or a `TIMESTAMP` column; without either, `predict` fails with `temporal query: pass params={'prediction_time': ...} or a TIMESTAMP column`. Binary classification needs `params={"binary_classification_threshold": ...}` unless the model's `scoring_defaults` set one. +- Serving needs an image with kumo-engine already installed; MLflow's stock serving containers cannot install it. Serve with `mlflow models serve -m --env-manager local` (on SageMaker endpoints, `-m /opt/ml/model -p 8080`). SageMaker batch transform takes JSON input (`dataframe_split` plus `params`), not CSV, which carries no params. +- GPU real-time endpoints on SageMaker need an inference image whose NVIDIA driver supports CUDA 12.6. The default GPU hosting image ships driver 470, so set a newer `InferenceAmiVersion` on the production variant, for example `al2023-ami-sagemaker-inference-gpu-4-1`. Not yet tested live. +- The `mlflow` extra installs `mlflow-skinny` (3.0.1 or later), the client without the MLflow server's dependencies. SageMaker managed MLflow servers run MLflow 3.0; a 3.17 client against a 3.0 server handles logging, registration, aliases and `load_model`. +- On Databricks, use Databricks Runtime 17.3 LTS ML or later: its torch already satisfies kumo-core, so installing kumo-engine keeps the runtime's torch, torchvision and MLflow. kumo-engine accepts pyarrow 16 through 24, so a runtime's pyarrow in that range (19 on 17.3 LTS ML) is kept too; a custom CA bundle for S3 needs pyarrow 21 or later. On 16.4 LTS ML pip replaces torch 2.6 and breaks the preinstalled torchvision. +- On Databricks, the parent folder of an MLflow experiment path must exist. Registering to Unity Catalog from outside Databricks also needs `boto3`, which the `mlflow` extra does not install. + +#### Which registry + +Pick one per registered model, since the two store different formats. On Databricks both work; use the MLflow model unless only kumo-engine jobs load it, since it also serves through `mlflow.pyfunc` (Model Serving, `spark_udf`). Both record only the `kumo_bundle_id` pointer; lineage stays in the bundle's `_provenance.json`. Both accept the same aliases (`kumoapi.plugin.check_model_alias`). + +| | Connector registry | MLflow model | +|---|---|---| +| Register | `kumo.MLOps(registry=ModelOutputConfig(connector, location=ModelRegistryLocation(name, alias=...)))` | `kumo.MLOps(tracking=MlflowTracking(), registry=ModelRegistryLocation(name, alias=...))` | +| Load | `Trainer.load_from_connector(connector, location=...)` | `mlflow.pyfunc.load_model(uri)` or `load_model(uri)` | +| Promote | the registry's own API (see [Promote a Kumo model](#promote-a-kumo-model)) | MLflow's `set_registered_model_alias` | +| Databricks | `DatabricksConnector` (Unity Catalog) | Unity Catalog through MLflow | +| SageMaker | not available | SageMaker managed MLflow (`MlflowTracking(target=SageMakerMLflow(arn), ...)`) | +| Snowflake | `SnowflakeConnector` (Snowflake Model Registry) | not available | +| Plain VM | not available (store under the artifact root instead) | any MLflow tracking server | + +#### SageMaker managed MLflow + +1. Create a tracking server once and reuse it; creation takes about 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 S3 prefix). +2. Give the identity that runs the job (notebook, VM role or pipeline execution role) `sagemaker-mlflow:*` on the tracking server ARN, and S3 read and write on the artifact store prefix and on the kumo artifact root. +3. Install `pip install "kumo-engine[mlflow,sagemaker]"` (the MLflow client, S3 and the `sagemaker-mlflow` plugin) and pass `MlflowTracking(target=SageMakerMLflow("arn:aws:sagemaker:::mlflow-tracking-server/"), experiment_name=...)`; the registry is the same server. Without an identity, the job's AWS identity (a role or `aws configure`) signs; `MLFLOW_TRACKING_URI` need not be set. The plugin signs and routes with the process-wide tracking URI, so each MLflow call kumo-engine makes sets it to the ARN and restores it afterwards: other threads that call MLflow during a fit see the SageMaker URI. +4. Model names use letters, digits and hyphens only (`users-churn`); `kumo.MLOps` checks this when it is built. +5. The run links MLflow prints 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. + +On S3 the artifact root, source tables and prediction outputs are plain `s3://` locations read and written with the job's AWS identity (instance profile, `AWS_PROFILE` or the job role): + +```python +from kumo.engine import PredictionOutputConfig, S3Connector + +data = S3Connector(root_dir="s3:///data") # data["orders.parquet"] reads s3:///data/orders.parquet; data["orders"] a directory of files +cfg = kumo.MLOps("s3:///kumo", + tracking=MlflowTracking(target=SageMakerMLflow(""), experiment_name="users-churn"), + registry=ModelRegistryLocation("users-churn", alias="candidate")) +output = PredictionOutputConfig(S3Connector(), output_paths={"predictions": "s3:///scores/"}) +store = S3Connector() # monitoring: load with absolute s3:// paths +``` + +### Running from a laptop + +`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 uses +the platform's ambient chain (a notebook, job or instance role, `aws +configure`, `databricks auth login`); one with credentials uses exactly +those. + +```python +import kumo.engine as kumo +from kumo.engine import ModelRegistryLocation +from kumo.engine.connector import DatabricksConnector, DatabricksIdentity, S3Identity, SnowflakeConnector, SnowflakeIdentity +from kumo.engine.integrations.mlflow import MlflowTracking, SageMakerMLflow +from kumo.engine.integrations.snowflake import SnowflakeTracking + +# Databricks MLflow and Unity Catalog, as a service principal (kumo-engine[mlflow,databricks]). +workspace = DatabricksConnector(catalog="main", schema_name="ml", warehouse_id="", + identity=DatabricksIdentity(host="https://", client_id="", client_secret="")) +cfg = kumo.MLOps("s3:///kumo", + tracking=MlflowTracking(target=workspace, experiment_name="/Users//users_churn"), + registry=ModelRegistryLocation("main.ml.users_churn", alias="candidate")) + +# SageMaker managed MLflow, with an access key or a role to assume (kumo-engine[mlflow,sagemaker]). +cfg = kumo.MLOps("s3:///kumo", + tracking=MlflowTracking(target=SageMakerMLflow("", identity=S3Identity(role_arn="")), + experiment_name="users-churn"), + registry=ModelRegistryLocation("users-churn", alias="candidate")) + +# Snowflake Experiment Tracking, as a key-pair user (kumo-engine[snowflake]). +snowflake = SnowflakeConnector(warehouse="", database="ML", schema_name="RUNS", + identity=SnowflakeIdentity(account="", username="", private_key="")) +cfg = kumo.MLOps("s3:///kumo", tracking=SnowflakeTracking(connector=snowflake, experiment_name="USERS_CHURN")) +``` + +- Each target is checked when `kumo.MLOps` is built, so wrong credentials or a missing catalog, schema or experiment folder parent fail before the fit starts. +- MLflow has no per-client credentials: its Databricks store reads `DATABRICKS_*` variables or a profile, and the SageMaker plugin signs with the process-wide tracking URI. So with explicit credentials, each MLflow call kumo-engine makes sets `DATABRICKS_HOST`, a short-lived `DATABRICKS_TOKEN` (never the client secret) and the process-wide tracking URI `databricks` (Unity Catalog registration downloads the logged model through it), or the process-wide tracking URI and the AWS session, and restores them when the call returns. With explicit credentials, don't call MLflow or Databricks yourself on other threads during a fit. Ambient identities change no credentials; a SageMaker target still sets the process-wide tracking URI for each call, so with a SageMaker target, other threads calling MLflow during a fit see the SageMaker URI. +- SageMaker model artifacts are uploaded and downloaded with the identity's credentials. `load_model(uri, tracking=MlflowTracking(target=..., experiment_name=...))` and `log_model(..., tracking=...)` take the same target. +- `examples/mlops/smoke.py databricks|sagemaker|snowflake` is a manual live check per platform: a small fit recorded and registered with an explicit identity, then loaded back, in a few minutes on a CPU. It is not run in CI. diff --git a/packages/kumo-engine/kumo/engine/__init__.py b/packages/kumo-engine/kumo/engine/__init__.py index db8c0743e..b9e467763 100644 --- a/packages/kumo-engine/kumo/engine/__init__.py +++ b/packages/kumo-engine/kumo/engine/__init__.py @@ -57,6 +57,7 @@ ModelOutputConfig, PredictionOutputConfig, PredictionResult, + ScoringDefaults, SnapshotOutputConfig, Trainer, TrainingResult, @@ -101,6 +102,7 @@ "Trainer", "TrainingResult", "PredictionResult", + "ScoringDefaults", "ModelOutputConfig", "PredictionOutputConfig", "TrainingTableOutputConfig", diff --git a/packages/kumo-engine/kumo/engine/_events.py b/packages/kumo-engine/kumo/engine/_events.py new file mode 100644 index 000000000..7309aab19 --- /dev/null +++ b/packages/kumo-engine/kumo/engine/_events.py @@ -0,0 +1,86 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +r"""Progress events reported by `Trainer.fit`.""" + +import hashlib +import logging +import math +import time +from dataclasses import asdict, dataclass, field +from typing import ( + Any, + Callable, + Dict, + Literal, + Mapping, + Optional, + Sequence, + cast, +) + +from kumo_core.instrumentation import NoOpTrainingInstrumentation + +logger = logging.getLogger(__name__) + +EventName = Literal['fit_started', 'trial_started', 'epoch_ended', + 'trial_ended', 'fit_ended'] + + +@dataclass(frozen=True) +class Event: + r"""One progress event. ``trial`` and ``epoch`` count from 1.""" + + name: EventName + run_id: str + time: float + trial: Optional[int] = None + epoch: Optional[int] = None + metrics: Mapping[str, float] = field(default_factory=dict) + params: Mapping[str, Any] = field(default_factory=dict) + + def to_dict(self) -> Dict[str, Any]: + return asdict(self) + + +def query_sha256(query: str) -> str: + return hashlib.sha256(query.encode()).hexdigest() + + +def _finite_floats(metrics: Mapping[str, Any]) -> Dict[str, float]: + result = {} + for key, value in metrics.items(): + try: + number = float(value) + except (TypeError, ValueError): + continue + if math.isfinite(number): + result[key] = number + return result + + +class EventEmitter(NoOpTrainingInstrumentation): + def __init__(self, run_id: str, + callbacks: Sequence[Callable[[Event], None]]) -> None: + self.run_id = run_id + self._callbacks = list(callbacks) + + def emit(self, name: EventName, *, trial: Optional[int] = None, + epoch: Optional[int] = None, + metrics: Optional[Mapping[str, float]] = None, + params: Optional[Mapping[str, Any]] = None) -> None: + event = Event(name=name, run_id=self.run_id, time=time.time(), + trial=trial, epoch=epoch, + metrics=_finite_floats(metrics or {}), + params=dict(params or {})) + for callback in list(self._callbacks): + try: + callback(event) + except Exception: + logger.warning( + 'Progress callback %r failed; it is disabled for the ' + 'rest of run %s', callback, self.run_id, exc_info=True) + self._callbacks.remove(callback) + + def event(self, name: str, **fields: Any) -> None: + self.emit(cast(EventName, name), **fields) diff --git a/packages/kumo-engine/kumo/engine/_manifest.py b/packages/kumo-engine/kumo/engine/_manifest.py index a6e8a871d..cbfb2bbba 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,6 +243,11 @@ class FitTiming(Record): seconds: float +@dataclass(frozen=True, kw_only=True) +class TrackingRef(Record): + run_id: str + + @dataclass(frozen=True, kw_only=True) class ModelManifest(Manifest): """How a model bundle was fitted: from which training table (embedded, @@ -252,6 +258,9 @@ class ModelManifest(Manifest): parents: Optional[ModelParents] = None fit: Optional[FitTiming] = None fit_mode: Optional[str] = None + tracking: Optional[TrackingRef] = omit_none() + + EMBED_OMITS = ('tracking', ) 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 120b71792..054577c90 100644 --- a/packages/kumo-engine/kumo/engine/_mlops.py +++ b/packages/kumo-engine/kumo/engine/_mlops.py @@ -7,8 +7,20 @@ import contextvars import os import posixpath +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 @@ -31,24 +43,38 @@ 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: where fits register models: a ``ModelOutputConfig`` with a + registry ``location``, or a ``ModelRegistryLocation`` registered + through the MLflow tracker. 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) @@ -60,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( @@ -83,6 +112,28 @@ def __post_init__(self) -> None: f'{list(KINDS)}.') if kinds: _require_writable(root) + from kumo.engine.trainer.output_config import ( + ModelOutputConfig, + check_model_location, + ) + + if isinstance(self.registry, ModelOutputConfig): + check_model_location(self.registry) + if isinstance(self.registry, ModelRegistryLocation): + if not callable(getattr(self.tracking, 'log_model', None)): + raise ValueError( + 'registry=ModelRegistryLocation(...) registers through ' + 'the MLflow tracker: pass tracking=MlflowTracking(...), ' + 'or registry=ModelOutputConfig(connector, location=...).') + if self.registry.version is not None: + raise ValueError('MLflow assigns registered model versions; ' + 'leave ModelRegistryLocation.version unset.') + if self.registry.alias is not None: + check_model_alias(self.registry.alias) + check = getattr(self.tracking, 'check', None) + if callable(check): + check(self.registry if isinstance( + self.registry, ModelRegistryLocation) else None) object.__setattr__(self, 'root', root) object.__setattr__(self, 'kinds', kinds) @@ -152,7 +203,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) @@ -168,6 +220,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 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 tracking or registration ' + f'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.""" 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/_tracking.py b/packages/kumo-engine/kumo/engine/_tracking.py new file mode 100644 index 000000000..4160e97aa --- /dev/null +++ b/packages/kumo-engine/kumo/engine/_tracking.py @@ -0,0 +1,482 @@ +# 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.artifacts import PROVENANCE_FILE +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.graph import GraphSnapshot + from kumo.engine.trainer.trainer import TrainingResult + +logger = logging.getLogger(__name__) + +DROPPED_PARAMS = 'kumo.dropped_params' +DROPPED_METRICS = 'kumo.dropped_metrics' +# Tracker calls run on one background thread so a slow tracking server +# never stalls training; these bound how long the fit waits for it. +QUEUE_SIZE = 1000 +FINISH_SECONDS = 30.0 +MODEL_SECONDS = 3600.0 + + +@dataclass(frozen=True) +class FitContext: + r"""What a fit trains, known once its model plan is resolved.""" + + run_id: str + query: str + model_plan: Mapping[str, Any] + inputs: Mapping[str, Mapping[str, Optional[str]]] = field( + default_factory=dict) + + +@dataclass(frozen=True) +class RunRecord: + r"""One tracked run; ``steps`` are per-epoch ``(epoch, value)`` + histories, ``artifacts`` map file names to text and ``inputs`` map + source table names to their ``connector``, ``location``, ``version`` and + ``version_kind``.""" + + name: Optional[str] = None + params: Mapping[str, str] = field(default_factory=dict) + metrics: Mapping[str, float] = field(default_factory=dict) + steps: Mapping[str, Tuple[Tuple[int, float], + ...]] = field(default_factory=dict) + tags: Mapping[str, str] = field(default_factory=dict) + artifacts: Mapping[str, str] = field(default_factory=dict) + inputs: Mapping[str, Mapping[str, Optional[str]]] = field( + default_factory=dict) + + +INPUT_FIELDS = ('connector', 'location', 'version', 'version_kind') + + +def fit_inputs(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(context: FitContext, result: Optional['TrainingResult'] = None, + fit_ended: Optional[Event] = None, *, + tags: Optional[Mapping[str, str]] = None) -> RunRecord: + r"""The fit run's record; ``tags`` are the user's extra tags.""" + metrics: Dict[str, float] = {} + pointers: Dict[str, str] = {} + artifacts = { + 'query.pql': context.query, + 'model_plan.json': json.dumps(dict(context.model_plan), indent=2, + default=str), + } + if result is not None: + metrics.update(result.metrics) + if fit_ended is not None and 'seconds' in fit_ended.metrics: + metrics['seconds'] = fit_ended.metrics['seconds'] + if result.id: + pointers['kumo.bundle_id'] = str(result.id) + provenance = Path(result.bundle_dir) / PROVENANCE_FILE + if provenance.is_file(): + artifacts[PROVENANCE_FILE] = provenance.read_text( + encoding='utf-8') + params = { + key: _param_value(value) + for key, value in _flatten(context.model_plan) + } + return RunRecord( + params=params, + metrics=_finite_floats(metrics), + tags={ + **_user_tags(tags, params), + 'kumo.kind': 'fit', + 'kumo.run_id': context.run_id, + **pointers, + **_error_tag(fit_ended), + }, + artifacts=artifacts, + inputs=dict(context.inputs), + ) + + +def _user_tags(tags: Optional[Mapping[str, Any]], + params: Mapping[str, str]) -> Dict[str, str]: + dropped = sorted(key for key in tags or {} + if key.startswith('kumo.') or key in params) + if dropped: + warnings.warn( + f'Tracking tags {dropped} are reserved and not recorded.', + UserWarning, stacklevel=3) + return { + key: str(value) + for key, value in (tags or {}).items() if key not in dropped + } + + +def _error_tag(ended: Optional[Event]) -> Dict[str, str]: + if ended is None or 'error' not in ended.params: + return {} + return {'kumo.error': str(ended.params['error'])} + + +def trial_record(context: FitContext, events: Sequence[Event]) -> RunRecord: + r"""One trial attempt's record from its events so far.""" + started = next(event for event in events if event.name == 'trial_started') + trial, attempt = started.trial, int(started.params['attempt']) + params: Dict[str, str] = { + key: _param_value(value) + for key, value in started.params.items() + } + metrics: Dict[str, float] = {} + steps: Dict[str, List[Tuple[int, float]]] = {} + ended: Dict[str, str] = {} + for event in events: + if event.name == 'epoch_ended' and event.epoch is not None: + for name, value in event.metrics.items(): + steps.setdefault(name, []).append((event.epoch, value)) + elif event.name == 'trial_ended': + ended = { + 'kumo.status': str(event.params.get('status')), + **_error_tag(event), + } + metrics.update({ + name if name == 'seconds' else f'best_{name}': value + for name, value in event.metrics.items() + }) + return RunRecord( + name=f'trial-{trial}-attempt-{attempt}', + params=params, + metrics=metrics, + steps={ + name: tuple(history) + for name, history in steps.items() + }, + tags={ + 'kumo.kind': 'trial', + 'kumo.run_id': f'{context.run_id}-trial-{trial}-attempt-{attempt}', + 'kumo.parent_run': context.run_id, + **ended, + }, + ) + + +class Tracking(Protocol): + r"""Records one `Trainer.fit` in an experiment tracker. + + Calls come from one background thread. `start` gets the fit run's + record without metrics, `on_trial` a trial attempt's record so far with + the event that updated it, and `finish` the fit run's final record + (`None` if tracking stopped early, which should end the run as failed). + Optional members: ``prepare()``, called first on the fit's own thread; + ``tags`` (extra fit-run tags); ``log_model(result, 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``. + """ + 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._context: Optional[FitContext] = None + self._trials: Dict[int, List[Event]] = {} + self._fit_ended: Optional[Event] = None + self._jobs: 'queue.Queue[Optional[Callable[[], Any]]]' = queue.Queue( + QUEUE_SIZE) + self._worker: Optional[threading.Thread] = None + self._behind = False + self.error: Optional[BaseException] = None + + @property + def enabled(self) -> bool: + return self._tracker is not None + + def start(self, context: Callable[[], FitContext]) -> None: + if self._tracker is None: + return + try: + value = context() + except Exception 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._context = value + self._call(lambda tracker: getattr(tracker, 'prepare', _none)()) + self._worker = threading.Thread(target=self._work, + name='kumo-tracking', daemon=True) + self._worker.start() + self._submit(lambda: self._call(lambda tracker: tracker.start( + run_record(value, tags=getattr(tracker, 'tags', None))))) + + def on_event(self, event: Event) -> None: + if self._context is None: + return + if event.name == 'fit_ended': + self._fit_ended = event + if event.trial is None: + return + if event.name == 'trial_started': + self._trials[event.trial] = [] + events = self._trials.setdefault(event.trial, []) + if 'attempt' not in event.params and events: + event = dataclasses.replace(event, params={ + **event.params, 'attempt': events[0].params['attempt'] + }) + events.append(event) + record = trial_record(self._context, events) + self._submit(lambda: self._call( + lambda tracker: tracker.on_trial(record, event))) + + def log_model( + self, + result: 'TrainingResult', + snapshot: 'GraphSnapshot', + 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 hasattr( + self._tracker, 'log_model')): + 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')(result, 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, result: Optional['TrainingResult'], + error: Optional[BaseException]) -> None: + worker = self._worker + if self._context is None or worker is None: + return + try: + record: Optional[RunRecord] = run_record(self._context, result, + self._fit_ended) + if error is not None and 'kumo.error' not in record.tags: + record = dataclasses.replace(record, tags={ + **record.tags, 'kumo.error': format_error(error)}) + 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/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..da763012f --- /dev/null +++ b/packages/kumo-engine/kumo/engine/integrations/_mlflow_model.py @@ -0,0 +1,217 @@ +# 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.trainer.trainer import Trainer + + self._trainer = Trainer.load(context.artifacts['bundle']) + self._columns = list(context.model_config['columns']) + self._snapshots: Dict[str, GraphSnapshot] = {} + self._default = context.model_config.get('snapshot') + if self._default: + self._snapshot(self._default) + + def predict(self, context: Any, model_input: pd.DataFrame, + params: Optional[Dict[str, Any]] = None) -> pd.DataFrame: + params = params or {} + if len(model_input) == 0 and not params.get('population'): + return pd.DataFrame(columns=self._columns) + uri = params.get('snapshot') or self._default + if not uri: + raise ValueError("Pass params={'snapshot': }.") + return predict(self._trainer, self._snapshot(uri), model_input, + params) + + def _snapshot(self, uri: str) -> GraphSnapshot: + if uri not in self._snapshots: + connector, path = snapshot_location(uri) + self._snapshots[uri] = GraphSnapshot.load_from_connector( + connector, path=path) + return self._snapshots[uri] + + +def predict(trainer: Any, snapshot: GraphSnapshot, rows: pd.DataFrame, + params: Mapping[str, Any]) -> pd.DataFrame: + from kumo.engine._scratch import scoped_scratch_dir + from kumo.engine.connector.file import FileConnector + + time = params.get('prediction_time') + time = None if time is None else pd.Timestamp(time) + options = dict( + binary_classification_threshold=params.get( + 'binary_classification_threshold'), + num_classes_to_return=params.get('num_classes_to_return'), + ) + _require_inputs(trainer, snapshot, rows, params, 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, + params: Mapping[str, Any], + time: Optional[pd.Timestamp]) -> None: + from kumo_core.pquery import query_is_temporal + from kumo_core.types import TaskType + + from kumo.engine.pquery import PredictiveQuery + + untimed = len(rows) > 0 and time is None and 'TIMESTAMP' not in rows + unthresholded = ( + params.get('binary_classification_threshold') is None and + trainer.result.scoring_defaults.binary_classification_threshold is None) + if not (untimed or unthresholded): + return + parsed = PredictiveQuery._from_snapshot(snapshot, + trainer._training_query)._parsed + if untimed and query_is_temporal(parsed, snapshot._native_graph): + raise ValueError("temporal query: pass params={'prediction_time': " + "...} or a TIMESTAMP column") + if unthresholded and (TaskType(parsed.task_type) + is TaskType.BINARY_CLASSIFICATION): + raise ValueError("binary classification: pass " + "params={'binary_classification_threshold': ...}") + + +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..95348b466 --- /dev/null +++ b/packages/kumo-engine/kumo/engine/integrations/_mlflow_target.py @@ -0,0 +1,230 @@ +# 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 nothing, +except that a 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) -> None: + self.tracking_uri = tracking_uri + self.registry_uri = registry_uri + self._swap = swap + self._aws_session = aws_session + + @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, tracking_uri: Optional[str]) -> Target: + mlflow = import_provider('mlflow', 'mlflow') + if target is not None and tracking_uri is not None: + raise ValueError('Pass target or tracking_uri, not both.') + if target is None: + if tracking_uri is None: + return Target(mlflow.get_tracking_uri(), mlflow.get_registry_uri()) + return Target(tracking_uri, registry_for(tracking_uri)) + 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'or a SageMakerMLflow, not {type(target).__name__}.') + + +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): + return Target('databricks', 'databricks-uc') + client = import_provider('kumo_databricks.session', + 'databricks')._client_for(principal) + return Target('databricks', 'databricks-uc', + swap=lambda: _databricks_token(client)) + + +@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..0eee9e20b --- /dev/null +++ b/packages/kumo-engine/kumo/engine/integrations/mlflow.py @@ -0,0 +1,859 @@ +# 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, + 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, +) +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.graph import GraphSnapshot + from kumo.engine.trainer.trainer import Trainer, TrainingResult + +logger = logging.getLogger(__name__) + +_MAX_PARAM_KEY = 250 +_MAX_PARAM_VALUE = 6000 +# MLflow's portable name charset (its Windows rule, which omits ':'). +_PARAM_KEY = re.compile(r'[/\w.\- ]+') +# MLflow's defaults (7 retries, 120 s timeout) can hold a call for minutes. +_HTTP_LIMITS = { + 'MLFLOW_HTTP_REQUEST_MAX_RETRIES': '2', + 'MLFLOW_HTTP_REQUEST_TIMEOUT': '30', +} +# Model upload only, and only where the limits above apply. +_UPLOAD_HTTP = {'MLFLOW_HTTP_REQUEST_TIMEOUT': '600'} +_PATH_CONNECTORS = ('LOCAL_FILE', 'S3', 'GCS') +_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 with + ``target`` or ``tracking_uri``. + 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. + serving_snapshot: exported snapshot URI the served model loads at + startup (see `log_model`). + target: where to record, with the credentials that reach it: a + ``DatabricksConnector`` (Databricks MLflow and Unity Catalog, + as its identity) or a `SageMakerMLflow`. `None` uses MLflow's + own tracking and registry URIs, read once here. + tracking_uri: another MLflow server, instead of ``target``. + + 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 or instance role) change nothing, + except that a SageMaker target still sets the process-wide tracking URI + for each call, since the plugin signs and routes with it. + + ``run_id`` is the MLflow id of the latest fit run, once it has started. + """ + + 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, + serving_snapshot: Optional[str] = None, + *, + target: Any = None, + tracking_uri: Optional[str] = None, + ) -> None: + self._mlflow = import_provider('mlflow', 'mlflow') + self.target = resolve_target(target, tracking_uri) + if experiment_name is None and (target is not None + or tracking_uri is not None): + raise ValueError('MlflowTracking(target=...) and ' + 'MlflowTracking(tracking_uri=...) need an ' + 'experiment_name.') + self._default_experiment: Optional[str] = ( + None if experiment_name is not None 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.serving_snapshot = serving_snapshot + 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() + + 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.""" + 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) + try: + with _limited_http(owned), self.target.scope(): + client = self.target.client() + if self.experiment_name is not None: + client.get_experiment_by_name(self.experiment_name) + 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 + + 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 start(self, record: RunRecord) -> None: + 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, + result: 'TrainingResult', + 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, result, snapshot, + name='model', run_id=run_id, + experiment_id=self._experiment_id, + serving_snapshot=self.serving_snapshot) + finally: + if self._parent == run_id: + os.environ.update({name: _HTTP_LIMITS[name] for name in upload}) + if self._parent != run_id: + logger.warning('The fit run %s ended before its model was ' + 'logged; the model is not registered', run_id) + return None + self._model_id = info.model_id + if registry is None: + return None + try: + return _register(self.target, self._client, info, registry, + registry_pointers(result)) + 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: Optional[str] = None, + tracking: Optional[MlflowTracking] = None, +) -> Any: + r"""Log a trainer's bundle as an MLflow pyfunc model in ``run_id`` (or + the active run, else a new one) with its source tables as run inputs; + returns its `ModelInfo`. ``tracking`` says where, with what credentials + (its experiment holds a new run); `None` uses MLflow's own tracking + URI. Scores with ``predict(entities, params={'snapshot': uri})``; with + ``serving_snapshot`` (an exported snapshot URI), each serving replica + downloads that snapshot once at startup and scores against it when no + ``snapshot`` param is given.""" + 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 = tracking or MlflowTracking() + target = tracking.target + 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, snapshot, name=name, + run_id=run_id, experiment_id=experiment_id, + serving_snapshot=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, + *, + tracking: Optional[MlflowTracking] = None, +) -> 'Trainer': + r"""Restore the `Trainer` of a model `log_model` logged, from the + server ``tracking`` reaches (`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 = (tracking or MlflowTracking()).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) + 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, result: 'TrainingResult', + 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.""" + from kumo_core.artifacts import PROVENANCE_FILE + + from kumo.engine.integrations import _mlflow_model, _mlflow_model_code + from kumo.engine.trainer._spec import read_training_spec + from kumo.engine.trainer.trainer import registry_pointers + + mlflow = import_provider('mlflow', 'mlflow') + if snapshot is None: + raise ValueError('log_model needs the snapshot the model trained on') + spec = read_training_spec(result.bundle_dir) + signature, example = _mlflow_model.signature(snapshot, + spec.predictive_query) + provenance = result.bundle_dir / PROVENANCE_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=registry_pointers(result), + 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': signature.outputs.input_names(), + 'snapshot': serving_snapshot, + }, + artifacts={ + 'bundle': str(result.bundle_dir), + **({'provenance': str(provenance)} + if provenance.is_file() else {}), + }, + signature=signature, + input_example=example, + metadata=registry_pointers(result), + pip_requirements=[_requirement(result)], + ) + 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 _requirement(result: 'TrainingResult') -> 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 = result.artifact_paths.get('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: + tags = version.tags or 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..dfa463060 --- /dev/null +++ b/packages/kumo-engine/kumo/engine/integrations/snowflake.py @@ -0,0 +1,281 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +r"""Snowflake Experiment Tracking for `Trainer.fit` +(``pip install 'kumo-engine[snowflake]'``).""" + +import logging +import re +import tempfile +from contextlib import ExitStack +from pathlib import Path +from typing import Any, Dict, List, Mapping, Optional, Set, Tuple + +from kumo_core.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: + session: Snowpark session; needs CREATE EXPERIMENT on the schema, or + MODIFY on an existing experiment. + connector: a ``SnowflakeConnector`` whose identity opens the session + instead. + experiment_name: experiment for the fit run, created if missing. + database: `None` uses the session's current database. + schema: `None` uses the session's current schema, or ``PUBLIC`` when + `database` is given or the session has none. + run_name: name of the fit run. `None` lets Snowflake pick one; a + taken name gets a suffix from the kumo run id. + tags: extra tags for the fit run. + native_model_lineage: keep the run open as a context for the whole + fit, so models registered meanwhile are linked to it. + + ``run_id`` is the latest fit run's name (Snowflake runs are identified + by name), once it has started. + """ + + def __init__( + self, + session: Any = None, + *, + connector: Any = None, + experiment_name: str, + database: Optional[str] = None, + schema: Optional[str] = None, + run_name: Optional[str] = None, + tags: Optional[Mapping[str, str]] = None, + native_model_lineage: bool = False, + ) -> None: + self._module = import_provider('snowflake.ml.experiment', 'snowflake') + if (session is None) == (connector is None): + raise ValueError('SnowflakeTracking takes a session or a ' + 'connector.') + self.session = session if session is not None else connector.session() + self.experiment_name = experiment_name + self.database = database + self.schema = schema + self.run_name = run_name + self.tags = dict(tags or {}) + self.native_model_lineage = native_model_lineage + self.run_id: Optional[str] = None + self._reset() + + def _reset(self) -> None: + self._tracking: Any = None + self._run: Any = None + self._run_context = ExitStack() + self._trial_tags: Dict[Tuple[int, int], Set[str]] = {} + self._metrics: Set[str] = set() + self._params = 0 + self._start_tags: Set[str] = set() + self._artifacts: Set[str] = set() + self._dropped_metrics: Set[str] = set() + self._dropped_params = 0 + + def 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, + record.tags['kumo.run_id'])) + self.run_id = str(self._run.name) + if self.native_model_lineage: + self._run_context.enter_context(self._run) + self._start_tags = set(record.tags) + self._log_params({**record.params, **record.tags}) + self._log_artifacts(record.artifacts) + + def on_trial(self, record: RunRecord, event: Event) -> None: + if self._run is None: + return + prefix = f'{record.name}/' + key = (event.trial or 0, int(event.params['attempt'])) + if event.name == 'trial_started': + self._trial_tags[key] = set(record.tags) + self._log_params({prefix + k: v + for k, v in {**record.params, + **record.tags}.items()}) + elif event.name == 'epoch_ended': + self._log_metrics({prefix + k: v + for k, v in event.metrics.items()}, + _EPOCH_METRICS, step=event.epoch or 0) + elif event.name == 'trial_ended': + self._log_metrics({prefix + k: v + for k, v in record.metrics.items()}, + _SUMMARY_METRICS) + logged = self._trial_tags.pop(key, set()) + self._log_params({prefix + k: v + for k, v in record.tags.items() + if k not in logged}) + + def finish(self, record: Optional[RunRecord], + error: Optional[BaseException]) -> None: + status = 'FAILED' + try: + if self._run is not None and record is not None: + self._log_metrics(record.metrics, _MAX_METRICS) + self._log_params({k: v + for k, v in record.tags.items() + if k not in self._start_tags}, _MAX_PARAMS) + self._log_artifacts(record.artifacts) + self._log_dropped() + if error is None: + status = 'FINISHED' + finally: + try: + if self._run is not None: + self._tracking.end_run(status=status) + except Exception 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, + run_id: str) -> Optional[str]: + if self.run_name is None: + return None + # start_run(name) resumes, or raises on, any run matching name as a + # SQL LIKE pattern and cannot create an existing run, so pick a name + # that is neither. + existing = [ + str(row['name']) + for row in tracking._sql_client.show_runs_in_experiment( + experiment_name=experiment) + ] + short = run_id.rsplit('-', 1)[-1] + for name in (self.run_name, _suffixed(self.run_name, short[:8]), + _suffixed(self.run_name, short)): + if not _taken(name, existing): + break + else: + name = None + if name != self.run_name: + logger.warning( + 'Snowflake run %s already exists in experiment %s; recording ' + 'this fit as %s', self.run_name, self.experiment_name, + name or 'a generated run name') + return name + + def _log_artifacts(self, artifacts: Mapping[str, str]) -> None: + with tempfile.TemporaryDirectory() as directory: + for name, text in artifacts.items(): + if name in self._artifacts: + continue + path = Path(directory) / name + path.write_text(text, encoding='utf-8') + self._tracking.log_artifact(str(path)) + self._artifacts.add(name) + + def _log_dropped(self) -> None: + dropped = { + DROPPED_METRICS: len(self._dropped_metrics), + DROPPED_PARAMS: self._dropped_params, + } + dropped = {name: count for name, count in dropped.items() if count} + if dropped: + self._tracking.log_params(dropped) + + def _log_params(self, params: Mapping[str, str], + limit: int = _MAX_PARAMS - _FINISH_PARAMS) -> None: + kept = {} + for name, value in params.items(): + if '$$' in name or '$$' in value or self._params >= limit: + self._dropped_params += 1 + else: + kept[name] = value + self._params += 1 + if kept: + self._tracking.log_params(kept) + + def _log_metrics(self, metrics: Mapping[str, float], limit: int, + step: int = 0) -> None: + kept = {} + for name, value in metrics.items(): + if '$$' in name: + self._dropped_metrics.add(name) + continue + if name not in self._metrics: + if len(self._metrics) >= limit: + self._dropped_metrics.add(name) + continue + self._metrics.add(name) + kept[name] = value + if kept: + self._tracking.log_metrics(kept, step=step) + + +def _suffixed(name: str, suffix: str) -> str: + if len(name) > 1 and name[0] == name[-1] == '"': + return f'{name[:-1]}_{suffix}"' + return f'{name}_{suffix}' + + +def _taken(name: str, existing: List[str]) -> bool: + resolved = resolve_snowflake_identifier(name) + regex = ''.join('.*' if c == '%' else '.' if c == '_' else re.escape(c) + for c in name) + return any( + other == resolved + or re.fullmatch(regex, other, re.IGNORECASE | re.DOTALL) is not None + for other in existing) diff --git a/packages/kumo-engine/kumo/engine/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..dfd6f151e 100644 --- a/packages/kumo-engine/kumo/engine/trainer/_provenance.py +++ b/packages/kumo-engine/kumo/engine/trainer/_provenance.py @@ -12,6 +12,7 @@ FitTiming, ModelManifest, ModelParents, + TrackingRef, embedded, read_manifest, stamp, @@ -30,6 +31,7 @@ def write_provenance( warm_start: Optional[ModelManifest], started_at: datetime, fit_mode: Optional[str], + tracked_run_id: Optional[str] = None, ) -> None: from kumo_core.artifacts import PROVENANCE_FILE, add_bundle_file @@ -50,6 +52,8 @@ def write_provenance( seconds=round((finished_at - started_at).total_seconds(), 3), ), fit_mode=fit_mode, + tracking=None if tracked_run_id is None else TrackingRef( + run_id=tracked_run_id), **stamp(), )) add_bundle_file(bundle_dir, PROVENANCE_ROLE, PROVENANCE_FILE) @@ -67,8 +71,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 +89,26 @@ 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)), + **stamp(), + )) + add_bundle_file(bundle_dir, PROVENANCE_ROLE, PROVENANCE_FILE) 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..43fc02c20 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, @@ -54,13 +61,17 @@ require_embedding_models, require_trained_embedding_models, ) +from kumo.engine._events import Event, EventEmitter, query_sha256 from kumo.engine._manifest import ModelManifest +from kumo.engine._mlops import full_error, warn_failure from kumo.engine._scratch import new_run_id, trainer_scratch_dir +from kumo.engine._tracking import FitContext, 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._scoring import ScoringDefaults, read_scoring_defaults from kumo.engine.trainer._spec import ( HOLDOUT_META_ROLE, HOLDOUT_ROLE, @@ -73,6 +84,7 @@ from kumo.engine.trainer.output_config import ( ModelOutputConfig, PredictionOutputConfig, + check_model_location, ) logger = logging.getLogger(__name__) @@ -84,6 +96,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 +114,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) @@ -127,6 +142,44 @@ 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: ModelOutputConfig) -> PublishedModel: + """Publish this model's bundle, as :meth:`Trainer.export`.""" + _require_publisher(config) + return _publish(self, 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``.""" + from kumo_core.artifacts import bundle_task_type + + from kumo.engine.trainer._provenance import record_derivation + from kumo.engine.trainer._scoring import ( + normalized_scoring_defaults, + write_scoring_defaults, + ) + + 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: @@ -330,6 +383,70 @@ 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 _mlflow_registry(settings: Any) -> Optional[ModelRegistryLocation]: + registry = None if settings is None else settings.registry + return registry if isinstance(registry, ModelRegistryLocation) else None + + +def _connector_registry(settings: Any) -> Optional[ModelOutputConfig]: + registry = None if settings is None else settings.registry + return registry if isinstance(registry, ModelOutputConfig) else None + + +def _publish(training_result: TrainingResult, + export_config: ModelOutputConfig, + errors: Optional[List[BaseException]] = None, + warn: bool = True) -> PublishedModel: + spec = read_training_spec(training_result.bundle_dir) + metadata = { + name: str(getattr(spec, name)) + for name in ('spec_version', 'predictive_query', 'snapshot_id') + } + metadata.update(registry_pointers(training_result)) + location = export_config.location + alias = location.alias if location else None + if location and alias: + location = replace(location, alias=None) + published = export_config.connector.sink_backend.publish_model( + export_config.connector, + str(training_result.bundle_dir), + path=export_config.path, + location=location, + metadata=metadata, + metrics=training_result.metrics, + ) + if location and alias: + try: + _set_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 +536,19 @@ 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: ModelOutputConfig) -> PublishedModel: + """Publish the fitted or loaded model bundle.""" + training_result = self._fitted('export') + _require_publisher(config) + return _publish(training_result, config) + @_without_torch_scatter_hint def fit( self, @@ -428,6 +558,7 @@ def fit( export_config: Optional[ModelOutputConfig] = None, warm_start: Optional[Self] = None, mlops: Optional[Any] = None, + callbacks: Sequence[Callable[[Event], None]] = (), ) -> TrainingResult: r"""Train a model against `train_table`. @@ -457,16 +588,101 @@ 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( + run_id, + [*callbacks, session.on_event] if session.enabled else callbacks) + fit_start = time.monotonic() + try: + started = { + 'snapshot_id': snapshot.id, + 'query_sha256': query_sha256(train_table.query), + } + except Exception: + started = {} + events.emit('fit_started', params=started) + try: + result = self._fit(train_table, snapshot=snapshot, + export_config=export_config, + warm_start=warm_start, run_id=run_id, + events=events, settings=settings, + tracking=session) + except BaseException as error: + events.emit('fit_ended', params={ + 'status': 'failed', + 'error': format_error(error), + }) + session.finish(None, error) + raise + events.emit('fit_ended', metrics={ + **result.metrics, 'seconds': time.monotonic() - fit_start + }, params={'status': 'succeeded'}) + registered, error = session.log_model(result, snapshot, + _mlflow_registry(settings)) + if registered is not None: + result = replace(result, registered_model=registered) + self._training_result = result + if registry is not None: + result, failure = self._register(result, registry, warn) + error = error or failure + session.finish(result, error) + raise_on_error(settings, session.error or error, + 'Trainer.fit (its model is in trainer.result)') + return result + + def _register( + self, result: TrainingResult, registry: ModelOutputConfig, warn: bool + ) -> Tuple[TrainingResult, Optional[BaseException]]: + errors: List[BaseException] = [] + try: + registered = _publish(result, registry, errors, warn) + except Exception as error: + if warn: + warn_failure( + f"Could not register the model: {full_error(error)}", + error, 3) + return result, error + result = replace(result, registered_model=registered) + self._training_result = result + return result, (errors or [None])[0] + + def _fit( + self, + train_table: TrainingTable, + *, + snapshot: GraphSnapshot, + export_config: Optional[ModelOutputConfig], + warm_start: Optional[Self], + run_id: str, + events: EventEmitter, + settings: Optional[Any], + tracking: TrackingSession, + ) -> TrainingResult: from kumo_core.automl.runtime import ( inherit_warm_start_model_plan, validate_warm_start_task_type, ) - from kumo_core.io.json import load_json + from kumo_core.io.json import load_json, to_json_dict from kumo_core.training.orchestration import ( CallableTrainingOptions, ResolvedTrainingInput, @@ -481,20 +697,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) @@ -529,11 +735,11 @@ def fit( search=AutoMLSearchOptions(), ) if warm_start_artifacts is not None: - source_schema = json.loads( - Path(warm_start_artifacts['prediction_schema']).read_text()) + from kumo_core.artifacts import bundle_task_type + validate_warm_start_task_type( TaskType(workflow.query.task_type), - TaskType(source_schema['task_type']), + bundle_task_type(warm_start_artifacts), ) from kumo_core.pquery import ( query_is_temporal, @@ -605,6 +811,12 @@ 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: FitContext( + run_id=run_id, + 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 +829,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(), ) @@ -669,7 +881,8 @@ def fit( 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), + started_at=started_at, fit_mode=fit_mode, + tracked_run_id=run_id if tracking.enabled else None), settings=settings, ) @@ -718,21 +931,9 @@ def _finish_fit( f"Publishing the model bundle without _provenance.json: " f"{error}", UserWarning, stacklevel=3) if export_config is not None: - spec = read_training_spec(training_result.bundle_dir) - published = export_config.connector.sink_backend.publish_model( - export_config.connector, - str(training_result.bundle_dir), - path=export_config.path, - location=export_config.location, - metadata={ - name: str(getattr(spec, name)) - for name in ('spec_version', 'predictive_query', - 'snapshot_id') - }, - metrics=training_result.metrics, - ) - training_result = replace(training_result, - published_model=published) + training_result = replace( + training_result, + published_model=_publish(training_result, export_config)) if settings is not None and settings.wants('model'): _store_model(training_result, settings) @@ -786,6 +987,8 @@ def predict( Binary classification requires an explicit threshold. Multilabel classification defaults to 0.5. Both accept thresholds from 0 to 1; thresholds change predicted labels without changing class scores. + + Omitted arguments fall back to the bundle's :attr:`scoring_defaults`. """ from kumo_core.prediction_options import ( validate_num_classes_to_return, @@ -852,11 +1055,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, @@ -890,6 +1094,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 +1122,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 +1186,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, @@ -1077,10 +1286,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( @@ -1108,5 +1325,30 @@ def _load_bundle( 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 _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(result: TrainingResult) -> Dict[str, str]: + return {BUNDLE_ID_KEY: result.bundle_id} if result.bundle_id else {} diff --git a/packages/kumo-engine/pyproject.toml b/packages/kumo-engine/pyproject.toml index fb5741d0b..b051564ae 100644 --- a/packages/kumo-engine/pyproject.toml +++ b/packages/kumo-engine/pyproject.toml @@ -71,6 +71,17 @@ text=[ gcs=[ "kumo-gcs>=0.1.0.dev1", ] +# nvidia-ml-py adds GPU system metrics. +mlflow=[ + "mlflow-skinny>=3.0.1,<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_fit_events.py b/packages/kumo-engine/test/test_fit_events.py new file mode 100644 index 000000000..c9efff26c --- /dev/null +++ b/packages/kumo-engine/test/test_fit_events.py @@ -0,0 +1,154 @@ +# 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 len({event.run_id for event in events}) == 1 + assert set(events[0].params) == {'snapshot_id', 'query_sha256'} + assert [(e.trial, e.epoch) for e in events if e.name == 'epoch_ended' + ] == [(1, 1), (1, 2), (2, 1), (2, 2)] + assert any(name.startswith('val_') for name in events[2].metrics) + assert events[1].params['attempt'] == 1 + assert events[4].params == {'status': 'succeeded', 'attempt': 1} + assert {'val_mae', 'seconds'} <= set(events[4].metrics) + assert events[8].params == { + 'status': 'failed', + 'attempt': 1, + 'error': 'RuntimeError: searcher down' + } + assert events[-1].metrics == {**result.metrics, + 'seconds': events[-1].metrics['seconds']} + assert events[-1].params == {'status': 'succeeded'} + assert [json.loads(json.dumps(event.to_dict()))['name'] + for event in events] == names + + +def test_failed_fit_reports_fit_ended(table_and_snapshot, monkeypatch): + import kumo.engine.trainer.trainer as trainer_module + + def fail(*args, **kwargs): + raise ValueError('stop') + + monkeypatch.setattr(trainer_module, 'build_generated_training_task_dataset', + fail) + table, snapshot = table_and_snapshot + events = [] + + with pytest.raises(ValueError, match='stop'): + kumo.Trainer(ModelPlan()).fit(table, snapshot=snapshot, + callbacks=[events.append]) + + assert [event.name for event in events] == ['fit_started', 'fit_ended'] + assert events[-1].params == { + 'status': 'failed', + 'error': 'ValueError: stop' + } + + +def test_event_metrics_keep_finite_numbers(): + from kumo.engine._events import EventEmitter + + received = [] + EventEmitter('run', [received.append]).emit( + 'epoch_ended', metrics={'val_loss': 'n/a', 'val_auc': float('nan'), + 'val_mae': 0.5}) + + assert received[0].metrics == {'val_mae': 0.5} + + +def test_epoch_event_without_recorded_metrics_does_not_raise(tmp_path): + from datetime import datetime, timezone + from types import SimpleNamespace + + from kumo_core.callbacks.train_progress_logger import ( + AutoTrainingProgressLogger, + ) + + received = [] + logger = AutoTrainingProgressLogger( + str(tmp_path), 1, on_event=lambda *a, **k: received.append(a)) + progress = SimpleNamespace(num_epochs_completed=1, metrics=[], + estimated_train_time=0.0) + logger.set_current_trainer( + SimpleNamespace(training_progress_logger=SimpleNamespace( + progress=progress)), trial=1) + logger.progress.start_time = datetime.now(timezone.utc) + + logger.on_train_epoch_end(None, None) + + assert received == [] diff --git a/packages/kumo-engine/test/test_mlflow_model.py b/packages/kumo-engine/test/test_mlflow_model.py new file mode 100644 index 000000000..8937403b3 --- /dev/null +++ b/packages/kumo-engine/test/test_mlflow_model.py @@ -0,0 +1,516 @@ +# 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.run_id': parent.data.tags['kumo.run_id'], + '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 + + 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()._trainer.scratch_dir + assert list(scratch.iterdir()) == [] + assert set(tables.glob('*')) == before | {tables / 'other-call'} + + +def test_pyfunc_names_missing_temporal_and_threshold_inputs( + fitted, monkeypatch): + from dataclasses import replace + + from kumo.engine.integrations._mlflow_model import predict + + trainer, table, snapshot = fitted + rows = pd.DataFrame({'ENTITY': range(5)}) + time = str(pd.Timestamp(table.timestamps[0])) + 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") + with pytest.raises(ValueError, match=r"^binary classification: pass " + r"params=\{'binary_classification_threshold': " + r"\.\.\.\}$"): + predict(trainer, snapshot, rows, {'prediction_time': time}) + + monkeypatch.setattr( + trainer, '_training_result', + replace(trainer.result, scoring_defaults=kumo.ScoringDefaults( + binary_classification_threshold=.5))) + assert len(predict(trainer, snapshot, rows, {'prediction_time': time + })) == 5 + + +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 FitContext, TrackingSession + + monkeypatch.setattr(_tracking, 'FINISH_SECONDS', 0.01) + + class Tracker: + def __init__(self, fail): + self.fail, self.calls = fail, [] + + def start(self, record): + self.calls.append('start') + + def on_trial(self, record, event): + pass + + def log_model(self, result, snapshot, 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: FitContext('r', 'q', {})) + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter('always') + _, error = session.log_model(result, None) + session.finish(result, error) + session._worker.join(5) + assert tracker.calls[:2] == ['start', 'log_model'] + assert tracker.calls[2] == ('finish', error) + assert (error is not None) is fail + assert any('upload failed' in str(w.message) for w in caught) is fail + + +@pytest.mark.parametrize('stall', ['slow upload', 'dead thread']) +def test_model_logging_wait_is_bounded(monkeypatch, stall): + import threading + + from kumo.engine import _tracking + from kumo.engine._tracking import FitContext, TrackingSession + + monkeypatch.setattr(_tracking, 'MODEL_SECONDS', 0.5) + release = threading.Event() + + class Tracker: + calls = [] + + def start(self, record): + pass + + def on_trial(self, record, event): + pass + + def log_model(self, result, snapshot, registry): + release.wait(30) + + def finish(self, record, error): + self.calls.append((record, error)) + + tracker = Tracker() + session = TrackingSession(tracker) + session.start(lambda: FitContext('r', 'q', {})) + if stall == 'dead thread': + session._submit(lambda: (_ for _ in ()).throw(SystemExit())) + session._worker.join(5) + with pytest.warns(UserWarning, match='Could not log the model'): + _, error = session.log_model(SimpleNamespace(), None) + release.set() + session.finish(None, error) + + assert type(error) is (TimeoutError + if stall == 'slow upload' else RuntimeError) + assert tracker.calls == [(None, None)] + + +def test_model_upload_gets_a_longer_http_timeout_and_a_late_one_is_dropped( + tracking_uri, monkeypatch, caplog): + import os + + from test_tracking_contract import CONTEXT + + from kumo.engine._tracking import TrackingSession + from kumo.engine.integrations import mlflow as integration + + for name in integration._HTTP_LIMITS: + monkeypatch.delenv(name, raising=False) + tracker = integration.MlflowTracking(log_system_metrics=False) + location = kumo.ModelRegistryLocation('m') + seen = [] + + def log_model(target, client, result, snapshot, **kwargs): + seen.append(os.environ.get('MLFLOW_HTTP_REQUEST_TIMEOUT')) + if len(seen) == 2: + tracker.finish(None, None) + return SimpleNamespace(model_id=None, model_uri='runs:/x/model') + + monkeypatch.setattr(integration, '_log_model', log_model) + monkeypatch.setattr(integration, '_register', lambda *args: 'registered') + for _ in range(2): + session = TrackingSession(tracker) + session.start(lambda: CONTEXT) + registered, error = session.log_model( + SimpleNamespace(bundle_id=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( + SimpleNamespace(artifact_paths={'encoders': str(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 FitContext, TrackingSession + + class Tracker: + def start(self, record): + raise RuntimeError('server down') + + def log_model(self, result, snapshot, registry): + raise AssertionError('not called') + + def finish(self, record, error): + pass + + session = TrackingSession(Tracker()) + session.start(lambda: FitContext('r', '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(SimpleNamespace(), 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.') + + monkeypatch.setattr(MlflowClient, 'create_experiment', no_folder) + table, snapshot = table_and_snapshot + cfg = kumo.MLOps(str(tmp_path / 'root'), + tracking=MlflowTracking('/NoSuchFolder/churn', + log_system_metrics=False), + on_error='raise') + 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 + + cfg = dataclasses.replace(cfg, on_error='warn') + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter('always') + trainer.fit(table, snapshot=snapshot, mlops=cfg) + 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..c18a2df68 --- /dev/null +++ b/packages/kumo-engine/test/test_mlflow_target.py @@ -0,0 +1,290 @@ +# 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, None) + assert target.tracking_uri == tracking_uri + assert not target.explicit + assert resolve(None, 'databricks').registry_uri == 'databricks-uc' + assert resolve(None, 'databricks://prod').registry_uri == ( + 'databricks-uc://prod') + assert resolve(None, 'http://mlflow:5000').registry_uri == ( + 'http://mlflow:5000') + with pytest.raises(ValueError, match='not both'): + resolve(SageMakerMLflow(ARN), 'http://x') + + +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()), None) + before = dict(os.environ) + with target.scope(): + assert dict(os.environ) == before + assert mlflow.get_tracking_uri() == tracking_uri + assert not target.explicit + + for target in (resolve(SageMakerMLflow(ARN, S3Identity()), None), + resolve(SageMakerMLflow(ARN), None)): + 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')), None) + + 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') + + +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')), None) + + 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_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(tracking_uri='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) + + +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, 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') 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..b28f30b4c --- /dev/null +++ b/packages/kumo-engine/test/test_mlflow_tracking.py @@ -0,0 +1,404 @@ +# 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 + +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.run_id', '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'] == { + 'run_id': parent.data.tags['kumo.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, + tracking_uri=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)) + + +@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(CONTEXT)) + tracker.target.tracking_uri = uri + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter('always') + tracker.finish(run_record(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(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 3d3af18d9..ea9974e7d 100644 --- a/packages/kumo-engine/test/test_mlops.py +++ b/packages/kumo-engine/test/test_mlops.py @@ -32,6 +32,30 @@ 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_a_connector_registry_checks_the_name_when_mlops_is_built(tmp_path): + def check_model_location(connector, location): + if '-' in location.model_name: + raise ValueError('bad name') + + connector = SimpleNamespace(sink_backend=SimpleNamespace( + check_model_location=check_model_location)) + good = kumo.ModelRegistryLocation('USERS_CHURN', alias='candidate') + kumo.MLOps(artifact_root=str(tmp_path), + registry=kumo.ModelOutputConfig(connector, location=good)) + with pytest.raises(ValueError, match='bad name'): + kumo.MLOps(artifact_root=str(tmp_path), + registry=kumo.ModelOutputConfig( + connector, + location=kumo.ModelRegistryLocation('users-churn'))) + with pytest.raises(ValueError, match='model alias'): + kumo.MLOps(artifact_root=str(tmp_path), + registry=kumo.ModelOutputConfig( + connector, location=kumo.ModelRegistryLocation( + 'USERS_CHURN', alias='LATEST'))) def test_a_root_is_a_uri_or_a_connector_and_path(tmp_path): 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..6ccfdf08e --- /dev/null +++ b/packages/kumo-engine/test/test_scoring_defaults.py @@ -0,0 +1,139 @@ +# 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() + with pytest.raises(ValueError, match='must be specified'): + ignored.predict(table, snapshot=snapshot) diff --git a/packages/kumo-engine/test/test_semantic_carry_forward.py b/packages/kumo-engine/test/test_semantic_carry_forward.py index fc5788a9b..5968ef805 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}) @@ -136,7 +137,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 +165,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 +275,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 +711,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 +724,8 @@ def run_predict() -> None: same_trainer = Trainer(ModelPlan()) monkeypatch.setattr(type(same_trainer), 'scratch_dir', property(lambda _: tmp_path)) - same_trainer._training_result = SimpleNamespace(artifact_paths={}) + same_trainer._training_result = SimpleNamespace( + artifact_paths={}, scoring_defaults=ScoringDefaults()) same_trainer._snapshot_id = snapshot.id same_trainer._training_query = query same_trainer.predict(table, snapshot=snapshot, output_config=output, diff --git a/packages/kumo-engine/test/test_snowflake_connector.py b/packages/kumo-engine/test/test_snowflake_connector.py index 1e4fcc687..b2b9254c8 100644 --- a/packages/kumo-engine/test/test_snowflake_connector.py +++ b/packages/kumo-engine/test/test_snowflake_connector.py @@ -108,3 +108,18 @@ def test_a_missing_target_field_with_no_override_is_reported( with pytest.raises(ValueError, match="names no warehouse"): SnowflakeConnector.from_spcs_env() + + +def test_session_is_the_readers_session(monkeypatch): + from kumo_snowflake import session as sessions + + connector = SnowflakeConnector(warehouse='WH', database='DB', + schema_name='S') + calls = [] + monkeypatch.setattr( + sessions._sessions, 'get', + lambda principal, target, fresh=False: calls.append(target) or 'live') + monkeypatch.setattr(type(connector.source_backend), '_principal', + lambda self, target: 'principal') + assert connector.session() == 'live' + assert calls == [connector] diff --git a/packages/kumo-engine/test/test_snowflake_tracking.py b/packages/kumo-engine/test/test_snowflake_tracking.py new file mode 100644 index 000000000..3ca74f771 --- /dev/null +++ b/packages/kumo-engine/test/test_snowflake_tracking.py @@ -0,0 +1,250 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import inspect +import threading +import time +from unittest import mock + +import _snowflake_experiment_fake +import pytest +from _snowflake_experiment_fake import FakeExperimentTracking, FakeSession +from test_tracking_contract import ( + CONTEXT, + EVENTS, + RESULT, + join_tracking_threads, +) + +from kumo.engine._tracking import TrackingSession + +FAKED = ('__init__', 'set_experiment', 'start_run', 'end_run', 'log_params', + 'log_metrics', 'log_artifact') + + +@pytest.mark.parametrize('method', FAKED) +def test_fake_matches_installed_experiment_tracking(method): + experiment = pytest.importorskip('snowflake.ml.experiment') + real = inspect.signature(getattr(experiment.ExperimentTracking, method)) + fake = inspect.signature(getattr(FakeExperimentTracking, method)) + assert [(p.name, p.kind, p.default) for p in fake.parameters.values() + ] == [(p.name, p.kind, p.default) for p in real.parameters.values()] + + +def test_installed_singleton_exposes_what_the_adapter_reads(monkeypatch): + experiment = pytest.importorskip('snowflake.ml.experiment') + cls = experiment.ExperimentTracking + monkeypatch.setattr(cls, '_instance', None) + session = mock.MagicMock() + session.get_current_database.return_value = 'DB' + session.get_current_schema.return_value = 'PUBLIC' + + tracking = cls(session) + + assert cls(mock.MagicMock()) is tracking + assert tracking._run is None + assert tracking._session is session + assert 'experiment_name' in inspect.signature( + tracking._sql_client.show_runs_in_experiment).parameters + + +@pytest.fixture +def fake(monkeypatch): + return _snowflake_experiment_fake.install(monkeypatch) + + +def _tracking(session=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: CONTEXT) + for event in EVENTS: + session.on_event(event) + if during is not None: + caught_up = threading.Event() + session._submit(caught_up.set) + caught_up.wait(30) + during() + session.finish(result, error) + + +def test_native_model_lineage_holds_the_run_open_during_the_fit(fake): + lineage = [] + _replay( + _tracking(database='ml', schema='runs', run_name='fit', + native_model_lineage=True), + during=lambda: lineage.extend(fake._instance.lineage_stack)) + + assert lineage == [('ENTS', 'FIT')] + assert fake._instance.lineage_stack == [] + assert fake._instance.calls[1] == ('set_experiment', 'ents', 'ml', 'runs') + + +def test_snowflake_errors_warn_once_and_end_the_run_failed(fake, caplog): + session = FakeSession() + fake(session).fail_on = {'log_metrics'} + + _replay(_tracking(session)) + + tracking = fake._instance + assert caplog.text.count('off for the rest of this run') == 1 + run, = tracking.runs.values() + assert run['status'] == 'FAILED' + assert tracking._run is None + + +def test_failed_fit_ends_the_run_failed(fake): + _replay(_tracking(), None, ValueError('stop')) + + run, = fake._instance.runs.values() + assert run['status'] == 'FAILED' + + +@pytest.mark.parametrize('busy', ['active run', 'other session']) +def test_an_instance_in_use_is_left_alone(fake, busy, caplog): + session = FakeSession() + user = fake(session, warnings_as_errors=True) + user.set_experiment('mine') + if busy == 'active run': + user.start_run('mine_1') + else: + session = FakeSession() + calls = len(user.calls) + + _replay(_tracking(session)) + + assert user.calls[calls:] == [] + assert user.warnings_as_errors is True + assert caplog.text.count('not tracked in Snowflake') == 1 + + +def test_location_defaults_to_the_session_not_the_last_fit(fake): + session = FakeSession(database='DEV', schema='WORK') + _replay(_tracking(session, database='ml')) + _replay(_tracking(session)) + + assert [ + call[1:] for call in fake._instance.calls if call[0] == 'set_experiment' + ] == [('ents', 'ml', 'PUBLIC'), ('ents', 'DEV', 'WORK')] + assert not any(call[0].startswith('run.') for call in fake._instance.calls) + + +@pytest.mark.parametrize('taken', ['FIT', 'MYXRUN']) +def test_an_existing_run_is_never_resumed(fake, taken, caplog): + session = FakeSession() + user = fake(session) + user.set_experiment('ents') + user.start_run(taken) + if taken == 'FIT': + user.end_run() + else: + user._run = None # still running, in another process + run_name = 'fit' if taken == 'FIT' else 'my_run' + + tracking = _tracking(session, run_name=run_name) + _replay(tracking) + + new = f"{run_name}_{CONTEXT.run_id.rsplit('-', 1)[-1][:8]}".upper() + assert tracking.run_id == new + assert set(user.runs) == {taken, new} + assert user.runs[taken]['params'] == {} + assert user.runs[new]['status'] == 'FINISHED' + assert 'recording this fit as' in caplog.text + + +def test_metric_and_param_limits(fake, monkeypatch): + import kumo.engine.integrations.snowflake as snowflake + + monkeypatch.setattr(snowflake, '_EPOCH_METRICS', 2) + monkeypatch.setattr(snowflake, '_SUMMARY_METRICS', 3) + monkeypatch.setattr(snowflake, '_MAX_METRICS', 5) + result = type(RESULT)(metrics={**RESULT.metrics, 'r2': float('nan')}, + 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 + + session = _SchemaSession([('RUNS', )]) + tracking = _tracking(None, connector=mock.Mock(session=lambda: session)) + tracking.check() + assert tracking.session is session + 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='a session or a connector'): + _tracking(None) 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..2efe3cb35 --- /dev/null +++ b/packages/kumo-engine/test/test_tracking_contract.py @@ -0,0 +1,400 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Every tracking adapter records the same run records as a recorded fit.""" + +import json +import threading +import time +from dataclasses import dataclass, field +from pathlib import Path +from types import SimpleNamespace +from typing import Any, Callable, Dict, List, Mapping, Optional, Tuple + +import pytest +from kumo_core.artifacts import PROVENANCE_FILE + +from kumo.engine._tracking import ( + DROPPED_METRICS, + DROPPED_PARAMS, + FitContext, + TrackingSession, + run_record, +) +from kumo.engine.tracking import Event, RunRecord + +RUN = 'trainer-fit-1' +USER_TAGS = {'team': 'ents', 'kumo.run_id': 'overridden'} +QUERY = 'PREDICT ents.Y FOR EACH ents.ID' +INPUTS = { + 'ents': {'connector': 'S3', 'location': 's3://bucket/ents.parquet', + 'version': '1790000000', 'version_kind': 'timestamp'}, + 'orders': {'connector': 'DATABRICKS', + 'location': '`main`.`sales`.`orders`', + 'version': '7_1790000000', 'version_kind': 'exact'}, + 'users': {'connector': 'SNOWFLAKE', 'location': '"DB"."SCH"."USERS"', + 'version': '1790000000', 'version_kind': 'timestamp'}, +} +CONTEXT = FitContext( + run_id=RUN, query=QUERY, inputs=INPUTS, + model_plan={ + 'optimization': {'base_lr': [0.01, 0.001], 'max_epochs': 2}, + 'training_job': {'tune_metric': 'mae'}, + 'column_processing': { + 'encoder_overrides': {'ents.amount ($)': 'numerical'} + }, + }) + + +def _event(name, trial=None, epoch=None, metrics=None, **params): + return Event(name=name, run_id=RUN, time=0.0, trial=trial, epoch=epoch, + metrics=metrics or {}, params=params) + + +def _epochs(trial, *values): + return [ + _event('epoch_ended', trial, epoch, {'train_loss': loss, + 'val_mae': mae}) + for epoch, (loss, mae) in enumerate(values, 1) + ] + + +# Trial 2's first attempt fails; epoch events carry no attempt, as in core. +EVENTS = [ + _event('trial_started', 1, attempt=1, **{'optimization.base_lr': 0.01}), + *_epochs(1, (0.5, 0.4), (0.3, 0.35)), + _event('trial_ended', 1, metrics={'val_mae': 0.35, 'seconds': 3.0}, + status='succeeded', attempt=1), + _event('trial_started', 2, attempt=1, **{'optimization.base_lr': 0.001}), + _event('trial_ended', 2, status='failed', attempt=1, + error='OutOfMemoryError: CUDA out of memory'), + _event('trial_started', 2, attempt=2, **{'optimization.base_lr': 0.001, + 'optimization.batch_size': 256}), + *_epochs(2, (0.6, 0.5), (0.4, 0.3)), + _event('trial_ended', 2, metrics={'val_mae': 0.3, 'seconds': 4.0}, + status='succeeded', attempt=2), + _event('fit_ended', metrics={'val_mae': 0.3, 'test_mae': 0.32, + 'seconds': 9.0}, status='succeeded'), +] +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: CONTEXT) + for event in events: + session.on_event(event) + session.finish(result, error) + records = { + record.tags['kumo.run_id']: record + for hook, record in recorder.calls if hook == 'on_trial' + } + records[RUN] = run_record(CONTEXT, result, events[-1], tags=USER_TAGS) + return records + + +@dataclass +class Observed: + name: Optional[str] + params: Mapping[str, str] + metrics: Mapping[str, float] + steps: Mapping[str, Tuple[Tuple[int, float], ...]] + tags: Mapping[str, str] + parent: Optional[str] = None + artifacts: Mapping[str, str] = field(default_factory=dict) + + +@dataclass(frozen=True) +class Adapter: + make: Callable[[], Any] + observe: Callable[[], Dict[str, Observed]] + kept_params: Callable[[Mapping[str, str]], Mapping[str, str]] = dict + nested_runs: bool = False + step_metrics: bool = False + artifacts: bool = False + + +def gaps(expected: Dict[str, RunRecord], observed: Dict[str, Observed], + adapter: Adapter) -> List[str]: + kept_params = adapter.kept_params + found = [] + if set(observed) != set(expected): + return [f'runs {sorted(observed)} != {sorted(expected)}'] + dropped = sum(len(record.params) - len(kept_params(record.params)) + for record in expected.values()) + counted = sum(int(seen.tags.get(DROPPED_PARAMS, 0)) + for seen in observed.values()) + if counted != dropped: + found.append(f'{DROPPED_PARAMS}: {counted} != {dropped}') + for run_id, record in expected.items(): + seen = observed[run_id] + + def differ(what, have, want): + if have != want: + found.append(f'{run_id} {what}: {have!r} != {want!r}') + + if record.name is not None: + differ('name', seen.name, record.name) + differ('params', dict(seen.params), dict(kept_params(record.params))) + differ('metrics', {k: seen.metrics.get(k) for k in record.metrics}, + dict(record.metrics)) + names = set(record.metrics) + if adapter.step_metrics: + names |= set(record.steps) + differ('steps', {k: seen.steps.get(k) for k in record.steps}, + dict(record.steps)) + differ('metric names', sorted(seen.metrics), sorted(names)) + differ('tags', {k: seen.tags.get(k) for k in record.tags}, + dict(record.tags)) + differ( + 'kumo tags', + sorted(k for k in seen.tags if k.startswith('kumo.') + and k not in (DROPPED_PARAMS, DROPPED_METRICS)), + sorted(k for k in record.tags if k.startswith('kumo.'))) + if adapter.nested_runs: + differ('parent', seen.parent, record.tags.get('kumo.parent_run')) + if adapter.artifacts: + differ('artifacts', dict(seen.artifacts), dict(record.artifacts)) + return found + + +@pytest.fixture +def mlflow_adapter(tmp_path, monkeypatch): + mlflow = pytest.importorskip('mlflow') + from kumo.engine.integrations.mlflow import MlflowTracking, _mlflow_params + + monkeypatch.chdir(tmp_path) + previous = mlflow.get_tracking_uri() + mlflow.set_tracking_uri(f'sqlite:///{tmp_path / "mlflow.db"}') + + def observe(): + client = mlflow.MlflowClient() + runs = client.search_runs( + [e.experiment_id for e in client.search_experiments()]) + kumo_ids = {r.info.run_id: r.data.tags['kumo.run_id'] for r in runs} + + def history(run_id, key): + return tuple((m.step, m.value) for m in sorted( + client.get_metric_history(run_id, key), + key=lambda m: (m.step, m.timestamp))) + + return { + kumo_ids[run.info.run_id]: Observed( + name=run.info.run_name, params=run.data.params, + metrics=run.data.metrics, tags=run.data.tags, + steps={k: history(run.info.run_id, k) + for k in run.data.metrics}, + parent=kumo_ids.get(run.data.tags.get('mlflow.parentRunId')), + artifacts={ + a.path: mlflow.artifacts.load_text( + f'{run.info.artifact_uri}/{a.path}') + for a in client.list_artifacts(run.info.run_id) + }) + for run in runs + } + + try: + yield Adapter( + make=lambda: MlflowTracking(log_system_metrics=False, + tags=USER_TAGS), observe=observe, + kept_params=_mlflow_params, nested_runs=True, step_metrics=True, + artifacts=True) + finally: + while mlflow.active_run() is not None: + mlflow.end_run() + mlflow.set_tracking_uri(previous) + + +@pytest.fixture +def snowflake_adapter(monkeypatch): + import _snowflake_experiment_fake + + fake = _snowflake_experiment_fake.install(monkeypatch) + from kumo.engine.integrations.snowflake import SnowflakeTracking + + def observe(): + observed = {} + for run in fake._instance.runs.values(): + names: Dict[Optional[str], Dict[str, Dict[str, Any]]] = {} + + def bucket(key): + prefix, _, name = key.rpartition('/') + return names.setdefault(prefix or None, { + 'params': {}, 'tags': {}, 'metrics': {}, 'steps': {} + }), name + + for key, value in run['params'].items(): + runs, name = bucket(key) + # Snowflake has no tags: they are params of the same name. + if name.startswith('kumo.') or name in USER_TAGS: + runs['tags'][name] = value + else: + runs['params'][name] = value + for key, value, step in run['metrics']: + runs, name = bucket(key) + runs['metrics'][name] = value + runs['steps'][name] = (*runs['steps'].get(name, ()), + (step, value)) + for prefix, runs in names.items(): + observed[runs['tags']['kumo.run_id']] = Observed( + name=prefix, + artifacts=run['artifacts'] if prefix is None else {}, + **runs) + return observed + + return Adapter( + make=lambda: SnowflakeTracking(_snowflake_experiment_fake.FakeSession(), + experiment_name='ents', tags=USER_TAGS), + observe=observe, step_metrics=True, artifacts=True) + + +@pytest.fixture(params=['mlflow', 'snowflake']) +def adapter(request) -> Adapter: + return request.getfixturevalue(f'{request.param}_adapter') + + +def test_records_of_the_recorded_fit(): + records = replay() + + fit = records[RUN] + assert fit.tags == {'team': 'ents', 'kumo.kind': 'fit', 'kumo.run_id': RUN, + '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(CONTEXT, tags={'kumo.bundle_id': 'x', + 'training_job.tune_metric': 'y', + 'a': 1}) + assert spoofed.tags.keys() - fit.tags.keys() == {'a'} + assert fit.metrics == {'val_mae': 0.3, 'test_mae': 0.32, 'seconds': 9.0} + assert fit.artifacts['query.pql'] == QUERY + assert json.loads(fit.artifacts['model_plan.json']) == CONTEXT.model_plan + assert fit.inputs == INPUTS + + oom = records[f'{RUN}-trial-2-attempt-1'] + assert oom == RunRecord( + name='trial-2-attempt-1', + params={'optimization.base_lr': '0.001', 'attempt': '1'}, tags={ + 'kumo.kind': 'trial', + 'kumo.run_id': f'{RUN}-trial-2-attempt-1', + 'kumo.parent_run': RUN, + 'kumo.status': 'failed', + 'kumo.error': 'OutOfMemoryError: CUDA out of memory', + }) + retry = records[f'{RUN}-trial-2-attempt-2'] + assert retry.params['attempt'] == '2' + assert retry.metrics == {'best_val_mae': 0.3, 'seconds': 4.0} + assert retry.steps == {'train_loss': ((1, 0.6), (2, 0.4)), + 'val_mae': ((1, 0.5), (2, 0.3))} + + +def test_adapter_records_the_same_runs(adapter, tmp_path): + (tmp_path / PROVENANCE_FILE).write_text('{"bundle": "b"}') + result = SimpleNamespace(metrics=RESULT.metrics, id=BUNDLE, + bundle_dir=tmp_path) + expected = replay(adapter.make(), result=result) + + assert expected[RUN].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[RUN].tags['kumo.error'] == 'ValueError: stop' + assert gaps(expected, adapter.observe(), adapter) == [] + + +def test_only_rank_zero_tracks(monkeypatch): + monkeypatch.setenv('LOCAL_RANK', '1') + recorder = _Recorder() + session = TrackingSession(recorder) + + session.start(lambda: CONTEXT) + session.finish(RESULT, None) + + assert recorder.calls == [] + + +def test_a_hung_tracker_does_not_stall_the_fit(monkeypatch, caplog): + import kumo.engine._tracking as tracking + + monkeypatch.setattr(tracking, 'FINISH_SECONDS', 0.5) + release = threading.Event() + + class Hung(_Recorder): + def start(self, record): + super().start(record) + release.wait(30) + + hung = Hung() + began = time.monotonic() + try: + replay(hung) + elapsed = time.monotonic() - began + finally: + release.set() + join_tracking_threads() + + assert elapsed < 1.5 + assert 'did not finish within' in caplog.text diff --git a/packages/kumo-engine/test/test_trainer_registry.py b/packages/kumo-engine/test/test_trainer_registry.py index 7ea9ba950..e20b5f387 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,182 @@ 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 == {} + + +@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')) + registry = kumo.ModelOutputConfig( + _FakeConnector(sink), + location=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'), 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'), 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) + registry = kumo.ModelOutputConfig( + _FakeConnector(_AliasingSink()), + location=ModelRegistryLocation('CHURN_GNN', alias='LATEST')) + with pytest.raises(ValueError, match='model alias'): + Trainer(ModelPlan()).fit( + SimpleNamespace(query='q'), snapshot=SimpleNamespace(id='s'), + mlops=kumo.MLOps(str(tmp_path), registry=registry, 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..79523d2dd 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.0.1,<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" From 7a97314e44e471b17b1bffd57139a4b69e81efd8 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Wed, 28 Oct 2026 18:41:42 +0000 Subject: [PATCH 02/16] Track, store, register and serve any bundle through internal seams Tracking builds each run record from a neutral RunSpec, and Trainer.fit adapts its context and result into one. One _store_and_register stores, logs and registers a bundle from its directory and ModelManifest. The manifest lists the model's dependencies: the predictor entry point that the MLflow pyfunc loads, the kumo-engine package, and content-addressed assets such as text-encoder word vectors, stored once under /assets/// and resolved with a sha256 check when the model loads. Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/docs/mlops.md | 29 +- packages/kumo-engine/kumo/engine/_deps.py | 301 ++++++++++++++++++ .../kumo-engine/kumo/engine/_embeddings.py | 90 +++++- packages/kumo-engine/kumo/engine/_manifest.py | 35 +- .../kumo-engine/kumo/engine/_predictor.py | 65 ++++ packages/kumo-engine/kumo/engine/_tracking.py | 182 +++++------ .../kumo/engine/integrations/_mlflow_model.py | 8 +- .../kumo/engine/integrations/mlflow.py | 65 ++-- .../kumo/engine/trainer/_provenance.py | 6 +- .../kumo/engine/trainer/_run_spec.py | 56 ++++ .../kumo/engine/trainer/trainer.py | 177 ++++++---- .../test/test_bundle_provenance.py | 9 +- .../kumo-engine/test/test_mlflow_model.py | 42 +-- .../kumo-engine/test/test_mlflow_tracking.py | 7 +- .../test/test_model_agnostic_mlops.py | 269 ++++++++++++++++ .../test/test_snowflake_tracking.py | 8 +- .../test/test_tracking_contract.py | 26 +- 17 files changed, 1144 insertions(+), 231 deletions(-) create mode 100644 packages/kumo-engine/kumo/engine/_deps.py create mode 100644 packages/kumo-engine/kumo/engine/_predictor.py create mode 100644 packages/kumo-engine/kumo/engine/trainer/_run_spec.py create mode 100644 packages/kumo-engine/test/test_model_agnostic_mlops.py diff --git a/packages/kumo-engine/docs/mlops.md b/packages/kumo-engine/docs/mlops.md index 0a79573eb..061a9a253 100644 --- a/packages/kumo-engine/docs/mlops.md +++ b/packages/kumo-engine/docs/mlops.md @@ -51,6 +51,33 @@ With MLOps on, each bundle is stored under `/models//`. A compl 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. +### Artifacts, bundles and dependencies + +- An **artifact** is anything with its own id and `_provenance.json`: a snapshot, a training table, a model, a prediction run, or an asset. +- A model's **bundle** holds only the model's own state: weights, encoders, the training spec and the holdout. +- Its manifest's **`parents`** are lineage: the training table and the warm-start model it came from. +- Its **`dependencies`** are what `predict` needs besides the bundle and a snapshot. The snapshot is an input to each call, never a dependency. Each entry has a `kind`: + - `predictor`: the entry point that scores the bundle; + - `package`: the `kumo-engine` version and its extras, with `[text]` only when the model encodes a text column; + - `asset`: a content-addressed file set, by `asset_kind`, `name`, `version` and `sha256`. +- Bundles fitted before dependencies were recorded list none. They load as before, with the built-in predictor. +- **Assets** are not copied into bundles. With MLOps on, each asset is stored once under `/assets///` with its own `_provenance.json`, and every model that uses it shares that copy. The dependency records where it was stored. +- **Resolving assets:** `Trainer.load`, the MLflow pyfunc and scoring jobs look for each asset at the location the dependency records, then under the artifact root of the enclosing `kumo.mlops` block, then in the embedding-model store. They check its sha256 and cache non-local copies under `~/.cache/kumo/assets//`. If an asset is stored but cannot be found, loading fails with `MissingDependencyError`, which names the asset kind, name and sha256. A bundle without stored assets resolves its models from the store, as before. +- **Text encoders** are the first assets (`asset_kind="text_encoder"`): the word vectors of each text column's embedding model. Word-vector models (`GloVeVectors`) are stored. Other `TokenEmbeddingModel`s are recorded by fingerprint and must be registered wherever the model is loaded. + +To train with your own word vectors and serve the model without registering them again, register them under a name, fit with MLOps on, and load the stored model anywhere that can reach the artifact root: + +```python +store = kumo.EmbeddingModelStore() +store.add("colors", kumo.GloVeVectors("/data/vectors")) # a .zip release or a directory of converted files +plan.column_processing.encoder_overrides = {"users.bio": TextEmbedding("colors")} +with kumo.mlops(kumo.MLOps(artifact_root="s3://bucket/kumo")): + result = kumo.Trainer(plan, embedding_models=store).fit(table, snapshot=snapshot) +# Elsewhere, with no store: the vectors come from /assets/text_encoder//. +trainer = kumo.Trainer.load_from_connector(kumo.S3Connector(root_dir="s3://bucket/kumo"), + path=f"models/{result.id}") +``` + ### Scoring defaults A model can carry the prediction arguments `predict` uses when a call omits them: @@ -290,7 +317,7 @@ scores = model.predict(pd.DataFrame({"ENTITY": [1, 2, 3]}), params={ - `fit` waits for the upload and registration before it returns, for up to an hour, with a 600 s MLflow HTTP timeout during the upload (unless `MLFLOW_HTTP_REQUEST_TIMEOUT` is set). If they fail, the tracking thread stops or the hour passes, `fit` warns once (`Could not log the model: ...` or `Could not register the model: ...`, with the full error) and still returns its result with `registered_model=None` (with `on_error="raise"`, it raises instead), the run ends `FAILED` with the error in `kumo.error`, and a model that finishes uploading later is not registered. - `MlflowTracking(log_model=True)` logs the model without registering it. `log_model(trainer, snapshot=...)` logs a fitted or loaded model in any run, with the bundle's source tables as run inputs. -- The model's metadata, the registered version's tags and `result.registered_model.metadata` hold `kumo_bundle_id`. The model's `provenance` artifact is the bundle's `_provenance.json`. Its pip requirement is `kumo-engine[s3,gcs]` of the same version, plus `[text]` when the model encodes a text column. A model trained with a custom embedding model also needs that model's package in the serving environment. +- The model's metadata, the registered version's tags and `result.registered_model.metadata` hold `kumo_bundle_id`. The model's `provenance` artifact is the bundle's `_provenance.json`. Its pip requirement is `kumo-engine[s3,gcs]` of the same version, plus `[text]` when the model encodes a text column. The pyfunc scores with the predictor the bundle's manifest names. It resolves the model's assets, such as non-default word vectors, as described in [Artifacts, bundles and dependencies](#artifacts-bundles-and-dependencies). - The pyfunc model loads a snapshot on its first `predict` with that URI and keeps it for the life of the process: each serving replica, and each Spark executor of `mlflow.pyfunc.spark_udf`, downloads the snapshot once and keeps every snapshot it has loaded (no eviction), so its disk must hold them. With `serving_snapshot=` (on `log_model` or `MlflowTracking`), each replica downloads that snapshot when it starts, and `predict` uses it when no `snapshot` param is given. `params={"population": True}` scores every entity and returns the whole frame in memory; for large populations use `Trainer.predict` with a `PredictionOutputConfig`. Each `predict` call works in its own scratch directory and removes it before returning. - Input rows have an `ENTITY` column and, optionally, a `TIMESTAMP` column. Scores come back in the input's row order, one per input row (repeated rows included); rows whose entity is not in the snapshot are left out. A temporal query scores entities at a time: pass `params={"prediction_time": ...}` or a `TIMESTAMP` column; without either, `predict` fails with `temporal query: pass params={'prediction_time': ...} or a TIMESTAMP column`. Binary classification needs `params={"binary_classification_threshold": ...}` unless the model's `scoring_defaults` set one. - Serving needs an image with kumo-engine already installed; MLflow's stock serving containers cannot install it. Serve with `mlflow models serve -m --env-manager local` (on SageMaker endpoints, `-m /opt/ml/model -p 8080`). SageMaker batch transform takes JSON input (`dataframe_split` plus `params`), not CSV, which carries no params. diff --git a/packages/kumo-engine/kumo/engine/_deps.py b/packages/kumo-engine/kumo/engine/_deps.py new file mode 100644 index 000000000..189614856 --- /dev/null +++ b/packages/kumo-engine/kumo/engine/_deps.py @@ -0,0 +1,301 @@ +# 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 +import warnings +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._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'): + location = _store(asset, sha256, settings) + 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, or None (with a + warning) if it failed.""" + 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) + filesystem.create_dir(root, recursive=True) + 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: + warnings.warn( + f"Could not store {asset.asset_kind} {asset.name!r} under the " + f"artifact root: {error}", UserWarning, stacklevel=4) + return None + + +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..1a51a42c8 100644 --- a/packages/kumo-engine/kumo/engine/_embeddings.py +++ b/packages/kumo-engine/kumo/engine/_embeddings.py @@ -4,7 +4,8 @@ import logging from contextlib import contextmanager -from typing import Any, Iterator, Optional +from pathlib import Path +from typing import Any, Dict, Iterator, List, Optional from kumo_core import embedding_models from kumo_core.embedding_models import EmbeddingModelStore @@ -48,3 +49,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/_manifest.py b/packages/kumo-engine/kumo/engine/_manifest.py index cbfb2bbba..e5aca170b 100644 --- a/packages/kumo-engine/kumo/engine/_manifest.py +++ b/packages/kumo-engine/kumo/engine/_manifest.py @@ -248,21 +248,54 @@ class TrackingRef(Record): run_id: str +@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: return (json.dumps(manifest.to_dict(), indent=2, sort_keys=True) + '\n').encode('utf-8') 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/_tracking.py b/packages/kumo-engine/kumo/engine/_tracking.py index 4160e97aa..562059341 100644 --- a/packages/kumo-engine/kumo/engine/_tracking.py +++ b/packages/kumo-engine/kumo/engine/_tracking.py @@ -27,7 +27,6 @@ Tuple, ) -from kumo_core.artifacts import PROVENANCE_FILE from kumo_core.instrumentation import format_error from kumo.engine._events import Event, _finite_floats @@ -36,8 +35,7 @@ if TYPE_CHECKING: from kumoapi.plugin import ModelRegistryLocation, PublishedModel - from kumo.engine.graph import GraphSnapshot - from kumo.engine.trainer.trainer import TrainingResult + from kumo.engine._manifest import ModelManifest logger = logging.getLogger(__name__) @@ -50,17 +48,6 @@ MODEL_SECONDS = 3600.0 -@dataclass(frozen=True) -class FitContext: - r"""What a fit trains, known once its model plan is resolved.""" - - run_id: str - query: str - model_plan: Mapping[str, Any] - inputs: Mapping[str, Mapping[str, Optional[str]]] = field( - default_factory=dict) - - @dataclass(frozen=True) class RunRecord: r"""One tracked run; ``steps`` are per-epoch ``(epoch, value)`` @@ -79,6 +66,28 @@ class RunRecord: 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 + run_id: 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') @@ -116,43 +125,25 @@ def _text(value: Any) -> Optional[str]: return None if value is None else str(value) -def run_record(context: FitContext, result: Optional['TrainingResult'] = None, - fit_ended: Optional[Event] = None, *, - tags: Optional[Mapping[str, str]] = None) -> RunRecord: - r"""The fit run's record; ``tags`` are the user's extra tags.""" - metrics: Dict[str, float] = {} - pointers: Dict[str, str] = {} - artifacts = { - 'query.pql': context.query, - 'model_plan.json': json.dumps(dict(context.model_plan), indent=2, - default=str), - } - if result is not None: - metrics.update(result.metrics) - if fit_ended is not None and 'seconds' in fit_ended.metrics: - metrics['seconds'] = fit_ended.metrics['seconds'] - if result.id: - pointers['kumo.bundle_id'] = str(result.id) - provenance = Path(result.bundle_dir) / PROVENANCE_FILE - if provenance.is_file(): - artifacts[PROVENANCE_FILE] = provenance.read_text( - encoding='utf-8') - params = { - key: _param_value(value) - for key, value in _flatten(context.model_plan) - } +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(metrics), + metrics=_finite_floats(spec.metrics), + steps={name: tuple(history) + for name, history in spec.steps.items()}, tags={ - **_user_tags(tags, params), - 'kumo.kind': 'fit', - 'kumo.run_id': context.run_id, - **pointers, - **_error_tag(fit_ended), + **_user_tags(spec.tags, params), + 'kumo.kind': spec.kind, + 'kumo.run_id': spec.run_id, + **{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=artifacts, - inputs=dict(context.inputs), + artifacts=dict(spec.artifacts), + inputs=dict(spec.inputs), ) @@ -170,50 +161,34 @@ def _user_tags(tags: Optional[Mapping[str, Any]], } -def _error_tag(ended: Optional[Event]) -> Dict[str, str]: - if ended is None or 'error' not in ended.params: - return {} - return {'kumo.error': str(ended.params['error'])} - - -def trial_record(context: FitContext, events: Sequence[Event]) -> RunRecord: - r"""One trial attempt's record from its events so far.""" +def trial_spec(run_id: str, events: Sequence[Event]) -> RunSpec: + r"""One trial attempt of run ``run_id``, from its events so far.""" started = next(event for event in events if event.name == 'trial_started') trial, attempt = started.trial, int(started.params['attempt']) - params: Dict[str, str] = { - key: _param_value(value) - for key, value in started.params.items() - } metrics: Dict[str, float] = {} steps: Dict[str, List[Tuple[int, float]]] = {} - ended: Dict[str, str] = {} + 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 = { - 'kumo.status': str(event.params.get('status')), - **_error_tag(event), - } + 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 RunRecord( + return RunSpec( + kind='trial', + run_id=f'{run_id}-trial-{trial}-attempt-{attempt}', name=f'trial-{trial}-attempt-{attempt}', - params=params, + params=dict(started.params), metrics=metrics, - steps={ - name: tuple(history) - for name, history in steps.items() - }, - tags={ - 'kumo.kind': 'trial', - 'kumo.run_id': f'{context.run_id}-trial-{trial}-attempt-{attempt}', - 'kumo.parent_run': context.run_id, - **ended, - }, + steps=steps, + parents={'parent_run': run_id}, + **ended, ) @@ -225,11 +200,12 @@ class Tracking(Protocol): 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(result, 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``. + ``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``. """ def start(self, record: RunRecord) -> None: ... @@ -252,9 +228,8 @@ def __init__(self, tracking: Optional[Tracking], warn: bool = True) -> None: self._tracker = tracking if _rank() == 0 else None self._warn = warn - self._context: Optional[FitContext] = None + self.spec: Optional[RunSpec] = None self._trials: Dict[int, List[Event]] = {} - self._fit_ended: Optional[Event] = None self._jobs: 'queue.Queue[Optional[Callable[[], Any]]]' = queue.Queue( QUEUE_SIZE) self._worker: Optional[threading.Thread] = None @@ -265,31 +240,29 @@ def __init__(self, tracking: Optional[Tracking], def enabled(self) -> bool: return self._tracker is not None - def start(self, context: Callable[[], FitContext]) -> None: + def start(self, spec: Callable[[], RunSpec]) -> None: if self._tracker is None: return try: - value = context() + 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._context = value + 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() self._submit(lambda: self._call(lambda tracker: tracker.start( - run_record(value, tags=getattr(tracker, 'tags', None))))) + run_record(dataclasses.replace(value, tags={ + **value.tags, **(getattr(tracker, 'tags', None) or {}) + }))))) def on_event(self, event: Event) -> None: - if self._context is None: - return - if event.name == 'fit_ended': - self._fit_ended = event - if event.trial is None: + if self.spec is None or event.trial is None: return if event.name == 'trial_started': self._trials[event.trial] = [] @@ -299,14 +272,15 @@ def on_event(self, event: Event) -> None: **event.params, 'attempt': events[0].params['attempt'] }) events.append(event) - record = trial_record(self._context, events) + record = run_record(trial_spec(self.spec.run_id, events)) self._submit(lambda: self._call( lambda tracker: tracker.on_trial(record, event))) def log_model( self, - result: 'TrainingResult', - snapshot: 'GraphSnapshot', + 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 @@ -327,7 +301,8 @@ def log() -> None: '' if self.error is None else f' ({full_error(self.error)})')) registered.append( - getattr(tracker, 'log_model')(result, snapshot, registry)) + getattr(tracker, 'log_model')(bundle_dir, manifest, + snapshot, registry)) except BaseException as error: failed.append(error) finally: @@ -363,17 +338,18 @@ def log() -> None: return None, error return (registered or [None])[0], None - def finish(self, result: Optional['TrainingResult'], + 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._context is None or worker is None: + if self.spec is None or worker is None: return try: - record: Optional[RunRecord] = run_record(self._context, result, - self._fit_ended) - if error is not None and 'kumo.error' not in record.tags: - record = dataclasses.replace(record, tags={ - **record.tags, 'kumo.error': format_error(error)}) + 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()) diff --git a/packages/kumo-engine/kumo/engine/integrations/_mlflow_model.py b/packages/kumo-engine/kumo/engine/integrations/_mlflow_model.py index da763012f..a1be11e94 100644 --- a/packages/kumo-engine/kumo/engine/integrations/_mlflow_model.py +++ b/packages/kumo-engine/kumo/engine/integrations/_mlflow_model.py @@ -29,9 +29,9 @@ class KumoModel(_pyfunc.PythonModel): def load_context(self, context: Any) -> None: - from kumo.engine.trainer.trainer import Trainer + from kumo.engine._predictor import load_predictor - self._trainer = Trainer.load(context.artifacts['bundle']) + 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') @@ -47,8 +47,8 @@ def predict(self, context: Any, model_input: pd.DataFrame, if not uri: raise ValueError("Pass params={'snapshot': }.") - return predict(self._trainer, self._snapshot(uri), model_input, - params) + return self._predictor.predict(model_input, self._snapshot(uri), + params) def _snapshot(self, uri: str) -> GraphSnapshot: if uri not in self._snapshots: diff --git a/packages/kumo-engine/kumo/engine/integrations/mlflow.py b/packages/kumo-engine/kumo/engine/integrations/mlflow.py index 0eee9e20b..e5020d97c 100644 --- a/packages/kumo-engine/kumo/engine/integrations/mlflow.py +++ b/packages/kumo-engine/kumo/engine/integrations/mlflow.py @@ -50,8 +50,9 @@ 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, TrainingResult + from kumo.engine.trainer.trainer import Trainer logger = logging.getLogger(__name__) @@ -262,7 +263,8 @@ def on_trial(self, record: RunRecord, event: Event) -> None: def log_model( self, - result: 'TrainingResult', + bundle_dir: Path, + manifest: Optional['ModelManifest'], snapshot: 'GraphSnapshot', registry: Optional[ModelRegistryLocation] = None, ) -> Optional[PublishedModel]: @@ -277,8 +279,8 @@ def log_model( } os.environ.update(upload) try: - info = _log_model(self.target, self._client, result, snapshot, - name='model', run_id=run_id, + info = _log_model(self.target, self._client, bundle_dir, manifest, + snapshot, name='model', run_id=run_id, experiment_id=self._experiment_id, serving_snapshot=self.serving_snapshot) finally: @@ -293,7 +295,7 @@ def log_model( return None try: return _register(self.target, self._client, info, registry, - registry_pointers(result)) + registry_pointers(_id(manifest))) except Exception as error: raise RegistrationError(full_error(error)) from error @@ -508,6 +510,7 @@ def log_model( ``serving_snapshot`` (an exported snapshot URI), each serving replica downloads that snapshot once at startup and scores against it when no ``snapshot`` param is given.""" + from kumo.engine._manifest import stub from kumo.engine._tracking import fit_inputs mlflow = import_provider('mlflow', 'mlflow') @@ -532,8 +535,10 @@ def log_model( 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, snapshot, name=name, - run_id=run_id, experiment_id=experiment_id, + 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=serving_snapshot) with target.scope(): if result.metrics: @@ -594,25 +599,35 @@ def _download(mlflow: Any, target: Any, model_uri: str, tracking_uri=target.tracking_uri, registry_uri=target.registry_uri) -def _log_model(target: Any, client: Any, result: 'TrainingResult', +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.""" + 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 read_training_spec + 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') - spec = read_training_spec(result.bundle_dir) - signature, example = _mlflow_model.signature(snapshot, - spec.predictive_query) - provenance = result.bundle_dir / PROVENANCE_FILE + 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, @@ -620,8 +635,7 @@ def _log_model(target: Any, client: Any, result: 'TrainingResult', 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=registry_pointers(result), - model_id=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: @@ -631,18 +645,21 @@ def _log_model(target: Any, client: Any, result: 'TrainingResult', local, mlflow_model=mlflow_model, python_model=_mlflow_model_code.__file__, model_config={ - 'columns': signature.outputs.input_names(), + 'columns': [] if signature is None else + signature.outputs.input_names(), 'snapshot': serving_snapshot, }, artifacts={ - 'bundle': str(result.bundle_dir), + 'bundle': str(bundle_dir), **({'provenance': str(provenance)} if provenance.is_file() else {}), }, signature=signature, input_example=example, - metadata=registry_pointers(result), - pip_requirements=[_requirement(result)], + metadata=pointers, + pip_requirements=[ + _requirement(listed_file(bundle_dir, 'encoders')) + ], ) with target.scope(): repository = target.artifact_repository( @@ -665,13 +682,17 @@ def _log_model(target: Any, client: Any, result: 'TrainingResult', return mlflow_model.get_model_info(logged) -def _requirement(result: 'TrainingResult') -> str: +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 = result.artifact_paths.get('encoders') + 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( diff --git a/packages/kumo-engine/kumo/engine/trainer/_provenance.py b/packages/kumo-engine/kumo/engine/trainer/_provenance.py index dfd6f151e..73e54e499 100644 --- a/packages/kumo-engine/kumo/engine/trainer/_provenance.py +++ b/packages/kumo-engine/kumo/engine/trainer/_provenance.py @@ -6,9 +6,10 @@ 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, @@ -32,6 +33,7 @@ def write_provenance( started_at: datetime, fit_mode: Optional[str], tracked_run_id: Optional[str] = None, + dependencies: Optional[List[Dependency]] = None, ) -> None: from kumo_core.artifacts import PROVENANCE_FILE, add_bundle_file @@ -54,6 +56,7 @@ def write_provenance( fit_mode=fit_mode, tracking=None if tracked_run_id is None else TrackingRef( run_id=tracked_run_id), + dependencies=dependencies, **stamp(), )) add_bundle_file(bundle_dir, PROVENANCE_ROLE, PROVENANCE_FILE) @@ -109,6 +112,7 @@ def record_derivation(bundle_dir: Path, parent_dir: Path, 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..7819797e2 --- /dev/null +++ b/packages/kumo-engine/kumo/engine/trainer/_run_spec.py @@ -0,0 +1,56 @@ +# 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.""" + + run_id: str + 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', + run_id=context.run_id, + 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/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index 43fc02c20..1539ddb24 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -56,21 +56,27 @@ 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._events import Event, EventEmitter, query_sha256 -from kumo.engine._manifest import ModelManifest +from kumo.engine._manifest import Dependency, ModelManifest, 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 FitContext, TrackingSession, fit_inputs +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, read_scoring_defaults from kumo.engine.trainer._spec import ( HOLDOUT_META_ROLE, @@ -145,7 +151,7 @@ def predict(self, *args: Any, **kwargs: Any) -> 'PredictionResult': def export(self, config: ModelOutputConfig) -> PublishedModel: """Publish this model's bundle, as :meth:`Trainer.export`.""" _require_publisher(config) - return _publish(self, config) + return _publish(self.bundle_dir, self.bundle_id, self.metrics, config) def with_scoring_defaults(self, defaults: ScoringDefaults) -> 'TrainingResult': @@ -329,7 +335,41 @@ 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, +) -> 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.""" + if settings is not None and settings.wants('model'): + _store_model(bundle_dir, manifest, settings) + registered, error = (None, None) if session is None else session.log_model( + bundle_dir, manifest, snapshot, _mlflow_registry(settings)) + registry = _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.""" from kumo_core.object_store_sink import holds_bundle, upload_bundle @@ -337,21 +377,21 @@ def _store_model(result: TrainingResult, settings: Any) -> None: 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"Could not store model {manifest.id} under the artifact " f"root: {error}", UserWarning, stacklevel=4) @@ -410,27 +450,29 @@ def _connector_registry(settings: Any) -> Optional[ModelOutputConfig]: return registry if isinstance(registry, ModelOutputConfig) else None -def _publish(training_result: TrainingResult, - export_config: ModelOutputConfig, +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: - spec = read_training_spec(training_result.bundle_dir) - metadata = { - name: str(getattr(spec, name)) - for name in ('spec_version', 'predictive_query', 'snapshot_id') - } - metadata.update(registry_pointers(training_result)) + 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(training_result.bundle_dir), + str(bundle_dir), path=export_config.path, location=location, metadata=metadata, - metrics=training_result.metrics, + metrics=metrics, ) if location and alias: try: @@ -547,7 +589,8 @@ def export(self, config: ModelOutputConfig) -> PublishedModel: """Publish the fitted or loaded model bundle.""" training_result = self._fitted('export') _require_publisher(config) - return _publish(training_result, config) + return _publish(training_result.bundle_dir, training_result.bundle_id, + training_result.metrics, config) @_without_torch_scatter_hint def fit( @@ -625,8 +668,8 @@ def fit( result = self._fit(train_table, snapshot=snapshot, export_config=export_config, warm_start=warm_start, run_id=run_id, - events=events, settings=settings, - tracking=session) + events=events, tracking=session, + deps=Dependencies(settings=settings)) except BaseException as error: events.emit('fit_ended', params={ 'status': 'failed', @@ -634,38 +677,25 @@ def fit( }) session.finish(None, error) raise + seconds = time.monotonic() - fit_start events.emit('fit_ended', metrics={ - **result.metrics, 'seconds': time.monotonic() - fit_start + **result.metrics, 'seconds': seconds }, params={'status': 'succeeded'}) - registered, error = session.log_model(result, snapshot, - _mlflow_registry(settings)) + 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) if registered is not None: result = replace(result, registered_model=registered) self._training_result = result - if registry is not None: - result, failure = self._register(result, registry, warn) - error = error or failure - session.finish(result, error) + 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 _register( - self, result: TrainingResult, registry: ModelOutputConfig, warn: bool - ) -> Tuple[TrainingResult, Optional[BaseException]]: - errors: List[BaseException] = [] - try: - registered = _publish(result, registry, errors, warn) - except Exception as error: - if warn: - warn_failure( - f"Could not register the model: {full_error(error)}", - error, 3) - return result, error - result = replace(result, registered_model=registered) - self._training_result = result - return result, (errors or [None])[0] - def _fit( self, train_table: TrainingTable, @@ -675,8 +705,8 @@ def _fit( warm_start: Optional[Self], run_id: str, events: EventEmitter, - settings: Optional[Any], tracking: TrackingSession, + deps: Dependencies, ) -> TrainingResult: from kumo_core.automl.runtime import ( inherit_warm_start_model_plan, @@ -811,12 +841,12 @@ 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: FitContext( + tracking.start(lambda: fit_spec(FitContext( run_id=run_id, 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, @@ -883,7 +913,7 @@ def _fit( warm_start_bundle_id), started_at=started_at, fit_mode=fit_mode, tracked_run_id=run_id if tracking.enabled else None), - settings=settings, + deps=deps, ) def _finish_fit( @@ -897,7 +927,7 @@ 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, ) -> TrainingResult: """Commit durable fit outputs before making this trainer usable.""" from kumo_core.artifacts import bundle_id_from_manifest @@ -923,19 +953,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: training_result = replace( - training_result, - published_model=_publish(training_result, export_config)) - if settings is not None and settings.wants('model'): - _store_model(training_result, settings) + 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 @@ -946,6 +978,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, @@ -1321,6 +1365,12 @@ 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, @@ -1330,6 +1380,15 @@ def _load_bundle( 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 @@ -1350,5 +1409,5 @@ def warn_if_other_bundle(metadata: Optional[Mapping[str, Any]], stacklevel=3) -def registry_pointers(result: TrainingResult) -> Dict[str, str]: - return {BUNDLE_ID_KEY: result.bundle_id} if result.bundle_id else {} +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/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_mlflow_model.py b/packages/kumo-engine/test/test_mlflow_model.py index 8937403b3..efd932dec 100644 --- a/packages/kumo-engine/test/test_mlflow_model.py +++ b/packages/kumo-engine/test/test_mlflow_model.py @@ -166,7 +166,7 @@ def predict_beside_another_call(self, *args, **kwargs): pyfunc.predict(rows, params={**params, 'population': population})), expected) - scratch = pyfunc.unwrap_python_model()._trainer.scratch_dir + scratch = pyfunc.unwrap_python_model()._predictor.trainer.scratch_dir assert list(scratch.iterdir()) == [] assert set(tables.glob('*')) == before | {tables / 'other-call'} @@ -261,7 +261,8 @@ def test_fit_waits_for_model_logging_and_fails_the_run_if_it_fails( import time from kumo.engine import _tracking - from kumo.engine._tracking import FitContext, TrackingSession + from kumo.engine._tracking import TrackingSession + from kumo.engine.trainer._run_spec import FitContext, fit_spec monkeypatch.setattr(_tracking, 'FINISH_SECONDS', 0.01) @@ -275,7 +276,7 @@ def start(self, record): def on_trial(self, record, event): pass - def log_model(self, result, snapshot, registry): + def log_model(self, bundle_dir, manifest, snapshot, registry): time.sleep(0.2) self.calls.append('log_model') if self.fail: @@ -288,11 +289,11 @@ def finish(self, record, error): for fail in (False, True): tracker = Tracker(fail) session = TrackingSession(tracker) - session.start(lambda: FitContext('r', 'q', {})) + session.start(lambda: fit_spec(FitContext('r', 'q', {}))) with warnings.catch_warnings(record=True) as caught: warnings.simplefilter('always') - _, error = session.log_model(result, None) - session.finish(result, error) + _, 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) @@ -305,7 +306,8 @@ def test_model_logging_wait_is_bounded(monkeypatch, stall): import threading from kumo.engine import _tracking - from kumo.engine._tracking import FitContext, TrackingSession + 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() @@ -319,7 +321,7 @@ def start(self, record): def on_trial(self, record, event): pass - def log_model(self, result, snapshot, registry): + def log_model(self, bundle_dir, manifest, snapshot, registry): release.wait(30) def finish(self, record, error): @@ -327,12 +329,12 @@ def finish(self, record, error): tracker = Tracker() session = TrackingSession(tracker) - session.start(lambda: FitContext('r', 'q', {})) + session.start(lambda: fit_spec(FitContext('r', 'q', {}))) if stall == 'dead thread': session._submit(lambda: (_ for _ in ()).throw(SystemExit())) session._worker.join(5) with pytest.warns(UserWarning, match='Could not log the model'): - _, error = session.log_model(SimpleNamespace(), None) + _, error = session.log_model(None, None, None) release.set() session.finish(None, error) @@ -349,6 +351,7 @@ def test_model_upload_gets_a_longer_http_timeout_and_a_late_one_is_dropped( 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) @@ -356,7 +359,7 @@ def test_model_upload_gets_a_longer_http_timeout_and_a_late_one_is_dropped( location = kumo.ModelRegistryLocation('m') seen = [] - def log_model(target, client, result, snapshot, **kwargs): + 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) @@ -366,9 +369,8 @@ def log_model(target, client, result, snapshot, **kwargs): monkeypatch.setattr(integration, '_register', lambda *args: 'registered') for _ in range(2): session = TrackingSession(tracker) - session.start(lambda: CONTEXT) - registered, error = session.log_model( - SimpleNamespace(bundle_id=None), None, location) + 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) @@ -391,8 +393,7 @@ def test_a_model_with_text_columns_requires_the_text_extra(tmp_path): def requirement(encoders): path = tmp_path / 'encoders.json' save_json(EncodersConfig(encoders={'users': encoders}), str(path)) - return _requirement( - SimpleNamespace(artifact_paths={'encoders': str(path)})) + return _requirement(path) assert requirement({'age': Numerical()}) == ( f'kumo-engine[s3,gcs]=={kumo.__version__}') @@ -403,27 +404,28 @@ def requirement(encoders): def test_a_tracker_dropped_earlier_says_the_model_was_not_logged(): import time - from kumo.engine._tracking import FitContext, TrackingSession + 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, result, snapshot, registry): + 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: FitContext('r', 'q', {})) + session.start(lambda: fit_spec(FitContext('r', '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(SimpleNamespace(), None) + 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' diff --git a/packages/kumo-engine/test/test_mlflow_tracking.py b/packages/kumo-engine/test/test_mlflow_tracking.py index b28f30b4c..9b459e264 100644 --- a/packages/kumo-engine/test/test_mlflow_tracking.py +++ b/packages/kumo-engine/test/test_mlflow_tracking.py @@ -22,6 +22,7 @@ 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') @@ -249,11 +250,11 @@ def test_a_tracked_fit_on_sagemaker_says_how_to_open_the_ui( from kumo.engine.integrations.mlflow import MlflowTracking tracker = MlflowTracking(log_system_metrics=False) - tracker.start(run_record(CONTEXT)) + 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(CONTEXT), None) + 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 " @@ -265,7 +266,7 @@ def test_user_disabled_autologging_stays_disabled(tracking_uri): tracker = MlflowTracking(log_system_metrics=False) with autologging_utils.disable_autologging(): - tracker.start(run_record(CONTEXT)) + tracker.start(run_record(fit_spec(CONTEXT))) tracker.finish(None, None) assert autologging_utils._AUTOLOGGING_GLOBALLY_DISABLED is True 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..648c1b66d --- /dev/null +++ b/packages/kumo-engine/test/test_model_agnostic_mlops.py @@ -0,0 +1,269 @@ +# 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 json +import os +import subprocess +import sys +import textwrap +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))) + + 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', run_id='loop-1', params={'lr': 0.1}, + inputs={'users': {'connector': 'S3'}}) + session.start(lambda: spec) + session.on_event(Event('trial_started', 'loop-1', 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.tags['kumo.parent_run'] == 'loop-1' + assert finished.metrics == {'auc': 0.9} + assert finished.tags == {'kumo.kind': 'custom', 'kumo.run_id': 'loop-1', + 'kumo.bundle_id': 'b'} + + +def test_any_bundle_is_stored_logged_and_registered(tmp_path): + 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')) + registry = kumo.ModelOutputConfig(SimpleNamespace(sink_backend=sink), + location=ModelRegistryLocation('TOY')) + settings = kumo.MLOps(str(tmp_path / 'root'), registry=registry) + tracker = _Tracker() + session = TrackingSession(tracker) + session.start(lambda: RunSpec(kind='custom', run_id='loop-1')) + + registered, error = _store_and_register( + bundle, manifest, settings, session=session, snapshot='snap', + metrics={'auc': 0.9}) + + assert error is None and registered.version == '1' + stored = tmp_path / 'root' / 'models' / manifest.id + assert (stored / 'weight.txt').read_text() == '2.0' + assert ('log_model', (bundle, manifest.id, 'snap')) in tracker.calls + 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 diff --git a/packages/kumo-engine/test/test_snowflake_tracking.py b/packages/kumo-engine/test/test_snowflake_tracking.py index 3ca74f771..aebe648c5 100644 --- a/packages/kumo-engine/test/test_snowflake_tracking.py +++ b/packages/kumo-engine/test/test_snowflake_tracking.py @@ -18,6 +18,7 @@ ) 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') @@ -63,7 +64,7 @@ def _tracking(session=FakeSession, **kwargs): def _replay(tracking, result=RESULT, error=None, during=None): session = TrackingSession(tracking) - session.start(lambda: CONTEXT) + session.start(lambda: fit_spec(CONTEXT)) for event in EVENTS: session.on_event(event) if during is not None: @@ -71,7 +72,10 @@ def _replay(tracking, result=RESULT, error=None, during=None): session._submit(caught_up.set) caught_up.wait(30) during() - session.finish(result, error) + 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): diff --git a/packages/kumo-engine/test/test_tracking_contract.py b/packages/kumo-engine/test/test_tracking_contract.py index 2efe3cb35..bcf9cf3f0 100644 --- a/packages/kumo-engine/test/test_tracking_contract.py +++ b/packages/kumo-engine/test/test_tracking_contract.py @@ -6,7 +6,7 @@ import json import threading import time -from dataclasses import dataclass, field +from dataclasses import dataclass, field, replace from pathlib import Path from types import SimpleNamespace from typing import Any, Callable, Dict, List, Mapping, Optional, Tuple @@ -17,11 +17,11 @@ from kumo.engine._tracking import ( DROPPED_METRICS, DROPPED_PARAMS, - FitContext, TrackingSession, run_record, ) from kumo.engine.tracking import Event, RunRecord +from kumo.engine.trainer._run_spec import FitContext, fit_spec, fitted_spec RUN = 'trainer-fit-1' USER_TAGS = {'team': 'ents', 'kumo.run_id': 'overridden'} @@ -118,15 +118,22 @@ def replay(tracker=None, events=EVENTS, result=RESULT, of it.""" recorder = _Recorder(tracker) session = TrackingSession(recorder) - session.start(lambda: CONTEXT) + session.start(lambda: fit_spec(CONTEXT)) for event in events: session.on_event(event) - session.finish(result, error) + session.finish( + None if result is None else lambda: fitted_spec( + session.spec, result, events[-1].metrics.get('seconds')), error) records = { record.tags['kumo.run_id']: record for hook, record in recorder.calls if hook == 'on_trial' } - records[RUN] = run_record(CONTEXT, result, events[-1], tags=USER_TAGS) + 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[RUN] = run_record(spec) return records @@ -304,9 +311,8 @@ def test_records_of_the_recorded_fit(): } with pytest.warns(UserWarning, match="'kumo.bundle_id', 'training_job.tune_metric'"): - spoofed = run_record(CONTEXT, tags={'kumo.bundle_id': 'x', - 'training_job.tune_metric': 'y', - 'a': 1}) + 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 @@ -370,8 +376,8 @@ def test_only_rank_zero_tracks(monkeypatch): recorder = _Recorder() session = TrackingSession(recorder) - session.start(lambda: CONTEXT) - session.finish(RESULT, None) + session.start(lambda: fit_spec(CONTEXT)) + session.finish(lambda: fitted_spec(session.spec, RESULT), None) assert recorder.calls == [] From 6ac846ac85f5289f37d4b5b7dc2f3a1ddb7ff97c Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Wed, 28 Oct 2026 21:13:47 +0000 Subject: [PATCH 03/16] Register through the tracker, record its own run, and set scoring defaults at fit kumo.MLOps(registry=) takes a ModelRegistryLocation only: MlflowTracking registers an MLflow pyfunc (Unity Catalog, SageMaker MLflow) and SnowflakeTracking registers in the Snowflake Model Registry through its connector. Trainer.load_from_connector(cfg, location=...) loads a registered version back the same way. Runs carry no Kumo run id: the model manifest records the tracker's own run as tracking (platform, run id, URL), read as TrainingResult.tracking, and trial runs are found by name. fit(scoring_defaults=...) stores the defaults in the bundle before it is stored and registered. MlflowTracking(target=) also takes a tracking URI and log_model/load_model take target=; SnowflakeTracking takes the connector or session first and schema_name=. Co-Authored-By: Claude Opus 5.5 --- examples/mlops/smoke.py | 23 ++-- packages/kumo-engine/kumo/engine/_events.py | 10 +- packages/kumo-engine/kumo/engine/_manifest.py | 6 ++ packages/kumo-engine/kumo/engine/_mlops.py | 71 +++++++++---- packages/kumo-engine/kumo/engine/_tracking.py | 48 ++++++--- .../engine/integrations/_mlflow_target.py | 55 ++++++++-- .../kumo/engine/integrations/mlflow.py | 100 ++++++++++++------ .../kumo/engine/integrations/snowflake.py | 58 ++++++---- .../kumo/engine/trainer/_provenance.py | 5 +- .../kumo/engine/trainer/_run_spec.py | 2 - .../kumo/engine/trainer/trainer.py | 73 ++++++++++--- packages/kumo-engine/test/test_fit_events.py | 3 +- .../kumo-engine/test/test_mlflow_model.py | 10 +- .../kumo-engine/test/test_mlflow_target.py | 37 ++++--- .../kumo-engine/test/test_mlflow_tracking.py | 7 +- packages/kumo-engine/test/test_mlops.py | 29 ++--- .../test/test_model_agnostic_mlops.py | 50 ++++++--- .../kumo-engine/test/test_scoring_defaults.py | 24 +++++ .../test/test_snowflake_tracking.py | 41 +++++-- .../test/test_tracking_contract.py | 37 ++++--- .../kumo-engine/test/test_trainer_registry.py | 36 +++++-- 21 files changed, 497 insertions(+), 228 deletions(-) diff --git a/examples/mlops/smoke.py b/examples/mlops/smoke.py index 47e152564..0d4b1cfcf 100644 --- a/examples/mlops/smoke.py +++ b/examples/mlops/smoke.py @@ -88,8 +88,8 @@ def tracking(args: argparse.Namespace, suffix: str): schema_name=args.schema, identity=SnowflakeIdentity( account=args.account, username=args.user, private_key=Path(args.private_key_file).read_text())) - return SnowflakeTracking(connector=connector, - experiment_name='KUMO_SMOKE'), None + return (SnowflakeTracking(connector, experiment_name='KUMO_SMOKE'), + f'KUMO_SMOKE_{suffix}') def main(argv: Optional[Sequence[str]] = None) -> None: @@ -109,10 +109,9 @@ def main(argv: Optional[Sequence[str]] = None) -> None: kumo.configure(scratch_dir=work / 'scratch') graph = data(work / 'data') tracker, model_name = tracking(args, time.strftime('%Y%m%d%H%M%S')) - cfg = kumo.MLOps( - args.artifact_root or str(work / 'root'), tracking=tracker, - registry=None if model_name is None else ModelRegistryLocation( - model_name, alias='candidate'), on_error='raise') + 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() @@ -122,15 +121,11 @@ def main(argv: Optional[Sequence[str]] = None) -> None: plan.training_job.num_experiments = 1 plan.optimization.max_epochs = 1 result = kumo.Trainer(plan).fit(table, snapshot=snapshot) - print(f'fit run {tracker.run_id}, bundle {result.id}, ' + print(f'fit run {result.tracking}, bundle {result.id}, ' f'registered {result.registered_model}') - if model_name is not None: - from kumo.engine.integrations.mlflow import load_model - - loaded = load_model(f'models:/{model_name}@candidate', - tracking=tracker) - assert loaded.result.id == result.id, (loaded.result.id, result.id) - print('loaded the registered model back') + 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') diff --git a/packages/kumo-engine/kumo/engine/_events.py b/packages/kumo-engine/kumo/engine/_events.py index 7309aab19..98e8c2394 100644 --- a/packages/kumo-engine/kumo/engine/_events.py +++ b/packages/kumo-engine/kumo/engine/_events.py @@ -32,7 +32,6 @@ class Event: r"""One progress event. ``trial`` and ``epoch`` count from 1.""" name: EventName - run_id: str time: float trial: Optional[int] = None epoch: Optional[int] = None @@ -60,17 +59,14 @@ def _finite_floats(metrics: Mapping[str, Any]) -> Dict[str, float]: class EventEmitter(NoOpTrainingInstrumentation): - def __init__(self, run_id: str, - callbacks: Sequence[Callable[[Event], None]]) -> None: - self.run_id = run_id + 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, run_id=self.run_id, time=time.time(), - trial=trial, epoch=epoch, + 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): @@ -79,7 +75,7 @@ def emit(self, name: EventName, *, trial: Optional[int] = None, except Exception: logger.warning( 'Progress callback %r failed; it is disabled for the ' - 'rest of run %s', callback, self.run_id, exc_info=True) + 'rest of this fit', callback, exc_info=True) self._callbacks.remove(callback) def event(self, name: str, **fields: Any) -> None: diff --git a/packages/kumo-engine/kumo/engine/_manifest.py b/packages/kumo-engine/kumo/engine/_manifest.py index e5aca170b..574e84836 100644 --- a/packages/kumo-engine/kumo/engine/_manifest.py +++ b/packages/kumo-engine/kumo/engine/_manifest.py @@ -245,7 +245,13 @@ class FitTiming(Record): @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) diff --git a/packages/kumo-engine/kumo/engine/_mlops.py b/packages/kumo-engine/kumo/engine/_mlops.py index 054577c90..a47175446 100644 --- a/packages/kumo-engine/kumo/engine/_mlops.py +++ b/packages/kumo-engine/kumo/engine/_mlops.py @@ -7,6 +7,7 @@ import contextvars import os import posixpath +import re import warnings from dataclasses import dataclass, field from typing import ( @@ -55,9 +56,10 @@ class MLOps: ``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: where fits register models: a ``ModelOutputConfig`` with a - registry ``location``, or a ``ModelRegistryLocation`` registered - through the MLflow tracker. + 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. @@ -112,30 +114,39 @@ def __post_init__(self) -> None: f'{list(KINDS)}.') if kinds: _require_writable(root) + 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, ) - if isinstance(self.registry, ModelOutputConfig): - check_model_location(self.registry) - if isinstance(self.registry, ModelRegistryLocation): - if not callable(getattr(self.tracking, 'log_model', None)): - raise ValueError( - 'registry=ModelRegistryLocation(...) registers through ' - 'the MLflow tracker: pass tracking=MlflowTracking(...), ' - 'or registry=ModelOutputConfig(connector, location=...).') - if self.registry.version is not None: + 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.') - if self.registry.alias is not None: - check_model_alias(self.registry.alias) - check = getattr(self.tracking, 'check', None) - if callable(check): - check(self.registry if isinstance( - self.registry, ModelRegistryLocation) else None) - object.__setattr__(self, 'root', root) - object.__setattr__(self, 'kinds', kinds) + 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 wants(self, kind: str) -> bool: r"""Whether artifacts of ``kind`` are stored under the root.""" @@ -265,3 +276,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/_tracking.py b/packages/kumo-engine/kumo/engine/_tracking.py index 562059341..159c9a83a 100644 --- a/packages/kumo-engine/kumo/engine/_tracking.py +++ b/packages/kumo-engine/kumo/engine/_tracking.py @@ -35,7 +35,7 @@ if TYPE_CHECKING: from kumoapi.plugin import ModelRegistryLocation, PublishedModel - from kumo.engine._manifest import ModelManifest + from kumo.engine._manifest import ModelManifest, TrackingRef logger = logging.getLogger(__name__) @@ -73,7 +73,6 @@ class RunSpec: the user's own.""" kind: str - run_id: str name: Optional[str] = None params: Mapping[str, Any] = field(default_factory=dict) metrics: Mapping[str, Any] = field(default_factory=dict) @@ -136,7 +135,6 @@ def run_record(spec: RunSpec) -> RunRecord: tags={ **_user_tags(spec.tags, params), 'kumo.kind': spec.kind, - 'kumo.run_id': spec.run_id, **{f'kumo.{name}': str(value) for name, value in spec.parents.items() if value}, **({} if spec.status is None else {'kumo.status': spec.status}), @@ -161,8 +159,8 @@ def _user_tags(tags: Optional[Mapping[str, Any]], } -def trial_spec(run_id: str, events: Sequence[Event]) -> RunSpec: - r"""One trial attempt of run ``run_id``, from its events so far.""" +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] = {} @@ -182,12 +180,10 @@ def trial_spec(run_id: str, events: Sequence[Event]) -> RunSpec: }) return RunSpec( kind='trial', - run_id=f'{run_id}-trial-{trial}-attempt-{attempt}', name=f'trial-{trial}-attempt-{attempt}', params=dict(started.params), metrics=metrics, steps=steps, - parents={'parent_run': run_id}, **ended, ) @@ -205,7 +201,10 @@ class Tracking(Protocol): `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``. + ``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: ... @@ -234,6 +233,7 @@ def __init__(self, tracking: Optional[Tracking], QUEUE_SIZE) self._worker: Optional[threading.Thread] = None self._behind = False + self._started = threading.Event() self.error: Optional[BaseException] = None @property @@ -256,10 +256,30 @@ def start(self, spec: Callable[[], RunSpec]) -> None: self._worker = threading.Thread(target=self._work, name='kumo-tracking', daemon=True) self._worker.start() - self._submit(lambda: self._call(lambda tracker: tracker.start( - run_record(dataclasses.replace(value, tags={ + + 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: @@ -272,7 +292,7 @@ def on_event(self, event: Event) -> None: **event.params, 'attempt': events[0].params['attempt'] }) events.append(event) - record = run_record(trial_spec(self.spec.run_id, events)) + record = run_record(trial_spec(events)) self._submit(lambda: self._call( lambda tracker: tracker.on_trial(record, event))) @@ -285,8 +305,8 @@ def log_model( ) -> 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 hasattr( - self._tracker, 'log_model')): + 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] = [] diff --git a/packages/kumo-engine/kumo/engine/integrations/_mlflow_target.py b/packages/kumo-engine/kumo/engine/integrations/_mlflow_target.py index 95348b466..8372eec35 100644 --- a/packages/kumo-engine/kumo/engine/integrations/_mlflow_target.py +++ b/packages/kumo-engine/kumo/engine/integrations/_mlflow_target.py @@ -56,11 +56,34 @@ class Target: def __init__(self, tracking_uri: str, registry_uri: Optional[str], swap: Optional[Callable[[], ContextManager[Any]]] = None, - aws_session: Any = None) -> 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: @@ -93,22 +116,33 @@ def artifact_repository(self, uri: str) -> Any: session_token=frozen.token) -def resolve(target: Any, tracking_uri: Optional[str]) -> Target: +def resolve(target: Any) -> Target: mlflow = import_provider('mlflow', 'mlflow') - if target is not None and tracking_uri is not None: - raise ValueError('Pass target or tracking_uri, not both.') + if isinstance(target, Target): + return target if target is None: - if tracking_uri is None: - return Target(mlflow.get_tracking_uri(), mlflow.get_registry_uri()) - return Target(tracking_uri, registry_for(tracking_uri)) + 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'or a SageMakerMLflow, not {type(target).__name__}.') + 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: @@ -135,7 +169,8 @@ def _databricks(connector: Any) -> Target: client = import_provider('kumo_databricks.session', 'databricks')._client_for(principal) return Target('databricks', 'databricks-uc', - swap=lambda: _databricks_token(client)) + swap=lambda: _databricks_token(client), + host=client.config.host) @contextmanager diff --git a/packages/kumo-engine/kumo/engine/integrations/mlflow.py b/packages/kumo-engine/kumo/engine/integrations/mlflow.py index e5020d97c..49d8a4b4e 100644 --- a/packages/kumo-engine/kumo/engine/integrations/mlflow.py +++ b/packages/kumo-engine/kumo/engine/integrations/mlflow.py @@ -24,6 +24,7 @@ FrozenSet, Iterator, List, + Literal, Mapping, Optional, Set, @@ -83,8 +84,8 @@ class MlflowTracking: Args: experiment_name: experiment for the fit run, created if missing. - `None` uses MLflow's current experiment; required with - ``target`` or ``tracking_uri``. + `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 @@ -94,13 +95,11 @@ class MlflowTracking: before the fit returns. ``kumo.MLOps(registry= ModelRegistryLocation(...))`` also registers it, which implies this. - serving_snapshot: exported snapshot URI the served model loads at - startup (see `log_model`). target: where to record, with the credentials that reach it: a ``DatabricksConnector`` (Databricks MLflow and Unity Catalog, - as its identity) or a `SageMakerMLflow`. `None` uses MLflow's - own tracking and registry URIs, read once here. - tracking_uri: another MLflow server, instead of ``target``. + 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 @@ -111,7 +110,8 @@ class MlflowTracking: except that a SageMaker target still sets the process-wide tracking URI for each call, since the plugin signs and routes with it. - ``run_id`` is the MLflow id of the latest fit run, once it has started. + ``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__( @@ -121,27 +121,20 @@ def __init__( log_system_metrics: bool = True, tags: Optional[Mapping[str, str]] = None, log_model: bool = False, - serving_snapshot: Optional[str] = None, *, target: Any = None, - tracking_uri: Optional[str] = None, ) -> None: self._mlflow = import_provider('mlflow', 'mlflow') - self.target = resolve_target(target, tracking_uri) - if experiment_name is None and (target is not None - or tracking_uri is not None): - raise ValueError('MlflowTracking(target=...) and ' - 'MlflowTracking(tracking_uri=...) need an ' - 'experiment_name.') + self.target = resolve_target(target) + self._named = target is not None self._default_experiment: Optional[str] = ( - None if experiment_name is not None else + 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.serving_snapshot = serving_snapshot self.logs_model = log_model self._outer: Optional[str] = None self.run_id: Optional[str] = None @@ -157,6 +150,16 @@ def _reset(self) -> 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) + def check_model_name(self, name: str) -> None: r"""Raises :class:`ValueError` unless this tracker's registry accepts ``name``.""" @@ -167,6 +170,7 @@ def check(self, registry: Optional[ModelRegistryLocation] = None) -> None: 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 @@ -195,7 +199,26 @@ def prepare(self) -> None: 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 start(self, record: RunRecord) -> None: + self._require_experiment() self._reset() self.run_id = None self._http_owned = frozenset(name for name in _HTTP_LIMITS @@ -282,7 +305,7 @@ def log_model( info = _log_model(self.target, self._client, bundle_dir, manifest, snapshot, name='model', run_id=run_id, experiment_id=self._experiment_id, - serving_snapshot=self.serving_snapshot) + serving_snapshot=None) finally: if self._parent == run_id: os.environ.update({name: _HTTP_LIMITS[name] for name in upload}) @@ -499,26 +522,32 @@ def log_model( snapshot: 'GraphSnapshot', name: str = 'model', run_id: Optional[str] = None, - serving_snapshot: Optional[str] = None, - tracking: Optional[MlflowTracking] = 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) with its source tables as run inputs; - returns its `ModelInfo`. ``tracking`` says where, with what credentials - (its experiment holds a new run); `None` uses MLflow's own tracking - URI. Scores with ``predict(entities, params={'snapshot': uri})``; with - ``serving_snapshot`` (an exported snapshot URI), each serving replica - downloads that snapshot once at startup and scores against it when no - ``snapshot`` param is given.""" + 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 = tracking or MlflowTracking() + 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 ( @@ -539,7 +568,7 @@ def log_model( (None if result.id is None else stub('model', result.id)), snapshot, name=name, run_id=run_id, experiment_id=experiment_id, - serving_snapshot=serving_snapshot) + serving_snapshot=snapshot_uri(serving_snapshot)) with target.scope(): if result.metrics: client.log_batch(run_id, metrics=[ @@ -557,22 +586,25 @@ def load_model( model_uri: str, embedding_models: Optional['EmbeddingModelStore'] = None, *, - tracking: Optional[MlflowTracking] = 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 ``tracking`` reaches (`None`: MLflow's own tracking URI).""" + 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 = (tracking or MlflowTracking()).target + 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) + trainer = Trainer.load(local / bundle, embedding_models, + accelerator=accelerator) warn_if_other_bundle(model.metadata, trainer) return trainer diff --git a/packages/kumo-engine/kumo/engine/integrations/snowflake.py b/packages/kumo-engine/kumo/engine/integrations/snowflake.py index dfa463060..78b3e23b6 100644 --- a/packages/kumo-engine/kumo/engine/integrations/snowflake.py +++ b/packages/kumo-engine/kumo/engine/integrations/snowflake.py @@ -7,6 +7,7 @@ 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 @@ -36,16 +37,17 @@ class SnowflakeTracking: parameters are not changed. Args: - session: Snowpark session; needs CREATE EXPERIMENT on the schema, or - MODIFY on an existing experiment. - connector: a ``SnowflakeConnector`` whose identity opens the session - instead. + 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: `None` uses the session's current schema, or ``PUBLIC`` when - `database` is given or the session has none. + 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 suffix from the kumo run id. + 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. @@ -54,26 +56,32 @@ class SnowflakeTracking: by name), once it has started. """ + platform = 'snowflake' + run_url = None + def __init__( self, - session: Any = None, + source: Any, *, - connector: Any = None, experiment_name: str, database: Optional[str] = None, - schema: 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 (session is None) == (connector is None): - raise ValueError('SnowflakeTracking takes a session or a ' - 'connector.') - self.session = session if session is not None else connector.session() + 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 + self.schema = schema_name self.run_name = run_name self.tags = dict(tags or {}) self.native_model_lineage = native_model_lineage @@ -92,6 +100,17 @@ def _reset(self) -> None: 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).""" @@ -135,8 +154,7 @@ def start(self, record: RunRecord) -> None: database_name=database, schema_name=schema or 'PUBLIC') self._run = tracking.start_run( - self._new_run_name(tracking, experiment.name, - record.tags['kumo.run_id'])) + 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) @@ -191,8 +209,8 @@ def finish(self, record: Optional[RunRecord], self._run_context.close() self._reset() - def _new_run_name(self, tracking: Any, experiment: Any, - run_id: str) -> Optional[str]: + 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 @@ -203,7 +221,7 @@ def _new_run_name(self, tracking: Any, experiment: Any, for row in tracking._sql_client.show_runs_in_experiment( experiment_name=experiment) ] - short = run_id.rsplit('-', 1)[-1] + 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): diff --git a/packages/kumo-engine/kumo/engine/trainer/_provenance.py b/packages/kumo-engine/kumo/engine/trainer/_provenance.py index 73e54e499..72c83129e 100644 --- a/packages/kumo-engine/kumo/engine/trainer/_provenance.py +++ b/packages/kumo-engine/kumo/engine/trainer/_provenance.py @@ -32,7 +32,7 @@ def write_provenance( warm_start: Optional[ModelManifest], started_at: datetime, fit_mode: Optional[str], - tracked_run_id: Optional[str] = None, + tracking: Optional[TrackingRef] = None, dependencies: Optional[List[Dependency]] = None, ) -> None: from kumo_core.artifacts import PROVENANCE_FILE, add_bundle_file @@ -54,8 +54,7 @@ def write_provenance( seconds=round((finished_at - started_at).total_seconds(), 3), ), fit_mode=fit_mode, - tracking=None if tracked_run_id is None else TrackingRef( - run_id=tracked_run_id), + tracking=tracking, dependencies=dependencies, **stamp(), )) diff --git a/packages/kumo-engine/kumo/engine/trainer/_run_spec.py b/packages/kumo-engine/kumo/engine/trainer/_run_spec.py index 7819797e2..1a916aad4 100644 --- a/packages/kumo-engine/kumo/engine/trainer/_run_spec.py +++ b/packages/kumo-engine/kumo/engine/trainer/_run_spec.py @@ -17,7 +17,6 @@ class FitContext: r"""What a fit trains, known once its model plan is resolved.""" - run_id: str query: str model_plan: Mapping[str, Any] inputs: Mapping[str, Mapping[str, Optional[str]]] = field( @@ -28,7 +27,6 @@ def fit_spec(context: FitContext, *, tags: Optional[Mapping[str, Any]] = None) -> RunSpec: return RunSpec( kind='fit', - run_id=context.run_id, params=dict(_flatten(context.model_plan)), tags=dict(tags or {}), inputs=dict(context.inputs), diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index 1539ddb24..9e12ff8f7 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -66,7 +66,7 @@ with_text_encoders, ) from kumo.engine._events import Event, EventEmitter, query_sha256 -from kumo.engine._manifest import Dependency, ModelManifest, stub +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 @@ -77,7 +77,12 @@ 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, read_scoring_defaults +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, @@ -144,6 +149,13 @@ 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) @@ -157,14 +169,11 @@ 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``.""" + ``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 - from kumo.engine.trainer._scoring import ( - normalized_scoring_defaults, - write_scoring_defaults, - ) if self._trainer is None: raise RuntimeError('This result has no model; use Trainer.load on ' @@ -441,13 +450,21 @@ def _require_publisher(export_config: ModelOutputConfig) -> None: def _mlflow_registry(settings: Any) -> Optional[ModelRegistryLocation]: - registry = None if settings is None else settings.registry - return registry if isinstance(registry, ModelRegistryLocation) else None + if settings is None or not _logs_models(settings.tracking): + return None + return settings.registry def _connector_registry(settings: Any) -> Optional[ModelOutputConfig]: - registry = None if settings is None else settings.registry - return registry if isinstance(registry, ModelOutputConfig) else None + 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], @@ -599,6 +616,7 @@ 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]] = (), @@ -622,6 +640,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; @@ -653,7 +674,6 @@ def fit( session = TrackingSession( None if settings is None else settings.tracking, warn=warn) events = EventEmitter( - run_id, [*callbacks, session.on_event] if session.enabled else callbacks) fit_start = time.monotonic() try: @@ -667,6 +687,7 @@ def fit( 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=Dependencies(settings=settings)) @@ -702,6 +723,7 @@ def _fit( *, snapshot: GraphSnapshot, export_config: Optional[ModelOutputConfig], + scoring_defaults: Optional[ScoringDefaults], warm_start: Optional[Self], run_id: str, events: EventEmitter, @@ -764,6 +786,9 @@ 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: from kumo_core.artifacts import bundle_task_type @@ -842,7 +867,6 @@ def _fit( raise ValueError('The test split is empty: shorten the query ' 'window or add data.') tracking.start(lambda: fit_spec(FitContext( - run_id=run_id, query=train_table.query, model_plan=to_json_dict(materialized.model_plan), inputs=fit_inputs(train_table.manifest), @@ -905,6 +929,7 @@ 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, @@ -912,7 +937,7 @@ def _fit( warm_start._training_result.bundle_dir, warm_start_bundle_id), started_at=started_at, fit_mode=fit_mode, - tracked_run_id=run_id if tracking.enabled else None), + tracking=tracking.pointer()), deps=deps, ) @@ -928,6 +953,7 @@ def _finish_fit( model_plan: Optional[ModelPlan] = None, provenance: Optional[Mapping[str, 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 @@ -943,6 +969,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( @@ -1315,12 +1346,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, diff --git a/packages/kumo-engine/test/test_fit_events.py b/packages/kumo-engine/test/test_fit_events.py index c9efff26c..08a22729c 100644 --- a/packages/kumo-engine/test/test_fit_events.py +++ b/packages/kumo-engine/test/test_fit_events.py @@ -78,7 +78,6 @@ def broken(event): 'trial_started', 'epoch_ended', 'epoch_ended', 'trial_ended', 'fit_ended', ] - assert len({event.run_id for event in events}) == 1 assert set(events[0].params) == {'snapshot_id', 'query_sha256'} assert [(e.trial, e.epoch) for e in events if e.name == 'epoch_ended' ] == [(1, 1), (1, 2), (2, 1), (2, 2)] @@ -124,7 +123,7 @@ def test_event_metrics_keep_finite_numbers(): from kumo.engine._events import EventEmitter received = [] - EventEmitter('run', [received.append]).emit( + EventEmitter([received.append]).emit( 'epoch_ended', metrics={'val_loss': 'n/a', 'val_auc': float('nan'), 'val_mae': 0.5}) diff --git a/packages/kumo-engine/test/test_mlflow_model.py b/packages/kumo-engine/test/test_mlflow_model.py index efd932dec..2f3cff81d 100644 --- a/packages/kumo-engine/test/test_mlflow_model.py +++ b/packages/kumo-engine/test/test_mlflow_model.py @@ -65,7 +65,6 @@ def test_fit_logs_registers_and_serves_the_model(table_and_snapshot, assert {k: v for k, v in parent.data.tags.items() if k.startswith('kumo.')} == { 'kumo.kind': 'fit', - 'kumo.run_id': parent.data.tags['kumo.run_id'], 'kumo.bundle_id': result.id, } assert result.registered_model == PublishedModel( @@ -95,6 +94,9 @@ def test_fit_logs_registers_and_serves_the_model(table_and_snapshot, assert list(served) == ['ENTITY', 'TARGET_PRED'] assert len(served) == len(range(0, 400, 7)) assert load_model('models:/m@champion').result.id == result.id + assert kumo.Trainer.load_from_connector( + cfg, location=kumo.ModelRegistryLocation('m', alias='champion') + ).result.id == result.id original = mlflow.models.Model.load @@ -289,7 +291,7 @@ def finish(self, record, error): for fail in (False, True): tracker = Tracker(fail) session = TrackingSession(tracker) - session.start(lambda: fit_spec(FitContext('r', 'q', {}))) + 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) @@ -329,7 +331,7 @@ def finish(self, record, error): tracker = Tracker() session = TrackingSession(tracker) - session.start(lambda: fit_spec(FitContext('r', 'q', {}))) + session.start(lambda: fit_spec(FitContext('q', {}))) if stall == 'dead thread': session._submit(lambda: (_ for _ in ()).throw(SystemExit())) session._worker.join(5) @@ -418,7 +420,7 @@ def finish(self, record, error): pass session = TrackingSession(Tracker()) - session.start(lambda: fit_spec(FitContext('r', 'q', {}))) + 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) diff --git a/packages/kumo-engine/test/test_mlflow_target.py b/packages/kumo-engine/test/test_mlflow_target.py index c18a2df68..27be764b0 100644 --- a/packages/kumo-engine/test/test_mlflow_target.py +++ b/packages/kumo-engine/test/test_mlflow_target.py @@ -40,30 +40,34 @@ def _connector(identity): def test_no_target_reads_mlflows_own_uris_once(tracking_uri): - target = resolve(None, None) + target = resolve(None) assert target.tracking_uri == tracking_uri assert not target.explicit - assert resolve(None, 'databricks').registry_uri == 'databricks-uc' - assert resolve(None, 'databricks://prod').registry_uri == ( + 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') - assert resolve(None, 'http://mlflow:5000').registry_uri == ( - 'http://mlflow:5000') - with pytest.raises(ValueError, match='not both'): - resolve(SageMakerMLflow(ARN), 'http://x') + 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()), None) + target = resolve(_connector(DatabricksIdentity())) before = dict(os.environ) with target.scope(): assert dict(os.environ) == before assert mlflow.get_tracking_uri() == tracking_uri assert not target.explicit - for target in (resolve(SageMakerMLflow(ARN, S3Identity()), None), - resolve(SageMakerMLflow(ARN), None)): + 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 @@ -85,7 +89,7 @@ def test_a_databricks_service_principal_lends_a_token_for_each_call( target = resolve( _connector( DatabricksIdentity(host='https://dbc.example', client_id='sp', - client_secret='secret')), None) + client_secret='secret'))) previous = mlflow.get_tracking_uri() with target.scope(): @@ -102,6 +106,8 @@ def test_a_databricks_service_principal_lends_a_token_for_each_call( 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( @@ -112,7 +118,7 @@ def test_an_aws_access_key_signs_sagemaker_calls_and_uploads( monkeypatch.setenv('SAGEMAKER_MLFLOW_ASSUME_ROLE_ARN', 'arn:aws:iam::x') target = resolve( SageMakerMLflow(ARN, S3Identity(access_key_id='AKIA', - secret_access_key='s3cret')), None) + secret_access_key='s3cret'))) with target.scope(): assert mlflow.get_tracking_uri() == ARN @@ -172,11 +178,14 @@ def test_an_unreachable_server_fails_when_mlops_is_built( from kumo.engine.integrations.mlflow import MlflowTracking monkeypatch.setenv('MLFLOW_HTTP_REQUEST_MAX_RETRIES', '0') - tracking = MlflowTracking(tracking_uri='http://127.0.0.1:9', + 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( @@ -255,7 +264,7 @@ def test_registering_a_registered_bundle_reuses_its_version( tracking_uri, caplog): from kumo.engine.integrations.mlflow import _register - target = resolve(None, None) + target = resolve(None) client = target.client() run = client.create_run('0') info = SimpleNamespace(model_uri=f'runs:/{run.info.run_id}/model', diff --git a/packages/kumo-engine/test/test_mlflow_tracking.py b/packages/kumo-engine/test/test_mlflow_tracking.py index 9b459e264..920e92272 100644 --- a/packages/kumo-engine/test/test_mlflow_tracking.py +++ b/packages/kumo-engine/test/test_mlflow_tracking.py @@ -116,11 +116,12 @@ def test_fit_logs_a_run_per_fit_and_per_trial(table_and_snapshot, tmp_path, assert client.get_experiment(parent.info.experiment_id).name == 'ents' assert parent.info.status == 'FINISHED' assert {k for k in parent.data.tags if k.startswith('kumo.')} == { - 'kumo.kind', 'kumo.run_id', 'kumo.bundle_id', 'kumo.dropped_params'} + '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'] == { - 'run_id': parent.data.tags['kumo.run_id']} + '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 @@ -231,7 +232,7 @@ def test_registry_model_names_are_checked_when_mlops_is_built( for name in good: check_model_name(name, uri) tracking = MlflowTracking(log_system_metrics=False, - tracking_uri=f'sqlite:///{tmp_path}/m.db', + target=f'sqlite:///{tmp_path}/m.db', experiment_name='e') tracking.target.registry_uri = uri for name in bad: diff --git a/packages/kumo-engine/test/test_mlops.py b/packages/kumo-engine/test/test_mlops.py index ea9974e7d..aeca7d9e9 100644 --- a/packages/kumo-engine/test/test_mlops.py +++ b/packages/kumo-engine/test/test_mlops.py @@ -36,26 +36,29 @@ def test_mlops_needs_an_artifact_root_it_can_write_files_to(tmp_path): kumo.MLOps(artifact_root=str(tmp_path), on_error='ignore') -def test_a_connector_registry_checks_the_name_when_mlops_is_built(tmp_path): +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') - connector = SimpleNamespace(sink_backend=SimpleNamespace( - check_model_location=check_model_location)) + tracker = SimpleNamespace(connector=SimpleNamespace( + sink_backend=SimpleNamespace( + check_model_location=check_model_location))) good = kumo.ModelRegistryLocation('USERS_CHURN', alias='candidate') - kumo.MLOps(artifact_root=str(tmp_path), - registry=kumo.ModelOutputConfig(connector, location=good)) + kumo.MLOps(str(tmp_path), tracking=tracker, registry=good) with pytest.raises(ValueError, match='bad name'): - kumo.MLOps(artifact_root=str(tmp_path), - registry=kumo.ModelOutputConfig( - connector, - location=kumo.ModelRegistryLocation('users-churn'))) + kumo.MLOps(str(tmp_path), tracking=tracker, + registry=kumo.ModelRegistryLocation('users-churn')) with pytest.raises(ValueError, match='model alias'): - kumo.MLOps(artifact_root=str(tmp_path), - registry=kumo.ModelOutputConfig( - connector, location=kumo.ModelRegistryLocation( - 'USERS_CHURN', alias='LATEST'))) + 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): diff --git a/packages/kumo-engine/test/test_model_agnostic_mlops.py b/packages/kumo-engine/test/test_model_agnostic_mlops.py index 648c1b66d..a73a30061 100644 --- a/packages/kumo-engine/test/test_model_agnostic_mlops.py +++ b/packages/kumo-engine/test/test_model_agnostic_mlops.py @@ -75,6 +75,8 @@ def on_trial(self, record, event): 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)) @@ -83,48 +85,62 @@ def finish(self, record, error): def test_any_run_is_tracked_from_its_spec(): tracker = _Tracker() session = TrackingSession(tracker) - spec = RunSpec(kind='custom', run_id='loop-1', params={'lr': 0.1}, + spec = RunSpec(kind='custom', params={'lr': 0.1}, inputs={'users': {'connector': 'S3'}}) session.start(lambda: spec) - session.on_event(Event('trial_started', 'loop-1', 0.0, trial=1, + 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.tags['kumo.parent_run'] == 'loop-1' + 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.run_id': 'loop-1', - 'kumo.bundle_id': 'b'} + assert finished.tags == {'kumo.kind': 'custom', 'kumo.bundle_id': 'b'} -def test_any_bundle_is_stored_logged_and_registered(tmp_path): +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')) - registry = kumo.ModelOutputConfig(SimpleNamespace(sink_backend=sink), - location=ModelRegistryLocation('TOY')) - settings = kumo.MLOps(str(tmp_path / 'root'), registry=registry) - tracker = _Tracker() + 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', run_id='loop-1')) + 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 and registered.version == '1' + assert error is None stored = tmp_path / 'root' / 'models' / manifest.id assert (stored / 'weight.txt').read_text() == '2.0' - assert ('log_model', (bundle, manifest.id, 'snap')) in tracker.calls - assert published == [(str(bundle), { - 'path': None, 'location': ModelRegistryLocation('TOY'), - 'metadata': {'kumo_bundle_id': manifest.id}, 'metrics': {'auc': 0.9} - })] + 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) diff --git a/packages/kumo-engine/test/test_scoring_defaults.py b/packages/kumo-engine/test/test_scoring_defaults.py index 6ccfdf08e..fd8812113 100644 --- a/packages/kumo-engine/test/test_scoring_defaults.py +++ b/packages/kumo-engine/test/test_scoring_defaults.py @@ -137,3 +137,27 @@ def test_unlisted_or_unreadable_defaults_are_ignored(fitted, tmp_path): assert ignored.result.scoring_defaults == kumo.ScoringDefaults() with pytest.raises(ValueError, match='must be specified'): ignored.predict(table, snapshot=snapshot) + + +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_snowflake_tracking.py b/packages/kumo-engine/test/test_snowflake_tracking.py index aebe648c5..7cbbacec9 100644 --- a/packages/kumo-engine/test/test_snowflake_tracking.py +++ b/packages/kumo-engine/test/test_snowflake_tracking.py @@ -3,6 +3,7 @@ # SPDX-License-Identifier: Apache-2.0 import inspect +import re import threading import time from unittest import mock @@ -81,7 +82,7 @@ def _replay(tracking, result=RESULT, error=None, during=None): def test_native_model_lineage_holds_the_run_open_during_the_fit(fake): lineage = [] _replay( - _tracking(database='ml', schema='runs', run_name='fit', + _tracking(database='ml', schema_name='runs', run_name='fit', native_model_lineage=True), during=lambda: lineage.extend(fake._instance.lineage_stack)) @@ -154,8 +155,8 @@ def test_an_existing_run_is_never_resumed(fake, taken, caplog): tracking = _tracking(session, run_name=run_name) _replay(tracking) - new = f"{run_name}_{CONTEXT.run_id.rsplit('-', 1)[-1][:8]}".upper() - assert tracking.run_id == new + 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' @@ -240,15 +241,43 @@ def sql(self, statement): 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', )]) - tracking = _tracking(None, connector=mock.Mock(session=lambda: session)) + connector = mock.Mock(spec=Connector, session=lambda: session) + tracking = _tracking(connector) tracking.check() - assert tracking.session is session + 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='a session or a connector'): + 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 index bcf9cf3f0..0d7bedbf1 100644 --- a/packages/kumo-engine/test/test_tracking_contract.py +++ b/packages/kumo-engine/test/test_tracking_contract.py @@ -23,8 +23,8 @@ from kumo.engine.tracking import Event, RunRecord from kumo.engine.trainer._run_spec import FitContext, fit_spec, fitted_spec -RUN = 'trainer-fit-1' -USER_TAGS = {'team': 'ents', 'kumo.run_id': 'overridden'} +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', @@ -36,7 +36,7 @@ 'version': '1790000000', 'version_kind': 'timestamp'}, } CONTEXT = FitContext( - run_id=RUN, query=QUERY, inputs=INPUTS, + query=QUERY, inputs=INPUTS, model_plan={ 'optimization': {'base_lr': [0.01, 0.001], 'max_epochs': 2}, 'training_job': {'tune_metric': 'mae'}, @@ -47,7 +47,7 @@ def _event(name, trial=None, epoch=None, metrics=None, **params): - return Event(name=name, run_id=RUN, time=0.0, trial=trial, epoch=epoch, + return Event(name=name, time=0.0, trial=trial, epoch=epoch, metrics=metrics or {}, params=params) @@ -125,7 +125,7 @@ def replay(tracker=None, events=EVENTS, result=RESULT, None if result is None else lambda: fitted_spec( session.spec, result, events[-1].metrics.get('seconds')), error) records = { - record.tags['kumo.run_id']: record + record.name: record for hook, record in recorder.calls if hook == 'on_trial' } spec = fit_spec(CONTEXT, tags=USER_TAGS) @@ -133,7 +133,7 @@ def replay(tracker=None, events=EVENTS, result=RESULT, 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[RUN] = run_record(spec) + records[FIT] = run_record(spec) return records @@ -196,7 +196,8 @@ def differ(what, have, want): and k not in (DROPPED_PARAMS, DROPPED_METRICS)), sorted(k for k in record.tags if k.startswith('kumo.'))) if adapter.nested_runs: - differ('parent', seen.parent, record.tags.get('kumo.parent_run')) + 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 @@ -215,7 +216,11 @@ def observe(): client = mlflow.MlflowClient() runs = client.search_runs( [e.experiment_id for e in client.search_experiments()]) - kumo_ids = {r.info.run_id: r.data.tags['kumo.run_id'] for r in runs} + 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( @@ -280,7 +285,7 @@ def bucket(key): runs['steps'][name] = (*runs['steps'].get(name, ()), (step, value)) for prefix, runs in names.items(): - observed[runs['tags']['kumo.run_id']] = Observed( + observed[prefix or FIT] = Observed( name=prefix, artifacts=run['artifacts'] if prefix is None else {}, **runs) @@ -300,8 +305,8 @@ def adapter(request) -> Adapter: def test_records_of_the_recorded_fit(): records = replay() - fit = records[RUN] - assert fit.tags == {'team': 'ents', 'kumo.kind': 'fit', 'kumo.run_id': RUN, + 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]', @@ -319,17 +324,15 @@ def test_records_of_the_recorded_fit(): assert json.loads(fit.artifacts['model_plan.json']) == CONTEXT.model_plan assert fit.inputs == INPUTS - oom = records[f'{RUN}-trial-2-attempt-1'] + 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.run_id': f'{RUN}-trial-2-attempt-1', - 'kumo.parent_run': RUN, 'kumo.status': 'failed', 'kumo.error': 'OutOfMemoryError: CUDA out of memory', }) - retry = records[f'{RUN}-trial-2-attempt-2'] + 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)), @@ -342,7 +345,7 @@ def test_adapter_records_the_same_runs(adapter, tmp_path): bundle_dir=tmp_path) expected = replay(adapter.make(), result=result) - assert expected[RUN].artifacts[PROVENANCE_FILE] == '{"bundle": "b"}' + assert expected[FIT].artifacts[PROVENANCE_FILE] == '{"bundle": "b"}' assert gaps(expected, adapter.observe(), adapter) == [] @@ -367,7 +370,7 @@ def test_adapter_records_a_failed_fit(adapter): expected = replay(adapter.make(), [*EVENTS[:-1], failed], None, ValueError('stop')) - assert expected[RUN].tags['kumo.error'] == 'ValueError: stop' + assert expected[FIT].tags['kumo.error'] == 'ValueError: stop' assert gaps(expected, adapter.observe(), adapter) == [] diff --git a/packages/kumo-engine/test/test_trainer_registry.py b/packages/kumo-engine/test/test_trainer_registry.py index e20b5f387..fc1fd6185 100644 --- a/packages/kumo-engine/test/test_trainer_registry.py +++ b/packages/kumo-engine/test/test_trainer_registry.py @@ -699,6 +699,20 @@ def test_invalid_export_alias_is_rejected_before_publishing( 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), @@ -716,17 +730,16 @@ def test_mlops_registry_registers_after_the_fit(tmp_path: Path, monkeypatch, monkeypatch.setattr( Trainer, '_fit', lambda self, *a, **k: TrainingResult( {}, {}, bundle, bundle_id='b1')) - registry = kumo.ModelOutputConfig( - _FakeConnector(sink), - location=ModelRegistryLocation('CHURN_GNN', alias='prod')) + 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'), registry=registry, - store=False)) + 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 @@ -740,8 +753,9 @@ def test_mlops_registry_registers_after_the_fit(tmp_path: Path, monkeypatch, trainer.fit( SimpleNamespace(query='PREDICT COUNT(*)'), snapshot=SimpleNamespace(id='s'), - mlops=kumo.MLOps(str(tmp_path / 'root'), registry=registry, - store=False, on_error='raise')) + mlops=kumo.MLOps(str(tmp_path / 'root'), tracking=tracker, + registry=registry, store=False, + on_error='raise')) assert isinstance(raised.value.__cause__, ZeroDivisionError) @@ -750,10 +764,10 @@ def test_mlops_registry_alias_is_checked_before_the_fit(tmp_path: Path, import kumo.engine as kumo monkeypatch.setattr(Trainer, '_fit', lambda *a, **k: 1 / 0) - registry = kumo.ModelOutputConfig( - _FakeConnector(_AliasingSink()), - location=ModelRegistryLocation('CHURN_GNN', alias='LATEST')) with pytest.raises(ValueError, match='model alias'): Trainer(ModelPlan()).fit( SimpleNamespace(query='q'), snapshot=SimpleNamespace(id='s'), - mlops=kumo.MLOps(str(tmp_path), registry=registry, store=False)) + mlops=kumo.MLOps( + str(tmp_path), tracking=_RegistryTracker(_AliasingSink()), + registry=ModelRegistryLocation('CHURN_GNN', alias='LATEST'), + store=False)) From c8490303b3069400a1a6a236d85e30787fffef3d Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Wed, 28 Oct 2026 22:11:47 +0000 Subject: [PATCH 04/16] Store, register and authenticate safely on every platform Store failures (snapshot, training table, model, asset) follow on_error like tracking and registration, and a model whose bundle or assets were not stored is never registered. Assets go to object stores without bucket-level calls, so prefix-scoped S3 rights suffice. A second MLflow registration into a Unity Catalog model reads search results whose tags are a method. result.export(cfg) registers a model through the MLOps tracker and returns the version that already holds its bundle. An ambient Databricks target points MLflow at the workspace for each call, so a profile works without environment variables, and a service principal's client ignores the token lent to MLflow on another thread. docs/mlops.md is restructured around a quickstart per platform, with the details in a reference section. Co-Authored-By: Claude Opus 5.5 --- .../src/kumo_databricks/session.py | 3 + packages/kumo-engine/docs/mlops.md | 497 +++++++----------- packages/kumo-engine/kumo/engine/_deps.py | 19 +- packages/kumo-engine/kumo/engine/_mlops.py | 10 +- .../kumo-engine/kumo/engine/_snapshot_io.py | 9 +- packages/kumo-engine/kumo/engine/graph.py | 6 +- .../engine/integrations/_mlflow_target.py | 10 +- .../kumo/engine/integrations/mlflow.py | 57 +- .../kumo/engine/pquery/training_table.py | 12 +- .../kumo/engine/trainer/trainer.py | 66 ++- .../kumo-engine/test/test_mlflow_model.py | 14 +- .../kumo-engine/test/test_mlflow_target.py | 48 +- .../test/test_model_agnostic_mlops.py | 81 +++ .../kumo-engine/test/test_snapshot_export.py | 14 + 14 files changed, 484 insertions(+), 362 deletions(-) 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-engine/docs/mlops.md b/packages/kumo-engine/docs/mlops.md index 061a9a253..1553fef8b 100644 --- a/packages/kumo-engine/docs/mlops.md +++ b/packages/kumo-engine/docs/mlops.md @@ -1,404 +1,289 @@ # MLOps with Kumo Engine -## Turning MLOps on +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 -`kumo.MLOps` says where artifacts go and what records them. Use it as a block, which notebooks and scheduled scripts alike can do, or pass it as `mlops=` on a single call: +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(artifact_root="s3://bucket/kumo") # or (connector, path), e.g. a volume or stage +cfg = kumo.MLOps("artifacts", tracking=MlflowTracking("churn"), + registry=kumo.ModelRegistryLocation("users-churn", alias="candidate")) with kumo.mlops(cfg): snapshot = graph.snapshot() + result = kumo.Trainer(pquery.suggest_model_plan(snapshot)).fit(pquery.generate_training_table(snapshot), + snapshot=snapshot, scoring_defaults=kumo.ScoringDefaults(binary_classification_threshold=0.5)) + scores = result.predict(snapshot).frame + print(result.registered_model.version, result.tracking.run_id, len(scores)) ``` -- **`artifact_root`** is required. It is 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. -- **Volumes and stages:** `(DatabricksConnector(...), "/Volumes////kumo")` and `(SnowflakeConnector(...), "@DB.SCHEMA.STAGE/kumo")` work as roots once the installed kumo-databricks or kumo-snowflake sink holds files at those paths; with an older plugin, building `kumo.MLOps` raises. -- **`export`** names the artifact kinds stored under the root when they are produced: `snapshot`, `training_table` and `model` (default: all). -- **`tracking`** and **`registry`** record fits and register models. -- **`on_error`** says what a tracking or registration failure in `fit` does. `"warn"` (default): the call warns, naming the error, and returns as usual. `"raise"`, for scheduled jobs: the call finishes and stores its result first, then raises a `RuntimeError` whose `__cause__` is the failure; after `fit`, `trainer.result` holds the fitted model. -- **Precedence:** a call's `mlops=` replaces the enclosing block, and an inner block replaces an outer one. A block covers the calling thread or asyncio task. -- **Off by default.** Without a block or `mlops=`, every call behaves as it always has. Plain exports (`export(config)`, `export_config=`, `output_config=`) never write to the artifact root. A path export that finds a complete copy of the same artifact at its destination reuses it instead of uploading again (logged at info): a snapshot with the same bytes, a training table or prediction directory with `_SUCCESS` and the same `_KUMO_EXPORT` identity, or a model bundle with the same files and `_provenance.json`. - -## Lineage - -Every artifact carries a `_provenance.json` manifest, metadata only. The leading underscore keeps Spark and Hive Parquet readers from treating it as data. - -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. - -### Training tables - -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. +### Databricks -`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. +In a notebook or job on Databricks Runtime 17.3 LTS ML or later: -With MLOps on, each table is stored once under `/training_tables//`. - -### Models +```python +import kumo.engine as kumo +from kumo.engine.integrations.mlflow import MlflowTracking -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. +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) +``` -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`). +### SageMaker -With MLOps on, each bundle is stored under `/models//`. A complete copy that is already there is kept; a partial one is replaced. +With the managed MLflow tracking server's ARN (setup: [SageMaker managed MLflow](#sagemaker-managed-mlflow)): -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. +```python +import kumo.engine as kumo +from kumo.engine.integrations.mlflow import MlflowTracking, SageMakerMLflow -### Artifacts, bundles and dependencies +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) +``` -- An **artifact** is anything with its own id and `_provenance.json`: a snapshot, a training table, a model, a prediction run, or an asset. -- A model's **bundle** holds only the model's own state: weights, encoders, the training spec and the holdout. -- Its manifest's **`parents`** are lineage: the training table and the warm-start model it came from. -- Its **`dependencies`** are what `predict` needs besides the bundle and a snapshot. The snapshot is an input to each call, never a dependency. Each entry has a `kind`: - - `predictor`: the entry point that scores the bundle; - - `package`: the `kumo-engine` version and its extras, with `[text]` only when the model encodes a text column; - - `asset`: a content-addressed file set, by `asset_kind`, `name`, `version` and `sha256`. -- Bundles fitted before dependencies were recorded list none. They load as before, with the built-in predictor. -- **Assets** are not copied into bundles. With MLOps on, each asset is stored once under `/assets///` with its own `_provenance.json`, and every model that uses it shares that copy. The dependency records where it was stored. -- **Resolving assets:** `Trainer.load`, the MLflow pyfunc and scoring jobs look for each asset at the location the dependency records, then under the artifact root of the enclosing `kumo.mlops` block, then in the embedding-model store. They check its sha256 and cache non-local copies under `~/.cache/kumo/assets//`. If an asset is stored but cannot be found, loading fails with `MissingDependencyError`, which names the asset kind, name and sha256. A bundle without stored assets resolves its models from the store, as before. -- **Text encoders** are the first assets (`asset_kind="text_encoder"`): the word vectors of each text column's embedding model. Word-vector models (`GloVeVectors`) are stored. Other `TokenEmbeddingModel`s are recorded by fingerprint and must be registered wherever the model is loaded. +### Snowflake -To train with your own word vectors and serve the model without registering them again, register them under a name, fit with MLOps on, and load the stored model anywhere that can reach the artifact root: +Tracking and the Snowflake Model Registry through one connector (`kumo-engine[snowflake,s3]`): ```python -store = kumo.EmbeddingModelStore() -store.add("colors", kumo.GloVeVectors("/data/vectors")) # a .zip release or a directory of converted files -plan.column_processing.encoder_overrides = {"users.bio": TextEmbedding("colors")} -with kumo.mlops(kumo.MLOps(artifact_root="s3://bucket/kumo")): - result = kumo.Trainer(plan, embedding_models=store).fit(table, snapshot=snapshot) -# Elsewhere, with no store: the vectors come from /assets/text_encoder//. -trainer = kumo.Trainer.load_from_connector(kumo.S3Connector(root_dir="s3://bucket/kumo"), - path=f"models/{result.id}") +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) ``` -### Scoring defaults +To serve only a better model, move its alias yourself: [Promote a model](#promote-a-model). -A model can carry the prediction arguments `predict` uses when a call omits them: +## How it fits together -```python -from kumo.engine import ModelOutputConfig, ScoringDefaults +- 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. -tuned = result.with_scoring_defaults( - ScoringDefaults(binary_classification_threshold=0.37)) -tuned.predict(snapshot) # thresholds at 0.37 -tuned.scoring_defaults # ScoringDefaults(binary_classification_threshold=0.37, ...) -``` +| id | names | where to find it | +|---|---|---| +| `snapshot.id` | a graph snapshot | `/snapshots///`; `result.snapshot_id` | +| `table.id` | a training table | `/training_tables//` | +| `result.id` | a model (its bundle id) | `/models//`; the `kumo.bundle_id` run tag; `kumo_bundle_id` on the registered version | +| `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 | -`with_scoring_defaults` returns a `TrainingResult` for a copy of the bundle with its own id. The copy's manifest embeds the original's in `parents.derived_from`. The threshold applies to binary and multilabel classification, `num_classes_to_return` to multiclass, multilabel and link prediction. Explicit `predict` arguments and a query's `TOP k` win; `ScoringDefaults()` removes the defaults. `Trainer.export(config)` publishes a fitted or loaded bundle like `fit(export_config=...)`. +Find runs with the platform's UI, or by tag: `kumo.kind` (`fit`, `trial`) and `kumo.bundle_id`. -### Registering models +## Turning MLOps on -With `registry=ModelOutputConfig(connector, location=...)`, every fit in the block registers its model in the connector's registry (Snowflake Model Registry, or Unity Catalog through a `DatabricksConnector`) after it is stored and tracked: +`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 -from kumo.engine import ModelOutputConfig, ModelRegistryLocation - -cfg = kumo.MLOps("s3://bucket/kumo", registry=ModelOutputConfig( - registry, location=ModelRegistryLocation("CHURN_GNN", alias="challenger"))) +cfg = kumo.MLOps("s3://bucket/kumo", tracking=..., registry=..., on_error="raise") with kumo.mlops(cfg): - result = trainer.fit(training_table, snapshot=snapshot) -result.registered_model # PublishedModel(name='CHURN_GNN', version='V3', uri=...) - -# Scoring jobs load whichever version an alias points at. -trainer = kumo.Trainer.load_from_connector(registry, location=ModelRegistryLocation( - "CHURN_GNN", alias="challenger")) + snapshot = graph.snapshot() +trainer.fit(table, snapshot=snapshot, mlops=cfg) ``` -- Registration is independent of `export`. If it fails, `fit` warns with the error and returns the result with `registered_model=None` (with `on_error="raise"`, it raises instead), and a tracked fit run ends failed with `kumo.error`; an invalid alias or a connector that cannot publish models raises before training, and an invalid alias or model name raises when `kumo.MLOps` is built. Snowflake model names are Snowflake identifiers (a letter or underscore, then letters, digits, underscores or `$`, e.g. `USERS_CHURN`; or a double-quoted name); Unity Catalog model names are at most 255 characters with no space, `/`, `.` or control character, optionally qualified with the connector's own catalog and schema. -- `result.export(ModelOutputConfig(connector, location=...))` on a `TrainingResult`, or `trainer.export(...)` for its latest result, publishes the model with or without a block and also accepts an `alias`; `Trainer.load_from_connector` loads it back. `fit(export_config=...)` is shorthand for fitting then exporting, and its version is `result.published_model`. -- An alias is a letter or underscore followed by letters, digits or underscores; `DEFAULT`, `FIRST`, `LAST`, `LATEST` and `v` are reserved. The alias is set on the new version after publishing, moving it from any version that held it; if that fails, the publish still stands, with a warning. Snowflake allows one alias per version; if moving an alias fails, it is put back on the version that held it. Snowflake's `DEFAULT`, `FIRST` and `LAST` can be loaded but not set. File, S3 and GCS connectors have no aliases. -- Each registered version records only pointers: the `kumo_engine` marker and `kumo_bundle_id`, as version metrics in the Snowflake Model Registry and as version tags in Unity Catalog through a `DatabricksConnector`. (A version registered through MLflow, below, gets a `kumo_bundle_id` version tag only.) The bundle's own `_provenance.json` holds its lineage. `list_models` and `fetch_model` return `kumo_bundle_id` in `PublishedModel.metadata`; `load_from_connector` warns if the fetched bundle differs from it. Through a `DatabricksConnector`, publishing a bundle that a ready version of the model already holds (same `kumo_bundle_id`) returns that version, logged at info, instead of adding another; an alias then moves to it. +- **`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; 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. -### Promote a Kumo model +### Reading artifacts back -To register every fit as a candidate and promote only a better one, register under a candidate alias (or none) and move the serving alias yourself with the registry's own API. Compare the new model's test metrics (`result.metrics`, also on its tracked fit run) with the current champion's (`load_from_connector(...).result.metrics`): +Pass the `kumo.MLOps` where a connector goes, with paths relative to the artifact root: ```python -champion = kumo.Trainer.load_from_connector( - registry, location=ModelRegistryLocation("CHURN_GNN", alias="champion")) -better = result.metrics["test_auroc"] > champion.result.metrics["test_auroc"] -version = result.registered_model.version +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) ``` -Unity Catalog (registered through MLflow or a `DatabricksConnector`) and SageMaker managed MLflow: +`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()`). -```python -from mlflow import MlflowClient +## Train and register -if better: - MlflowClient().set_registered_model_alias("main.ml.churn_gnn", "champion", version) -``` +`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. -`set_registered_model_alias` moves the alias from the version that held it. For models registered through MLflow, the champion's metrics are also on its run: `client.get_run(client.get_model_version_by_alias(name, "champion").run_id).data.metrics`. +- `result.registered_model` is the version `MLOps(registry=)` registered; `result.published_model` the one `fit(export_config=...)` published. +- `result.tracking` is the tracker's run: `TrackingRef(platform="databricks", run_id="...", url="https://.../ml/experiments/.../runs/...")`. The model's `_provenance.json` records the same. +- `fit(scoring_defaults=kumo.ScoringDefaults(...))` stores prediction defaults in the bundle before it is stored and registered: `binary_classification_threshold` (binary and multilabel classification) and `num_classes_to_return` (multiclass, multilabel, link prediction). `predict` and the MLflow pyfunc apply them when a call omits them; explicit arguments and a query's `TOP k` win. `result.with_scoring_defaults(...)` retunes a fitted model as a new bundle whose manifest names the original in `parents.derived_from`. +- `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. -Snowflake Model Registry, where a version holds at most one alias, so the alias is unset first: +**Model names** follow the registry's rules, checked when `kumo.MLOps` is built: -```sql -ALTER MODEL DB.SCHEMA.CHURN_GNN VERSION champion UNSET ALIAS; -ALTER MODEL DB.SCHEMA.CHURN_GNN VERSION V3 SET ALIAS = champion; -``` +| 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` | -Scoring jobs that load `alias="champion"` (or `models:/@champion`) pick up the new version on their next run. +An alias is a letter or underscore followed by letters, digits or underscores; `DEFAULT`, `FIRST`, `LAST`, `LATEST` and `v` are reserved. -### Moving snapshots between machines +### Promote a model -To build an index on one machine and train or score on another, export it to a file destination and load it there: +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: ```python -from kumo.engine import GraphSnapshot, S3Connector, SnapshotOutputConfig +champion = kumo.Trainer.load_from_connector(cfg, location=kumo.ModelRegistryLocation(name, alias="champion")) +if result.metrics["test_auroc"] > champion.result.metrics["test_auroc"]: + mlflow.MlflowClient().set_registered_model_alias(name, "champion", result.registered_model.version) +``` + +In the Snowflake Model Registry a version holds at most one alias, so unset it first: -location = graph.snapshot().export(SnapshotOutputConfig(S3Connector(), "s3://bucket/snapshots")) -# On the training machine: -snapshot = GraphSnapshot.load_from_connector(S3Connector(), path=location) +```sql +ALTER MODEL ML.MODELS.CHURN_GNN VERSION champion UNSET ALIAS; +ALTER MODEL ML.MODELS.CHURN_GNN VERSION V3 SET ALIAS = champion; ``` -- **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. +Jobs that load `alias="champion"` pick up the new version on their next run. -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. +### Trackers -### Scoring without the graph +`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. -A scoring job does not need the training job's `Graph` code. It loads the snapshot the training job exported and the model, then scores: +`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`. -```python -snapshot = GraphSnapshot.load_from_connector(store, path=location) -trainer = Trainer.load_from_connector(store, path="models/") -result = trainer.predict(snapshot, binary_classification_threshold=0.5) -``` +## Score -- `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, ...)`. +`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. -## Progress events +### MLflow pyfunc -`fit(callbacks=[...])` calls each callback with every progress -`kumo.engine.tracking.Event`: `fit_started`, then per search trial -`trial_started`, `epoch_ended` and `trial_ended`, then `fit_ended`. Ended events -carry `params["status"]` and, on failure, a short `params["error"]`. A callback -that raises is logged and not called again; it never fails the fit. For JSON -lines on stdout: +A model registered through `MlflowTracking` is an MLflow pyfunc: it serves with `mlflow.pyfunc.load_model(uri)`, `spark_udf` and Model Serving. ```python -import json - -trainer.fit(training_table, snapshot=snapshot, - callbacks=[lambda e: print(json.dumps(e.to_dict()))]) +model = mlflow.pyfunc.load_model("models:/main.ml.users_churn@champion") +scores = model.predict(pd.DataFrame({"ENTITY": [1, 2, 3]}), params={ + "snapshot": "/Volumes/main/ml/kumo/snapshots//", + "prediction_time": "2024-06-01"}) # temporal queries ``` -## Experiment tracking +- Input rows have an `ENTITY` column and optionally a `TIMESTAMP`; scores come back in input order, one per scored row. Binary classification needs `params={"binary_classification_threshold": ...}` unless the model carries scoring defaults. +- 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). -`kumo.MLOps(tracking=...)` records each fit in an experiment tracker: one run -per fit, with each hyperparameter-search trial attempt as a child run (MLflow) -or inside the fit run (Snowflake). Tracking is on only -inside a `kumo.mlops(cfg)` block or for a `fit` given `mlops=cfg`. +## Credentials and running off-platform + +`MlflowTracking(target=...)` and `SnowflakeTracking(connector)` take the same identities as the connectors, so a fit on a laptop or any VM can record to Databricks, SageMaker or Snowflake. An identity with no credentials (or no identity) uses the platform's ambient chain (a notebook, job or instance role, `aws configure`, `databricks auth login`); one with credentials uses exactly those. ```python -from kumo.engine.integrations.mlflow import MlflowTracking # kumo-engine[mlflow] +from kumo.engine import DatabricksConnector, DatabricksIdentity, S3Identity, SnowflakeConnector, SnowflakeIdentity -cfg = kumo.MLOps(artifact_root="s3://bucket/kumo", - tracking=MlflowTracking(experiment_name="users-spend")) -with kumo.mlops(cfg): - result = trainer.fit(training_table, snapshot=snapshot) +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", + identity=SnowflakeIdentity(account="", username="", private_key="")), + experiment_name="USERS_CHURN") ``` -- A tracking or registration failure does not fail the fit: it is logged, and reported once as a `UserWarning` naming the full error, so both notebooks and job logs show it. With `on_error="raise"` the call raises once instead, after storing its result, with the error as the cause. Only global rank 0 tracks. -- Tracker calls run on a background thread, so a slow or unreachable tracking server does not hold up training. At the end of the fit, `fit` waits at most 30 s for tracking to finish; the rest continues in the background. -- During the fit, MLflow HTTP calls use 2 retries and a 30 s timeout unless `MLFLOW_HTTP_REQUEST_MAX_RETRIES` / `MLFLOW_HTTP_REQUEST_TIMEOUT` are set. -- `MlflowTracking()` records where MLflow's own configuration points when it is built (`MLFLOW_TRACKING_URI`, or `mlflow.set_tracking_uri`): on Databricks this is automatic, and with nothing configured MLflow writes a local store in the working directory. `MlflowTracking(target=...)` names the server and the credentials instead; see [Running from a laptop](#running-from-a-laptop). -- `kumo.MLOps(...)` checks the tracker when it is built: one experiment lookup, plus one registry listing and (Unity Catalog) the model's catalog and schema when it registers. An unreachable server, refused credentials or a missing catalog or schema raise there, before any work. -- The `View run ... at` links MLflow prints do not open on SageMaker; after a tracked fit kumo-engine warns how to open the UI: from SageMaker Studio, or at the URL `aws sagemaker create-presigned-mlflow-tracking-server-url --tracking-server-name ` prints. -- MLflow rejects `@` in metric names, so link-prediction metrics such as `test_map@10` are logged as `test_map_at_10`. -- Per-epoch metrics use step = epoch. After the fit, `tracking.run_id` is the fit run's MLflow run id. -- The fit run is named `fit-` (for example `fit-20261027T061838Z`) unless `run_name` is given; trial runs are `trial--attempt-`. -- The fit's source tables are logged as MLflow run inputs, read from the training table's embedded snapshot manifest (plus a custom training table's source). Databricks tables are logged as Delta tables at the snapshot's Delta version, so Unity Catalog links each table to a model registered from the run; file and object-store sources by URI. Other warehouse tables and Custom SQL sources are not logged as inputs; they stay in the model's `_provenance.json`. +`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. `examples/mlops/smoke.py databricks|sagemaker|snowflake` is a manual live check per platform. -Snowflake records the fit in Snowflake Experiment Tracking (Snowsight: AI & ML -» Experiments), using a Snowpark session: +### SageMaker managed MLflow -```python -from kumo.engine.integrations.snowflake import SnowflakeTracking # kumo-engine[snowflake] +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. +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`. -cfg = kumo.MLOps(artifact_root="s3://bucket/kumo", - tracking=SnowflakeTracking(session, experiment_name="users_spend")) -with kumo.mlops(cfg): - result = trainer.fit(training_table, snapshot=snapshot) -``` +## Reference -- Snowflake has no child runs or tags, so trial attempts are logged in the fit run as `trial-1-attempt-1/` params and metrics, and tags are params of the same name. -- `native_model_lineage=True` links a model the fit publishes through Snowflake lineage. -- Names beyond Snowflake's 200 metrics and 1000 params per run are counted in `kumo.dropped_metrics` / `kumo.dropped_params`. -- If a Snowflake run is already active in the process, the fit is not tracked there; an existing run is never resumed (a taken `run_name` gets a suffix). -- Snowflake Experiment Tracking has no dataset inputs, so source tables are not recorded there. -- The Snowpark session's parameters are left as they are. After the fit, `tracking.run_id` is the fit run's name. +### Install -Runs carry pointers, not lineage: +- 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 (cu130) 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, MLflow 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`. +- **MLflow:** the `mlflow` extra installs `mlflow-skinny` 3.0.1 or later; 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`). -| tag | on | value | -|---|---|---| -| `kumo.kind` | every run | `fit` or `trial` | -| `kumo.run_id` | every run | the kumo run id (trials: `-trial--attempt-`) | -| `kumo.bundle_id` | a finished fit run | the model's bundle id (`models//` under the artifact root) | -| `kumo.parent_run`, `kumo.status`, `kumo.error` | trial runs; `kumo.error` also on a failed fit (including a failed model logging or registration) | outcome | +### Exporting and loading artifacts -The model's `_provenance.json` records the tracked fit as `tracking.run_id` -(the `kumo.run_id` tag) and is logged on the fit run with `query.pql` and -`model_plan.json`. +Every artifact has one way out and one way back. A `kumo.MLOps` stands in for the connector to read the artifact root. -Which id is which: - -| id | what it names | where to read it | +| Artifact | Export | Load | |---|---|---| -| kumo run id | one `fit` call | the `kumo.run_id` tag; `tracking.run_id` in the model's `_provenance.json` | -| bundle id | the fitted model's content | `result.id`; the `kumo.bundle_id` tag; `kumo_bundle_id` in the registry; `models//` under the artifact root | -| tracker run id | the run in the tracker | the tracker's `run_id` attribute (`tracking.run_id` on the `MlflowTracking` object: the MLflow run id; Snowflake: the run name) | -| registry version | the registered model version | `result.registered_model.version` | +| 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=...)` | -To find a bundle's MLflow fit run, search for its tag: ``mlflow.search_runs(filter_string="tags.`kumo.bundle_id` = '' AND tags.`kumo.kind` = 'fit'", search_all_experiments=True)``. +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. -Other systems plug in through the `kumo.engine.tracking.Tracking` protocol, -which receives ready-built `RunRecord`s so every tracker records the same -names: `start(record)` with the fit run, `on_trial(record, event)` with a trial -attempt's run after each of its events, and `finish(record, error)` with the -fit run's final record (`None` when tracking stopped early). A tracker that -raises is not called again for that fit and its run is ended as failed. -Optional members: `prepare()`, `tags`, `check(registry)` (called when `kumo.MLOps` is built) and `log_model(result, snapshot, registry)` (see MLflow models below). Params are the resolved model plan -(fit) or searched hyperparameters and `attempt` (trial). `RunRecord.inputs` -maps each source table to its `connector`, `location`, `version` and -`version_kind`. A tracker's `tags=` are added as given, except `kumo.*` names -and param names. +### Lineage and manifests -### MLflow models +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. -With `registry=ModelRegistryLocation(name, alias=...)` and an `MlflowTracking` tracker, every fit in the block logs its model on the fit run as an MLflow pyfunc model and registers it as a new version of `name` (Unity Catalog with `MLFLOW_REGISTRY_URI=databricks-uc`, SageMaker managed MLflow, or any MLflow registry). One alias is set on the new version. A bundle the registry already holds under `name` (same `kumo_bundle_id` version tag) is not registered again: its version is returned, the alias moves to it, and an INFO log says it was reused. A bare `ModelRegistryLocation` without such a tracker, a fixed `version`, an invalid alias, or a name the registry (`mlflow.get_registry_uri()`) rejects raises when `kumo.MLOps` is built: +- **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. -- SageMaker managed MLflow (a tracking server or MLflow app ARN): letters, digits and hyphens, starting and ending with a letter or digit, at most 57 letters and digits (`^[a-zA-Z0-9](-*[a-zA-Z0-9]){0,56}$`), so no `_` or `.`: `users-churn`. -- Unity Catalog (`databricks-uc`): `catalog.schema.model`, each part at most 255 characters with no space, `/` or control character: `main.ml.users_spend`. On a cluster a shorter name takes the current catalog and schema. -- Other MLflow registries: no extra rule. +### Artifacts, bundles and dependencies -The rules conflict, so no one name fits every registry: SageMaker forbids the `.` Unity Catalog needs between its parts and the `_` Snowflake needs instead of `-`. The examples use `users-churn` for SageMaker and other MLflow registries, `..users_churn` for Unity Catalog and `USERS_CHURN` for Snowflake. +- 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 -import mlflow -import pandas as pd -from kumo.engine import ModelRegistryLocation -from kumo.engine.integrations.mlflow import MlflowTracking, load_model - -cfg = kumo.MLOps("s3://bucket/kumo", - tracking=MlflowTracking(experiment_name="/Shared/users_spend"), - registry=ModelRegistryLocation("main.ml.users_spend", alias="champion")) +store = kumo.EmbeddingModelStore() +store.add("colors", kumo.GloVeVectors("/data/vectors")) +plan.column_processing.encoder_overrides = {"users.bio": TextEmbedding("colors")} with kumo.mlops(cfg): - result = trainer.fit(training_table, snapshot=snapshot) -result.registered_model # PublishedModel(name='main.ml.users_spend', version='4', uri='models:/main.ml.users_spend/4') - -trainer = load_model("models:/main.ml.users_spend@champion") # the Trainer back -model = mlflow.pyfunc.load_model("models:/main.ml.users_spend@champion") -scores = model.predict(pd.DataFrame({"ENTITY": [1, 2, 3]}), params={ - "snapshot": "s3://bucket/kumo/snapshots//", - "prediction_time": "2024-06-01", # temporal queries - "binary_classification_threshold": 0.5, # binary classification only -}) + 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 ``` -- `fit` waits for the upload and registration before it returns, for up to an hour, with a 600 s MLflow HTTP timeout during the upload (unless `MLFLOW_HTTP_REQUEST_TIMEOUT` is set). If they fail, the tracking thread stops or the hour passes, `fit` warns once (`Could not log the model: ...` or `Could not register the model: ...`, with the full error) and still returns its result with `registered_model=None` (with `on_error="raise"`, it raises instead), the run ends `FAILED` with the error in `kumo.error`, and a model that finishes uploading later is not registered. -- `MlflowTracking(log_model=True)` logs the model without registering it. `log_model(trainer, snapshot=...)` logs a fitted or loaded model in any run, with the bundle's source tables as run inputs. -- The model's metadata, the registered version's tags and `result.registered_model.metadata` hold `kumo_bundle_id`. The model's `provenance` artifact is the bundle's `_provenance.json`. Its pip requirement is `kumo-engine[s3,gcs]` of the same version, plus `[text]` when the model encodes a text column. The pyfunc scores with the predictor the bundle's manifest names. It resolves the model's assets, such as non-default word vectors, as described in [Artifacts, bundles and dependencies](#artifacts-bundles-and-dependencies). -- The pyfunc model loads a snapshot on its first `predict` with that URI and keeps it for the life of the process: each serving replica, and each Spark executor of `mlflow.pyfunc.spark_udf`, downloads the snapshot once and keeps every snapshot it has loaded (no eviction), so its disk must hold them. With `serving_snapshot=` (on `log_model` or `MlflowTracking`), each replica downloads that snapshot when it starts, and `predict` uses it when no `snapshot` param is given. `params={"population": True}` scores every entity and returns the whole frame in memory; for large populations use `Trainer.predict` with a `PredictionOutputConfig`. Each `predict` call works in its own scratch directory and removes it before returning. -- Input rows have an `ENTITY` column and, optionally, a `TIMESTAMP` column. Scores come back in the input's row order, one per input row (repeated rows included); rows whose entity is not in the snapshot are left out. A temporal query scores entities at a time: pass `params={"prediction_time": ...}` or a `TIMESTAMP` column; without either, `predict` fails with `temporal query: pass params={'prediction_time': ...} or a TIMESTAMP column`. Binary classification needs `params={"binary_classification_threshold": ...}` unless the model's `scoring_defaults` set one. -- Serving needs an image with kumo-engine already installed; MLflow's stock serving containers cannot install it. Serve with `mlflow models serve -m --env-manager local` (on SageMaker endpoints, `-m /opt/ml/model -p 8080`). SageMaker batch transform takes JSON input (`dataframe_split` plus `params`), not CSV, which carries no params. -- GPU real-time endpoints on SageMaker need an inference image whose NVIDIA driver supports CUDA 12.6. The default GPU hosting image ships driver 470, so set a newer `InferenceAmiVersion` on the production variant, for example `al2023-ami-sagemaker-inference-gpu-4-1`. Not yet tested live. -- The `mlflow` extra installs `mlflow-skinny` (3.0.1 or later), the client without the MLflow server's dependencies. SageMaker managed MLflow servers run MLflow 3.0; a 3.17 client against a 3.0 server handles logging, registration, aliases and `load_model`. -- On Databricks, use Databricks Runtime 17.3 LTS ML or later: its torch already satisfies kumo-core, so installing kumo-engine keeps the runtime's torch, torchvision and MLflow. kumo-engine accepts pyarrow 16 through 24, so a runtime's pyarrow in that range (19 on 17.3 LTS ML) is kept too; a custom CA bundle for S3 needs pyarrow 21 or later. On 16.4 LTS ML pip replaces torch 2.6 and breaks the preinstalled torchvision. -- On Databricks, the parent folder of an MLflow experiment path must exist. Registering to Unity Catalog from outside Databricks also needs `boto3`, which the `mlflow` extra does not install. +### Moving snapshots between machines -#### Which registry +- **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. -Pick one per registered model, since the two store different formats. On Databricks both work; use the MLflow model unless only kumo-engine jobs load it, since it also serves through `mlflow.pyfunc` (Model Serving, `spark_udf`). Both record only the `kumo_bundle_id` pointer; lineage stays in the bundle's `_provenance.json`. Both accept the same aliases (`kumoapi.plugin.check_model_alias`). +### Custom trackers -| | Connector registry | MLflow model | -|---|---|---| -| Register | `kumo.MLOps(registry=ModelOutputConfig(connector, location=ModelRegistryLocation(name, alias=...)))` | `kumo.MLOps(tracking=MlflowTracking(), registry=ModelRegistryLocation(name, alias=...))` | -| Load | `Trainer.load_from_connector(connector, location=...)` | `mlflow.pyfunc.load_model(uri)` or `load_model(uri)` | -| Promote | the registry's own API (see [Promote a Kumo model](#promote-a-kumo-model)) | MLflow's `set_registered_model_alias` | -| Databricks | `DatabricksConnector` (Unity Catalog) | Unity Catalog through MLflow | -| SageMaker | not available | SageMaker managed MLflow (`MlflowTracking(target=SageMakerMLflow(arn), ...)`) | -| Snowflake | `SnowflakeConnector` (Snowflake Model Registry) | not available | -| Plain VM | not available (store under the artifact root instead) | any MLflow tracking server | - -#### SageMaker managed MLflow - -1. Create a tracking server once and reuse it; creation takes about 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 S3 prefix). -2. Give the identity that runs the job (notebook, VM role or pipeline execution role) `sagemaker-mlflow:*` on the tracking server ARN, and S3 read and write on the artifact store prefix and on the kumo artifact root. -3. Install `pip install "kumo-engine[mlflow,sagemaker]"` (the MLflow client, S3 and the `sagemaker-mlflow` plugin) and pass `MlflowTracking(target=SageMakerMLflow("arn:aws:sagemaker:::mlflow-tracking-server/"), experiment_name=...)`; the registry is the same server. Without an identity, the job's AWS identity (a role or `aws configure`) signs; `MLFLOW_TRACKING_URI` need not be set. The plugin signs and routes with the process-wide tracking URI, so each MLflow call kumo-engine makes sets it to the ARN and restores it afterwards: other threads that call MLflow during a fit see the SageMaker URI. -4. Model names use letters, digits and hyphens only (`users-churn`); `kumo.MLOps` checks this when it is built. -5. The run links MLflow prints 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. - -On S3 the artifact root, source tables and prediction outputs are plain `s3://` locations read and written with the job's AWS identity (instance profile, `AWS_PROFILE` or the job role): +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. -```python -from kumo.engine import PredictionOutputConfig, S3Connector - -data = S3Connector(root_dir="s3:///data") # data["orders.parquet"] reads s3:///data/orders.parquet; data["orders"] a directory of files -cfg = kumo.MLOps("s3:///kumo", - tracking=MlflowTracking(target=SageMakerMLflow(""), experiment_name="users-churn"), - registry=ModelRegistryLocation("users-churn", alias="candidate")) -output = PredictionOutputConfig(S3Connector(), output_paths={"predictions": "s3:///scores/"}) -store = S3Connector() # monitoring: load with absolute s3:// paths -``` - -### Running from a laptop +### Progress events -`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 uses -the platform's ambient chain (a notebook, job or instance role, `aws -configure`, `databricks auth login`); one with credentials uses exactly -those. +`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()))]`. -```python -import kumo.engine as kumo -from kumo.engine import ModelRegistryLocation -from kumo.engine.connector import DatabricksConnector, DatabricksIdentity, S3Identity, SnowflakeConnector, SnowflakeIdentity -from kumo.engine.integrations.mlflow import MlflowTracking, SageMakerMLflow -from kumo.engine.integrations.snowflake import SnowflakeTracking - -# Databricks MLflow and Unity Catalog, as a service principal (kumo-engine[mlflow,databricks]). -workspace = DatabricksConnector(catalog="main", schema_name="ml", warehouse_id="", - identity=DatabricksIdentity(host="https://", client_id="", client_secret="")) -cfg = kumo.MLOps("s3:///kumo", - tracking=MlflowTracking(target=workspace, experiment_name="/Users//users_churn"), - registry=ModelRegistryLocation("main.ml.users_churn", alias="candidate")) - -# SageMaker managed MLflow, with an access key or a role to assume (kumo-engine[mlflow,sagemaker]). -cfg = kumo.MLOps("s3:///kumo", - tracking=MlflowTracking(target=SageMakerMLflow("", identity=S3Identity(role_arn="")), - experiment_name="users-churn"), - registry=ModelRegistryLocation("users-churn", alias="candidate")) - -# Snowflake Experiment Tracking, as a key-pair user (kumo-engine[snowflake]). -snowflake = SnowflakeConnector(warehouse="", database="ML", schema_name="RUNS", - identity=SnowflakeIdentity(account="", username="", private_key="")) -cfg = kumo.MLOps("s3:///kumo", tracking=SnowflakeTracking(connector=snowflake, experiment_name="USERS_CHURN")) -``` +### Limits -- Each target is checked when `kumo.MLOps` is built, so wrong credentials or a missing catalog, schema or experiment folder parent fail before the fit starts. -- MLflow has no per-client credentials: its Databricks store reads `DATABRICKS_*` variables or a profile, and the SageMaker plugin signs with the process-wide tracking URI. So with explicit credentials, each MLflow call kumo-engine makes sets `DATABRICKS_HOST`, a short-lived `DATABRICKS_TOKEN` (never the client secret) and the process-wide tracking URI `databricks` (Unity Catalog registration downloads the logged model through it), or the process-wide tracking URI and the AWS session, and restores them when the call returns. With explicit credentials, don't call MLflow or Databricks yourself on other threads during a fit. Ambient identities change no credentials; a SageMaker target still sets the process-wide tracking URI for each call, so with a SageMaker target, other threads calling MLflow during a fit see the SageMaker URI. -- SageMaker model artifacts are uploaded and downloaded with the identity's credentials. `load_model(uri, tracking=MlflowTracking(target=..., experiment_name=...))` and `log_model(..., tracking=...)` take the same target. -- `examples/mlops/smoke.py databricks|sagemaker|snowflake` is a manual live check per platform: a small fit recorded and registered with an explicit identity, then loaded back, in a few minutes on a CPU. It is not run in CI. +- 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/_deps.py b/packages/kumo-engine/kumo/engine/_deps.py index 189614856..f8d6b5713 100644 --- a/packages/kumo-engine/kumo/engine/_deps.py +++ b/packages/kumo-engine/kumo/engine/_deps.py @@ -10,7 +10,6 @@ import os import shutil import tempfile -import warnings from pathlib import Path from typing import ( Any, @@ -126,6 +125,7 @@ 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) @@ -144,7 +144,10 @@ def require(self, asset: Asset) -> Dependency: location = None settings = self._settings if settings is not None and settings.wants('model'): - location = _store(asset, sha256, settings) + 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, @@ -214,8 +217,9 @@ def asset_cache() -> Path: def _store(asset: Asset, sha256: str, settings: Any) -> Optional[str]: r"""Writes ``asset`` once to ``/assets///``, - its ``_provenance.json`` last; returns where, or None (with a - warning) if it failed.""" + 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}' @@ -226,7 +230,7 @@ def _store(asset: Asset, sha256: str, settings: Any) -> Optional[str]: marker = f'{root}/{PROVENANCE_FILE}' if filesystem.get_file_info(marker).is_file: return settings.root.uri(relative) - filesystem.create_dir(root, recursive=True) + _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: @@ -235,10 +239,9 @@ def _store(asset: Asset, sha256: str, settings: Any) -> Optional[str]: writer.write(manifest_bytes(_manifest(asset, sha256))) return settings.root.uri(relative) except Exception as error: - warnings.warn( + raise RuntimeError( f"Could not store {asset.asset_kind} {asset.name!r} under the " - f"artifact root: {error}", UserWarning, stacklevel=4) - return None + f"artifact root: {error}") from error def _manifest(asset: Asset, sha256: str) -> AssetManifest: diff --git a/packages/kumo-engine/kumo/engine/_mlops.py b/packages/kumo-engine/kumo/engine/_mlops.py index a47175446..aa108d828 100644 --- a/packages/kumo-engine/kumo/engine/_mlops.py +++ b/packages/kumo-engine/kumo/engine/_mlops.py @@ -233,13 +233,13 @@ def resolve(settings: Optional[MLOps] = None) -> Optional[MLOps]: def raise_on_error(settings: Optional[MLOps], error: Optional[BaseException], call: str, stacklevel: int = 2) -> None: - r"""Reports ``error``, a tracking or registration failure of a finished - ``call``: raises when ``settings.on_error`` is ``'raise'``, else warns - unless `warn_failure` already warned about it.""" + 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 tracking or registration ' - f'failed: {full_error(error)}') + 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) 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/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/_mlflow_target.py b/packages/kumo-engine/kumo/engine/integrations/_mlflow_target.py index 8372eec35..62ecac2aa 100644 --- a/packages/kumo-engine/kumo/engine/integrations/_mlflow_target.py +++ b/packages/kumo-engine/kumo/engine/integrations/_mlflow_target.py @@ -9,8 +9,9 @@ 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 nothing, -except that a SageMaker target still sets the process-wide tracking URI. +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 @@ -165,7 +166,10 @@ def _databricks(connector: Any) -> Target: principal = _principal(connector.identity) if (principal is None or principal.config.authentication is DatabricksAuthenticationKind.AMBIENT): - return Target('databricks', 'databricks-uc') + # 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', diff --git a/packages/kumo-engine/kumo/engine/integrations/mlflow.py b/packages/kumo-engine/kumo/engine/integrations/mlflow.py index 49d8a4b4e..502e1d3f1 100644 --- a/packages/kumo-engine/kumo/engine/integrations/mlflow.py +++ b/packages/kumo-engine/kumo/engine/integrations/mlflow.py @@ -106,9 +106,10 @@ class MlflowTracking: 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 or instance role) change nothing, - except that a SageMaker target still sets the process-wide tracking URI - for each call, since the plugin signs and routes with it. + 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. @@ -217,6 +218,50 @@ def load_model(self, location: ModelRegistryLocation, 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() @@ -804,8 +849,10 @@ def _registered_version(client: Any, name: str, return None raise for version in page: - tags = version.tags or client.get_model_version( - name, version.version).tags + # 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 diff --git a/packages/kumo-engine/kumo/engine/pquery/training_table.py b/packages/kumo-engine/kumo/engine/pquery/training_table.py index db3888293..82d616e11 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/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index 9e12ff8f7..e39170f02 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -160,8 +160,14 @@ 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: ModelOutputConfig) -> PublishedModel: - """Publish this model's bundle, as :meth:`Trainer.export`.""" + 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) @@ -353,15 +359,25 @@ def _store_and_register( 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.""" - if settings is not None and settings.wants('model'): - _store_model(bundle_dir, manifest, settings) + 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, _mlflow_registry(settings)) - registry = _connector_registry(settings) + 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] = [] @@ -380,7 +396,7 @@ def _store_and_register( 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 @@ -399,9 +415,8 @@ def _store_model(bundle_dir: Path, manifest: Optional[ModelManifest], filesystem.delete_dir_contents(root, missing_dir_ok=True) upload_bundle(filesystem, root, str(bundle_dir)) except Exception as error: - warnings.warn( - f"Could not store model {manifest.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]]: @@ -449,6 +464,19 @@ def _require_publisher(export_config: ModelOutputConfig) -> None: 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 @@ -602,12 +630,10 @@ def _fitted(self, method: str) -> TrainingResult: "this same Trainer instance.") return self._training_result - def export(self, config: ModelOutputConfig) -> PublishedModel: - """Publish the fitted or loaded model bundle.""" - training_result = self._fitted('export') - _require_publisher(config) - return _publish(training_result.bundle_dir, training_result.bundle_id, - training_result.metrics, config) + 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( @@ -675,6 +701,7 @@ def fit( 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 = { @@ -690,7 +717,7 @@ def fit( scoring_defaults=scoring_defaults, warm_start=warm_start, run_id=run_id, events=events, tracking=session, - deps=Dependencies(settings=settings)) + deps=deps) except BaseException as error: events.emit('fit_ended', params={ 'status': 'failed', @@ -706,7 +733,8 @@ def fit( 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) + 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 diff --git a/packages/kumo-engine/test/test_mlflow_model.py b/packages/kumo-engine/test/test_mlflow_model.py index 2f3cff81d..a62f90694 100644 --- a/packages/kumo-engine/test/test_mlflow_model.py +++ b/packages/kumo-engine/test/test_mlflow_model.py @@ -94,9 +94,17 @@ def test_fit_logs_registers_and_serves_the_model(table_and_snapshot, assert list(served) == ['ENTITY', 'TARGET_PRED'] assert len(served) == len(range(0, 400, 7)) assert load_model('models:/m@champion').result.id == result.id - assert kumo.Trainer.load_from_connector( - cfg, location=kumo.ModelRegistryLocation('m', alias='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 diff --git a/packages/kumo-engine/test/test_mlflow_target.py b/packages/kumo-engine/test/test_mlflow_target.py index 27be764b0..80aeb006d 100644 --- a/packages/kumo-engine/test/test_mlflow_target.py +++ b/packages/kumo-engine/test/test_mlflow_target.py @@ -62,9 +62,13 @@ def test_ambient_identities_set_no_credentials(monkeypatch, tracking_uri): target = resolve(_connector(DatabricksIdentity())) before = dict(os.environ) with target.scope(): - assert dict(os.environ) == before - assert mlflow.get_tracking_uri() == tracking_uri - assert not target.explicit + # 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))): @@ -297,3 +301,41 @@ def test_the_smoke_script_compiles(tmp_path): 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_model_agnostic_mlops.py b/packages/kumo-engine/test/test_model_agnostic_mlops.py index a73a30061..4bfa258ad 100644 --- a/packages/kumo-engine/test/test_model_agnostic_mlops.py +++ b/packages/kumo-engine/test/test_model_agnostic_mlops.py @@ -4,11 +4,13 @@ """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 @@ -283,3 +285,82 @@ def test_a_model_serves_from_its_bundle_snapshot_and_assets_alone( 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_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) From ff396edfc2c6fd396fa1cec62959deb59c8fcbc2 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Thu, 29 Oct 2026 00:36:27 +0000 Subject: [PATCH 05/16] Give your own MLflow calls the tracker's identity and fix the promotion recipes MlflowTracking.client() yields an MlflowClient for the tracker's server with its credentials and tracking URI set for the block, so promotion and run lookups reach the server as the identity kumo-engine uses. The promotion recipes handle the first promotion, and the Snowflake one moves candidate off the new version. kumo.TextEmbedding is exported, and a text column left unencoded says when no word-vector model is installed. Co-Authored-By: Claude Opus 5.5 --- .../src/kumo_core/encoder_selection.py | 15 +++++++- .../kumo-core/tests/test_encoder_selection.py | 14 ++++++- packages/kumo-engine/docs/mlops.md | 37 +++++++++++++------ packages/kumo-engine/kumo/engine/__init__.py | 2 + .../kumo/engine/integrations/mlflow.py | 11 ++++++ .../kumo-engine/test/test_embedding_models.py | 7 +++- .../kumo-engine/test/test_mlflow_target.py | 22 +++++++++++ 7 files changed, 92 insertions(+), 16 deletions(-) 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/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-engine/docs/mlops.md b/packages/kumo-engine/docs/mlops.md index 1553fef8b..a430f6caf 100644 --- a/packages/kumo-engine/docs/mlops.md +++ b/packages/kumo-engine/docs/mlops.md @@ -58,6 +58,16 @@ with kumo.mlops(cfg): 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]`): @@ -151,15 +161,18 @@ An alias is a letter or underscore followed by letters, digits or underscores; ` 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: ```python -champion = kumo.Trainer.load_from_connector(cfg, location=kumo.ModelRegistryLocation(name, alias="champion")) -if result.metrics["test_auroc"] > champion.result.metrics["test_auroc"]: - mlflow.MlflowClient().set_registered_model_alias(name, "champion", result.registered_model.version) +with cfg.tracking.client() as client: # the tracker's MLflow server, as its identity + first = "champion" not in client.get_registered_model(name).aliases + if first or result.metrics["test_auroc"] > kumo.Trainer.load_from_connector( + cfg, location=kumo.ModelRegistryLocation(name, alias="champion")).result.metrics["test_auroc"]: + client.set_registered_model_alias(name, "champion", result.registered_model.version) ``` -In the Snowflake Model Registry a version holds at most one alias, so unset it first: +In the Snowflake Model Registry a version holds at most one alias: take `champion` off the old version (there is none on the first promotion) and `candidate` off the new one, then set it: ```sql -ALTER MODEL ML.MODELS.CHURN_GNN VERSION champion UNSET ALIAS; +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; ``` @@ -206,16 +219,16 @@ 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", +snowflake = SnowflakeTracking(SnowflakeConnector(warehouse="", database="ML", schema_name="RUNS", role="", identity=SnowflakeIdentity(account="", username="", private_key="")), experiment_name="USERS_CHURN") ``` -`load_model(uri, target=...)` and `log_model(..., target=...)` take the same targets. A `DatabricksConnector` without an identity uses a notebook, a job or a `~/.databrickscfg` profile with no environment variables to set. MLflow has no per-client credentials, so with explicit credentials each MLflow call kumo-engine makes sets `DATABRICKS_HOST` and a short-lived `DATABRICKS_TOKEN` (never the secret), or the SageMaker tracking URI and AWS session, and restores them afterwards: don't call MLflow or Databricks yourself on other threads during a fit. `examples/mlops/smoke.py databricks|sagemaker|snowflake` is a manual live check per platform. +`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. +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`. @@ -225,7 +238,7 @@ snowflake = SnowflakeTracking(SnowflakeConnector(warehouse="", database="ML" ### Install - A bare `kumo-engine[mlflow]` with `--find-links ` can resolve to an unrelated `kumo-engine` on PyPI; install the engine wheel by path: `pip install --find-links "/[mlflow]"`. -- **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 (cu130) 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()`. +- **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, MLflow 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`. - **MLflow:** the `mlflow` extra installs `mlflow-skinny` 3.0.1 or later; 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`). @@ -258,13 +271,15 @@ Every artifact carries `_provenance.json`, metadata only (the underscore keeps S ```python store = kumo.EmbeddingModelStore() -store.add("colors", kumo.GloVeVectors("/data/vectors")) -plan.column_processing.encoder_overrides = {"users.bio": TextEmbedding("colors")} +store.add("glove100", kumo.GloVeVectors("glove.6B.zip", dim=100)) # the official GloVe release .zip, unchanged +plan.column_processing.encoder_overrides = {"users.bio": kumo.TextEmbedding("glove100")} with kumo.mlops(cfg): result = kumo.Trainer(plan, embedding_models=store).fit(table, snapshot=snapshot) trainer = kumo.Trainer.load_from_connector(cfg, path=f"models/{result.id}") # anywhere, no store ``` +A text column `infer_metadata()` reads as categorical (short values from a small vocabulary) takes a text encoder once its stype is text: `users.with_column("bio", stype="text")`. + ### Moving snapshots between machines - **Layout:** `export` writes `///`, holding `kumo_snapshot.json` and one directory per export generation (`index/` and `_provenance.json`). diff --git a/packages/kumo-engine/kumo/engine/__init__.py b/packages/kumo-engine/kumo/engine/__init__.py index b9e467763..642cdece9 100644 --- a/packages/kumo-engine/kumo/engine/__init__.py +++ b/packages/kumo-engine/kumo/engine/__init__.py @@ -21,6 +21,7 @@ TokenEmbeddingModel, ) from kumo_core.utils.nofile import ensure_nofile +from kumoapi.encoder import TextEmbedding from kumoapi.graph import Edge from kumoapi.plugin import ModelRegistryLocation from kumoapi.table import Column @@ -69,6 +70,7 @@ "configure", "EmbeddingModelStore", "GloVeVectors", + "TextEmbedding", "TokenEmbeddingModel", "ModelRegistryLocation", "Connector", diff --git a/packages/kumo-engine/kumo/engine/integrations/mlflow.py b/packages/kumo-engine/kumo/engine/integrations/mlflow.py index 502e1d3f1..96c616790 100644 --- a/packages/kumo-engine/kumo/engine/integrations/mlflow.py +++ b/packages/kumo-engine/kumo/engine/integrations/mlflow.py @@ -43,6 +43,7 @@ from kumo.engine.integrations._mlflow_target import ( SAGEMAKER_ARN, SageMakerMLflow, + global_tracking_uri, ) from kumo.engine.integrations._mlflow_target import resolve as resolve_target @@ -161,6 +162,16 @@ def run_url(self) -> Optional[str]: 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``.""" diff --git a/packages/kumo-engine/test/test_embedding_models.py b/packages/kumo-engine/test/test_embedding_models.py index ee2eb68da..a51475c8a 100644 --- a/packages/kumo-engine/test/test_embedding_models.py +++ b/packages/kumo-engine/test/test_embedding_models.py @@ -14,7 +14,6 @@ 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.typing import Stype from kumo.engine import ( @@ -24,6 +23,7 @@ Graph, PredictiveQuery, Table, + TextEmbedding, Trainer, ) @@ -148,7 +148,10 @@ 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], diff --git a/packages/kumo-engine/test/test_mlflow_target.py b/packages/kumo-engine/test/test_mlflow_target.py index 80aeb006d..4b9b360a8 100644 --- a/packages/kumo-engine/test/test_mlflow_target.py +++ b/packages/kumo-engine/test/test_mlflow_target.py @@ -139,6 +139,28 @@ def test_an_aws_access_key_signs_sagemaker_calls_and_uploads( 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( From 91f86944798860601a6187fdb7654a7912f843cb Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Thu, 29 Oct 2026 00:58:59 +0000 Subject: [PATCH 06/16] Require MLflow 3.5 for the mlflow extra and test against that floor Loading a model by alias passes download_artifacts(registry_uri=), which MLflow added in 3.5, so on Databricks Runtime 17.3 (mlflow-skinny 3.0.1) it failed with a TypeError. The extra now requires mlflow-skinny 3.5, which %pip installs over the runtime's, and a CI job runs the MLflow suites at that floor. WriteMode is exported from kumo.engine, and the Databricks quickstart shows its install line. Co-Authored-By: Claude Opus 5.5 --- .github/workflows/testing.yml | 35 ++++++++++++++++++++ packages/kumo-engine/docs/mlops.md | 6 ++-- packages/kumo-engine/kumo/engine/__init__.py | 2 ++ packages/kumo-engine/pyproject.toml | 5 +-- uv.lock | 2 +- 5 files changed, 44 insertions(+), 6 deletions(-) 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/packages/kumo-engine/docs/mlops.md b/packages/kumo-engine/docs/mlops.md index a430f6caf..5b9dccab0 100644 --- a/packages/kumo-engine/docs/mlops.md +++ b/packages/kumo-engine/docs/mlops.md @@ -26,7 +26,7 @@ with kumo.mlops(cfg): ### Databricks -In a notebook or job on Databricks Runtime 17.3 LTS ML or later: +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 @@ -239,8 +239,8 @@ snowflake = SnowflakeTracking(SnowflakeConnector(warehouse="", database="ML" - 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, MLflow 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`. -- **MLflow:** the `mlflow` extra installs `mlflow-skinny` 3.0.1 or later; a 3.x client works against SageMaker's MLflow 3.0 servers. +- **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`. +- **MLflow:** the `mlflow` extra installs `mlflow-skinny` 3.5 or later (Databricks Runtime 17.3 ships 3.0.1, which `%pip` upgrades); a 3.x client works against SageMaker's MLflow 3.0 servers. - **SageMaker serving:** MLflow's stock serving containers cannot install kumo-engine; on SageMaker endpoints serve with `-m /opt/ml/model -p 8080`. GPU endpoints need a driver for CUDA 12.6 (set `InferenceAmiVersion`, e.g. `al2023-ami-sagemaker-inference-gpu-4-1`); batch transform takes JSON (`dataframe_split` plus `params`). ### Exporting and loading artifacts diff --git a/packages/kumo-engine/kumo/engine/__init__.py b/packages/kumo-engine/kumo/engine/__init__.py index 642cdece9..ca5a6b00f 100644 --- a/packages/kumo-engine/kumo/engine/__init__.py +++ b/packages/kumo-engine/kumo/engine/__init__.py @@ -23,6 +23,7 @@ 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 @@ -71,6 +72,7 @@ "EmbeddingModelStore", "GloVeVectors", "TextEmbedding", + "WriteMode", "TokenEmbeddingModel", "ModelRegistryLocation", "Connector", diff --git a/packages/kumo-engine/pyproject.toml b/packages/kumo-engine/pyproject.toml index b051564ae..c55d4f8f5 100644 --- a/packages/kumo-engine/pyproject.toml +++ b/packages/kumo-engine/pyproject.toml @@ -71,9 +71,10 @@ text=[ gcs=[ "kumo-gcs>=0.1.0.dev1", ] -# nvidia-ml-py adds GPU system metrics. +# 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.0.1,<4", + "mlflow-skinny>=3.5,<4", "nvidia-ml-py", ] # SageMaker managed MLflow: 0.5.0 added the per-thread boto3 session diff --git a/uv.lock b/uv.lock index 79523d2dd..321bd266a 100644 --- a/uv.lock +++ b/uv.lock @@ -3541,7 +3541,7 @@ requires-dist = [ { name = "kumo-glove-6b-50d", marker = "extra == 'text'", editable = "tools/glove-6b-50d" }, { name = "kumo-s3", marker = "extra == 's3'", editable = "packages/kumo-s3" }, { name = "kumo-snowflake", extras = ["sdk"], marker = "extra == 'snowflake'", editable = "packages/kumo-snowflake" }, - { name = "mlflow-skinny", marker = "extra == 'mlflow'", specifier = ">=3.0.1,<4" }, + { 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" }, From e4c5cfa30faa214a8b78da114a9c1035a4774d41 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Thu, 29 Oct 2026 00:59:46 +0000 Subject: [PATCH 07/16] Give a stored snapshot's location: MLOps.snapshot_uri The pyfunc's snapshot param takes the snapshot's location under the artifact root, which users rebuilt by hand from the snapshot id. Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/docs/mlops.md | 2 +- packages/kumo-engine/kumo/engine/_mlops.py | 5 +++++ packages/kumo-engine/test/test_mlops.py | 3 +++ 3 files changed, 9 insertions(+), 1 deletion(-) diff --git a/packages/kumo-engine/docs/mlops.md b/packages/kumo-engine/docs/mlops.md index 5b9dccab0..de9348704 100644 --- a/packages/kumo-engine/docs/mlops.md +++ b/packages/kumo-engine/docs/mlops.md @@ -200,7 +200,7 @@ A model registered through `MlflowTracking` is an MLflow pyfunc: it serves with ```python model = mlflow.pyfunc.load_model("models:/main.ml.users_churn@champion") scores = model.predict(pd.DataFrame({"ENTITY": [1, 2, 3]}), params={ - "snapshot": "/Volumes/main/ml/kumo/snapshots//", + "snapshot": cfg.snapshot_uri(snapshot), # "/Volumes/main/ml/kumo/snapshots//" "prediction_time": "2024-06-01"}) # temporal queries ``` diff --git a/packages/kumo-engine/kumo/engine/_mlops.py b/packages/kumo-engine/kumo/engine/_mlops.py index aa108d828..c436a63db 100644 --- a/packages/kumo-engine/kumo/engine/_mlops.py +++ b/packages/kumo-engine/kumo/engine/_mlops.py @@ -148,6 +148,11 @@ def _check_registry(self) -> None: '...).') 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 diff --git a/packages/kumo-engine/test/test_mlops.py b/packages/kumo-engine/test/test_mlops.py index aeca7d9e9..53f929086 100644 --- a/packages/kumo-engine/test/test_mlops.py +++ b/packages/kumo-engine/test/test_mlops.py @@ -75,6 +75,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(): From a43acd07590bcea85dbde5145b8cfc9c6ee79496 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Thu, 29 Oct 2026 00:59:58 +0000 Subject: [PATCH 08/16] Say where Python's temporary files go next to the scratch directory Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/docs/mlops.md | 1 + 1 file changed, 1 insertion(+) diff --git a/packages/kumo-engine/docs/mlops.md b/packages/kumo-engine/docs/mlops.md index de9348704..7db3f3f77 100644 --- a/packages/kumo-engine/docs/mlops.md +++ b/packages/kumo-engine/docs/mlops.md @@ -240,6 +240,7 @@ snowflake = SnowflakeTracking(SnowflakeConnector(warehouse="", database="ML" - 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`). From c5b02a76a10b0f8f51f492c78e848d9b306e65ce Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Thu, 29 Oct 2026 03:53:39 +0000 Subject: [PATCH 09/16] Stay silent about inferred encoders for columns the model plan overrides Fit logged "Text column 'users.BIO' will not be encoded by default" for a column whose encoder the plan set to TextEmbedding, because encoder inference runs for every column before the overrides apply. Fit now drops inference warnings for overridden columns and keeps them for the rest. Co-Authored-By: Claude Opus 5.5 --- .../kumo-engine/kumo/engine/_embeddings.py | 15 ++++++++--- .../kumo/engine/trainer/trainer.py | 26 ++++++++++--------- .../kumo-engine/test/test_embedding_models.py | 6 +++-- 3 files changed, 29 insertions(+), 18 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/_embeddings.py b/packages/kumo-engine/kumo/engine/_embeddings.py index 1a51a42c8..ffeb42ff7 100644 --- a/packages/kumo-engine/kumo/engine/_embeddings.py +++ b/packages/kumo-engine/kumo/engine/_embeddings.py @@ -5,18 +5,25 @@ import logging from contextlib import contextmanager from pathlib import Path -from typing import Any, Dict, Iterator, List, Optional +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) diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index e39170f02..b1c21e075 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -863,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) diff --git a/packages/kumo-engine/test/test_embedding_models.py b/packages/kumo-engine/test/test_embedding_models.py index a51475c8a..681fef6c9 100644 --- a/packages/kumo-engine/test/test_embedding_models.py +++ b/packages/kumo-engine/test/test_embedding_models.py @@ -14,6 +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 Null from kumoapi.typing import Stype from kumo.engine import ( @@ -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')) @@ -155,7 +157,7 @@ def test_each_fit_reports_what_inference_leaves_unencoded(fitted): assert fitted[-1] == { 'tables': [], 'failed fit': [unencoded], - 'fit': [unencoded], + 'fit': [], } From a435d8d69ff3296100f72f3ef16127ee25ce4936 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Thu, 29 Oct 2026 07:37:27 +0000 Subject: [PATCH 10/16] Score binary models without a threshold: probabilities, labels only with one A model fitted without scoring defaults (the notebook flow) failed every default scoring path: predict, the pyfunc and the score job raised "binary_classification_threshold must be specified". Binary predict now always writes True_PROB and sets TARGET_PRED only from a threshold the call passes or the bundle stores; without one TARGET_PRED is null, so no threshold is invented and the output schema stays the same. Verified: test_scoring_defaults, test_mlflow_model and test_semantic_carry_forward, kumo-core test_prediction. Co-Authored-By: Claude Opus 5.5 --- .../kumo-core/src/kumo_core/prediction.py | 1 + .../train/diskgraph_nested_trainer.py | 5 ++++ packages/kumo-engine/docs/mlops.md | 4 ++-- .../kumo/engine/integrations/_mlflow_model.py | 16 +++---------- .../kumo/engine/trainer/trainer.py | 13 ++++++---- .../kumo-engine/test/test_mlflow_model.py | 24 +++++++------------ .../kumo-engine/test/test_scoring_defaults.py | 4 ++-- .../test/test_semantic_carry_forward.py | 1 - 8 files changed, 29 insertions(+), 39 deletions(-) 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-engine/docs/mlops.md b/packages/kumo-engine/docs/mlops.md index 7db3f3f77..1a7ab98c8 100644 --- a/packages/kumo-engine/docs/mlops.md +++ b/packages/kumo-engine/docs/mlops.md @@ -142,7 +142,7 @@ table = kumo.TrainingTable.load_from_connector(cfg, path=f"training_tables/{tabl - `result.registered_model` is the version `MLOps(registry=)` registered; `result.published_model` the one `fit(export_config=...)` published. - `result.tracking` is the tracker's run: `TrackingRef(platform="databricks", run_id="...", url="https://.../ml/experiments/.../runs/...")`. The model's `_provenance.json` records the same. -- `fit(scoring_defaults=kumo.ScoringDefaults(...))` stores prediction defaults in the bundle before it is stored and registered: `binary_classification_threshold` (binary and multilabel classification) and `num_classes_to_return` (multiclass, multilabel, link prediction). `predict` and the MLflow pyfunc apply them when a call omits them; explicit arguments and a query's `TOP k` win. `result.with_scoring_defaults(...)` retunes a fitted model as a new bundle whose manifest names the original in `parents.derived_from`. +- `fit(scoring_defaults=kumo.ScoringDefaults(...))` stores prediction defaults in the bundle before it is stored and registered: `binary_classification_threshold` (binary and multilabel classification) and `num_classes_to_return` (multiclass, multilabel, link prediction). `predict` 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. @@ -204,7 +204,7 @@ scores = model.predict(pd.DataFrame({"ENTITY": [1, 2, 3]}), params={ "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 needs `params={"binary_classification_threshold": ...}` unless the model carries scoring defaults. +- 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). diff --git a/packages/kumo-engine/kumo/engine/integrations/_mlflow_model.py b/packages/kumo-engine/kumo/engine/integrations/_mlflow_model.py index a1be11e94..75d063c23 100644 --- a/packages/kumo-engine/kumo/engine/integrations/_mlflow_model.py +++ b/packages/kumo-engine/kumo/engine/integrations/_mlflow_model.py @@ -70,7 +70,7 @@ def predict(trainer: Any, snapshot: GraphSnapshot, rows: pd.DataFrame, 'binary_classification_threshold'), num_classes_to_return=params.get('num_classes_to_return'), ) - _require_inputs(trainer, snapshot, rows, params, time) + _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'): @@ -138,28 +138,18 @@ def _seconds(values: pd.Series) -> pd.Series: def _require_inputs(trainer: Any, snapshot: GraphSnapshot, rows: pd.DataFrame, - params: Mapping[str, Any], time: Optional[pd.Timestamp]) -> None: from kumo_core.pquery import query_is_temporal - from kumo_core.types import TaskType from kumo.engine.pquery import PredictiveQuery - untimed = len(rows) > 0 and time is None and 'TIMESTAMP' not in rows - unthresholded = ( - params.get('binary_classification_threshold') is None and - trainer.result.scoring_defaults.binary_classification_threshold is None) - if not (untimed or unthresholded): + if len(rows) == 0 or time is not None or 'TIMESTAMP' in rows: return parsed = PredictiveQuery._from_snapshot(snapshot, trainer._training_query)._parsed - if untimed and query_is_temporal(parsed, snapshot._native_graph): + if query_is_temporal(parsed, snapshot._native_graph): raise ValueError("temporal query: pass params={'prediction_time': " "...} or a TIMESTAMP column") - if unthresholded and (TaskType(parsed.task_type) - is TaskType.BINARY_CLASSIFICATION): - raise ValueError("binary classification: pass " - "params={'binary_classification_threshold': ...}") def snapshot_location(uri: str) -> Tuple[Connector, str]: diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index b1c21e075..04219f639 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -1089,9 +1089,10 @@ 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`. """ @@ -1183,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( diff --git a/packages/kumo-engine/test/test_mlflow_model.py b/packages/kumo-engine/test/test_mlflow_model.py index a62f90694..cb93defd8 100644 --- a/packages/kumo-engine/test/test_mlflow_model.py +++ b/packages/kumo-engine/test/test_mlflow_model.py @@ -181,32 +181,24 @@ def predict_beside_another_call(self, *args, **kwargs): assert set(tables.glob('*')) == before | {tables / 'other-call'} -def test_pyfunc_names_missing_temporal_and_threshold_inputs( - fitted, monkeypatch): - from dataclasses import replace - +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)}) - time = str(pd.Timestamp(table.timestamps[0])) 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") - with pytest.raises(ValueError, match=r"^binary classification: pass " - r"params=\{'binary_classification_threshold': " - r"\.\.\.\}$"): - predict(trainer, snapshot, rows, {'prediction_time': time}) - - monkeypatch.setattr( - trainer, '_training_result', - replace(trainer.result, scoring_defaults=kumo.ScoringDefaults( - binary_classification_threshold=.5))) - assert len(predict(trainer, snapshot, rows, {'prediction_time': time - })) == 5 + + 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): diff --git a/packages/kumo-engine/test/test_scoring_defaults.py b/packages/kumo-engine/test/test_scoring_defaults.py index fd8812113..3ba5d2537 100644 --- a/packages/kumo-engine/test/test_scoring_defaults.py +++ b/packages/kumo-engine/test/test_scoring_defaults.py @@ -135,8 +135,8 @@ def test_unlisted_or_unreadable_defaults_are_ignored(fitted, tmp_path): with pytest.warns(UserWarning, match='Ignoring the scoring defaults'): ignored = kumo.Trainer.load(bundle) assert ignored.result.scoring_defaults == kumo.ScoringDefaults() - with pytest.raises(ValueError, match='must be specified'): - ignored.predict(table, snapshot=snapshot) + 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): diff --git a/packages/kumo-engine/test/test_semantic_carry_forward.py b/packages/kumo-engine/test/test_semantic_carry_forward.py index 5968ef805..81da3228e 100644 --- a/packages/kumo-engine/test/test_semantic_carry_forward.py +++ b/packages/kumo-engine/test/test_semantic_carry_forward.py @@ -129,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, From 8feb640d8a32dc373c5d5895d79f7f2ace77c743 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Thu, 29 Oct 2026 09:59:40 +0000 Subject: [PATCH 11/16] Create the tracker's experiment when kumo.MLOps is built, or say why not An experiment under a Databricks folder that did not exist made every tracked fit fail to record, only with a warning under on_error='warn': result.tracking was None and nothing pointed at the folder. MlflowTracking.check, which runs when kumo.MLOps is built, now creates a missing experiment. If MLflow refuses, it raises with MLflow's error and says the parent folder must exist and how to create it. Verified: test_mlflow_tracking, test_mlflow_target. Co-Authored-By: Claude Opus 5.5 --- .../kumo/engine/integrations/mlflow.py | 20 ++++++++++++++++++- .../kumo-engine/test/test_mlflow_model.py | 8 +++++--- .../kumo-engine/test/test_mlflow_tracking.py | 19 ++++++++++++++++++ 3 files changed, 43 insertions(+), 4 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/integrations/mlflow.py b/packages/kumo-engine/kumo/engine/integrations/mlflow.py index 96c616790..774eb1bfb 100644 --- a/packages/kumo-engine/kumo/engine/integrations/mlflow.py +++ b/packages/kumo-engine/kumo/engine/integrations/mlflow.py @@ -188,11 +188,13 @@ def check(self, registry: Optional[ModelRegistryLocation] = None) -> None: 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: - client.get_experiment_by_name(self.experiment_name) + missing = client.get_experiment_by_name( + self.experiment_name) is None else: client.search_experiments(max_results=1) if registry is not None: @@ -204,6 +206,22 @@ def check(self, registry: Optional[ModelRegistryLocation] = None) -> None: 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() diff --git a/packages/kumo-engine/test/test_mlflow_model.py b/packages/kumo-engine/test/test_mlflow_model.py index cb93defd8..51eb51fb7 100644 --- a/packages/kumo-engine/test/test_mlflow_model.py +++ b/packages/kumo-engine/test/test_mlflow_model.py @@ -490,12 +490,15 @@ def no_folder(self, name, *args, **kwargs): raise MlflowException(f'RESOURCE_DOES_NOT_EXIST: Parent of {name} ' f'does not exist.') - monkeypatch.setattr(MlflowClient, 'create_experiment', no_folder) 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) @@ -507,10 +510,9 @@ def no_folder(self, name, *args, **kwargs): assert 'MlflowException: RESOURCE_DOES_NOT_EXIST' in off.getMessage() assert not off.exc_info - cfg = dataclasses.replace(cfg, on_error='warn') with warnings.catch_warnings(record=True) as caught: warnings.simplefilter('always') - trainer.fit(table, snapshot=snapshot, mlops=cfg) + 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 == [ diff --git a/packages/kumo-engine/test/test_mlflow_tracking.py b/packages/kumo-engine/test/test_mlflow_tracking.py index 920e92272..709b9885f 100644 --- a/packages/kumo-engine/test/test_mlflow_tracking.py +++ b/packages/kumo-engine/test/test_mlflow_tracking.py @@ -241,6 +241,25 @@ def test_registry_model_names_are_checked_when_mlops_is_built( 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'), From 3dc527a6a27bad67ed74a8a4b45978b8cfb96c5c Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Thu, 29 Oct 2026 12:36:41 +0000 Subject: [PATCH 12/16] Keep column types when a table export stamps its id columns prepared_artifact_file added PREDICTION_ID and BUNDLE_ID through a pandas round trip, which turns an all-null bool column into Arrow's null type. A binary run scored without a threshold (TARGET_PRED null) therefore loaded into Snowflake with TARGET_PRED as NUMBER, while a thresholded run or a baseline loaded it as BOOLEAN, and CREATE MODEL MONITOR refused the pair. The id columns are now added in Arrow, so every source column keeps its type. Verified: kumo-core test_artifact_export (new: an all-null bool and a UTC timestamp keep their types), kumo-snowflake and kumo-databricks sink tests. Co-Authored-By: Claude Opus 5.5 --- .../kumo-core/src/kumo_core/artifact_export.py | 12 +++++++++--- packages/kumo-core/tests/test_artifact_export.py | 15 +++++++++++++++ 2 files changed, 24 insertions(+), 3 deletions(-) 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/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): From ca64a2cb69b2f821932a8d1f8702e0cc59c0bce9 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Thu, 29 Oct 2026 16:37:20 +0000 Subject: [PATCH 13/16] Report snapshot progress through the progress events Graph.snapshot printed nothing kumo-side while it staged tables and built the index (75 s of silence on a 32k-row graph in a usability run). snapshot(callbacks=[...]) now takes the same callbacks as fit and reports snapshot_started, table_staged per table and snapshot_ended (snapshot id, reused), each with the seconds since the start and logged at INFO. mlops.md documents them, and its TrainingTable.load_from_connector examples say pquery= after #50's rename. Verified: test_fit_events (new: a snapshot's events, then a reused one). Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/docs/mlops.md | 6 +-- packages/kumo-engine/kumo/engine/_events.py | 7 +-- packages/kumo-engine/kumo/engine/graph.py | 51 ++++++++++++++++++-- packages/kumo-engine/test/test_fit_events.py | 20 ++++++++ 4 files changed, 74 insertions(+), 10 deletions(-) diff --git a/packages/kumo-engine/docs/mlops.md b/packages/kumo-engine/docs/mlops.md index 1a7ab98c8..e209eeaed 100644 --- a/packages/kumo-engine/docs/mlops.md +++ b/packages/kumo-engine/docs/mlops.md @@ -131,7 +131,7 @@ Pass the `kumo.MLOps` where a connector goes, with paths relative to the artifac snapshot = kumo.GraphSnapshot.load_from_connector(cfg, id=result.snapshot_id) trainer = kumo.Trainer.load_from_connector(cfg, path=f"models/{result.id}") champion = kumo.Trainer.load_from_connector(cfg, location=kumo.ModelRegistryLocation("users-churn", alias="champion")) -table = kumo.TrainingTable.load_from_connector(cfg, path=f"training_tables/{table_id}", pq=pquery, snapshot=snapshot) +table = kumo.TrainingTable.load_from_connector(cfg, path=f"training_tables/{table_id}", pquery=pquery, snapshot=snapshot) ``` `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()`). @@ -251,7 +251,7 @@ Every artifact has one way out and one way back. A `kumo.MLOps` stands in for th | 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)` | +| Training table | `table.export(TrainingTableOutputConfig(connector, name))` | `TrainingTable.load_from_connector(connector, path=name, pquery=pquery, 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. @@ -295,7 +295,7 @@ Other systems plug in through `kumo.engine.tracking.Tracking`, which receives re ### 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()))]`. +`fit(callbacks=[...])` calls each callback with every `kumo.engine.tracking.Event`: `fit_started`, per trial `trial_started`, `epoch_ended` and `trial_ended`, then `fit_ended`. Ended events carry `params["status"]` and, on failure, `params["error"]`; a callback that raises is logged and not called again. For JSON lines: `callbacks=[lambda e: print(json.dumps(e.to_dict()))]`. `graph.snapshot(callbacks=[...])` reports `snapshot_started`, `table_staged` per table and `snapshot_ended` (with `snapshot_id` and `reused`), each with `params["seconds"]` since the start, and logs each at INFO. ### Limits diff --git a/packages/kumo-engine/kumo/engine/_events.py b/packages/kumo-engine/kumo/engine/_events.py index 98e8c2394..ecf1c2a79 100644 --- a/packages/kumo-engine/kumo/engine/_events.py +++ b/packages/kumo-engine/kumo/engine/_events.py @@ -1,7 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. # All rights reserved. # SPDX-License-Identifier: Apache-2.0 -r"""Progress events reported by `Trainer.fit`.""" +r"""Progress events reported by `Trainer.fit` and `Graph.snapshot`.""" import hashlib import logging @@ -24,7 +24,8 @@ logger = logging.getLogger(__name__) EventName = Literal['fit_started', 'trial_started', 'epoch_ended', - 'trial_ended', 'fit_ended'] + 'trial_ended', 'fit_ended', 'snapshot_started', + 'table_staged', 'snapshot_ended'] @dataclass(frozen=True) @@ -75,7 +76,7 @@ def emit(self, name: EventName, *, trial: Optional[int] = None, except Exception: logger.warning( 'Progress callback %r failed; it is disabled for the ' - 'rest of this fit', callback, exc_info=True) + 'rest of this call', callback, exc_info=True) self._callbacks.remove(callback) def event(self, name: str, **fields: Any) -> None: diff --git a/packages/kumo-engine/kumo/engine/graph.py b/packages/kumo-engine/kumo/engine/graph.py index 982ee731b..528d1c4a9 100644 --- a/packages/kumo-engine/kumo/engine/graph.py +++ b/packages/kumo-engine/kumo/engine/graph.py @@ -60,17 +60,29 @@ """ import json +import logging import threading +import time import warnings from dataclasses import dataclass, field, replace from functools import cached_property, partial from pathlib import Path -from typing import TYPE_CHECKING, Any, Dict, List, Mapping, Optional, Sequence +from typing import ( + TYPE_CHECKING, + Any, + Callable, + Dict, + List, + Mapping, + Optional, + Sequence, +) from kumo_core.artifacts import PROVENANCE_FILE from kumo_core.pquery import NativeGraph, normalize_graph from kumoapi.graph import Edge +from kumo.engine._events import Event, EventEmitter from kumo.engine._manifest import ( RECONSTRUCTED, SnapshotManifest, @@ -95,6 +107,7 @@ SNAPSHOT_MANIFEST_FILE = PROVENANCE_FILE +logger = logging.getLogger(__name__) _warned_scratch: set[Path] = set() _warned_backfill: set[Path] = set() @@ -268,6 +281,21 @@ def report_other_graph(message: str, before: Optional[Mapping[str, Any]], stacklevel=stacklevel + 1) +class _SnapshotEvents(EventEmitter): + r"""Snapshot progress: each event goes to the callbacks and the log, + with the seconds since the snapshot started.""" + + def __init__(self, callbacks: Sequence[Callable[[Event], None]]): + super().__init__(callbacks) + self._started = time.monotonic() + + def emit(self, name: Any, **fields: Any) -> None: + params = {**fields.pop('params', {}), + 'seconds': round(time.monotonic() - self._started, 1)} + logger.info('Graph.snapshot: %s %s', name, params) + super().emit(name, params=params, **fields) + + def _store() -> Any: from kumo_core.diskgraph.snapshot import ( DiskgraphLocalConfig, @@ -671,7 +699,9 @@ def validate(self) -> None: self.native_graph def snapshot(self, *, record_sources: bool = True, - mlops: Optional[Any] = None) -> GraphSnapshot: + mlops: Optional[Any] = None, + callbacks: Sequence[Callable[[Event], None]] = () + ) -> GraphSnapshot: r"""Materialize this graph: read each table's connector, run the result through `kumo-core` to build column stats and the on-disk index. The index is stored under a directory keyed by both `self.id` @@ -688,6 +718,9 @@ def snapshot(self, *, record_sources: bool = True, record_sources: if False, the snapshot manifest records each source as connector and table name only (no bucket, path or account) and Custom SQL as a hash only (no SQL text). + callbacks: called with each progress event, + `snapshot_started`, `table_staged` (per table) and + `snapshot_ended`; each is also logged at INFO. """ self.validate() for table in self._named_tables.values(): @@ -716,9 +749,12 @@ def snapshot(self, *, record_sources: bool = True, from kumo.engine._mlops import resolve settings = resolve(mlops) + events = _SnapshotEvents(callbacks) + events.emit('snapshot_started', params={ + 'graph': graph_key, 'tables': list(self._named_tables)}) with _lock(graph_key): snapshot = self._snapshot_no_lock(graph_key, persist_dir, - record_sources) + record_sources, events) if settings is not None and settings.wants('snapshot'): from kumo.engine._mlops import raise_on_error from kumo.engine._snapshot_io import store @@ -730,7 +766,9 @@ def snapshot(self, *, record_sources: bool = True, return snapshot def _snapshot_no_lock(self, graph_key: str, persist_dir: Path, - record_sources: bool = True) -> GraphSnapshot: + record_sources: bool = True, + events: Optional['_SnapshotEvents'] = None + ) -> GraphSnapshot: r"""`snapshot()`'s body, run under its `FileLock`. Not meant to be called on its own -- `persist_dir`, unlike everything else here, isn't a pure function of `self` alone (it's `snapshot()`'s lock key, @@ -792,6 +830,8 @@ def _snapshot_no_lock(self, graph_key: str, persist_dir: Path, "limitation in kumoai/backend/local/graph.py's " "_stage_sources, which raises the same way.") table_data_versions[alias] = prepared.version + if events is not None: + events.emit('table_staged', params={'table': alias}) source_tables[alias] = source_table( connector, name, sql, prepared.version, record_sources=record_sources) @@ -848,6 +888,9 @@ def _snapshot_no_lock(self, graph_key: str, persist_dir: Path, record_sources, fresh=existing is None) _warn_ingest_data_loss(result, graph_config) + if events is not None: + events.emit('snapshot_ended', params={ + 'snapshot_id': snapshot_id, 'reused': existing is not None}) return GraphSnapshot( id=snapshot_id, index_path=Path(result.index_path), diff --git a/packages/kumo-engine/test/test_fit_events.py b/packages/kumo-engine/test/test_fit_events.py index 08a22729c..682a68a81 100644 --- a/packages/kumo-engine/test/test_fit_events.py +++ b/packages/kumo-engine/test/test_fit_events.py @@ -151,3 +151,23 @@ def test_epoch_event_without_recorded_metrics_does_not_raise(tmp_path): logger.on_train_epoch_end(None, None) assert received == [] + + +def test_a_snapshot_reports_progress_events(graph, tmp_path): + previous = get_scratch_dir() + kumo.configure(scratch_dir=tmp_path / 'scratch') + try: + events, again = [], [] + snapshot = graph.snapshot(callbacks=[events.append]) + graph.snapshot(callbacks=[again.append]) + finally: + kumo.configure(scratch_dir=previous) + + assert [event.name for event in events] == [ + 'snapshot_started', 'table_staged', 'table_staged', 'snapshot_ended'] + assert [event.params['table'] for event in events[1:3]] == [ + 'users', 'events'] + assert events[-1].params['snapshot_id'] == snapshot.id + assert events[-1].params['reused'] is False + assert again[-1].params['reused'] is True + assert all(event.params['seconds'] >= 0 for event in events) From c4ff3aaf88ab0589703b0618458c9d0f385e6c55 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Thu, 29 Oct 2026 16:43:07 +0000 Subject: [PATCH 14/16] Keep PyG's import warning and Lightning's tips out of kumo calls MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit With torch 2.15, import kumo.engine printed PyTorch Geometric's FutureWarning about its import-time torch.jit.script calls, and every fit or predict logged Lightning's '💡 Tip' lines advertising its cloud products; users can act on neither. kumo.engine now imports torch_geometric.nn first with that one warning filtered, and fit and predict drop Lightning's Tip records while they run. Other warnings and Lightning log lines, and the user's own Lightning runs, are unchanged. Verified: test_fit_events (new: Tip lines dropped inside a kumo call only), test_optional_imports; import kumo.engine with -W error::FutureWarning on torch 2.15. Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/__init__.py | 1 + packages/kumo-engine/kumo/engine/_quiet.py | 12 ++++++++ .../kumo/engine/trainer/trainer.py | 28 +++++++++++++------ packages/kumo-engine/test/test_fit_events.py | 23 +++++++++++++++ 4 files changed, 55 insertions(+), 9 deletions(-) create mode 100644 packages/kumo-engine/kumo/engine/_quiet.py diff --git a/packages/kumo-engine/kumo/engine/__init__.py b/packages/kumo-engine/kumo/engine/__init__.py index ca5a6b00f..3b94059a7 100644 --- a/packages/kumo-engine/kumo/engine/__init__.py +++ b/packages/kumo-engine/kumo/engine/__init__.py @@ -27,6 +27,7 @@ from kumoapi.plugin import ModelRegistryLocation from kumoapi.table import Column +from kumo.engine import _quiet # noqa: F401 from kumo.engine._manifest import ( ModelManifest, SnapshotManifest, diff --git a/packages/kumo-engine/kumo/engine/_quiet.py b/packages/kumo-engine/kumo/engine/_quiet.py new file mode 100644 index 000000000..3a407b1b2 --- /dev/null +++ b/packages/kumo-engine/kumo/engine/_quiet.py @@ -0,0 +1,12 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +r"""Imports PyTorch Geometric without the FutureWarning torch 2.15+ raises +for its import-time ``torch.jit.script`` calls, which users cannot act on.""" + +import warnings + +with warnings.catch_warnings(): + warnings.filterwarnings('ignore', r'`torch\.jit\.script` is deprecated', + FutureWarning) + import torch_geometric.nn # noqa: F401 diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index 04219f639..e7a3446f5 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -544,7 +544,8 @@ def _stage_model_bundle(artifacts: Mapping[str, str], dest: Path) -> Path: def _without_torch_scatter_hint(method: _F) -> _F: r"""Runs ``method`` without PyTorch Geometric's GPU hint that min/max - aggregation "can be accelerated via the 'torch-scatter' package". + aggregation "can be accelerated via the 'torch-scatter' package", and + without Lightning's "Tip" log lines advertising its cloud products. kumo-engine uses PyTorch's own scatter instead of that optional extension (see ``test_wheel_uses_torch_fallbacks_without_native_extensions``), so @@ -552,18 +553,27 @@ def _without_torch_scatter_hint(method: _F) -> _F: """ @functools.wraps(method) def wrapper(*args: Any, **kwargs: Any) -> Any: - with warnings.catch_warnings(): - warnings.filterwarnings( - 'ignore', - message=r"The usage of `scatter\(reduce='\w+'\)` can be " - r"accelerated via the 'torch-scatter' package", - category=UserWarning, - ) - return method(*args, **kwargs) + tips = logging.getLogger('pytorch_lightning.utilities.rank_zero') + tips.addFilter(_no_lightning_tip) + try: + with warnings.catch_warnings(): + warnings.filterwarnings( + 'ignore', + message=r"The usage of `scatter\(reduce='\w+'\)` can be " + r"accelerated via the 'torch-scatter' package", + category=UserWarning, + ) + return method(*args, **kwargs) + finally: + tips.removeFilter(_no_lightning_tip) return cast(_F, wrapper) +def _no_lightning_tip(record: logging.LogRecord) -> bool: + return not str(record.msg).startswith('💡 Tip') + + class Trainer: r"""Trains and runs a GNN model in-process. diff --git a/packages/kumo-engine/test/test_fit_events.py b/packages/kumo-engine/test/test_fit_events.py index 682a68a81..fa34faae4 100644 --- a/packages/kumo-engine/test/test_fit_events.py +++ b/packages/kumo-engine/test/test_fit_events.py @@ -3,6 +3,7 @@ # SPDX-License-Identifier: Apache-2.0 import json +import logging import numpy as np import pandas as pd @@ -171,3 +172,25 @@ def test_a_snapshot_reports_progress_events(graph, tmp_path): assert events[-1].params['reused'] is False assert again[-1].params['reused'] is True assert all(event.params['seconds'] >= 0 for event in events) + + +def test_kumo_calls_drop_lightning_tips_and_only_those(caplog): + from pytorch_lightning.utilities.rank_zero import rank_zero_info + + from kumo.engine.trainer.trainer import _without_torch_scatter_hint + + @_without_torch_scatter_hint + def fit(): + rank_zero_info('💡 Tip: inside') + rank_zero_info('GPU available: False') + + lightning = logging.getLogger('pytorch_lightning') + lightning.addHandler(caplog.handler) # it does not propagate + try: + with caplog.at_level('INFO', logger='pytorch_lightning'): + fit() + rank_zero_info('💡 Tip: outside') + finally: + lightning.removeHandler(caplog.handler) + assert 'inside' not in caplog.text + assert 'GPU available' in caplog.text and 'outside' in caplog.text From 7308c25f956dab2f5dfd64ee9d198b658ef6215b Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Thu, 29 Oct 2026 16:45:20 +0000 Subject: [PATCH 15/16] Link Databricks runs through the workspace URL on a cluster On a cluster, result.tracking.url used the SDK's host, the regional endpoint, whose run links do not open the workspace's page. It now uses the workspace URL the notebook context names (MLflow's get_workspace_info_from_dbutils), and the SDK's host elsewhere. Verified: test_mlflow_target. Co-Authored-By: Claude Opus 5.5 --- .../kumo/engine/integrations/_mlflow_target.py | 11 +++++++++++ packages/kumo-engine/test/test_mlflow_target.py | 13 +++++++++++++ 2 files changed, 24 insertions(+) diff --git a/packages/kumo-engine/kumo/engine/integrations/_mlflow_target.py b/packages/kumo-engine/kumo/engine/integrations/_mlflow_target.py index 62ecac2aa..7f3483377 100644 --- a/packages/kumo-engine/kumo/engine/integrations/_mlflow_target.py +++ b/packages/kumo-engine/kumo/engine/integrations/_mlflow_target.py @@ -139,6 +139,17 @@ def resolve(target: Any) -> Target: def _workspace_host() -> Optional[str]: + # On a cluster the SDK's host is the regional one; the notebook context + # names the workspace's own URL, which run links need. + try: + from mlflow.utils.databricks_utils import ( + get_workspace_info_from_dbutils, + ) + host, _ = get_workspace_info_from_dbutils() + if host: + return host + except Exception: + pass try: return import_provider('databricks.sdk.config', 'databricks').Config().host diff --git a/packages/kumo-engine/test/test_mlflow_target.py b/packages/kumo-engine/test/test_mlflow_target.py index 4b9b360a8..e883cd238 100644 --- a/packages/kumo-engine/test/test_mlflow_target.py +++ b/packages/kumo-engine/test/test_mlflow_target.py @@ -114,6 +114,19 @@ def test_a_databricks_service_principal_lends_a_token_for_each_call( 'https://dbc.example/ml/experiments/7/runs/r') +def test_a_cluster_run_links_to_the_workspace_url_not_the_regional_host( + monkeypatch): + from mlflow.utils import databricks_utils + + from kumo.engine.integrations._mlflow_target import resolve + + target = resolve('databricks') + monkeypatch.setattr(databricks_utils, 'get_workspace_info_from_dbutils', + lambda: ('https://dbc-1.example', '42')) + assert target.run_url('7', 'r') == ( + 'https://dbc-1.example/ml/experiments/7/runs/r') + + def test_an_aws_access_key_signs_sagemaker_calls_and_uploads( monkeypatch, tracking_uri): boto3 = pytest.importorskip('boto3') From a86f651efeade724021e4efab23d8b7a1e21209c Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Thu, 29 Oct 2026 17:01:11 +0000 Subject: [PATCH 16/16] Give the prefix-scoped S3 test fake a delete and keep a URI test offline Building kumo.MLOps now writes and deletes a probe under its root. The prefix-scoped S3 fake gains delete_file (real prefix rights include s3:DeleteObject), and the snapshot_uri test, which only formats a URI, builds its MLOps with store=False so it makes no S3 call. Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/test/test_mlops.py | 2 +- packages/kumo-engine/test/test_model_agnostic_mlops.py | 3 +++ 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/packages/kumo-engine/test/test_mlops.py b/packages/kumo-engine/test/test_mlops.py index 53f929086..e258f3eb8 100644 --- a/packages/kumo-engine/test/test_mlops.py +++ b/packages/kumo-engine/test/test_mlops.py @@ -76,7 +76,7 @@ def test_a_root_is_a_uri_or_a_connector_and_path(tmp_path): 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( + assert kumo.MLOps('s3://bucket/kumo', store=False).snapshot_uri( SimpleNamespace(id='ab12@cd34')) == 's3://bucket/kumo/snapshots/ab12/cd34' diff --git a/packages/kumo-engine/test/test_model_agnostic_mlops.py b/packages/kumo-engine/test/test_model_agnostic_mlops.py index 4bfa258ad..7cd604a04 100644 --- a/packages/kumo-engine/test/test_model_agnostic_mlops.py +++ b/packages/kumo-engine/test/test_model_agnostic_mlops.py @@ -311,6 +311,9 @@ def close(self): return Writer() + def delete_file(self, path): + del self.objects[path] + def test_an_asset_is_stored_on_s3_with_prefix_scoped_rights( tmp_path, monkeypatch):