Extend KD cli commands.

PiperOrigin-RevId: 885246572

Author: The kauldron Authors

kauldron/cli/README.md

@@ -35,3 +35,15 @@ kd.cli.data.ElementSpec(cfg).execute()
 
 * `data`
   - `element_spec`: Display the element spec of the training data pipeline.
+  - `batch`: Display the stats (shapes, dtype, min, max, mean) of an actual fetched batch.
+
+* `run`
+  - `eval_shape`: Run a train step through `jax.eval_shape` (shapes only, no compute).
+  - `train`: Run `trainer.train()` directly to test the training loop.
+  - `eval`: Run `trainer.eval()` and print the metrics output.
+
+* `inspect`
+  - `model_overview`: Display the model overview (parameters, inputs, shapes) similar to the one used in the colab environments.
+
+* `multi`
+  - `execute`: Run multiple commands sequentially. For example, use `--cmds="data batch, run eval, inspect model_overview"` for layered execution.

kauldron/cli/data.py

@@ -19,6 +19,7 @@
 import dataclasses
 from typing import Union
 
+from kauldron import inspect as kd_inspect
 from kauldron import kontext
 from kauldron.cli import cmd_utils as cu
 import tensorflow_datasets as tfds
@@ -48,8 +49,39 @@ def __call__(self):
       print(f"  {k}: {v.dtype}{list(v.shape)}")
 
 
+@dataclasses.dataclass(frozen=True, kw_only=True)
+class Batch(cu.SubCommand):
+  """Display the batch statistics (shapes, dtype, min, max, mean)."""
+
+  ds_path: str = "train_ds"
+
+  def __call__(self):
+    self.print_config_origin()
+    trainer = self.trainer  # trigger config resolution
+    ds = kontext.get_by_path(trainer, self.ds_path)
+
+    with cu.timed("Getting real batch"):
+      batch = next(iter(ds))
+
+    print("")
+    print(f"Dataset: {self.ds_path}")
+    # NOTE: kd_inspect is a konfig.import
+    stats_df = kd_inspect.get_batch_stats(batch)
+    if hasattr(stats_df, "to_markdown"):
+      print(stats_df.to_markdown())
+    else:
+      print(stats_df.to_string(max_rows=None, max_cols=None))
+
+_SUBCOMMANDS = {
+    "element_spec": ElementSpec,
+    "batch": Batch,
+}
+
+
 @dataclasses.dataclass(frozen=True, kw_only=True)
 class Data(cu.CommandGroup):
   """Data commands."""
 
-  sub_command: Union[ElementSpec]  # Union required for simple_parsing
+  sub_command: Union[ElementSpec, Batch] = dataclasses.field(
+      metadata={"subparsers": _SUBCOMMANDS}
+  )

kauldron/cli/inspect_cli.py

@@ -0,0 +1,64 @@
+# Copyright 2026 The kauldron Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Inspect-related CLI commands."""
+
+from __future__ import annotations
+
+import dataclasses
+from typing import Union
+
+from kauldron import inspect as kd_inspect
+from kauldron.cli import cmd_utils as cu
+
+@dataclasses.dataclass(frozen=True, kw_only=True)
+class ModelOverview(cu.SubCommand):
+  """Display the model overview (parameters, inputs, shapes, etc.)."""
+
+  def __call__(self):
+    self.print_config_origin()
+    trainer = self.trainer  # trigger config resolution
+
+    with cu.timed("Getting model overview"):
+      df = kd_inspect.get_colab_model_overview(
+          model=trainer.model,
+          train_ds=trainer.train_ds,
+          ds_sharding=trainer.sharding.batch,
+          model_config=trainer.raw_cfg.model if trainer.raw_cfg else None,
+          rngs=trainer.rng_streams.init_rngs(),
+      )
+
+    print("\n======== Model Overview ========")
+    if hasattr(df, "data"):  # Extract raw DataFrame from StyledDataFrame
+      raw_df = df.data
+    else:
+      raw_df = df
+
+    if hasattr(raw_df, "to_markdown"):
+      print(raw_df.to_markdown(index=False))
+    else:
+      print(raw_df.to_string(index=False))
+
+
+_SUBCOMMANDS = {
+    "model_overview": ModelOverview,
+}
+
+@dataclasses.dataclass(frozen=True, kw_only=True)
+class Inspect(cu.CommandGroup):
+  """Inspect commands."""
+
+  sub_command: Union[ModelOverview] = dataclasses.field(
+      metadata={"subparsers": _SUBCOMMANDS}
+  )

kauldron/cli/main.py

@@ -31,6 +31,8 @@
 from kauldron.cli import cmd_utils as cu
 from kauldron.cli import config
 from kauldron.cli import data
+from kauldron.cli import inspect_cli
+from kauldron.cli import multi
 from kauldron.cli import patch_config
 from kauldron.cli import run
 import simple_parsing
@@ -48,7 +50,8 @@
 @dataclasses.dataclass(frozen=True, kw_only=True)
 class Args:
 
