Return the published model's holdout, independent of explain

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,29 +1517,23 @@ 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
 
+        refit = not original and fit_mode.is_refit
         # 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 path[-1] == '/':
-                    path = path[:-1]
-                path = f"{path}_{self.fit_mode.value}"
-                fsutil.makedirs(path, exist_ok=True)
+            if not refit:
                 return path
-            else:
-                return path
-
-        if self.name is not None:
-            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 path[-1] == '/':
-                path = path[:-1]
-            path = f"{path}_{self.fit_mode.value}"
+        else:
+            if self.name is not None:
+                path = osp.join(path, self.name)
+            if self.track_version:
+                path = osp.join(path, self.version_dir_name)
+        if refit:
+            path = f"{path.removesuffix('/')}_{fit_mode.value}"
         fsutil.makedirs(path, exist_ok=True)
         return path
 
@@ -2531,16 +2526,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)

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):

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

@@ -34,6 +34,7 @@
     resolve_query_forecast,
 )
 from kumo_core.train.artifact_manifest import PredictionSelectionRef
+from kumo_core.training.config import FitMode
 from kumo_core.training.explain_aggregation import aggregate_subgraph_features
 from kumo_core.training.workflow import TrainingWorkflowRequest
 from kumo_core.types import ExplainMethod, PredictionArtifactType
@@ -77,11 +78,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 +189,15 @@ def _find_holdout(log_dir: str) -> tuple[Optional[Path], Optional[bool]]:
     return path, truncated
 
 
+def _published_holdout(trainer: Any) -> tuple[Optional[Path], Optional[bool]]:
+    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 +505,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),

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,24 @@ def test_holdout_df_after_fit(graph):
     assert not (result.bundle_dir / 'test_df.parquet').exists()
 
 
+@pytest.mark.parametrize('refit_trainval,refit_full,explain,model', [
+    (False, False, False, 'fit'),
+    (True, False, True, 'refit_trainval'),
+    (True, True, True, None),
+])
+def test_holdout_is_the_published_models(graph, refit_trainval, refit_full,
+                                         explain, model):
+    result = _fit(graph, QUERY, refit_trainval, refit_full, explain)
+
+    if model is None:
+        assert result.holdout_path is None
+    else:
+        suffix = '' if model == 'fit' else '_refit_trainval'
+        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) '

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)
@@ -869,10 +874,9 @@ def fake_train(training_input, **kwargs):
         both_workers.wait(timeout=5)
         time.sleep(.05)
         run_name = Path(training_input.home_dir).name
-        return SimpleNamespace(
-            metrics={'run_name':
-                     run_name}, artifact_paths={}, explain_disabled=True,
-            best_trainer=SimpleNamespace(log_dir=str(tmp_path)))
+        return SimpleNamespace(metrics={'run_name': run_name},
+                               artifact_paths={}, explain_disabled=True,
+                               best_trainer=_best_trainer(tmp_path))
 
     monkeypatch.setattr(
         'kumo.engine.trainer.trainer.build_generated_training_task_dataset',

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):