From 7cdc1effbf520756b6119e0a7e017dbf60d9b7be Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Sat, 24 Oct 2026 18:05:09 +0000 Subject: [PATCH] Report fit progress events to callbacks fit(callbacks=[...]) receives each kumo.engine.tracking.Event: fit_started, trial_started, epoch_ended, trial_ended and fit_ended, with status and a short error on ended events. A callback that raises is logged and disabled; it never fails the fit. fit resolves its MLOps settings once on the caller's thread and passes them down. Co-Authored-By: Claude Opus 5.5 --- .../src/kumo_core/automl/auto_trainer.py | 37 ++++- .../src/kumo_core/automl/optuna_searcher.py | 1 + .../src/kumo_core/automl/searcher.py | 15 +- .../callbacks/train_progress_logger.py | 25 ++- .../src/kumo_core/instrumentation.py | 10 ++ .../src/kumo_core/training/runtime.py | 4 +- .../kumo-core/tests/test_searched_params.py | 38 +++++ packages/kumo-engine/docs/mlops.md | 16 ++ packages/kumo-engine/kumo/engine/_events.py | 86 ++++++++++ packages/kumo-engine/kumo/engine/tracking.py | 10 ++ .../kumo/engine/trainer/_runtime.py | 6 +- .../kumo/engine/trainer/trainer.py | 68 ++++++-- packages/kumo-engine/test/test_fit_events.py | 154 ++++++++++++++++++ 13 files changed, 447 insertions(+), 23 deletions(-) create mode 100644 packages/kumo-core/tests/test_searched_params.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/test/test_fit_events.py diff --git a/packages/kumo-core/src/kumo_core/automl/auto_trainer.py b/packages/kumo-core/src/kumo_core/automl/auto_trainer.py index 5dd818966..3abe36170 100644 --- a/packages/kumo-core/src/kumo_core/automl/auto_trainer.py +++ b/packages/kumo-core/src/kumo_core/automl/auto_trainer.py @@ -12,6 +12,7 @@ import math import os import os.path as osp +import time import traceback import warnings from dataclasses import asdict, dataclass, fields, is_dataclass @@ -39,7 +40,7 @@ AutoTrainerProgress, AutoTrainingProgressLogger, ) -from kumo_core.instrumentation import TrainingInstrumentation +from kumo_core.instrumentation import TrainingInstrumentation, format_error from kumo_core.io import filesystem from kumo_core.io.json import load_json, save_json from kumo_core.job_config import JobConfig, TrainStatus @@ -458,6 +459,7 @@ def _train( metrics_client_job_name: str, pruning: Callback | None = None, status: TrainStatus | None = None, + trial: int | None = None, ) -> list[float]: r"""Config-driven training function. @@ -713,7 +715,7 @@ def _train( trainer.pl_trainer.callbacks.append(pruning) # Add progress logging callback: - self.training_progress_logger.set_current_trainer(trainer) + self.training_progress_logger.set_current_trainer(trainer, trial) trainer.pl_trainer.callbacks.append(self.training_progress_logger) # Save AutoTrainer YAML (Trainer YAML will be saved within @@ -847,7 +849,7 @@ def _run( logger.info(f"Initialized '{self.searcher}'") self.training_progress_logger = AutoTrainingProgressLogger( - self.log_dir, self.num_experiments) + self.log_dir, self.num_experiments, on_event=self._emit) tune_mode = MetricStore.tune_mode(self.tune_metric) metrics_job_name = self.runtime.metrics_job_name() @@ -872,12 +874,17 @@ def _run( logger.info(f"Starting AutoML experiment " f"{i + 1}/{self.num_experiments}") + self._emit('trial_started', trial=i + 1, params={ + **self.searcher.sampled_params, + 'attempt': j + 1, + }) batch_size = cfg.optimization.batch_size[0] with self.instrumentation.timing( "auto_trainer.run.status_add_trial", trial=i + 1): status.add_trial(batch_size) status.save(self.home_dir) + trial_start = time.monotonic() with self.instrumentation.timing( "auto_trainer.run.train_trial", trial=i + 1): val_performance = self._train( @@ -885,7 +892,9 @@ def _run( metrics_client_job_name=metrics_job_name, pruning=pruning, status=status, + trial=i + 1, ) + trial_seconds = time.monotonic() - trial_start best_performance = _eval_fn[tune_mode](val_performance) logger.info(f"AutoML experiment " f"{i + 1}/{self.num_experiments} finished " @@ -902,9 +911,21 @@ def _run( "auto_trainer.run.status_save_success", trial=i + 1): status.save(self.home_dir) + self._emit('trial_ended', trial=i + 1, metrics={ + self.val_tune_metric: best_performance, + 'seconds': trial_seconds, + }, params={ + 'status': 'succeeded', + 'attempt': j + 1, + }) break except RuntimeError as e: + self._emit('trial_ended', trial=i + 1, params={ + 'status': 'failed', + 'attempt': j + 1, + 'error': format_error(e), + }) stacktrace = ''.join( traceback.format_exception(e, e, e.__traceback__)) logger.warning(stacktrace) @@ -957,6 +978,16 @@ def _run( should_init=init_best_trainer): return self.get_best_trainer(should_init=init_best_trainer) + def _emit(self, name: str, **fields: Any) -> None: + event = getattr(self.instrumentation, 'event', None) + if event is None: + return + try: + event(name, **fields) + except Exception: + logger.warning(f"Could not report progress event '{name}'", + exc_info=True) + def _handle_oom_retry( self, cfg: ModelPlan, diff --git a/packages/kumo-core/src/kumo_core/automl/optuna_searcher.py b/packages/kumo-core/src/kumo_core/automl/optuna_searcher.py index 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/searcher.py b/packages/kumo-core/src/kumo_core/automl/searcher.py index f8810e7b6..9116bf389 100644 --- a/packages/kumo-core/src/kumo_core/automl/searcher.py +++ b/packages/kumo-core/src/kumo_core/automl/searcher.py @@ -39,6 +39,7 @@ def __init__(self, space: PlanT, max_attempts: int = 50) -> None: self.space = space self.max_attempts = max_attempts self.state: list[tuple[PlanT, list[float]]] = [] + self.sampled_params: dict[str, Any] = {} @property def tune_metric(self) -> str: @@ -68,6 +69,7 @@ def sample_value(self, key: str, value: list[Any]) -> Any: def sample_instance(self) -> PlanT: instance = copy.deepcopy(self.space) + self.sampled_params = {} for field in fields(instance): section = getattr(instance, field.name) @@ -76,6 +78,7 @@ def sample_instance(self) -> PlanT: continue assert isinstance(value, list) + searched = len(value) > 1 if isinstance(value[0], PlanMixin): index_of_plan = self.sample_value(key, list(range(len(value)))) @@ -85,17 +88,25 @@ def sample_instance(self) -> PlanT: subvalue, MissingType): continue assert isinstance(subvalue, list) + searched = searched or len(subvalue) > 1 unique_key = f"{key}{index_of_plan}_{subkey}" setattr( plan, subkey, [self.sample_value(unique_key, subvalue)], ) - setattr(section, key, [plan]) + chosen = plan else: - setattr(section, key, [self.sample_value(key, value)]) + chosen = self.sample_value(key, value) + setattr(section, key, [chosen]) + if searched: + self.sampled_params[f"{field.name}.{key}"] = ( + chosen if chosen is None + or isinstance(chosen, (bool, int, float, str)) else + str(chosen)) return instance def __repr__(self) -> str: return f"{self.__class__.__name__}()" + diff --git a/packages/kumo-core/src/kumo_core/callbacks/train_progress_logger.py b/packages/kumo-core/src/kumo_core/callbacks/train_progress_logger.py index 45f124db5..abdbd7803 100644 --- a/packages/kumo-core/src/kumo_core/callbacks/train_progress_logger.py +++ b/packages/kumo-core/src/kumo_core/callbacks/train_progress_logger.py @@ -5,7 +5,7 @@ import logging import os.path as osp from datetime import datetime, timezone -from typing import Any +from typing import Any, Callable import kumoapi.jobs as jobs_api from kumo_core.io.json import load_json, save_json @@ -140,17 +140,22 @@ class AutoTrainingProgressLogger(MonitoringCallback): # pragma: no cover # Constant to account for overhead not captured in trial estimation: OVERHEAD_CONSTANT = 1.30 - def __init__(self, log_dir: str, num_experiments: int): + def __init__(self, log_dir: str, num_experiments: int, + on_event: Callable[..., None] | None = None): super().__init__() self.path = osp.join(log_dir, self.FILE_NAME) + self.on_event = on_event + self.trial: int | None = None self.progress = load_json(self.path, AutoTrainerProgress, raise_on_missing=False) if self.progress is None: self.progress = AutoTrainerProgress(total_trials=num_experiments) - def set_current_trainer(self, trainer: Any | None): + def set_current_trainer(self, trainer: Any | None, + trial: int | None = None): self.trainer = trainer + self.trial = trial def on_train_start(self, pl_trainer, pl_module): if self.progress.start_time is None: @@ -184,6 +189,20 @@ def on_train_epoch_end(self, pl_trainer, pl_module): save_json(self.progress, self.path) + if self.on_event is not None and self.trial is not None: + try: + progress = self.trainer.training_progress_logger.progress + epoch = progress.num_epochs_completed + metrics = progress.metrics[epoch - 1] + self.on_event('epoch_ended', trial=self.trial, epoch=epoch, + metrics={ + **metrics.train_metrics, + **metrics.validation_metrics + }) + except Exception: + logger.warning('Could not report epoch_ended', + exc_info=True) + def on_validation_end(self, pl_trainer, pl_module): if pl_trainer.state.fn == TrainerFn.VALIDATING: save_json(self.progress, self.path) diff --git a/packages/kumo-core/src/kumo_core/instrumentation.py b/packages/kumo-core/src/kumo_core/instrumentation.py index ab887f8aa..141bcd748 100644 --- a/packages/kumo-core/src/kumo_core/instrumentation.py +++ b/packages/kumo-core/src/kumo_core/instrumentation.py @@ -24,6 +24,9 @@ def span(self, name: str, **fields: Any) -> ContextManager[None]: def timing(self, name: str, **fields: Any) -> ContextManager[None]: ... + def event(self, name: str, **fields: Any) -> None: + return None + class NoOpTrainingInstrumentation: def record_metrics_job_name(self, job_name: str) -> None: @@ -37,6 +40,9 @@ def span(self, name: str, **fields: Any) -> Iterator[None]: def timing(self, name: str, **fields: Any) -> Iterator[None]: yield + def event(self, name: str, **fields: Any) -> None: + return None + _instrumentation: TrainingInstrumentation = NoOpTrainingInstrumentation() @@ -54,6 +60,10 @@ def timing(name: str, **fields: Any) -> ContextManager[None]: return _instrumentation.timing(name, **fields) +def format_error(error: BaseException) -> str: + return f'{type(error).__name__}: {error}'[:200] + + def trace() -> Callable[[Callable[..., T]], Callable[..., T]]: def decorate(function: Callable[..., T]) -> Callable[..., T]: @wraps(function) diff --git a/packages/kumo-core/src/kumo_core/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_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-engine/docs/mlops.md b/packages/kumo-engine/docs/mlops.md index bc2b84697..f8f21b41d 100644 --- a/packages/kumo-engine/docs/mlops.md +++ b/packages/kumo-engine/docs/mlops.md @@ -84,3 +84,19 @@ result = trainer.predict(snapshot, binary_classification_threshold=0.5) - `predict(prediction_table, snapshot=...)` keeps working as before. - `fit` results score too: `result.predict(snapshot, ...)`. - `PredictiveQuery.from_snapshot(snapshot, query)` binds a query to a snapshot without its `Graph`. Its graph check only means the snapshot belongs to the query's graph, not that it is compatible with a trained model. + +## 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()))]) +``` 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/tracking.py b/packages/kumo-engine/kumo/engine/tracking.py new file mode 100644 index 000000000..ef29cd7f5 --- /dev/null +++ b/packages/kumo-engine/kumo/engine/tracking.py @@ -0,0 +1,10 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +r"""Progress events and experiment-tracking types for `Trainer.fit`.""" + +from kumo.engine._events import Event + +__all__ = [ + 'Event', +] 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/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index 6e903e98d..085f2d8fb 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -20,6 +20,7 @@ Literal, Mapping, Optional, + Sequence, TypeVar, Union, cast, @@ -27,6 +28,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 @@ -53,6 +55,7 @@ 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._scratch import new_run_id, trainer_scratch_dir from kumo.engine.connector.base import Connector @@ -384,6 +387,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`. @@ -414,10 +418,60 @@ def fit( 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. + 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 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 + + settings = resolve(mlops) + run_id = f'{self._id}-fit-{new_run_id()}' + events = EventEmitter(run_id, 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) + except BaseException as error: + events.emit('fit_ended', params={ + 'status': 'failed', + 'error': format_error(error), + }) + raise + events.emit('fit_ended', metrics={ + **result.metrics, 'seconds': time.monotonic() - fit_start + }, params={'status': 'succeeded'}) + return result + + def _fit( + self, + train_table: TrainingTable, + *, + snapshot: GraphSnapshot, + export_config: Optional[ModelOutputConfig], + warm_start: Optional[Self], + run_id: str, + events: EventEmitter, + settings: Optional[Any], + ) -> TrainingResult: from kumo_core.automl.runtime import ( inherit_warm_start_model_plan, validate_warm_start_task_type, @@ -437,20 +491,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) @@ -569,7 +613,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(), ) 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 == []