Via tools/grade_run.py against LM Studio (load -> all prompts -> unload -> verify).
kat-coder: lfu/tts/webhook/automation/rust (lfu has a real SyntaxError: 'self._ bump_freq')
qwen3-coder-30b: lfu/tts/webhook/automation/rust (all 4 py parse clean; terse ~2k tok outputs)
Grading deferred to a later session (capture-only mode).
Script hardening:
- --resident / --unload-all commands (read loaded_instances, POST unload, verify)
- MEMORY GUARD: aborts if a different model is resident (never stack 30GB models)
- TTFT self-discovery dump (LM Studio returns empty stats non-streaming;
fix next session: switch to /api/v1/chat stream:true -> chat.end result.stats)
Co-Authored-By: Claude <noreply@anthropic.com>
338 lines
12 KiB
Python
338 lines
12 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Async batch processor with checkpointing, idempotency, retries, and graceful shutdown.
|
|
|
|
Usage:
|
|
python batch_processor.py --items job-a job-b ... [--checkpoint checkpoint.json]
|
|
|
|
Or import and use the BatchProcessor class directly.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import json
|
|
import os
|
|
import signal
|
|
import sys
|
|
import tempfile
|
|
import time
|
|
from dataclasses import asdict, dataclass, field
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Public exception
|
|
# ---------------------------------------------------------------------------
|
|
|
|
class ProcessingError(Exception):
|
|
"""Raised by ``process()`` when an external call fails."""
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Simulated flaky worker (the "external call")
|
|
# ---------------------------------------------------------------------------
|
|
|
|
async def process(item: str) -> str:
|
|
"""
|
|
Simulate a flaky external call.
|
|
|
|
Succeeds ~80 % of the time, raises ``ProcessingError`` otherwise.
|
|
Latency is uniform random in [10 ms, 80 ms].
|
|
"""
|
|
import random
|
|
await asyncio.sleep(random.uniform(0.010, 0.080))
|
|
if random.random() < 0.20: # 20 % failure rate
|
|
raise ProcessingError(f"transient failure for {item!r}")
|
|
return f"done:{item}"
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Checkpoint data model
|
|
# ---------------------------------------------------------------------------
|
|
|
|
@dataclass
|
|
class ItemResult:
|
|
status: str # "succeeded" | "failed"
|
|
attempts: int = 0
|
|
error: str = ""
|
|
|
|
|
|
@dataclass
|
|
class Checkpoint:
|
|
"""Persisted progress. All fields are JSON-serialisable."""
|
|
completed: dict[str, ItemResult] = field(default_factory=dict)
|
|
|
|
# Convenience accessors
|
|
@property
|
|
def succeeded(self) -> list[str]:
|
|
return [k for k, v in self.completed.items() if v.status == "succeeded"]
|
|
|
|
@property
|
|
def failed(self) -> list[str]:
|
|
return [k for k, v in self.completed.items() if v.status == "failed"]
|
|
|
|
@property
|
|
def summary(self) -> dict[str, int]:
|
|
return {
|
|
"succeeded": len(self.succeeded),
|
|
"failed": len(self.failed),
|
|
}
|
|
|
|
def is_done(self, item: str) -> bool:
|
|
return item in self.completed
|
|
|
|
def mark_succeeded(self, item: str) -> None:
|
|
self.completed[item] = ItemResult(status="succeeded")
|
|
|
|
def mark_failed(self, item: str, error: str = "") -> None:
|
|
self.completed[item] = ItemResult(status="failed", error=error)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Checkpoint I/O (atomic writes)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def load_checkpoint(path: Path) -> Checkpoint:
|
|
"""Load checkpoint from disk, or return an empty one if absent."""
|
|
if not path.exists():
|
|
return Checkpoint()
|
|
raw = json.loads(path.read_text(encoding="utf-8"))
|
|
completed: dict[str, ItemResult] = {}
|
|
for k, v in raw.get("completed", {}).items():
|
|
completed[k] = ItemResult(**v)
|
|
return Checkpoint(completed=completed)
|
|
|
|
|
|
async def save_checkpoint(cp: Checkpoint, path: Path) -> None:
|
|
"""
|
|
Atomically persist *cp* to *path*.
|
|
|
|
Writes to a temp file in the same directory then ``os.replace()`` —
|
|
this is atomic on every POSIX filesystem and on Windows when the dest
|
|
is on the same volume. A crash mid-write can never corrupt the target.
|
|
"""
|
|
parent = path.parent
|
|
parent.mkdir(parents=True, exist_ok=True)
|
|
loop = asyncio.get_event_loop()
|
|
|
|
def _write() -> None:
|
|
fd, tmp = tempfile.mkstemp(suffix=".tmp", dir=parent)
|
|
try:
|
|
payload = json.dumps(asdict(cp), indent=2, ensure_ascii=False) + "\n"
|
|
os.write(fd, payload.encode("utf-8"))
|
|
os.close(fd)
|
|
# Atomic rename — the only point where the target is visible.
|
|
os.replace(tmp, str(path))
|
|
except BaseException:
|
|
os.close(fd) if not None else None
|
|
try:
|
|
os.unlink(tmp)
|
|
except OSError:
|
|
pass
|
|
raise
|
|
|
|
await loop.run_in_executor(None, _write)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Batch processor
|
|
# ---------------------------------------------------------------------------
|
|
|
|
class BatchProcessor:
|
|
"""
|
|
Process a list of items with bounded concurrency, retries, checkpointing,
|
|
and graceful SIGINT handling.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
items: list[str],
|
|
*,
|
|
checkpoint_path: str | Path = "checkpoint.json",
|
|
max_concurrency: int = 8,
|
|
max_retries: int = 3,
|
|
base_backoff: float = 0.1,
|
|
) -> None:
|
|
self.items = list(items)
|
|
self.checkpoint_path = Path(checkpoint_path)
|
|
self.max_concurrency = max_concurrency
|
|
self.max_retries = max_retries
|
|
self.base_backoff = base_backoff
|
|
|
|
# Load any existing checkpoint for idempotency
|
|
self.cp = load_checkpoint(self.checkpoint_path)
|
|
|
|
# Semaphore for bounded concurrency
|
|
self._semaphore: asyncio.Semaphore | None = None
|
|
|
|
# Graceful-shutdown flag — set by the signal handler
|
|
self._shutdown_requested = asyncio.Event()
|
|
|
|
# Track in-flight tasks so we can wait for them on shutdown
|
|
self._in_flight: set[asyncio.Task[Any]] = set()
|
|
|
|
# Wall-clock start time (set in ``run``)
|
|
self._start_time: float = 0.0
|
|
|
|
# ------------------------------------------------------------------
|
|
# Signal handling
|
|
# ------------------------------------------------------------------
|
|
|
|
def _install_signal_handlers(self) -> None:
|
|
loop = asyncio.get_event_loop()
|
|
for sig in (signal.SIGINT, signal.SIGTERM):
|
|
loop.add_signal_handler(sig, self._on_signal, sig)
|
|
|
|
def _remove_signal_handlers(self) -> None:
|
|
loop = asyncio.get_event_loop()
|
|
for sig in (signal.SIGINT, signal.SIGTERM):
|
|
try:
|
|
loop.remove_signal_handler(sig)
|
|
except ValueError:
|
|
pass
|
|
|
|
def _on_signal(self, sig: signal.Signals) -> None:
|
|
# First signal: request graceful shutdown.
|
|
# Subsequent signals during the drain window force immediate exit.
|
|
if not self._shutdown_requested.is_set():
|
|
print(f"\n[BatchProcessor] {sig.name} received — draining in-flight tasks …", file=sys.stderr)
|
|
self._shutdown_requested.set()
|
|
else:
|
|
print(f"\n[BatchProcessor] {sig.name} received again — forcing exit.", file=sys.stderr)
|
|
sys.exit(1)
|
|
|
|
# ------------------------------------------------------------------
|
|
# Core processing
|
|
# ------------------------------------------------------------------
|
|
|
|
async def _process_one(self, item: str) -> None:
|
|
"""Retry *item* up to max_retries with exponential backoff."""
|
|
attempts = 0
|
|
last_error = ""
|
|
while attempts < self.max_retries:
|
|
attempts += 1
|
|
try:
|
|
result = await process(item)
|
|
self.cp.mark_succeeded(item)
|
|
await save_checkpoint(self.cp, self.checkpoint_path)
|
|
print(f" ✓ {item} (attempt {attempts})")
|
|
return
|
|
except ProcessingError as exc:
|
|
last_error = str(exc)
|
|
if attempts < self.max_retries:
|
|
backoff = self.base_backoff * (2 ** (attempts - 1))
|
|
print(f" ✗ {item} attempt {attempts}/{self.max_retries} failed — retry in {backoff:.2f}s", file=sys.stderr)
|
|
try:
|
|
await asyncio.sleep(backoff)
|
|
except asyncio.CancelledError:
|
|
# We were cancelled mid-backoff — do NOT mark failed.
|
|
raise
|
|
|
|
# Exhausted retries → terminal failure
|
|
self.cp.mark_failed(item, error=last_error)
|
|
await save_checkpoint(self.cp, self.checkpoint_path)
|
|
print(f" ☠ {item} failed after {attempts} attempts", file=sys.stderr)
|
|
|
|
async def _run_item(self, item: str) -> None:
|
|
"""Wrap ``_process_one`` with concurrency tracking and shutdown guard."""
|
|
task = asyncio.current_task()
|
|
assert task is not None
|
|
self._in_flight.add(task)
|
|
try:
|
|
# Check shutdown before starting work on this item.
|
|
if self._shutdown_requested.is_set():
|
|
return
|
|
await self._process_one(item)
|
|
finally:
|
|
self._in_flight.discard(task)
|
|
|
|
# ------------------------------------------------------------------
|
|
# Public entry point
|
|
# ------------------------------------------------------------------
|
|
|
|
async def run(self) -> dict[str, Any]:
|
|
"""
|
|
Run the batch. Returns a summary dict suitable for JSON output.
|
|
"""
|
|
self._start_time = time.monotonic()
|
|
self._install_signal_handlers()
|
|
|
|
# Items already in the checkpoint are skipped.
|
|
pending = [it for it in self.items if not self.cp.is_done(it)]
|
|
skipped = len(self.items) - len(pending)
|
|
|
|
self._semaphore = asyncio.Semaphore(self.max_concurrency)
|
|
|
|
try:
|
|
async def _bounded(item: str) -> None:
|
|
async with self._semaphore:
|
|
await self._run_item(item)
|
|
|
|
workers = [asyncio.create_task(_bounded(it)) for it in pending]
|
|
|
|
# Wait until either all workers finish or shutdown is requested.
|
|
if workers:
|
|
done, pending_tasks = await asyncio.wait(
|
|
workers,
|
|
return_when=asyncio.FIRST_COMPLETED,
|
|
)
|
|
# Check if we should drain or abort.
|
|
if self._shutdown_requested.is_set():
|
|
# Cancel any still-pending workers (they won't start new items).
|
|
for t in pending_tasks:
|
|
t.cancel()
|
|
# Wait for all to finish/cancel.
|
|
if pending_tasks:
|
|
await asyncio.gather(*pending_tasks, return_exceptions=True)
|
|
else:
|
|
# Some finished; continue with the rest.
|
|
if pending_tasks:
|
|
done2, _ = await asyncio.wait(pending_tasks)
|
|
await asyncio.gather(*done2, return_exceptions=True)
|
|
|
|
# Final flush (idempotent — same content).
|
|
await save_checkpoint(self.cp, self.checkpoint_path)
|
|
|
|
finally:
|
|
self._remove_signal_handlers()
|
|
|
|
elapsed_ms = int((time.monotonic() - self._start_time) * 1000)
|
|
summary = {
|
|
"succeeded": len(self.cp.succeeded),
|
|
"failed": len(self.cp.failed),
|
|
"skipped": skipped,
|
|
"total": len(self.items),
|
|
"elapsed_ms": elapsed_ms,
|
|
}
|
|
return summary
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# CLI entry point
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def main(argv: list[str] | None = None) -> int:
|
|
import argparse
|
|
|
|
parser = argparse.ArgumentParser(description="Async batch processor with checkpointing")
|
|
parser.add_argument("items", nargs="+", help="Item identifiers to process")
|
|
parser.add_argument("--checkpoint", default="checkpoint.json", help="Checkpoint file path")
|
|
parser.add_argument("--concurrency", type=int, default=8, help="Max concurrent items")
|
|
parser.add_argument("--retries", type=int, default=3, help="Max retries per item")
|
|
args = parser.parse_args(argv)
|
|
|
|
processor = BatchProcessor(
|
|
args.items,
|
|
checkpoint_path=args.checkpoint,
|
|
max_concurrency=args.concurrency,
|
|
max_retries=args.retries,
|
|
)
|
|
|
|
summary = asyncio.run(processor.run())
|
|
print(json.dumps(summary))
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main()) |