Report fit progress events to callbacks

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,20 +874,27 @@ 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(
                             cfg=cfg,
                             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,

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

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__}()"
+

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)

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)

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.

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)

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()))])
+```