Add API retries and concurrent batch processing

src/swe_data_processing/cli.py

@@ -5,7 +5,8 @@ from __future__ import annotations
 import argparse
 import json
 import sys
-from collections.abc import Callable
+from collections.abc import Callable, Iterator
+from concurrent.futures import FIRST_COMPLETED, Future, ThreadPoolExecutor, wait
 from pathlib import Path
 from typing import Any
 
@@ -39,30 +40,79 @@ def _run_streaming_stage(
     errors_path: Path | None,
     limit: int | None,
     resume: bool,
+    workers: int,
     stage_name: str,
     processor: Callable[[dict[str, Any]], dict[str, Any] | None],
 ) -> int:
     """Run one append-only stage with resume support and progress reporting."""
 
+    if workers < 1:
+        raise ValueError("workers must be at least 1")
     completed = load_completed_ids(output_path) if resume else set()
     processed = 0
     failures = 0
-    for record in iter_records(input_path):
+
+    def selected_records() -> Iterator[dict[str, Any]]:
+        selected = 0
+        for record in iter_records(input_path):
+            sample_id = get_sample_id(record)
+            if sample_id in completed:
+                continue
+            if limit is not None and selected >= limit:
+                break
+            selected += 1
+            yield record
+
+    def process_one(
+        record: dict[str, Any],
+    ) -> tuple[str, dict[str, Any] | None, Exception | None]:
         sample_id = get_sample_id(record)
-        if sample_id in completed:
-            continue
-        if limit is not None and processed >= limit:
-            break
         try:
             result = processor(record)
-            if result is not None:
-                append_jsonl(output_path, result)
         except Exception as exc:  # noqa: BLE001 - each sample must fail independently.
+            return sample_id, None, exc
+        return sample_id, result, None
+
+    def record_outcome(
+        outcome: tuple[str, dict[str, Any] | None, Exception | None],
+    ) -> None:
+        nonlocal processed, failures
+        sample_id, result, error = outcome
+        if error is not None:
             failures += 1
-            _write_error(errors_path, sample_id, stage_name, exc)
+            _write_error(errors_path, sample_id, stage_name, error)
+        elif result is not None:
+            append_jsonl(output_path, result)
         processed += 1
         if processed % 10 == 0:
             print(f"{stage_name}: processed={processed} failures={failures}", file=sys.stderr)
+
+    records = iter(selected_records())
+    if workers == 1:
+        for record in records:
+            record_outcome(process_one(record))
+    else:
+        with ThreadPoolExecutor(max_workers=workers, thread_name_prefix=stage_name) as executor:
+            pending: set[
+                Future[tuple[str, dict[str, Any] | None, Exception | None]]
+            ] = set()
+            for _ in range(workers):
+                try:
+                    record = next(records)
+                except StopIteration:
+                    break
+                pending.add(executor.submit(process_one, record))
+
+            while pending:
+                finished, pending = wait(pending, return_when=FIRST_COMPLETED)
+                for future in finished:
+                    record_outcome(future.result())
+                    try:
+                        record = next(records)
+                    except StopIteration:
+                        continue
+                    pending.add(executor.submit(process_one, record))
+
     print(f"{stage_name}: completed={processed} failures={failures}", file=sys.stderr)
     return 1 if failures else 0
 
@@ -76,6 +126,7 @@ def command_features(args: argparse.Namespace) -> int:
         errors_path=args.errors,
         limit=args.limit,
         resume=args.resume,
+        workers=args.workers,
         stage_name="features",
         processor=lambda record: {
             "sample_id": get_sample_id(record),
@@ -96,6 +147,7 @@ def command_classify(args: argparse.Namespace) -> int:
             errors_path=args.errors,
             limit=args.limit,
             resume=args.resume,
+            workers=args.workers,
             stage_name="classification",
             processor=lambda record: classify_record(record, client),
         )
@@ -112,6 +164,7 @@ def command_audit(args: argparse.Namespace) -> int:
             errors_path=args.errors,
             limit=args.limit,
             resume=args.resume,
+            workers=args.workers,
             stage_name="trajectory_audit",
             processor=lambda record: audit_trajectory(record, client),
         )
@@ -136,6 +189,7 @@ def command_repair_plan(args: argparse.Namespace) -> int:
             errors_path=args.errors,
             limit=args.limit,
             resume=args.resume,
+            workers=args.workers,
             stage_name="repair_plan",
             processor=processor,
         )
@@ -195,6 +249,7 @@ def command_review(args: argparse.Namespace) -> int:
             errors_path=args.errors,
             limit=args.limit,
             resume=args.resume,
+            workers=args.workers,
             stage_name="repair_review",
             processor=processor,
         )
@@ -246,6 +301,12 @@ def _add_stream_arguments(parser: argparse.ArgumentParser) -> None:
     parser.add_argument("--errors", type=Path, help="Optional error JSONL")
     parser.add_argument("--limit", type=int, help="Maximum number of new records")
     parser.add_argument("--resume", action="store_true", help="Skip IDs already present in output")
+    parser.add_argument(
+        "--workers",
+        type=int,
+        default=1,
+        help="Concurrent record workers; API stages can use this for parallel requests",
+    )
 
 
 def build_parser() -> argparse.ArgumentParser: