Dynamic batch should be write back to config if it's inferred.

tests/cli/grpo_main_test.py

@@ -20,6 +20,8 @@
 import os
 import pathlib
 import tempfile
+from typing import Any
+from typing import cast
 from unittest import mock
 
 from absl.testing import absltest
@@ -577,6 +579,26 @@ def test_standard_grpo_kv_cache(self):
     self.assertEqual(cfg.kv_cache_size, 256 + 512 + 256)
 
 
+class ComputeParamsTest(absltest.TestCase):
+
+  def test_compute_params_persists_dynamic_num_batches(self):
+    pipeline = _make_pipeline("")
+    pipeline.config["batch_size"] = 8
+    pipeline.config["num_batches"] = 0
+    pipeline.config["num_train_epochs"] = 1
+    pipeline.config["train_fraction"] = 0.8
+    rl_training_config = cast(dict[str, Any], pipeline.config["rl_training_config"])
+    rl_training_config["max_steps"] = 0
+
+    raw_dataset = mock.Mock()
+    raw_dataset.__len__ = mock.Mock(return_value=7473)
+
+    pipeline.compute_params(raw_dataset)
+
+    self.assertEqual(pipeline.config["num_batches"], 934)
+    self.assertEqual(rl_training_config["max_steps"], 747)
+
+
 # ---------------------------------------------------------------------------
 # GRPOConfig construction
 # ---------------------------------------------------------------------------

tests/cli/utils/data_test.py

@@ -205,6 +205,40 @@ def test_filters_by_prompt_length(self):
     self.assertLen(batches, 1)
     self.assertEqual(batches[0], [{"prompts": "short", "answer": 1}])
 
+  def test_raises_when_prompt_length_filter_removes_all_examples(self):
+    tokenizer = _FakeTokenizer()
+    dataset = _BaseDataset([
+        {"prompts": "this is too long", "answer": 1},
+        {"prompts": "also too long", "answer": 2},
+    ])
+
+    with self.assertRaisesRegex(
+        ValueError, "empty after post_init_dataset filtering"
+    ):
+      data_lib.post_init_dataset(
+          dataset,
+          tokenizer=tokenizer,  # pytype: disable=wrong-arg-types
+          batch_size=2,
+          num_batches=None,
+          max_prompt_length=2,
+      )
+
+  def test_raises_when_fraction_makes_training_split_empty(self):
+    tokenizer = _FakeTokenizer()
+    dataset = _BaseDataset([
+        {"prompts": "short", "answer": 1},
+    ])
+
+    with self.assertRaisesRegex(ValueError, "empty after post_init_dataset split"):
+      data_lib.post_init_dataset(
+          dataset,
+          tokenizer=tokenizer,  # pytype: disable=wrong-arg-types
+          batch_size=1,
+          num_batches=None,
+          max_prompt_length=None,
+          fraction=0.5,
+      )
+
   def test_limits_num_batches(self):
     tokenizer = _FakeTokenizer()
     dataset = _BaseDataset(

tunix/cli/config.py

@@ -15,6 +15,8 @@
 import ast
 import collections
 from collections.abc import Callable
+from collections.abc import Mapping
+from collections.abc import MutableMapping
 import copy
 import importlib
 import inspect
@@ -242,6 +244,60 @@ def __init__(self, argv: list[str], **kwargs):
     self.check_supported_workflow()
     self._validate_perf_metrics(entry_point=argv[0])
 
+  def _config_mapping(self, key: str) -> dict[str, Any]:
+    """Returns a config section as a plain dictionary.
+
+    This narrows nested config sections that may otherwise be inferred as broad
+    unions of scalars and mappings by static type checkers.
+    """
+    value = self.config.get(key)
+    if value is None:
+      return {}
+    if not isinstance(value, Mapping):
+      raise TypeError(
+          f"Expected config section {key!r} to be a mapping, got"
+          f" {type(value).__name__}."
+      )
+    return dict(value)
+
+  def _mutable_config_mapping(self, key: str) -> MutableMapping[str, Any]:
+    """Returns a mutable config section for in-place updates."""
+    value = self.config.get(key)
+    if value is None:
+      section: dict[str, Any] = {}
+      self.config[key] = section
+      return section
+    if not isinstance(value, MutableMapping):
+      raise TypeError(
+          f"Expected config section {key!r} to be a mutable mapping, got"
+          f" {type(value).__name__}."
+      )
+    return value
+
+  def _config_string(self, key: str, default: str = "") -> str:
+    """Returns a string config value with validation."""
+    value = self.config.get(key, default)
+    if value is None:
+      return default
+    if not isinstance(value, str):
+      raise TypeError(
+          f"Expected config value {key!r} to be a string, got"
+          f" {type(value).__name__}."
+      )
+    return value
+
+  def _config_bool(self, key: str, default: bool = False) -> bool:
+    """Returns a boolean config value with validation."""
+    value = self.config.get(key, default)
+    if value is None:
+      return default
+    if not isinstance(value, bool):
+      raise TypeError(
+          f"Expected config value {key!r} to be a bool, got"
+          f" {type(value).__name__}."
+      )
+    return value
+
   def _validate_perf_metrics(self, entry_point: str):
     """Validates that perf metrics are only enabled for GRPO.