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
This commit is contained in:
aniruddhaadak80
2026-08-24 00:09:37 -07:00
committed by Teknium
parent 316d52faf2
commit a26154aceb
3 changed files with 234 additions and 32 deletions

View File

@@ -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:")

View File

@@ -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

View File

@@ -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") == ""