From 0e6ea1cb5a05e49c877a0040e8706633e76f5438 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Wed, 21 Oct 2026 22:14:31 +0000 Subject: [PATCH 1/2] Read the holdout of the published model, not the last-evaluated one fit() read test_df.parquet from best_trainer.log_dir, which is whichever refit stage ran last, not necessarily the model the bundle publishes. With refit_full the published model has seen the test split, so its holdout is meaningless. Read the holdout from the published model's log dir (Trainer.published_fit_mode, also used to pick the prediction model) and return none for refit_full. get_log_dir takes an optional fit_mode override. Verified: test_holdout_df (fit / refit_trainval with explain / refit_full), test_training_explanations, test_semantic_carry_forward: 35 passed. Co-Authored-By: Claude Opus 5.5 --- .../kumo-core/src/kumo_core/train/trainer.py | 33 +++++++++++-------- .../kumo/engine/trainer/trainer.py | 26 +++++++++++---- packages/kumo-engine/test/test_holdout_df.py | 28 ++++++++++++++-- .../test/test_semantic_carry_forward.py | 15 ++++++--- .../test/test_training_explanations.py | 3 +- 5 files changed, 76 insertions(+), 29 deletions(-) diff --git a/packages/kumo-core/src/kumo_core/train/trainer.py b/packages/kumo-core/src/kumo_core/train/trainer.py index b18f552e8..efe6968f2 100644 --- a/packages/kumo-core/src/kumo_core/train/trainer.py +++ b/packages/kumo-core/src/kumo_core/train/trainer.py @@ -1507,6 +1507,7 @@ def get_log_dir( self, local: bool = False, original: bool = False, + fit_mode: FitMode | None = None, ) -> str: r"""Switch the log directory based on `self.fit_mode`. `version_xx/`: fit mode. @@ -1516,16 +1517,17 @@ def get_log_dir( If `original=True`, then we get the original (non-refitting) log directory. This way, we can update `config.json` in the original - directory""" + directory. `fit_mode` overrides `self.fit_mode`.""" + fit_mode = self.fit_mode if fit_mode is None else fit_mode path = self.root_dir if not local else self.local_root_dir # In case the `root_dir` contains a trainer config, it is a log dir # for a previous model version. if fsutil.exists(osp.join(path, 'config.json')): - if not original and self.fit_mode.is_refit: + if not original and fit_mode.is_refit: if path[-1] == '/': path = path[:-1] - path = f"{path}_{self.fit_mode.value}" + path = f"{path}_{fit_mode.value}" fsutil.makedirs(path, exist_ok=True) return path else: @@ -1535,10 +1537,10 @@ def get_log_dir( path = osp.join(path, self.name) if self.track_version: path = osp.join(path, self.version_dir_name) - if not original and self.fit_mode.is_refit: + if not original and fit_mode.is_refit: if path[-1] == '/': path = path[:-1] - path = f"{path}_{self.fit_mode.value}" + path = f"{path}_{fit_mode.value}" fsutil.makedirs(path, exist_ok=True) return path @@ -2531,16 +2533,21 @@ def _preprocess_top_k(self, top_k: int | None): # Batch prediction ######################################################## + def published_fit_mode(self) -> FitMode | None: + r"""The fit mode predictions use by default: the model trained on the + most data, or `None` before fitting.""" + if self.refit_cfg.refitted_full: + return FitMode.refit_full + if self.refit_cfg.refitted_trainval: + return FitMode.refit_trainval + if self.refit_cfg.fitted: + return FitMode.fit + return None + def _set_fit_mode_to_use(self, fit_mode_to_use: FitMode | None = None): - # Infer `fit_mode_to_use` - Use model that's trained on most data: if fit_mode_to_use is None: - if self.refit_cfg.refitted_full: - fit_mode_to_use = FitMode.refit_full - elif self.refit_cfg.refitted_trainval: - fit_mode_to_use = FitMode.refit_trainval - elif self.refit_cfg.fitted: - fit_mode_to_use = FitMode.fit - else: + fit_mode_to_use = self.published_fit_mode() + if fit_mode_to_use is None: raise RuntimeError("Trainer is not fitted yet.") else: self._check_readiness_for_prediction(fit_mode_to_use) diff --git a/packages/kumo-engine/kumo/engine/trainer/trainer.py b/packages/kumo-engine/kumo/engine/trainer/trainer.py index 1b525b010..28c3156a3 100644 --- a/packages/kumo-engine/kumo/engine/trainer/trainer.py +++ b/packages/kumo-engine/kumo/engine/trainer/trainer.py @@ -77,11 +77,12 @@ class TrainingResult: `published_model` records the destination's assigned name, version and URI when `fit()` exports the bundle. It is `None` without an export. - `holdout_path` is the holdout Parquet file written by `fit()`, or `None` - when the task produced none (also for results of `Trainer.load`). It is - not part of the published model bundle. The file is capped at 1 GiB in - memory; `holdout_truncated` is true when rows were dropped to fit, and - `None` when the size record could not be read. + `holdout_path` is the published model's holdout Parquet file written by + `fit()`, or `None` when it has none (no holdout task, a model refit on the + full data, or results of `Trainer.load`). It is not part of the published + model bundle. The file is capped at 1 GiB in memory; `holdout_truncated` is + true when rows were dropped to fit, and `None` when the size record could + not be read. """ metrics: Dict[str, float] @@ -187,6 +188,17 @@ def _find_holdout(log_dir: str) -> tuple[Optional[Path], Optional[bool]]: return path, truncated +def _published_holdout(trainer: Any) -> tuple[Optional[Path], Optional[bool]]: + from kumo_core.training.config import FitMode + + published = trainer.published_fit_mode() + if published is FitMode.refit_full: + logger.info('The published model is refit on the full data, including ' + 'the test split, so it has no holdout predictions.') + return None, False + return _find_holdout(trainer.get_log_dir(fit_mode=published)) + + def _collect_garbage(stage: str) -> None: # Lightning links each trainer to its data module and model in reference # cycles; free a finished run now rather than at a later collection. @@ -494,8 +506,8 @@ def fit( 'failed for run %s. The trained model remains available.', run_id) - holdout_path, holdout_truncated = _find_holdout( - result.best_trainer.log_dir) + holdout_path, holdout_truncated = _published_holdout( + result.best_trainer) training_result = TrainingResult( metrics=dict(result.metrics), artifact_paths=dict(result.artifact_paths), diff --git a/packages/kumo-engine/test/test_holdout_df.py b/packages/kumo-engine/test/test_holdout_df.py index 8c03af648..70bfe82c6 100644 --- a/packages/kumo-engine/test/test_holdout_df.py +++ b/packages/kumo-engine/test/test_holdout_df.py @@ -2,6 +2,8 @@ # All rights reserved. # SPDX-License-Identifier: Apache-2.0 +import re + import numpy as np import pandas as pd import pytest @@ -65,13 +67,16 @@ def graph(tmp_path_factory): QUERY = 'PREDICT COUNT(events.*, 0, 7, days) > 0 FOR EACH ents.ID' -def _fit(graph, query): +def _fit(graph, query, refit_trainval=True, refit_full=False, explain=False): graph, snapshot = graph pquery = kumo.PredictiveQuery(graph=graph, query=query) plan = ModelPlan() - plan.training_job.disable_explain = True + plan.training_job.disable_explain = not explain + plan.training_job.refit_trainval = refit_trainval + plan.training_job.refit_full = refit_full plan.optimization.max_epochs = 1 - plan.optimization.max_steps_per_epoch = 2 + # refit_full splits each epoch's batches across 3 stages. + plan.optimization.max_steps_per_epoch = 3 if refit_full else 2 trainer = kumo.Trainer(model_plan=plan) return trainer.fit(pquery.generate_training_table(snapshot), snapshot=snapshot) @@ -90,6 +95,23 @@ def test_holdout_df_after_fit(graph): assert not (result.bundle_dir / 'test_df.parquet').exists() +@pytest.mark.parametrize('refit_trainval,refit_full,explain,suffix', [ + (False, False, False, ''), + (True, False, True, '_refit_trainval'), + (True, True, False, None), +]) +def test_holdout_is_the_published_models(graph, refit_trainval, refit_full, + explain, suffix): + result = _fit(graph, QUERY, refit_trainval, refit_full, explain) + + if suffix is None: + assert result.holdout_path is None + else: + assert re.fullmatch(rf'version_\d+{suffix}', + result.holdout_path.parent.name) + assert len(result.holdout_df()) > 0 + + def test_no_holdout(graph): result = _fit( graph, 'PREDICT LIST_DISTINCT(events.ITEM_ID, 0, 7, days) ' diff --git a/packages/kumo-engine/test/test_semantic_carry_forward.py b/packages/kumo-engine/test/test_semantic_carry_forward.py index 15a4e2330..a1533b131 100644 --- a/packages/kumo-engine/test/test_semantic_carry_forward.py +++ b/packages/kumo-engine/test/test_semantic_carry_forward.py @@ -44,6 +44,11 @@ def _stub_generated_task_dataset_builder(monkeypatch): ) +def _best_trainer(log_dir): + return SimpleNamespace(get_log_dir=lambda fit_mode: str(log_dir), + published_fit_mode=lambda: None) + + def _write_prediction_shard(options, frame): from pathlib import Path @@ -406,7 +411,7 @@ def fake_materialize(request, task_dataset, generated, **kwargs): metrics={}, artifact_paths={}, explain_disabled=True, - best_trainer=SimpleNamespace(log_dir=str(tmp_path)), + best_trainer=_best_trainer(tmp_path), )) monkeypatch.setattr('kumo.engine.trainer.trainer.write_training_spec', lambda *args, **kwargs: None) @@ -536,7 +541,7 @@ def fake_materialize(request, task_dataset, generated, **kwargs): metrics={}, artifact_paths={}, explain_disabled=True, - best_trainer=SimpleNamespace(log_dir=str(tmp_path)), + best_trainer=_best_trainer(tmp_path), )) monkeypatch.setattr('kumo.engine.trainer.trainer.write_training_spec', lambda *args, **kwargs: None) @@ -612,7 +617,7 @@ def fake_materialize(request, task_dataset, generated, **kwargs): metrics={}, artifact_paths={}, explain_disabled=True, - best_trainer=SimpleNamespace(log_dir=str(tmp_path)), + best_trainer=_best_trainer(tmp_path), )) monkeypatch.setattr('kumo.engine.trainer.trainer.write_training_spec', lambda *args, **kwargs: None) @@ -791,7 +796,7 @@ def fake_materialize(request, task_dataset, generated, **kwargs): metrics={}, artifact_paths={}, explain_disabled=True, - best_trainer=SimpleNamespace(log_dir=str(tmp_path)), + best_trainer=_best_trainer(tmp_path), )) monkeypatch.setattr('kumo.engine.trainer.trainer.write_training_spec', lambda *args, **kwargs: None) @@ -872,7 +877,7 @@ def fake_train(training_input, **kwargs): return SimpleNamespace( metrics={'run_name': run_name}, artifact_paths={}, explain_disabled=True, - best_trainer=SimpleNamespace(log_dir=str(tmp_path))) + best_trainer=_best_trainer(tmp_path)) monkeypatch.setattr( 'kumo.engine.trainer.trainer.build_generated_training_task_dataset', diff --git a/packages/kumo-engine/test/test_training_explanations.py b/packages/kumo-engine/test/test_training_explanations.py index 5e8ed9df1..399dd1eed 100644 --- a/packages/kumo-engine/test/test_training_explanations.py +++ b/packages/kumo-engine/test/test_training_explanations.py @@ -51,7 +51,8 @@ def explain(**kwargs): best_trainer = SimpleNamespace( explain=Mock(side_effect=explain), subgraph_collector_dir=lambda: str(tmp_path / 'subgraph_collector'), - log_dir=str(tmp_path), + get_log_dir=lambda fit_mode: str(tmp_path), + published_fit_mode=lambda: None, ) def aggregate_features(*args): From 86fa3c6c46ba5c0dc521ad4a52547fd50bd92d18 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Thu, 22 Oct 2026 00:34:55 +0000 Subject: [PATCH 2/2] Expose published_fit_mode on the artifact test's fake trainer test_artifacts.py's fake trainer borrows Trainer methods; select_fit_mode now calls published_fit_mode, so the fake needs it too. Co-Authored-By: Claude Opus 5.5 --- packages/kumo-core/tests/test_artifacts.py | 1 + 1 file changed, 1 insertion(+) diff --git a/packages/kumo-core/tests/test_artifacts.py b/packages/kumo-core/tests/test_artifacts.py index 35c2cef65..071f4339b 100644 --- a/packages/kumo-core/tests/test_artifacts.py +++ b/packages/kumo-core/tests/test_artifacts.py @@ -199,6 +199,7 @@ def test_vital_source_resolution_is_core_owned(tmp_path: Path, refit) -> None: class SourceTrainer(SimpleNamespace): _set_fit_mode_to_use = Trainer._set_fit_mode_to_use + published_fit_mode = Trainer.published_fit_mode @property def best_ckpt(self):