fix(sft): tolerate trace-shape tool calls; opt-in skip for unrenderable samples

packages/prime-rl-configs/src/prime_rl/configs/sft.py

@@ -98,6 +98,9 @@ class SFTDataConfig(BaseDataConfig):
     loss_mask: LossMaskConfig = LossMaskConfig()
     """Which message types contribute to the loss."""
 
+    skip_invalid_samples: bool = False
+    """Skip (and log) samples the renderer cannot render instead of crashing the run. Off by default so data problems fail loudly."""
+
     @model_validator(mode="after")
     def validate_subsets_and_splits(self):
         if self.subsets is not None or self.splits is not None:

src/prime_rl/trainer/sft/data.py

@@ -143,12 +143,14 @@ def __init__(
         loss_mask_config: LossMaskConfig = LossMaskConfig(),
         max_examples: int | None = None,
         max_epochs: int | None = None,
+        skip_invalid_samples: bool = False,
     ):
         super().__init__()
         self.logger = get_logger()
         self.dataset = dataset
         self.num_examples = len(self.dataset)
         self.renderer = renderer
+        self.skip_invalid_samples = skip_invalid_samples
         self.shuffle = shuffle
         self.seed = seed
         self.seq_len = seq_len
@@ -307,7 +309,15 @@ def __iter__(self):
             example = dataset[(self.step - 1) % self.num_examples]
 
             # Process example
-            processed_example = self._process(cast(dict, example))
+            try:
+                processed_example = self._process(cast(dict, example))
+            except ValueError as e:
+                if not self.skip_invalid_samples:
+                    raise
+                self.logger.warning(
+                    f"Skipping example {cast(dict, example).get('__index', '')} because it could not be rendered: {e}"
+                )
+                continue
 
             # If processed example is None, skip it (e.g. if tokenized sample exceeds context window)
             if processed_example is None:
@@ -514,6 +524,7 @@ def setup_dataset(
             loss_mask_config=config.loss_mask,
             non_dp_size=non_dp_size,
             max_epochs=max_epochs,
+            skip_invalid_samples=config.skip_invalid_samples,
         )
     else:
         raise ValueError(f"Invalid dataset type: {config.type}")

src/prime_rl/utils/chat_template.py

@@ -23,7 +23,16 @@ def normalize_messages(messages: Any, default_role: str) -> list[dict[str, Any]]
 
 
 def deserialize_tool_calls(messages: list[dict[str, Any]]) -> list[dict[str, Any]]:
-    def _deserialize_tool_call(tool_call: dict[str, Any]) -> dict[str, Any]:
+    def _deserialize_tool_call(tool_call: dict[str, Any] | str) -> dict[str, Any]:
+        # Verifiers traces store tool calls as JSON strings in a flat
+        # {"id", "name", "arguments"} shape; normalize to the OAI form.
+        if isinstance(tool_call, str):
+            tool_call = json.loads(tool_call)
+        if "function" not in tool_call:
+            tool_call = {
+                **{k: tool_call[k] for k in ("id", "type") if k in tool_call},
+                "function": {k: v for k, v in tool_call.items() if k not in ("id", "type")},
+            }
         function = tool_call.get("function", {})
         arguments = function.get("arguments")
         if isinstance(arguments, str):

tests/unit/train/sft/test_sft_dataset.py

@@ -399,3 +399,41 @@ def test_null_messages_falls_back_to_prompt_and_completion():
     )
 
     assert next(iter(mixed_row_dataset)) == next(iter(expected_dataset))
+
+
+def test_deserialize_tool_calls_accepts_trace_shapes():
+    """Verifiers traces store tool calls as flat JSON strings; both they and
+    the OAI dict shape must deserialize to the OAI form."""
+    from prime_rl.utils.chat_template import deserialize_tool_calls
+
+    flat_string = '{"id": "t1", "name": "ipython", "arguments": "{\\"code\\": \\"print(1)\\"}"}'
+    oai_dict = {"id": "t2", "type": "function", "function": {"name": "ipython", "arguments": '{"code": "print(2)"}'}}
+    [message] = deserialize_tool_calls([{"role": "assistant", "content": "", "tool_calls": [flat_string, oai_dict]}])
+
+    first, second = message["tool_calls"]
+    assert first["id"] == "t1"
+    assert first["function"] == {"name": "ipython", "arguments": {"code": "print(1)"}}
+    assert second["function"] == {"name": "ipython", "arguments": {"code": "print(2)"}}
+
+
+class _RaisingRenderer(_DummyRenderer):
+    def render(self, messages, **kwargs):
+        if "bad" in messages[-1]["content"]:
+            raise ValueError("unrenderable sample")
+        return super().render(messages, **kwargs)
+
+
+def test_skip_invalid_samples_knob():
+    """A raising sample crashes by default and is skipped (with the good
+    samples still yielded) when skip_invalid_samples is on."""
+    dataset = Dataset.from_list(
+        [{"messages": [{"role": "assistant", "content": content}]} for content in ("a0", "bad", "a1")]
+    )
+
+    crashing = SFTDataset(dataset, _RaisingRenderer(), shuffle=False, max_epochs=1)
+    with pytest.raises(ValueError, match="unrenderable sample"):
+        list(crashing)
+
+    skipping = SFTDataset(dataset, _RaisingRenderer(), shuffle=False, max_epochs=1, skip_invalid_samples=True)
+    samples = list(skipping)
+    assert len(samples) == 2