Add API retries and concurrent batch processing

.env.example

@@ -8,7 +8,7 @@ GLM_MODEL=glm-5.2
 
 # Conservative defaults for structured quality-control output.
 GLM_TIMEOUT_SECONDS=300
-GLM_MAX_RETRIES=3
+GLM_MAX_RETRIES=5
 GLM_MAX_TOKENS=8192
 GLM_TEMPERATURE=0.0
 GLM_REASONING_EFFORT=high

README.md

@@ -127,13 +127,18 @@ Optional settings and their defaults:
 
 ```bash
 export GLM_TIMEOUT_SECONDS=300
-export GLM_MAX_RETRIES=3
+export GLM_MAX_RETRIES=5
 export GLM_MAX_TOKENS=8192
 export GLM_TEMPERATURE=0.0
 export GLM_REASONING_EFFORT=high
 export GLM_THINKING_ENABLED=true
 ```
 
+`GLM_MAX_RETRIES=5` means one initial request plus at most five retries. The
+client retries timeouts, connection failures, HTTP 429/5xx responses, malformed
+JSON, and schema-invalid model output with bounded exponential backoff. HTTP
+401/403 authentication failures are never retried.
+
 If the gateway rejects GLM-specific `thinking` or `reasoning_effort` fields, the
 client automatically retries with the portable OpenAI-compatible request subset.
 
@@ -209,9 +214,15 @@ swe-qc audit \
   --input samples/sample_20_seed_20260805.jsonl \
   --output qc_outputs/sample20.audits.jsonl \
   --errors qc_outputs/sample20.audit.errors.jsonl \
