fix(prm): PRM producer conformance — save tokenizer + standard train-result shape (fixes soup train task=prm KeyError; surfaced by v0.71.30 live smoke)

src/soup_cli/trainer/prm.py

@@ -16,6 +16,7 @@ per project policy — ``python -m soup_cli.cli --help`` must not pull torch.
 from __future__ import annotations
 
 import logging
+import time
 from functools import lru_cache
 from pathlib import Path
 from typing import Any, Optional
@@ -85,6 +86,45 @@ def make_prm_trainer_class(base_cls: type) -> type:
     return _PRMTrainer
 
 
+def build_prm_train_result(
+    *,
+    log_history: list,
+    metrics: Any,
+    global_step: int,
+    duration_secs: float,
+    output_dir: str,
+) -> dict:
+    """Build the standard trainer-result dict for the PRM path (v0.71.30).
+
+    Mirrors the shape every other trainer wrapper returns so ``commands/train.py``
+    can render its summary — previously the PRM wrapper returned a bespoke dict
+    missing ``initial_loss`` / ``final_loss`` / ``duration`` / ``total_steps``,
+    crashing the CLI with a ``KeyError`` right after ``save_model``.
+    """
+    train_losses = [e["loss"] for e in log_history if isinstance(e, dict) and "loss" in e]
+    fallback = 0.0
+    if isinstance(metrics, dict):
+        try:
+            fallback = float(metrics.get("train_loss", 0.0))
+        except (TypeError, ValueError):
+            fallback = 0.0
+    initial = train_losses[0] if train_losses else fallback
+    final = train_losses[-1] if train_losses else fallback
+    hours = int(duration_secs // 3600)
+    minutes = int((duration_secs % 3600) // 60)
+    duration = f"{hours}h {minutes}m" if hours > 0 else f"{minutes}m"
+    return {
+        "status": "ok",
+        "initial_loss": initial,
+        "final_loss": final,
+        "duration": duration,
+        "duration_secs": duration_secs,
+        "total_steps": global_step,
+        "output_dir": output_dir,
+        "metrics": metrics,
+    }
+
+
 def _prepare_prm_dataset(raw_rows: list[dict], tokenizer: Any, max_length: int) -> list[dict]:
     """Tokenise PRM rows into (input_ids, attention_mask, step_positions, labels).
 
@@ -294,6 +334,18 @@ class PRMTrainerWrapper:
             data_collator=collator,
         )
         console.print("[green]Starting PRM training...[/]")
+        start = time.time()
         result = self.trainer.train()
         self.trainer.save_model(str(output_dir))
-        return {"status": "ok", "output_dir": str(output_dir), "metrics": result.metrics}
+        # v0.71.30 — save the tokenizer alongside the model so the PRM
+        # checkpoint is loadable standalone (soup shrink / PRMScorer /
+        # `soup train prm_reward=<dir>` all call AutoTokenizer.from_pretrained
+        # on the dir). Previously the tokenizer was never persisted.
+        self.tokenizer.save_pretrained(str(output_dir))
+        return build_prm_train_result(
+            log_history=self.trainer.state.log_history,
+            metrics=result.metrics,
+            global_step=self.trainer.state.global_step,
+            duration_secs=time.time() - start,
+            output_dir=str(output_dir),
+        )

tests/test_v07130.py

@@ -727,6 +727,60 @@ class TestRecipes:
         assert len(RECIPES) == 137
 
 
+# ---------------------------------------------------------------------------
+# PRM producer fixes surfaced by the live smoke (train-result shape + tokenizer)
+# ---------------------------------------------------------------------------
+class TestBuildPrmTrainResult:
+    def test_has_all_summary_keys(self):
+        from soup_cli.trainer.prm import build_prm_train_result
+
+        out = build_prm_train_result(
+            log_history=[{"loss": 3.2}, {"loss": 1.1}],
+            metrics={"train_loss": 2.0},
+            global_step=6,
+            duration_secs=125.0,
+            output_dir="./out",
+        )
+        # These are exactly the keys commands/train.py's summary indexes.
+        for key in (
+            "initial_loss",
+            "final_loss",
+            "duration",
+            "duration_secs",
+            "total_steps",
+            "output_dir",
+        ):
+            assert key in out, key
+        assert out["initial_loss"] == pytest.approx(3.2)
+        assert out["final_loss"] == pytest.approx(1.1)
+        assert out["total_steps"] == 6
+        assert out["duration"] == "2m"
+
+    def test_empty_log_history_falls_back_to_metrics(self):
+        from soup_cli.trainer.prm import build_prm_train_result
+
+        out = build_prm_train_result(
+            log_history=[],
+            metrics={"train_loss": 2.5},
+            global_step=0,
+            duration_secs=3700.0,
+            output_dir="./out",
+        )
+        assert out["initial_loss"] == pytest.approx(2.5)
+        assert out["final_loss"] == pytest.approx(2.5)
+        assert out["duration"] == "1h 1m"
+
+    def test_train_saves_tokenizer(self):
+        # Guard: PRMTrainerWrapper.train() must persist the tokenizer so the
+        # PRM checkpoint is loadable standalone by PRMScorer.
+        import inspect
+
+        from soup_cli.trainer.prm import PRMTrainerWrapper
+
+        src = inspect.getsource(PRMTrainerWrapper.train)
+        assert "self.tokenizer.save_pretrained" in src
+
+
 class TestNoTopLevelTorch:
     def test_prm_reward_has_no_top_level_torch(self):
         import soup_cli.utils.prm_reward as mod