From f1a090e8c4d84e72aa3b81586970ba34b9f29677 Mon Sep 17 00:00:00 2001 From: Tom Lu Date: Thu, 6 Aug 2026 23:21:46 +0800 Subject: [PATCH] Add API retries and concurrent batch processing --- .env.example | 2 +- README.md | 13 ++++- src/swe_data_processing/cli.py | 79 +++++++++++++++++++++++++++---- src/swe_data_processing/client.py | 11 +++-- src/swe_data_processing/config.py | 4 +- tests/test_cli.py | 54 +++++++++++++++++++++ tests/test_client.py | 36 +++++++++++++- tests/test_config.py | 4 +- 8 files changed, 183 insertions(+), 20 deletions(-) create mode 100644 tests/test_cli.py diff --git a/.env.example b/.env.example index 78d75d5..4891ed8 100644 --- a/.env.example +++ b/.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 diff --git a/README.md b/README.md index 0e27093..4e039ea 100644 --- a/README.md +++ b/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. diff --git a/src/swe_data_processing/cli.py b/src/swe_data_processing/cli.py index 893b74d..36cf6bf 100644 --- a/src/swe_data_processing/cli.py +++ b/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: diff --git a/src/swe_data_processing/client.py b/src/swe_data_processing/client.py index 8da8b1a..eb40a26 100644 --- a/src/swe_data_processing/client.py +++ b/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}" ) diff --git a/src/swe_data_processing/config.py b/src/swe_data_processing/config.py index a27151f..00f4c36 100644 --- a/src/swe_data_processing/config.py +++ b/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, diff --git a/tests/test_cli.py b/tests/test_cli.py new file mode 100644 index 0000000..dfa3db1 --- /dev/null +++ b/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) + } diff --git a/tests/test_client.py b/tests/test_client.py index 7fe318b..e691ce3 100644 --- a/tests/test_client.py +++ b/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"} diff --git a/tests/test_config.py b/tests/test_config.py index e9b054d..f0762a1 100644 --- a/tests/test_config.py +++ b/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