Initial Open-SWE-Traces cleanup pipeline

scripts/split_jsonl_remaining.py

@@ -0,0 +1,74 @@
+#!/usr/bin/env python3
+"""Split unfinished JSONL records into balanced temporary retry shards."""
+
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+
+
+def record_id(value: dict) -> str:
+    """Return the stable trajectory identifier used by audit manifests."""
+
+    sample_id = value.get("trajectory_id") or value.get("sample_id")
+    if not isinstance(sample_id, str) or not sample_id:
+        raise ValueError("JSONL record has no trajectory_id/sample_id")
+    return sample_id
+
+
+def main() -> None:
+    """Write unfinished input records round-robin across the requested shards."""
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("input", type=Path)
+    parser.add_argument("completed", type=Path)
+    parser.add_argument("output_dir", type=Path)
+    parser.add_argument("--shards", type=int, default=3)
+    args = parser.parse_args()
+    if args.shards < 1:
+        raise SystemExit("--shards must be positive")
+
+    completed_ids: set[str] = set()
+    with args.completed.open(encoding="utf-8") as handle:
+        for line in handle:
+            if line.strip():
+                completed_ids.add(record_id(json.loads(line)))
+
+    args.output_dir.mkdir(parents=True, exist_ok=True)
+    handles = [
+        (args.output_dir / f"part_{index}.input.jsonl").open("w", encoding="utf-8")
+        for index in range(args.shards)
+    ]
+    counts = [0] * args.shards
+    try:
+        unfinished = 0
+        with args.input.open(encoding="utf-8") as handle:
+            for line in handle:
+                if not line.strip():
+                    continue
+                value = json.loads(line)
+                if record_id(value) in completed_ids:
+                    continue
+                shard = unfinished % args.shards
+                handles[shard].write(line)
+                counts[shard] += 1
+                unfinished += 1
+    finally:
+        for handle in handles:
+            handle.close()
+
+    print(
+        json.dumps(
+            {
+                "completed": len(completed_ids),
+                "unfinished": sum(counts),
+                "shard_counts": counts,
+            },
+            indent=2,
+        )
+    )
+
+
+if __name__ == "__main__":
+    main()