Allow None for max_val_steps and max_test_steps

packages/kumo-api/kumoapi/model_plan.py

@@ -1070,14 +1070,17 @@ class OptimizationPlan(PlanMixin):
         mini-batch (*default:* ``run_mode``-dependent).
         **Supported Task Types:** All
 
-    :ivar max_val_steps: (``int``) The maximum number of steps to be included
-        in a validation pass; one step corresponds to one forward pass of a
-        mini-batch (*default:* ``run_mode``-dependent).
+    :ivar max_val_steps: (``int | None``) The maximum number of steps to be
+        included in a validation pass; one step corresponds to one forward pass
+        of a mini-batch (*default:* ``run_mode``-dependent). ``None`` evaluates
+        the whole validation split every epoch; time and memory then grow
+        with the split size, since the split's logits are held in memory.
         **Supported Task Types:** All
 
-    :ivar max_test_steps: (``int``) The maximum number of steps to be included
-        in a test pass; one step corresponds to one forward pass of a
-        mini-batch (*default:* ``run_mode``-dependent).
+    :ivar max_test_steps: (``int | None``) The maximum number of steps to be
+        included in a test pass; one step corresponds to one forward pass of a
+        mini-batch (*default:* ``run_mode``-dependent). ``None`` evaluates the
+        whole test split; memory then grows with the split size.
         **Supported Task Types:** All
 
     :ivar loss: (``list[str]``) The loss type to use in the model optimizer
@@ -1166,15 +1169,17 @@ class OptimizationPlan(PlanMixin):
         metadata=Metadata(),
     )
 
-    max_val_steps: Union[conint(ge=1, le=8000), MissingType] = compat_field(
-        default_factory=lambda: MissingType.VALUE,
-        metadata=Metadata(),
-    )
+    max_val_steps: Union[conint(ge=1, le=8000), MissingType,
+                         None] = compat_field(
+                             default_factory=lambda: MissingType.VALUE,
+                             metadata=Metadata(),
+                         )
 
-    max_test_steps: Union[conint(ge=1, le=8000), MissingType] = compat_field(
-        default_factory=lambda: MissingType.VALUE,
-        metadata=Metadata(),
-    )
+    max_test_steps: Union[conint(ge=1, le=8000), MissingType,
+                          None] = compat_field(
+                              default_factory=lambda: MissingType.VALUE,
+                              metadata=Metadata(),
+                          )
 
     loss: Union[compat_conlist(
         Union[

test/config/test_model_plan.py

@@ -1608,6 +1608,30 @@ def test_model_plan_serialization():
     assert "Link Prediction" in str(errors[0])
 
 
+@pytest.mark.parametrize('max_steps, expected', [
+    (None, (None, None)),
+    (MissingType.VALUE, (2000, 4000)),
+])
+def test_model_plan_eval_steps(max_steps, expected):
+    task = (TaskType.BINARY_CLASSIFICATION, QueryType.STATIC)
+    plan = ModelPlan()
+    plan.optimization.max_val_steps = max_steps
+    plan.optimization.max_test_steps = max_steps
+    plan = model_plan_from_yaml(model_plan_to_yaml(plan, *task))
+    plan = from_json(to_json(plan), ModelPlan)
+
+    errors, _ = plan.validate(
+        *task,
+        has_train_table_time_col=False,
+        run_mode=RunMode.NORMAL,
+        raise_on_error=False,
+    )
+
+    assert errors == []
+    assert (plan.optimization.max_val_steps,
+            plan.optimization.max_test_steps) == expected
+
+
 def test_model_plan_metric_assignment():
     plan = ModelPlan.from_default(
         task_type=TaskType.BINARY_CLASSIFICATION,

test/train/datamodule/test_diskgraph_lightning_data.py

@@ -904,24 +904,38 @@ def test_val_dataloader_does_not_replay_wheel_iterator(
     assert len(calls["history"]) == 1
 
 
-def test_val_dataloader_honors_max_val_steps_without_replay(
+@pytest.mark.parametrize("stage", ["val", "test"])
+@pytest.mark.parametrize("max_steps, num_batches, num_rows", [
+    (3, 3, 6),
+    (None, 10, 20),
+])
+def test_eval_dataloader_honors_max_steps_without_replay(
     monkeypatch: pytest.MonkeyPatch,
     tmp_path,
+    stage: str,
+    max_steps: int | None,
+    num_batches: int,
+    num_rows: int,
 ) -> None:
-    calls = _install_fake_diskgraph(monkeypatch, batches_per_iterator=10)
+    input_ids = [[2 * i, 2 * i + 1] for i in range(10)]
+    calls = _install_fake_diskgraph(monkeypatch,
+                                    batch_input_ids_by_stage={stage: input_ids})
     dm = DiskgraphLightningDataModule(
-        _manifest(tmp_path=tmp_path, stage_counts={"val": 20}),
+        _manifest(tmp_path=tmp_path, stage_counts={stage: 20}),
         _mapper_bank(),
         sampling_plan={"layers": []},
         train_batch_size=2,
         eval_batch_size=2,
-        max_val_steps=3,
+        **{f"max_{stage}_steps": max_steps},
     )
+    loader = getattr(dm, f"{stage}_dataloader")
 
-    assert len(dm.val_dataloader()) == 3
-    batches = list(dm.val_dataloader())
+    assert len(loader()) == num_batches
+    batches = list(loader())
 
-    assert len(batches) == 3
+    assert len(batches) == num_batches
+    assert sum(batch["train_table"].input_id.numel()
+               for batch in batches) == num_rows
     assert len(calls["history"]) == 1
 
 

test/train/test_diskgraph_nested_trainer.py

@@ -849,6 +849,17 @@ def test_init_passes_eval_limits_and_runtime_options(
     }
 
 
+def test_init_passes_uncapped_eval_limits(
+        monkeypatch: pytest.MonkeyPatch) -> None:
+    stub = _StubTrainer(max_val_steps=None, max_test_steps=None)
+
+    _run_init(stub, monkeypatch=monkeypatch)
+
+    dm_kwargs = stub.datamodule_kwargs  # type: ignore[attr-defined]
+    assert dm_kwargs["max_val_steps"] is None
+    assert dm_kwargs["max_test_steps"] is None
+
+
 def test_init_applies_diskgraph_auto_config_to_ttng_runtime(
         monkeypatch: pytest.MonkeyPatch) -> None:
     """TT-NG must forward the legacy policy to the wheel, not bypass it."""