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:
@@ -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:")
|
||||
|
||||
@@ -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
|
||||
|
||||
175
tests/test_batch_runner_discard_resume.py
Normal file
175
tests/test_batch_runner_discard_resume.py
Normal 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") == ""
|
||||
Reference in New Issue
Block a user