+  --workers 20 \
   --resume
 ```
 
+`--workers` bounds the number of records processed concurrently. JSONL writes
+remain serialized in the main thread, so each completed record is appended
+atomically even when API requests run in parallel. Output order follows request
+completion order; `sample_id` remains the stable join key.
+
 The deterministic score combines weighted process dimensions with penalties
 for minor, major, and critical behavior issues. Failed exploratory calls are
 not penalized when the agent interprets them correctly and recovers.

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:

src/swe_data_processing/client.py

@@ -104,7 +104,8 @@ class GLMClient:
         """Request one JSON object and validate it against ``schema``.
 
         Invalid JSON and schema violations are retried because model output can
-        occasionally be malformed. Authentication errors are never retried.
+        occasionally be malformed. ``max_retries`` counts retries after the
+        initial request. Authentication errors are never retried.
         """
 
         schema_text = json.dumps(schema, ensure_ascii=False, separators=(",", ":"))
@@ -116,7 +117,8 @@ class GLMClient:
         compatibility_fallback = False
         last_error: Exception | None = None
 
-        for attempt in range(self.settings.max_retries):
+        total_attempts = self.settings.max_retries + 1
+        for attempt in range(total_attempts):
             request_payload = self._build_request(
                 full_system_prompt,
                 payload,
@@ -161,14 +163,15 @@ class GLMClient:
                 raise
             except (httpx.HTTPError, json.JSONDecodeError, ValidationError, GLMClientError) as exc:
                 last_error = exc
-                if attempt + 1 < self.settings.max_retries:
+                if attempt + 1 < total_attempts:
                     # Bounded exponential backoff avoids synchronized retries
                     # without making a single failed sample block indefinitely.
                     delay = min(20.0, 2.0**attempt) + random.uniform(0.0, 0.5)
                     time.sleep(delay)
 
         raise GLMClientError(
-            f"GLM failed to return valid structured output after {self.settings.max_retries} attempts: "
+            f"GLM failed after {total_attempts} attempts "
+            f"({self.settings.max_retries} retries): "
             f"{type(last_error).__name__}: {last_error}"
         )
 

src/swe_data_processing/config.py

@@ -47,7 +47,7 @@ class Settings:
     api_path: str = "/v1/chat/completions"
     model: str = "glm-5.2"
     timeout_seconds: float = 300.0
-    max_retries: int = 3
+    max_retries: int = 5
     max_tokens: int = 8192
     temperature: float = 0.0
     reasoning_effort: str = "high"
@@ -85,7 +85,7 @@ class Settings:
             api_path=os.getenv("GLM_API_PATH", "/v1/chat/completions"),
             model=os.getenv("GLM_MODEL", "glm-5.2"),
             timeout_seconds=_get_float("GLM_TIMEOUT_SECONDS", 300.0, 1.0),
-            max_retries=_get_int("GLM_MAX_RETRIES", 3, 1),
+            max_retries=_get_int("GLM_MAX_RETRIES", 5, 0),
             max_tokens=_get_int("GLM_MAX_TOKENS", 8192, 1),
             temperature=_get_float("GLM_TEMPERATURE", 0.0, 0.0),
             reasoning_effort=reasoning_effort,

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)
+    }

tests/test_client.py

@@ -5,6 +5,7 @@ from __future__ import annotations
 import json
 
 import httpx
+import pytest
 
 from swe_data_processing.client import GLMClient
 from swe_data_processing.config import Settings
@@ -33,7 +34,7 @@ def test_client_parses_fenced_json_and_never_sends_key_in_body() -> None:
         "required": ["status"],
         "properties": {"status": {"const": "ok"}},
     }
-    settings = Settings(api_key="test-secret", max_retries=1)
+    settings = Settings(api_key="test-secret", max_retries=0)
     with GLMClient(settings, transport=httpx.MockTransport(handler)) as client:
         response = client.invoke_json(system_prompt="Return JSON.", payload={"ping": True}, schema=schema)
 
@@ -62,10 +63,41 @@ def test_client_falls_back_when_vendor_extensions_are_rejected() -> None:
         "required": ["status"],
         "properties": {"status": {"const": "ok"}},
     }
-    settings = Settings(api_key="test-secret", max_retries=2)
+    settings = Settings(api_key="test-secret", max_retries=1)
     with GLMClient(settings, transport=httpx.MockTransport(handler)) as client:
         response = client.invoke_json(system_prompt="Return JSON.", payload={}, schema=schema)
 
     assert "thinking" in calls[0]
     assert "thinking" not in calls[1]
     assert response.compatibility_fallback_used is True
+
+
+def test_client_retries_five_failures_before_succeeding(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """Five retryable failures are followed by a sixth and final attempt."""
+
+    calls = 0
+
+    def handler(_: httpx.Request) -> httpx.Response:
+        nonlocal calls
+        calls += 1
+        if calls <= 5:
+            return httpx.Response(503, json={"error": {"message": "temporarily unavailable"}})
+        return httpx.Response(
+            200,
+            json={"choices": [{"message": {"content": "{\"status\":\"ok\"}"}}]},
+        )
+
+    monkeypatch.setattr("swe_data_processing.client.time.sleep", lambda _: None)
+    schema = {
+        "type": "object",
+        "required": ["status"],
+        "properties": {"status": {"const": "ok"}},
+    }
+    settings = Settings(api_key="test-secret", max_retries=5)
+    with GLMClient(settings, transport=httpx.MockTransport(handler)) as client:
+        response = client.invoke_json(system_prompt="Return JSON.", payload={}, schema=schema)
+
+    assert calls == 6
+    assert response.data == {"status": "ok"}

tests/test_config.py

@@ -24,4 +24,6 @@ def test_missing_api_key_is_rejected(monkeypatch: pytest.MonkeyPatch) -> None:
 
 def test_feature_only_settings_can_omit_key(monkeypatch: pytest.MonkeyPatch) -> None:
     monkeypatch.delenv("GLM_API_KEY", raising=False)
-    assert Settings.from_env(require_api_key=False).api_key == ""
+    settings = Settings.from_env(require_api_key=False)
+    assert settings.api_key == ""
+    assert settings.max_retries == 5