From a26154aceb74abf5ea9243f0e1c1a128c19b6e63 Mon Sep 17 00:00:00 2001
From: aniruddhaadak80 <aniruddhaadak80@gmail.com>
Date: Mon, 24 Aug 2026 00:09:37 -0700
Subject: [PATCH] fix(batch_runner): teach the resume content scan to honor
 discard tombstones

Complete the #93527 fix: the tombstone now carries the human prompt
text via _entry_prompt_text (handling flat prompt, ShareGPT, and
chat-style shapes), _scan_completed_prompts_by_content counts
discarded rows as completed instead of only reading ShareGPT
conversations, and the merge step reports excluded tombstones in the
combined-count summary. Adds a dedicated regression suite covering the
tombstone round-trip, the all-discarded-batch resume path, and merge
exclusion.

Salvaged from #93579 (issue reporter's PR), building on #93542.
Fixes #93527
---
 batch_runner.py                           |  88 +++++++----
 tests/test_batch_runner_checkpoint.py     |   3 +
 tests/test_batch_runner_discard_resume.py | 175 ++++++++++++++++++++++
 3 files changed, 234 insertions(+), 32 deletions(-)
 create mode 100644 tests/test_batch_runner_discard_resume.py

diff --git a/batch_runner.py b/batch_runner.py
index 5eccb1a53a..5d3c38fc16 100644
--- a/batch_runner.py
+++ b/batch_runner.py
@@ -457,15 +457,19 @@ def _process_batch_worker(args: Tuple) -> Dict[str, Any]:
                 print(f"   🚫 Prompt {prompt_index} discarded (no reasoning in any turn)")
                 discarded_no_reasoning += 1
                 completed_in_batch.append(prompt_index)
-                # Write a tombstone row so the content-based resume scan (which
-                # only reads batch_*.jsonl) can see this prompt was already
-                # processed and discarded, not just left unprocessed.
+                # Tombstone row (#93527): resume filters exclusively by
+                # scanning batch_*.jsonl rows for prompt content, so a
+                # discarded sample without a row is invisible to --resume
+                # and gets re-run at full cost on every restart. The
+                # tombstone carries just enough for the scan; the merge
+                # step excludes it from trajectories.jsonl.
+                tombstone = {
+                    "prompt_index": prompt_index,
+                    "discarded": "no_reasoning",
+                    "prompt": _entry_prompt_text(prompt_data),
+                }
                 with open(batch_output_file, 'a', encoding='utf-8') as f:
-                    f.write(json.dumps({
-                        "prompt_index": prompt_index,
-                        "conversations": result["trajectory"],
-                        "discarded": "no_reasoning",
-                    }, ensure_ascii=False) + "\n")
+                    f.write(json.dumps(tombstone, ensure_ascii=False) + "\n")
                     f.flush()
                     os.fsync(f.fileno())
                 continue
@@ -537,6 +541,31 @@ def _process_batch_worker(args: Tuple) -> Dict[str, Any]:
     }
 
 