-  command: config.Config | data.Data | run.Run
+  command: config.Config | data.Data | run.Run | inspect_cli.Inspect | multi.Multi
+
 
   patch: patch_config.PatchConfig = dataclasses.field(
       default_factory=patch_config.PatchConfig,

kauldron/cli/multi.py

@@ -0,0 +1,97 @@
+# Copyright 2026 The kauldron Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Multi-command execution CLI."""
+
+from __future__ import annotations
+
+import dataclasses
+from typing import Union
+
+from kauldron.cli import cmd_utils as cu
+from kauldron.cli import data
+from kauldron.cli import inspect_cli
+from kauldron.cli import run
+
+_COMMANDS = {
+    "data": {
+        "element_spec": data.ElementSpec,
+        "batch": data.Batch,
+    },
+    "run": {
+        "eval_shape": run.EvalShape,
+        "train": run.Train,
+        "eval": run.Eval,
+    },
+    "inspect": {
+        "model_overview": inspect_cli.ModelOverview,
+    }
+}
+
+@dataclasses.dataclass(frozen=True, kw_only=True)
+class Execute(cu.SubCommand):
+  """Execute multiple commands sequentially."""
+
+  cmds: str = ""
+
+  def __call__(self):
+    self.print_config_origin()
+    trainer = None
+
+    cmds_list = self.cmds.split(",")
+    if not cmds_list or (len(cmds_list) == 1 and not cmds_list[0].strip()):
+      print(
+          'No subcommands provided to multi. Use --multi.execute.cmds="data'
+          ' batch, run train"'
+      )
+      return
+
+    for cmd_str in cmds_list:
+      cmd_str = cmd_str.strip()
+      if not cmd_str:
+        continue
+      parts = cmd_str.split()
+      if len(parts) != 2:
+        print(f"Invalid command format: {cmd_str}. Expected 'group subcommand'")
+        continue
+      group_name, sub_name = parts
+
+      if group_name not in _COMMANDS or sub_name not in _COMMANDS[group_name]:
+        print(
+            f"Unknown command: '{cmd_str}'. Available commands: "
+            + ", ".join(
+                [f"{g} {s}" for g, subs in _COMMANDS.items() for s in subs]
+            )
+        )
+        continue
+
+      cmd_cls = _COMMANDS[group_name][sub_name]
+
+      print(f"\\n{'='*40}\\nExecuting: {cmd_str}\\n{'='*40}")
+
+      cmd = cmd_cls(cfg=self.cfg, origin=None)  # pytype: disable=wrong-keyword-args
+
+      # Share resolved trainer to avoid redundant recompilations/resolutions
+      if trainer is None:
+        trainer = cmd.trainer
+      else:
+        cmd.__dict__["trainer"] = trainer
+
+      cmd()
+
+@dataclasses.dataclass(frozen=True, kw_only=True)
+class Multi(cu.CommandGroup):
+  """Multi commands."""
+
+  sub_command: Union[Execute]

kauldron/cli/run.py

@@ -20,7 +20,9 @@
 import functools
 from typing import Union
 
+from etils import epy
 import jax
+import kauldron as kd
 from kauldron.cli import cmd_utils as cu
 from kauldron.data import utils as data_utils
 import tensorflow_datasets as tfds
@@ -71,17 +73,66 @@ def __call__(self) -> None:
     # TODO(klausg): could try to run the metrics computation too.
 
 
+@dataclasses.dataclass(frozen=True, kw_only=True)
+class Train(cu.SubCommand):
+  """Run trainer.train()."""
+
+  def __call__(self) -> None:
+    self.print_config_origin()
+
+    # Ensure exactly one training step is performed to match kd_test.ipynb
+    self.cfg.stop_after_steps = 1
+
+    if hasattr(self.cfg, 'evals'):
+      kd.kontext.set_by_path(self.cfg, 'evals.**.num_batches', 1)
+
+    trainer = self.trainer  # trigger config resolution
+
+    with cu.timed('trainer.train()'):
+      train_state, aux = trainer.train()
+      del train_state, aux
+    # We don't print the output of train() as it might be very verbose
+    # but we can print that it succeeded.
+    print('Successfully completed trainer.train()')
+
+@dataclasses.dataclass(frozen=True, kw_only=True)
+class Eval(cu.SubCommand):
+  """Run trainer.eval()."""
+
+  def __call__(self) -> None:
+    self.print_config_origin()
+
+    if hasattr(self.cfg, 'evals'):
+      kd.kontext.set_by_path(self.cfg, 'evals.**.num_batches', 1)
+
+    trainer = self.trainer  # trigger config resolution
+
+    with cu.timed('trainer.init_state()'):
+      state = trainer.init_state()
+
+    eval_metrics = {}
+    for name, evaluator in trainer.evals.items():
+      with cu.timed(f'evaluator.evaluate({name})'):
+        eval_metrics[name] = evaluator.evaluate(state=state, step=0)
+
+    if eval_metrics:
+      print('Evaluator metrics:')
+      epy.pprint(eval_metrics)
+
+
 _SUBCOMMANDS = {
     # Manually name the subcommand 'eval_shape' beacause simple-parsing
     # would turn this into 'evalshape' instead.
     'eval_shape': EvalShape,
+    'train': Train,
+    'eval': Eval,
 }
 
 
 @dataclasses.dataclass(frozen=True, kw_only=True)
 class Run(cu.CommandGroup):
   """Run commands for local training validation."""
 
-  sub_command: Union[EvalShape] = dataclasses.field(
+  sub_command: Union[EvalShape, Train, Eval] = dataclasses.field(
       metadata={'subparsers': _SUBCOMMANDS}
   )