Add API retries and concurrent batch processing

tests/test_cli.py

@@ -0,0 +1,54 @@
+"""Tests for the streaming CLI runner."""
+
+from __future__ import annotations
+
+import json
+import threading
+import time
+from pathlib import Path
+
+from swe_data_processing.cli import _run_streaming_stage
+
+
+def test_streaming_stage_processes_records_concurrently(tmp_path: Path) -> None:
+    """Multiple workers run processors in parallel while one thread writes JSONL."""
+
+    input_path = tmp_path / "input.jsonl"
+    output_path = tmp_path / "output.jsonl"
+    input_path.write_text(
+        "".join(
+            json.dumps({"trajectory_id": f"sample-{index}"}) + "\n" for index in range(20)
+        ),
+        encoding="utf-8",
+    )
+    lock = threading.Lock()
+    active = 0
+    peak_active = 0
+
+    def processor(record: dict) -> dict:
+        nonlocal active, peak_active
+        with lock:
+            active += 1
+            peak_active = max(peak_active, active)
+        time.sleep(0.01)
+        with lock:
+            active -= 1
+        return {"sample_id": record["trajectory_id"]}
+
+    status = _run_streaming_stage(
+        input_path=input_path,
+        output_path=output_path,
+        errors_path=None,
+        limit=None,
+        resume=False,
+        workers=4,
+        stage_name="test",
+        processor=processor,
+    )
+
+    written = [json.loads(line) for line in output_path.read_text().splitlines()]
+    assert status == 0
+    assert peak_active == 4
+    assert {value["sample_id"] for value in written} == {
+        f"sample-{index}" for index in range(20)
+    }