+def _entry_prompt_text(entry: Dict) -> str:
+    """Extract the human prompt text from a dataset or trajectory entry.
+
+    Handles the shapes that appear across batch_runner: a flat
+    ``entry["prompt"]``, ShareGPT-style ``conversations`` (``from``/``value``),
+    chat-style ``conversations``/``messages`` (``role``/``content``), and the
+    discard tombstones written by the no-reasoning discard path.
+    """
+    if not isinstance(entry, dict):
+        return ""
+    text = str(entry.get("prompt") or "").strip()
+    if text:
+        return text
+    for key in ("conversations", "messages"):
+        for msg in entry.get(key, []) or []:
+            if not isinstance(msg, dict):
+                continue
+            role = msg.get("role") or msg.get("from")
+            if role in {"user", "human"}:
+                text = str(msg.get("content") or msg.get("value") or "").strip()
+                if text:
+                    return text
+    return ""
+
+
 class BatchRunner:
     """
     Manages batch processing of agent prompts with checkpointing and statistics.
@@ -766,19 +795,17 @@ class BatchRunner:
                     for line in f:
                         try:
                             entry = json.loads(line.strip())
-                            
+
                             # Skip failed entries - we want to retry these
                             if entry.get("failed", False):
                                 continue
-                            
-                            # Extract the human/user prompt from conversations
-                            conversations = entry.get("conversations", [])
-                            for msg in conversations:
-                                if msg.get("from") == "human":
-                                    prompt_text = msg.get("value", "").strip()
-                                    if prompt_text:
-                                        completed_prompts.add(prompt_text)
-                                    break  # Only need the first human message
+
+                            # Discard tombstones count as completed — the
+                            # prompt was processed and deliberately dropped
+                            # (#93527); re-running it would just re-discard.
+                            prompt_text = _entry_prompt_text(entry)
+                            if prompt_text:
+                                completed_prompts.add(prompt_text)
                         except json.JSONDecodeError:
                             continue
             except Exception as e:
@@ -1006,11 +1033,8 @@ class BatchRunner:
         
         # Aggregate all batch statistics and update checkpoint
         total_reasoning_stats = {"total_assistant_turns": 0, "turns_with_reasoning": 0, "turns_without_reasoning": 0}
-        total_discarded_no_reasoning = 0
 
         for batch_result in results:
-            total_discarded_no_reasoning += batch_result.get("discarded_no_reasoning", 0)
-
             # Aggregate tool stats
             for tool_name, stats in batch_result.get("tool_stats", {}).items():
                 if tool_name not in total_tool_stats:
@@ -1057,7 +1081,7 @@ class BatchRunner:
         
         total_entries = 0
         filtered_entries = 0
-        discarded_tombstones = 0
+        tombstone_entries = 0
         batch_files_found = 0
         
         # Find ALL batch files in the output directory (handles resume merging old + new)
@@ -1074,15 +1098,15 @@ class BatchRunner:
                         try:
                             data = json.loads(line)
 
-                            # Discard tombstones exist only so resume can see
-                            # these prompts as done; they carry no full
-                            # trajectory and must not enter the training file.
+                            # Discard tombstones are resume bookkeeping, not
+                            # training data (#93527) — never enter the merged
+                            # trajectories file.
                             if data.get("discarded"):
-                                discarded_tombstones += 1
+                                tombstone_entries += 1
                                 continue
 
                             tool_stats = data.get('tool_stats', {})
-
+                            
                             # Check for invalid tool names (model hallucinations)
                             invalid_tools = [k for k in tool_stats if k not in VALID_TOOLS]
                             
@@ -1099,9 +1123,7 @@ class BatchRunner:
         
         if filtered_entries > 0:
             print(f"⚠️  Filtered {filtered_entries} corrupted entries out of {total_entries} total")
-        if discarded_tombstones > 0:
-            print(f"ℹ️  Excluded {discarded_tombstones} discarded (no-reasoning) tombstone rows out of {total_entries} total")
-        print(f"✅ Combined {batch_files_found} batch files into trajectories.jsonl ({total_entries - filtered_entries - discarded_tombstones} entries)")
+        print(f"✅ Combined {batch_files_found} batch files into trajectories.jsonl ({total_entries - filtered_entries - tombstone_entries} entries)")
         
         # Save final statistics
         final_stats = {
@@ -1115,7 +1137,9 @@ class BatchRunner:
             "duration_seconds": round(time.time() - start_time, 2),
             "tool_statistics": total_tool_stats,
             "reasoning_statistics": total_reasoning_stats,
-            "discarded_no_reasoning": total_discarded_no_reasoning,
+            "discarded_no_reasoning": sum(
+                r.get("discarded_no_reasoning", 0) for r in results
+            ),
         }
         
         with open(self.stats_file, 'w', encoding='utf-8') as f:
@@ -1126,7 +1150,7 @@ class BatchRunner:
         print("📊 BATCH PROCESSING COMPLETE")
         print("=" * 70)
         print(f"✅ Prompts processed this run: {sum(r.get('processed', 0) for r in results)}")
-        print(f"✅ Total trajectories in merged file: {total_entries - filtered_entries}")
+        print(f"✅ Total trajectories in merged file: {total_entries - filtered_entries - tombstone_entries}")
         print(f"✅ Total batch files merged: {batch_files_found}")
         print(f"⏱️  Total duration: {round(time.time() - start_time, 2)}s")
         print("\n📈 Tool Usage Statistics:")
diff --git a/tests/test_batch_runner_checkpoint.py b/tests/test_batch_runner_checkpoint.py
index 0a9096b3a3..899d35afe4 100644
--- a/tests/test_batch_runner_checkpoint.py
+++ b/tests/test_batch_runner_checkpoint.py
@@ -183,6 +183,9 @@ class TestBatchWorkerResumeBehavior:
         assert len(lines) == 1
         entry = json.loads(lines[0])
         assert entry["discarded"] == "no_reasoning"
+        # The lightweight tombstone carries the human prompt text so the
+        # content scan can match it without a full trajectory payload.
+        assert entry["prompt"] == "hi"
 
     def test_resume_after_all_discarded_batch_reruns_zero_prompts(self, tmp_path, monkeypatch):
         """Regression for the issue: a resumed run must not re-execute
diff --git a/tests/test_batch_runner_discard_resume.py b/tests/test_batch_runner_discard_resume.py
new file mode 100644
index 0000000000..7f8c725e0f
--- /dev/null
+++ b/tests/test_batch_runner_discard_resume.py
@@ -0,0 +1,175 @@
+"""Discarded samples must be visible to --resume (#93527).
+
+The no-reasoning discard path used to mark prompts completed only in the
+checkpoint index while writing no batch_*.jsonl row. Resume filters
+exclusively by scanning those files for prompt content, so every
+discarded sample was re-run at full cost on each restart. The fix writes
+a tombstone row that the content scan treats as completed and the
+trajectories.jsonl merge excludes.
+"""
+
+import json
+from unittest.mock import MagicMock, patch
+
+import batch_runner
+from batch_runner import (
+    BatchRunner,
+    _entry_prompt_text,
+    _process_batch_worker,
+)
+
+
+# ─────────────────────────────────────────────────────────────────────
+# Worker: discard leaves a tombstone row
+# ─────────────────────────────────────────────────────────────────────
+
+
+def _discarded_result():
+    return {
+        "success": True,
+        "trajectory": [{"role": "assistant", "content": "x"}],
+        "reasoning_stats": {"has_any_reasoning": False},
+        "tool_stats": {},
+        "metadata": {},
+        "completed": True,
+        "api_calls": 1,
+        "toolsets_used": [],
+    }
+
+
+def test_discard_writes_tombstone_row(tmp_path, monkeypatch):
+    monkeypatch.setattr(
+        "batch_runner._process_single_prompt", lambda *a, **kw: _discarded_result()
+    )
+
+    _process_batch_worker((1, [(0, {"prompt": "hi"})], tmp_path, set(), {"verbose": False}))
+
+    batch_file = tmp_path / "batch_1.jsonl"
+    rows = [json.loads(line) for line in batch_file.read_text(encoding="utf-8").splitlines() if line.strip()]
+    assert len(rows) == 1
+    assert rows[0]["discarded"] == "no_reasoning"
+    assert rows[0]["prompt_index"] == 0
+    assert rows[0]["prompt"] == "hi"
+
+
+# ─────────────────────────────────────────────────────────────────────
+# Content scan: tombstones count as completed
+# ─────────────────────────────────────────────────────────────────────
+
+
+def _scan_runner(tmp_path):
+    runner = BatchRunner.__new__(BatchRunner)
+    runner.output_dir = tmp_path
+    return runner
+
+
+def test_content_scan_treats_tombstone_as_completed(tmp_path):
+    (tmp_path / "batch_1.jsonl").write_text(
+        json.dumps({"prompt_index": 0, "discarded": "no_reasoning", "prompt": "tombstoned q"})
+        + "\n"
+        + json.dumps({
+            "conversations": [{"from": "human", "value": "normal q"}],
+            "completed": True,
+        })
+        + "\n",
+        encoding="utf-8",
+    )
+
+    completed = _scan_runner(tmp_path)._scan_completed_prompts_by_content()
+
+    assert completed == {"tombstoned q", "normal q"}
+
+
+def test_content_scan_still_skips_failed_rows(tmp_path):
+    (tmp_path / "batch_1.jsonl").write_text(
+        json.dumps({"failed": True, "conversations": [{"from": "human", "value": "retry me"}]})
+        + "\n",
+        encoding="utf-8",
+    )
+
+    assert _scan_runner(tmp_path)._scan_completed_prompts_by_content() == set()
+
+
+# ─────────────────────────────────────────────────────────────────────
+# Merge: tombstones never enter trajectories.jsonl
+# ─────────────────────────────────────────────────────────────────────
+
+
+def _make_real_runner(tmp_path, monkeypatch):
+    dataset = tmp_path / "dataset.jsonl"
+    dataset.write_text(json.dumps({"prompt": "hi"}) + "\n", encoding="utf-8")
+    monkeypatch.chdir(tmp_path)
+    return BatchRunner(
+        dataset_file=str(dataset),
+        batch_size=1,
+        run_name="discard-resume-test",
+        num_workers=1,
+    )
+
+
+def _fake_pool(batch_results):
+    pool = MagicMock()
+    pool.imap_unordered.return_value = iter(batch_results)
+    pool_cm = MagicMock()
+    pool_cm.__enter__ = MagicMock(return_value=pool)
+    pool_cm.__exit__ = MagicMock(return_value=False)
+    return pool_cm
+
+
+def test_merge_excludes_tombstones_from_trajectories(tmp_path, monkeypatch):
+    # Pre-existing output from an earlier session: one real trajectory +
+    # one tombstone. run(resume=False) re-processes its batches through the
+    # patched Pool, then merges ALL batch files on disk.
+    out_dir = tmp_path / "data" / "discard-resume-test"
+    out_dir.mkdir(parents=True)
+    (out_dir / "batch_1.jsonl").write_text(
+        json.dumps({
+            "prompt_index": 0,
+            "conversations": [{"from": "human", "value": "real q"}],
+            "completed": True,
+            "tool_stats": {},
+        })
+        + "\n"
+        + json.dumps({"prompt_index": 1, "discarded": "no_reasoning", "prompt": "dropped q"})
+        + "\n",
+        encoding="utf-8",
+    )
+    monkeypatch.setattr("sys.argv", ["batch_runner.py"])
+
+    runner = _make_real_runner(tmp_path, monkeypatch)
+    # Point the runner at the pre-populated directory instead of cwd/data.
+    runner.output_dir = out_dir
+
+    batch_result = {
+        "batch_num": 1,
+        "processed": 0,
+        "skipped": 0,
+        "tool_stats": {},
+        "reasoning_stats": {},
+        "discarded_no_reasoning": 0,
+        "completed_prompts": [],
+    }
+    with patch.object(batch_runner, "Pool", return_value=_fake_pool([batch_result])):
+        runner.run()
+
+    merged = (out_dir / "trajectories.jsonl").read_text(encoding="utf-8").splitlines()
+    parsed = [json.loads(line) for line in merged if line.strip()]
+    assert len(parsed) == 1
+    assert "discarded" not in parsed[0]
+    stats = json.loads((out_dir / "statistics.json").read_text(encoding="utf-8"))
+    assert "discarded_no_reasoning" in stats
+
+
+# ─────────────────────────────────────────────────────────────────────
+# Prompt-text extraction shapes
+# ─────────────────────────────────────────────────────────────────────
+
+
+def test_entry_prompt_text_shapes():
+    assert _entry_prompt_text({"prompt": "flat"}) == "flat"
+    assert _entry_prompt_text({"conversations": [{"from": "human", "value": "sharegpt"}]}) == "sharegpt"
+    assert _entry_prompt_text({"conversations": [{"role": "user", "content": "chat"}]}) == "chat"
+    assert _entry_prompt_text({"messages": [{"role": "user", "content": "msgs"}]}) == "msgs"
+    assert _entry_prompt_text({"prompt": "  padded  ", "discarded": "x"}) == "padded"
+    assert _entry_prompt_text({}) == ""
+    assert _entry_prompt_text("not-a-dict") == ""