| """Review one shared case queue using independently limited SJTU accounts.""" |
|
|
| import argparse |
| import os |
| import re |
| import threading |
| from contextlib import contextmanager |
| from pathlib import Path |
|
|
| from vimeml.review.run import MODEL, ROOT, QuotaLimiter, run |
|
|
|
|
| @contextmanager |
| def output_lock(output): |
| """Cross-platform, process-scoped exclusive lock; no stale PID cleanup.""" |
| output.mkdir(parents=True, exist_ok=True) |
| with (output / ".review.lock").open("a+b") as handle: |
| if handle.seek(0, 2) == 0: |
| handle.write(b"0") |
| handle.flush() |
| handle.seek(0) |
| try: |
| if os.name == "nt": |
| import msvcrt |
| msvcrt.locking(handle.fileno(), msvcrt.LK_NBLCK, 1) |
| else: |
| import fcntl |
| fcntl.flock(handle.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB) |
| except OSError: |
| raise RuntimeError("Review output is already in use; stop the other reviewer before resuming.") from None |
| try: |
| yield |
| finally: |
| handle.seek(0) |
| if os.name == "nt": |
| msvcrt.locking(handle.fileno(), msvcrt.LK_UNLCK, 1) |
| else: |
| fcntl.flock(handle.fileno(), fcntl.LOCK_UN) |
|
|
|
|
| class AccountPool: |
| """Choose an available account, with separate quotas and in-flight caps.""" |
|
|
| def __init__(self, api_keys, labels=None, workers_per_key=3, rpm=8, tpm=80000, |
| window=60.0, dry_run=False): |
| self.secrets = tuple(api_keys) |
| self.size = len(self.secrets) |
| self.labels = tuple(labels or [f"account_{index + 1}" for index in range(self.size)]) |
| if not 1 <= self.size <= 8 or len(self.labels) != self.size or len(set(self.labels)) != self.size: |
| raise ValueError("Provide 1..8 accounts with unique labels.") |
| if not dry_run and (any(not isinstance(key, str) or not key.strip() for key in self.secrets) |
| or len(set(self.secrets)) != self.size): |
| raise ValueError("Account keys must be nonempty and distinct; duplicate keys do not add quota.") |
| if not 1 <= workers_per_key <= 8 or rpm < 1 or tpm < 1: |
| raise ValueError("Invalid concurrency or quota.") |
| self.workers_per_key, self.rpm, self.tpm = workers_per_key, rpm, tpm |
| self.dry_run = dry_run |
| self.limiters = [QuotaLimiter(rpm, tpm, window) for _ in self.secrets] |
| self.lock, self.cursor = threading.Lock(), 0 |
| self.states = [{"account": label, "enabled": True, "disabled_http_status": None, |
| "in_flight": 0, "requests_started": 0, "successful_requests": 0, |
| "http_429": 0, "requests_with_usage": 0, "total_tokens_reported": 0} |
| for label in self.labels] |
|
|
| def plan(self): |
| return {"count": self.size, "labels": list(self.labels), |
| "workers_per_key": self.workers_per_key, "rpm_per_key": self.rpm, |
| "tpm_per_key": self.tpm, "aggregate_rpm": self.rpm * self.size, |
| "aggregate_tpm": self.tpm * self.size} |
|
|
| def acquire(self, tokens, stop): |
| if self.dry_run: |
| raise RuntimeError("A dry-run pool cannot send requests.") |
| if tokens > self.tpm: |
| raise ValueError("Single request exceeds per-account token quota.") |
| while not stop.is_set(): |
| waits = [] |
| with self.lock: |
| if not any(state["enabled"] for state in self.states): |
| return None |
| for offset in range(self.size): |
| index = (self.cursor + offset) % self.size |
| state = self.states[index] |
| if not state["enabled"]: |
| continue |
| if state["in_flight"] >= self.workers_per_key: |
| waits.append(0.05) |
| continue |
| ticket, wait = self.limiters[index].try_acquire(tokens) |
| if ticket is not None: |
| state["in_flight"] += 1 |
| state["requests_started"] += 1 |
| self.cursor = (index + 1) % self.size |
| return {"index": index, "quota_ticket": ticket} |
| waits.append(wait) |
| stop.wait(min(max(min(waits, default=0.05), 0.001), 0.05)) |
| return None |
|
|
| def key(self, ticket): |
| return self.secrets[ticket["index"]] |
|
|
| def label(self, ticket): |
| return self.labels[ticket["index"]] |
|
|
| def release(self, ticket): |
| with self.lock: |
| self.states[ticket["index"]]["in_flight"] -= 1 |
|
|
| def finish(self, ticket, usage): |
| self.limiters[ticket["index"]].finish(ticket["quota_ticket"], usage) |
| if isinstance(usage, dict) and type(usage.get("total_tokens")) is int and usage["total_tokens"] >= 0: |
| with self.lock: |
| state = self.states[ticket["index"]] |
| state["requests_with_usage"] += 1 |
| state["total_tokens_reported"] += usage["total_tokens"] |
|
|
| def success(self, ticket): |
| with self.lock: |
| self.states[ticket["index"]]["successful_requests"] += 1 |
|
|
| def pause(self, ticket, seconds): |
| self.limiters[ticket["index"]].pause(seconds) |
| with self.lock: |
| self.states[ticket["index"]]["http_429"] += 1 |
|
|
| def disable(self, ticket, status): |
| with self.lock: |
| state = self.states[ticket["index"]] |
| state["enabled"], state["disabled_http_status"] = False, status |
|
|
| def snapshot(self): |
| with self.lock: |
| return [dict(state) for state in self.states] |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument("--input-dir", type=Path, default=ROOT / "outputs/corpus-review-v1") |
| parser.add_argument("--output-dir", type=Path, help="Default: INPUT/review (compatible with single-key caches)") |
| parser.add_argument("--key-env", nargs="+", default=["SJTU_API_KEY", "SJTU_API_KEY_2"], |
| metavar="ENV_NAME", help="Names of distinct accounts' environment variables; never pass actual keys") |
| parser.add_argument("--workers-per-key", type=int, default=3) |
| parser.add_argument("--rpm-per-key", type=int, default=8) |
| parser.add_argument("--tpm-per-key", type=int, default=80000) |
| parser.add_argument("--batch-size", type=int, default=10) |
| parser.add_argument("--max-batch-tokens", type=int, default=10000) |
| parser.add_argument("--max-output-tokens", type=int, default=2048) |
| parser.add_argument("--timeout", type=int, default=180) |
| parser.add_argument("--retries", type=int, default=2) |
| parser.add_argument("--model", default=MODEL) |
| parser.add_argument("--dry-run", action="store_true", help="Read caches and estimate independent quotas; no keys, writes or requests") |
| args = parser.parse_args() |
| if (not 1 <= len(args.key_env) <= 8 or len(set(args.key_env)) != len(args.key_env) |
| or any(not re.fullmatch(r"[A-Za-z_][A-Za-z0-9_]*", name) for name in args.key_env)): |
| parser.error("--key-env requires 1..8 unique environment variable names, not API keys.") |
| if not 1 <= args.workers_per_key <= 8 or not 1 <= args.rpm_per_key <= 10 or not 1 <= args.tpm_per_key <= 100000: |
| parser.error("workers-per-key must be 1..8; rpm-per-key 1..10; tpm-per-key 1..100000.") |
| if (min(args.batch_size, args.timeout, args.max_output_tokens) < 1 |
| or not args.max_output_tokens < args.max_batch_tokens <= args.tpm_per_key |
| or not 0 <= args.retries <= 5): |
| parser.error("Invalid batch, output budget, timeout or retries.") |
| keys = [None] * len(args.key_env) if args.dry_run else [os.environ.get(name, "").strip() for name in args.key_env] |
| if not args.dry_run: |
| missing = [name for name, key in zip(args.key_env, keys) if not key] |
| if missing: |
| parser.error("Set these environment variables before starting: " + ", ".join(missing)) |
| if len(set(keys)) != len(keys): |
| parser.error("Two environment variables contain the same key; provide independent account keys.") |
| pool = AccountPool(keys, labels=args.key_env, workers_per_key=args.workers_per_key, |
| rpm=args.rpm_per_key, tpm=args.tpm_per_key, dry_run=args.dry_run) |
| result = run(args.input_dir, args.output_dir or args.input_dir / "review", |
| workers=len(keys) * args.workers_per_key, batch_size=args.batch_size, |
| rpm=args.rpm_per_key, tpm=args.tpm_per_key, max_batch_tokens=args.max_batch_tokens, |
| max_output_tokens=args.max_output_tokens, timeout=args.timeout, retries=args.retries, |
| dry_run=args.dry_run, model=args.model, request_pool=pool) |
| if not args.dry_run and result["pending_cases"]: |
| raise SystemExit(1) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
