[Vedaant-J]

Series: kumo-engine MLOps series: every snapshot, training table, model and prediction run carries a self-contained _provenance.json; with kumo.mlops(...) stores each artifact under one artifact root, tracks fits and prediction runs in MLflow or Snowflake and registers models; volumes and stages can hold files; prediction runs are monitored for drift and realized performance.
Stack: Experiment tracking: #52 → #60 → #68
Merge order: one linear stack; merge bottom to top (#53 first)
This PR: Adds progress events for Trainer.fit: fit(callbacks=[...]) calls each callback with a structured Event for the fit, every search trial and every epoch, so a job's logs or an experiment tracker can follow a fit while it runs.

Summary

This PR adds progress events for Trainer.fit: fit(callbacks=[...]) calls each callback with a structured Event for the fit, every search trial and every epoch, so a job's logs or an experiment tracker can follow a fit while it runs.

Problem. Teams running fit() as a Databricks, SageMaker or Snowflake job need to answer questions like these from the job's logs or tracker while it runs:

• Is my training job still progressing, or is it stuck?
• Which trial is best so far, and with which hyperparameters?
• Which epoch is a trial on, and are its validation metrics improving?
• Did a trial fail or get retried after an OOM, and with what error?

Today none of these can be answered without parsing progress-bar text.

Cause. fit() reports progress only as Lightning/tqdm progress-bar text on stdout, for example:

Epoch 1: 100%|██████████| 2/2 [00:01<00:00,  1.33it/s, val_loss=0.272, val_mae=0.628, val_mse=0.554, val_rmse=0.744, train_loss=0.758, ...]

That text has no trial number or hyperparameters, is overwritten in place with carriage returns, and rounds metrics. There is no programmatic hook either: the engine's training bindings hard-wire NoOpTrainingInstrumentation, and fit() takes no callbacks.

Change.

• kumo.engine.tracking.Event, a frozen dataclass with name, run_id, time, trial, epoch, metrics and params, plus to_dict().
• Trainer.fit(..., callbacks=[...]), where a callback is any Callable[[Event], None]. The default is no callbacks.
• A "Progress events" section in docs/mlops.md. JSON lines on stdout, for platform log capture, are one callback:

import json

trainer.fit(training_table, snapshot=snapshot,
            callbacks=[lambda e: print(json.dumps(e.to_dict()))])

Events are emitted in this order: fit_started; then per trial trial_started, epoch_ended once per epoch, and trial_ended; then fit_ended.

Example: the output of the callback above for a real two-trial, two-epoch CPU fit on a synthetic 400-row users table (PREDICT users.spend_next_30d FOR EACH users.user_id). The lines are trimmed to the first trial's events, the second trial's trial_ended, and fit_ended with its metrics cut to the *_mae keys (the real line also has loss/mse/rmse per split). Hashes are shortened and floats rounded to 4 digits:

{"name": "fit_started", "run_id": "f1070de2…-fit-50fe1fe6…", "time": 1792572589.95, "trial": null, "epoch": null, "metrics": {}, "params": {"snapshot_id": "77d088e9…@647e1e0794fba55c", "query_sha256": "cbb5b05e…"}}
{"name": "trial_started", "run_id": "f1070de2…-fit-50fe1fe6…", "time": 1792572590.21, "trial": 1, "epoch": null, "metrics": {}, "params": {"optimization.base_lr": 0.0001, "optimization.weight_decay": 5e-06, "optimization.batch_size": 1024, "optimization.lr_scheduler": "LRScheduler(\n  name=cosine_with_warmup_restarts,\n  interval=step,\n  kwargs={warmup_ratio_or_steps=0.1},\n)", "model_architecture.model": "GNNModelPlan(name='GNN', channels=[256], aggregation=[[sum, mean, min, max, std]], dropout=[0.0])", "model_architecture.num_pre_message_passing_layers": 2, "model_architecture.num_post_message_passing_layers": 1, "model_architecture.activation": "relu", "attempt": 1}}
{"name": "epoch_ended", "run_id": "f1070de2…-fit-50fe1fe6…", "time": 1792572593.15, "trial": 1, "epoch": 1, "metrics": {"val_loss": 1.0675, "val_mae": 1.4688, "val_mse": 3.5513, "val_rmse": 1.8845, "train_loss": 1.0454, "train_mae": 1.4242, "train_mse": 3.8349, "train_rmse": 1.9583}, "params": {}}
{"name": "epoch_ended", "run_id": "f1070de2…-fit-50fe1fe6…", "time": 1792572594.80, "trial": 1, "epoch": 2, "metrics": {"val_loss": 0.8225, "val_mae": 1.2265, "val_mse": 2.6122, "val_rmse": 1.6162, "train_loss": 0.8367, "train_mae": 1.2248, "train_mse": 2.7899, "train_rmse": 1.6703}, "params": {}}
{"name": "trial_ended", "run_id": "f1070de2…-fit-50fe1fe6…", "time": 1792572595.49, "trial": 1, "epoch": null, "metrics": {"val_mae": 1.2265, "seconds": 5.2858}, "params": {"status": "succeeded", "attempt": 1}}
{"name": "trial_ended", "run_id": "f1070de2…-fit-50fe1fe6…", "time": 1792572600.44, "trial": 2, "epoch": null, "metrics": {"val_mae": 0.9712, "seconds": 4.6662}, "params": {"status": "succeeded", "attempt": 1}}
{"name": "fit_ended", "run_id": "f1070de2…-fit-50fe1fe6…", "time": 1792572605.80, "trial": null, "epoch": null, "metrics": {"val_mae": 0.9712, "test_mae": 0.913, "test_refit_mae": 1.0872, "train_mae": 1.142, "seconds": 15.8578}, "params": {"status": "succeeded"}}

Here trial 2 (val_mae 0.971) beats trial 1 (1.227), and the fit_ended metrics are the final TrainingResult.metrics.

A fit that raises ends with a failed fit_ended before the exception propagates. From the same script, with the training-dataset build step patched to raise ValueError('training table has no rows for split "train"'):

{"name": "fit_ended", "run_id": "d602ea2b…-fit-8ba42ea7…", "time": 1792572605.81, "trial": null, "epoch": null, "metrics": {}, "params": {"status": "failed", "error": "ValueError: training table has no rows for split \"train\""}}

What each field records:

• run_id is the fit's run id, the same one that names the run's scratch directory. It is shared by every event of one fit.
• time is wall-clock epoch seconds. trial and epoch count from 1 and are null where they don't apply.
• fit_started.params holds the snapshot id and a SHA-256 of the query. The query text itself is never emitted. If these can't be read (malformed fit() arguments), the event has empty params and the fit then raises its usual error.
• trial_started.params holds the hyperparameters searched for that trial, keyed by their model-plan path, plus attempt. Values that are not plain numbers or strings appear as their string form.
• epoch_ended.metrics holds the train_* and val_* metrics logged at the end of that epoch.
• trial_ended gives the tune metric, seconds and status. It is emitted once per attempt: a succeeded trial_ended only after the trial is recorded, so a trial that trains and then fails during bookkeeping reports one failed trial_ended, not both. A failed attempt also gives error, the error type and message (at most 200 characters). An OOM retry is reported as a new attempt of the same trial.
• fit_ended gives the final metrics plus seconds, or status: failed and error (same format) if the fit raised.
• Metrics are finite floats only. Non-numeric or NaN/inf values are dropped without a warning.

In kumo-core, the training instrumentation gains an additive event(name, **fields) hook (a no-op by default) and local_training_bindings accepts an instrumentation. The engine passes an emitter that implements event and fans each event out to the callbacks. AutoTrainer emits the trial events, the existing AutoTrainingProgressLogger emits epoch_ended from the epoch it just recorded, and the searcher records the values it sampled. With no callbacks the emitter has nothing to call, so no extra work is done.

Back-compat

This PR only reports information. It adds no new checks or errors on any existing path.

Situation	Behaviour
Existing fit() call, no callbacks (the default)	Unchanged. Nothing is attached to kumo-core and nothing is printed; the bundle files and metric keys are the same.
fit(callbacks=...)	Callbacks receive the events above, in order. Bundle files and metric keys are the same as without callbacks.
A callback raises	A warning is logged and that callback is not called again for the rest of the fit. The fit continues, and other callbacks still run.
kumo-core cannot build a trial or epoch event	A warning is logged and that event is dropped. The fit continues.
Malformed fit() arguments (e.g. not a training table)	fit_started is emitted with empty params, then fit() raises the same documented error as without callbacks.
The fit raises	fit_ended with status: failed is delivered, then the same exception propagates.
fit() with an export_config whose connector cannot publish models	Still raises NotImplementedError, now before fit_started is emitted.

Known limitations

• Events cover fit only. predict emits none; with MLOps tracking on, prediction runs are recorded by the tracker instead (Log prediction runs to MLflow and Snowflake through the tracking adapters #73).
• Events are delivered synchronously on the training thread; a slow callback slows the fit. Record fits in experiment tracking and publish models to registries #60's tracking session moves tracker calls to a background thread.

Testing Done

• New test_fit_events.py (4 tests):
  • a two-trial CPU fit checks event order and numbering, metrics, the to_dict() JSON round trip, that a raising callback is called once and does not fail the fit, and that a trial whose searcher update fails after training reports exactly one failed trial_ended with the error;
  • a failed fit emits fit_started then fit_ended with status: failed and error: "ValueError: stop";
  • non-numeric and NaN metrics are dropped;
  • an epoch whose metrics were not recorded produces no epoch_ended and does not raise.
• New kumo-core test_searched_params.py: only searched fields are reported, non-scalar values as strings, JSON-serialisable.
• Re-run on the current head (rebased onto Score a snapshot with Trainer.predict(snapshot) and TrainingResult.predict #51, fit now also takes mlops=): test_fit_events.py + kumo-core test_searched_params.py: 5 passed; test_bundle_provenance.py (MLOps model store through the same fit): 7 passed.
• Not re-run after the last commits: test_fit_validation_metrics.py, test_trainer_registry.py, kumo-core test_training.py, test_local_training_runtime.py, make check-kumo-core.
• The example above was captured on an earlier head; event fields and order did not change, except that failure events now carry error as <Type>: <message>.

Checklist

• No breaking changes: fit() without callbacks behaves as before.
• Unit tests added.

🤖 Generated with Claude Code

[Vedaant-J]

Change map

File	Function / scope	What changed	Why
packages/kumo-core/src/kumo_core/automl/auto_trainer.py	AutoTrainer.run, _train, _emit	emits trial_started (searched params, attempt) and trial_ended once per attempt (succeeded only after the trial is recorded; failed with error); passes the trial number to the progress logger	a tracker sees each trial's start, outcome and duration
packages/kumo-core/src/kumo_core/automl/optuna_searcher.py	OptunaSearcher.sample_instance	keeps sampled_params from the fallback sampler	trial params are reported with the Optuna searcher too
packages/kumo-core/src/kumo_core/automl/searcher.py	Searcher.sample_instance, sampled_params	records the values sampled for fields that have more than one choice; non-scalar values as strings	trial_started params without re-walking the search space
packages/kumo-core/src/kumo_core/callbacks/train_progress_logger.py	AutoTrainingProgressLogger.on_train_epoch_end, set_current_trainer	emits epoch_ended with the train_* / val_* metrics of the epoch it just recorded; a failed lookup only warns	reuses the existing epoch callback instead of adding a second one
packages/kumo-core/src/kumo_core/instrumentation.py	TrainingInstrumentation.event, format_error	additive no-op event hook; format_error builds <Type>: <message> cut to 200 chars	one channel for events, one error format
packages/kumo-core/src/kumo_core/training/runtime.py	local_training_bindings	accepts instrumentation (default no-op)	lets the engine receive events
packages/kumo-core/tests/test_searched_params.py	new	only searched fields are reported; non-scalars as strings; JSON-safe	covers sampled_params
packages/kumo-engine/docs/mlops.md	"Progress events"	new section: event order, status/error, raising callbacks, JSON-lines example	user docs for fit(callbacks=)
packages/kumo-engine/kumo/engine/_events.py	Event, EventEmitter, query_sha256, _finite_floats	new: event type with to_dict(); emitter implements the event hook; a raising callback is disabled for the rest of the fit; finite-float metrics only	public event API
packages/kumo-engine/kumo/engine/tracking.py	module	new: public home of Event (later PRs add the tracking adapter types)	Event is not a top-level kumo.engine name
packages/kumo-engine/kumo/engine/trainer/_runtime.py	local_training_bindings	forwards instrumentation	engine-to-core adapter
packages/kumo-engine/kumo/engine/trainer/trainer.py	Trainer.fit, _fit	callbacks=; fit_started / fit_ended (status, error) around the previous body (moved to _fit); fit_started params are best-effort; the publish-capability check and the mlops= settings resolution run before fit_started and the settings are passed to _fit	fit-level events without changing the fit path
packages/kumo-engine/test/test_fit_events.py	new	event order and payloads, failed fit, late trial failure reports one failed trial_ended, raising callback, metric filtering, epoch without recorded metrics does not raise	coverage

Requirement: training progress must reach experiment trackers and platform log capture as structured data. Existing behaviour: no hook; instrumentation is a no-op. Smallest change: one additive instrumentation hook plus fit(callbacks=); JSON lines are a one-line callback, not a library setting. Stopping condition: fit/trial/epoch events delivered in order with status and errors; behaviour unchanged when no callback is passed.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[coderabbitai]

Note

Reviews paused

It looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the reviews.auto_review.auto_pause_after_reviewed_commits setting.

• Last reviewed commit: 17329c9
• Latest commit, not reviewed: 7cdc1ef

Use the following commands to manage reviews:

• @coderabbitai resume to resume automatic reviews.
• @coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

• ▶️ Resume reviews
• 🔍 Trigger review

No actionable comments were generated in the recent review.

ℹ️ Recent review info

⚙️ Run configuration

• Configuration used: Repository: NVIDIA/kumo-engine/.coderabbit.yaml
• Review profile: QUIET
• Plan: Enterprise
• Run ID: 3645eaae-4326-4a3c-89d3-90bf106173bb

📥 Commits

Reviewing files that changed from the base of the PR and between d858fbc and 17329c9.

📒 Files selected for processing: 1

• packages/kumo-engine/kumo/engine/trainer/trainer.py

Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 11 remain after this review.

Walkthrough

Trainer.fit now accepts callbacks and emits fit lifecycle events. Core training reports trial and epoch events with searched parameters, metrics, duration, attempt, and failure details. Event payloads include a run ID and timestamp. Callback failures are logged and do not interrupt fitting. Tests cover event ordering, payloads, failure paths, metric filtering, and training result parity. Documentation includes callback usage and JSON-lines output.

Sequence Diagram(s)

sequenceDiagram
  participant Trainer.fit
  participant EventEmitter
  participant local_training_bindings
  participant AutoTrainer
  participant Callbacks
  Trainer.fit->>EventEmitter: Emit fit_started
  Trainer.fit->>local_training_bindings: Pass core_sink
  local_training_bindings->>AutoTrainer: Pass on_event
  AutoTrainer->>EventEmitter: Send trial and epoch events
  EventEmitter->>Callbacks: Deliver events
  Trainer.fit->>EventEmitter: Emit fit_ended

Loading

Priority: Normal

Merge Risk: Low · up to 17329

Training failures still propagate, but some failed trials may lack a final progress event, and one callback may alter data received by another. These bounded event-reporting risks should be accepted or fixed before relying on callbacks for tracking.

Pre-merge checks: 2 · 1

Failed checks: 1 inconclusive

Check name	Status	Explanation	Resolution
Change Map Design Review		The required marked PR conversation comment is not available. The repository contains only the change-map template and validation code; the supplied context provides summaries, not the exact `<!-- kum…	Provide the complete PR conversation comment beginning with <!-- kumo-pr-change-map -->, or enable retrieval of the PR comments, then reassess the design against the changed-file diff.

Passed checks: 2 passed

Check name	Status	Explanation
Pr Size		The pull request changes 12 files. The limit is 40 files, so the check passes.
Public Resource Access		The pull request adds no external URL, repository location, or package source location. The new README section contains only local API names and a JSON callback example. Existing external links in the…

Full details: Public Resource Access

Explanation

The pull request adds no external URL, repository location, or package source location. The new README section contains only local API names and a JSON callback example. Existing external links in the README are unchanged. Local fixture paths in the new tests are excluded by the check.

Full details: Change Map Design Review

Explanation

The required marked PR conversation comment is not available. The repository contains only the change-map template and validation code; the supplied context provides summaries, not the exact <!-- kumo-pr-change-map --> comment or its table. This check requires reviewing that table row by row, including its stated failure, existing behavior, smallest change, verification, and stopping condition. The changed-file diff is available, but it cannot establish what the proposed design claimed or what material the map omitted.

Resolution

Provide the complete PR conversation comment beginning with <!-- kumo-pr-change-map -->, or enable retrieval of the PR comments, then reassess the design against the changed-file diff.

✨ Finishing Touches: 💡 1

🛠️ Fix failing CI checks: 💡

• Commit to this branch
• Create a new PR

• Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

Actionable comments posted: 1

---

ℹ️ Review info

⚙️ Run configuration

• Configuration used: Repository: NVIDIA/kumo-engine/.coderabbit.yaml
• Review profile: QUIET
• Plan: Enterprise
• Run ID: e8b43b2e-69aa-4d3f-bea4-e1f3db823fbc

📥 Commits

Reviewing files that changed from the base of the PR and between 9bcbaca and d45e5ed.

📒 Files selected for processing: 12

• packages/kumo-core/src/kumo_core/automl/auto_trainer.py
• packages/kumo-core/src/kumo_core/automl/searcher.py
• packages/kumo-core/src/kumo_core/callbacks/train_progress_logger.py
• packages/kumo-core/src/kumo_core/training/orchestration.py
• packages/kumo-core/src/kumo_core/training/runtime.py
• packages/kumo-core/tests/test_searched_params.py
• packages/kumo-engine/kumo/engine/__init__.py
• packages/kumo-engine/kumo/engine/_events.py
• packages/kumo-engine/kumo/engine/config.py
• packages/kumo-engine/kumo/engine/trainer/_runtime.py
• packages/kumo-engine/kumo/engine/trainer/trainer.py
• packages/kumo-engine/test/test_fit_events.py

Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 6 remain after this review.

[coderabbitai]

Caution

Some comments are outside the diff and can’t be posted inline due to GitHub limitations.

⚠️ Outside diff range comments (1)

🟡 Minor · Keep nonfinite parameters from disabling JSON progress. · _events.py:76-83

packages/kumo-engine/kumo/engine/_events.py:76-83
🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win

Keep nonfinite parameters from disabling JSON progress.

A positive infinity can reach trial_started.params through the tunable base_lr search value. _print_json passes that value to json.dumps(..., allow_nan=False), which raises. EventEmitter.emit then removes _print_json, so that event has no JSON line and later progress events are no longer printed. Normalize nonfinite parameter values to JSON null before serialization.

Suggested fix

+def _json_safe(value: Any) -> Any:
+    if isinstance(value, float) and not math.isfinite(value):
+        return None
+    if isinstance(value, Mapping):
+        return {key: _json_safe(item) for key, item in value.items()}
+    if isinstance(value, (list, tuple)):
+        return [_json_safe(item) for item in value]
+    return value
+
+
 def _print_json(event: Event) -> None:
-    line = json.dumps({'event': 'kumo.engine', **event.to_dict()},
+    payload = event.to_dict()
+    payload['params'] = _json_safe(payload['params'])
+    line = json.dumps({'event': 'kumo.engine', **payload},
                       allow_nan=False, default=str)

---

ℹ️ Review info

⚙️ Run configuration

• Configuration used: Repository: NVIDIA/kumo-engine/.coderabbit.yaml
• Review profile: QUIET
• Plan: Enterprise
• Run ID: bd991a3d-f0da-47b5-91f0-ab8b380c97a4

📥 Commits

Reviewing files that changed from the base of the PR and between d45e5ed and 096fa30.

📒 Files selected for processing: 1

• packages/kumo-core/src/kumo_core/automl/auto_trainer.py

Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 8 remain after this review.

[coderabbitai]

Note

Quiet mode is enabled, so only the most important comments were posted inline. Other review comments are grouped below.

🟡 Other comments: 1

packages/kumo-core/src/kumo_core/automl/auto_trainer.py: 1

942-942: 🗄️ Data Integrity & Integration | 🟡 Minor | ⚡ Quick win

End trials when status persistence fails.

If status.save() raises OSError after trial_started, this RuntimeError handler does not emit trial_ended. The outer fit reports failure, but the callback still has an open trial. Add a separate handler that emits a failed trial_ended for other exceptions and re-raises them. Keep the existing retry behavior for RuntimeError.

---

ℹ️ Review info

⚙️ Run configuration

• Configuration used: Repository: NVIDIA/kumo-engine/.coderabbit.yaml
• Review profile: QUIET
• Plan: Enterprise
• Run ID: f8fa803b-a4ee-4cfb-afc6-93d6677fb3ba

📥 Commits

Reviewing files that changed from the base of the PR and between 096fa30 and 8007f0e.

📒 Files selected for processing: 6

• packages/kumo-core/src/kumo_core/automl/auto_trainer.py
• packages/kumo-core/src/kumo_core/callbacks/train_progress_logger.py
• packages/kumo-engine/README.md
• packages/kumo-engine/kumo/engine/_events.py
• packages/kumo-engine/kumo/engine/trainer/trainer.py
• packages/kumo-engine/test/test_fit_events.py

💤 Files with no reviewable changes: 2

• packages/kumo-core/src/kumo_core/callbacks/train_progress_logger.py
• packages/kumo-engine/kumo/engine/_events.py

Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 7 remain after this review.

[coderabbitai]

Note

Quiet mode is enabled, so only the most important comments were posted inline. Other review comments are grouped below.

🟡 Other comments: 1

packages/kumo-engine/kumo/engine/_events.py: 1

58-59: 🗄️ Data Integrity & Integration | 🟡 Minor | ⚡ Quick win

Isolate event data between callbacks.

Event is frozen, but its metrics and params dictionaries are writable. If one callback changes either dictionary, later callbacks receive the changed values. This can make a JSON-lines callback or tracker record incorrect progress without any callback raising an error. Make the event data immutable, or give each callback an independent copy.

---

ℹ️ Review info

⚙️ Run configuration

• Configuration used: Repository: NVIDIA/kumo-engine/.coderabbit.yaml
• Review profile: QUIET
• Plan: Enterprise
• Run ID: 9993e071-1343-408f-9450-23db06591684

📥 Commits

Reviewing files that changed from the base of the PR and between 8007f0e and d858fbc.

📒 Files selected for processing: 3

• packages/kumo-engine/kumo/engine/_events.py
• packages/kumo-engine/kumo/engine/trainer/trainer.py
• packages/kumo-engine/test/test_fit_events.py

Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 10 remain after this review.

[Vedaant-J]

Folded into #60 (tracking & registry), which now contains this PR's final code: the consolidated stack shows only the end state. Review continues there.