Coverage for orchestrant/benchmark/client.py: 96%
192 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-18 00:48 +0000
1"""The front end the bench tools share: plumbing, so kept out of `tool_sha256` and grading."""
3import json
4import os
5import re
6import sys
7import time
8import urllib.error
9import urllib.request
10from datetime import UTC, datetime
12from orchestrant.benchmark.gateway import note_reply
15def _row(entry_label, backend, base_url, model, url):
16 """One candidate as a dict, before disambiguation."""
17 label = entry_label or model or backend or url or "unnamed"
18 return {
19 "label": label,
20 "explicit_label": bool(entry_label),
21 "backend": backend,
22 "base_url": url,
23 "raw_base_url": base_url,
24 "model": model,
25 "entry": {},
26 }
29def disambiguate(rows):
30 """Make every candidate label unique, in place, or refuse loudly: consumers key on it."""
31 seen = {}
32 for row in rows:
33 seen.setdefault(row["label"], []).append(row)
34 for label, group in seen.items():
35 if len(group) < 2:
36 continue
37 explicit = [r for r in group if r["explicit_label"]]
38 if explicit:
39 raise SystemExit(
40 f"two candidates share the label {label!r}. A label names one "
41 f"measured endpoint; give each an explicit distinct --label / "
42 f'"label" field.'
43 )
44 for row in group:
45 suffix = row["backend"] or row["base_url"]
46 if suffix:
47 row["label"] = f"{label} ({suffix})"
48 final = [r["label"] for r in rows]
49 if len(set(final)) != len(final):
50 dupes = sorted({lbl for lbl in final if final.count(lbl) > 1})
51 raise SystemExit(
52 f"candidate labels still collide after disambiguation: {dupes}. "
53 f'Give each candidate a distinct "label".'
54 )
55 return rows
58def load_candidates(path, resolve_backend, resolve_entry=None):
59 """Read a candidates/--compare file into rows; a `_comment`-only element is not one."""
60 with open(path) as f:
61 entries = json.load(f)
62 if not isinstance(entries, list):
63 raise SystemExit(f"{path}: expected a JSON list of candidates")
64 rows = []
65 for entry in entries:
66 if isinstance(entry, dict) and set(entry) <= {"_comment"}:
67 continue
68 backend, base_url = entry.get("backend"), entry.get("base_url")
69 url, model, _ = resolve_backend(backend, base_url)
70 row = _row(
71 entry.get("label"), backend, base_url, entry.get("model") or model, url
72 )
73 if resolve_entry is not None:
74 row["entry"] = resolve_entry(backend, base_url) or {}
75 rows.append(row)
76 return disambiguate(rows)
79def candidate_rows(args, resolve_backend, resolve_entry=None):
80 """The same candidates as resolve_candidates, as whole rows that keep the registry NAME."""
81 compare_file = getattr(args, "compare", None)
82 if compare_file:
83 return load_candidates(compare_file, resolve_backend, resolve_entry)
84 backend, base_url = getattr(args, "backend", None), getattr(args, "base_url", None)
85 url, model, _ = resolve_backend(backend, base_url)
86 row = _row(
87 getattr(args, "label", None),
88 backend,
89 base_url,
90 getattr(args, "model", None) or model,
91 url,
92 )
93 if resolve_entry is not None:
94 row["entry"] = resolve_entry(backend, base_url) or {}
95 return disambiguate([row])
98def resolve_candidates(args, resolve_backend, resolve_entry=None):
99 """Return [(label, base_url, model)] from --compare or the single-run flags.
101 With `resolve_entry`, 4-tuples ending in the backends.json entry.
102 """
103 rows = candidate_rows(args, resolve_backend, resolve_entry)
104 if resolve_entry is None:
105 return [(r["label"], r["base_url"], r["model"]) for r in rows]
106 return [(r["label"], r["base_url"], r["model"], r["entry"]) for r in rows]
109def run_start(tool_files, base_url=None, seconds=3):
110 """The start time, source hash and host load a tool records before its first request.
112 Pass the same `tool_files` to write_report, or the start hash is not compared.
113 """
114 from orchestrant.benchmark import hostload
115 from orchestrant.benchmark.provenance import tool_fingerprint
117 record = {
118 "started_utc": datetime.now(UTC).isoformat(),
119 "tool_files": list(tool_files),
120 "tool_sha256": tool_fingerprint(*tool_files) if tool_files else None,
121 }
122 try:
123 record["host_load"] = hostload.load_snapshot(seconds, lane=base_url)
124 except Exception as e: # a load reading must never cost the run
125 record["host_load"] = {
126 "other_cores": None,
127 "note": f"{type(e).__name__}: {e}"[:160],
128 }
129 print(hostload.load_line(record["host_load"]), flush=True)
130 return record
133def _served_model(reports):
134 """The one model id every report row served, or None when they differ."""
135 models = {r.get("model") for r in reports or () if isinstance(r, dict)}
136 return models.pop() if len(models) == 1 else None
139def write_report(
140 path,
141 benchmark,
142 config,
143 reports,
144 base_url,
145 tool_files,
146 extra=None,
147 *,
148 run_start=None,
149):
150 """Write one report in the shared envelope, BEFORE anything that merely prints."""
151 from orchestrant.benchmark.provenance import collect
153 start, same_files = run_start or {}, True
154 started = {}
155 if start:
156 # Hashes of different file sets always differ; comparing them would invent an edit.
157 same_files = list(start.get("tool_files") or ()) == list(tool_files)
158 started = {
159 "tool_sha256_at_start": start.get("tool_sha256") if same_files else None,
160 "host_load": start.get("host_load"),
161 "run_started_utc": start.get("started_utc"),
162 }
163 provenance = collect(base_url, tool_files, model=_served_model(reports), **started)
164 if not same_files:
165 provenance.setdefault("incomplete", []).append("tool_sha256_at_start")
166 for key, value in (extra or {}).items():
167 if key == "incomplete":
168 provenance["incomplete"] = provenance.get("incomplete", []) + list(value)
169 else:
170 provenance[key] = value
171 payload = {
172 "benchmark": benchmark,
173 "provenance": provenance,
174 "config": config,
175 "reports": reports,
176 }
177 tmp = f"{path}.tmp"
178 with open(tmp, "w") as f:
179 json.dump(payload, f, indent=2)
180 # Atomic: a Ctrl-C mid-write must not leave a truncated JSON.
181 os.replace(tmp, path)
182 return payload
185# One request path, one API-key read: see benchmarks/docs/llm-benchmark-review-2026-09-05.md § R10
188class Response:
189 """What post_json returns: `json()` for a body; `lines()`, then `gave_up`, for a stream.
191 Use it as a context manager; urllib does not close the socket for you.
192 """
194 def __init__(self, raw, started, deadline):
195 self.raw = raw
196 self.status = getattr(raw, "status", None)
197 self.headers = getattr(raw, "headers", {})
198 self.started = started
199 self.gave_up = False
200 self._deadline = deadline
202 def json(self):
203 return json.load(self.raw)
205 def lines(self):
206 for chunk in self.raw:
207 if self._deadline and time.monotonic() - self.started > self._deadline:
208 self.gave_up = True
209 return
210 yield chunk.decode("utf-8", "replace").strip()
212 def close(self):
213 self.raw.close()
215 def __enter__(self):
216 return self
218 def __exit__(self, *exc):
219 self.close()
220 return False
223def request_headers(entry):
224 """Content-Type, the entry's own headers, then Bearer auth read here from api_key_env."""
225 entry = entry or {}
226 headers = {"Content-Type": "application/json"}
227 for key, value in (entry.get("headers") or {}).items():
228 headers[str(key)] = str(value)
229 var = entry.get("api_key_env")
230 if var:
231 if not os.environ.get(var):
232 raise SystemExit(
233 f"this backend reads its API key from the environment variable "
234 f"{var}, which is unset or empty. Export it before benchmarking "
235 f"(the value is never printed or written to a report)."
236 )
237 headers["Authorization"] = f"Bearer {os.environ[var]}"
238 return headers
241def request_extras(entry):
242 """The entry's request_extra body keys — safe to record in a report."""
243 return dict((entry or {}).get("request_extra") or {})
246def utf8_stdio():
247 """Write stdout/stderr as UTF-8: a piped Windows stdout is cp1252 and raises on '→'."""
248 for stream in (sys.stdout, sys.stderr):
249 reconfigure = getattr(stream, "reconfigure", None)
250 if reconfigure is not None:
251 reconfigure(encoding="utf-8", errors="replace")
254def spacer(base_url, model, entry=None):
255 """One throwaway request: GenieX answers an identical follow-up along a cache path."""
256 body = {
257 "model": model,
258 "messages": [{"role": "user", "content": "Reply with the single word: ok"}],
259 "max_tokens": 1,
260 "stream": False,
261 }
262 try:
263 with post_json(
264 f"{base_url}/v1/chat/completions", body, entry=entry, timeout=120
265 ) as r:
266 r.json()
267 except Exception: # nosec B110 -- best effort, see the docstring
268 pass
271def entry_config(entry):
272 """What a committed report may say about a backends.json entry: names, never values."""
273 entry = entry or {}
274 # Cython 3.3's sorted() optimisation crashes on `sorted(x or {})`, so the fallback is bound first.
275 headers = entry.get("headers") or {}
276 return {
277 "request_extra": request_extras(entry),
278 "headers": sorted(headers),
279 "api_key_env": entry.get("api_key_env"),
280 "probe": bool(entry.get("probe", True)),
281 }
284REDACTED = "<redacted>"
285# Whole words only: --api-key loses its value, --max-tokens keeps its own.
286_SECRET_WORDS = frozenset(
287 ("key", "apikey", "token", "secret", "password", "passwd", "auth", "bearer")
288)
289# Credential-shaped values under any flag; hf_/ghp_ need one alphanumeric run, sparing paths.
290_SECRET_SHAPES = (
291 (re.compile(r"\b(?:sk|pk|rk|glpat)-[\w-]{16,}"), REDACTED),
292 (re.compile(r"\b(?:hf|gh[pousr])_[A-Za-z0-9]{16,}"), REDACTED),
293 (re.compile(r"\bgithub_pat_\w{16,}"), REDACTED),
294 (re.compile(r"(?i)\b(bearer\s+)\S+"), rf"\1{REDACTED}"),
295 (
296 re.compile(r"(?i)([?&](?:api[-_]?key|key|token|access[-_]?token)=)[^&#\s]+"),
297 rf"\1{REDACTED}",
298 ),
299 (re.compile(r"(://[^/:@\s]+:)[^@/\s]+@"), rf"\1{REDACTED}@"),
300)
303def _secret_flag(flag):
304 """Is `flag` (an argument up to any '=') named for a credential?"""
305 if not flag.startswith("-"):
306 return False
307 return bool(_SECRET_WORDS & set(re.split(r"[-_.]", flag.lstrip("-").lower())))
310def redact_argv(argv):
311 """A command line as a report may record it: every credential value replaced."""
312 out, hide_next = [], False
313 for arg in map(str, argv):
314 if hide_next and not arg.startswith("-"):
315 out.append(REDACTED)
316 hide_next = False
317 continue
318 flag, has_value, _ = arg.partition("=")
319 secret = _secret_flag(flag)
320 hide_next = secret and not has_value
321 kept = f"{flag}={REDACTED}" if secret and has_value else arg
322 for pattern, replacement in _SECRET_SHAPES:
323 kept = pattern.sub(replacement, kept)
324 out.append(kept)
325 return out
328def post_json(url, body, entry=None, stream=False, timeout=300, deadline=None):
329 """POST one JSON body, with the entry's auth, headers and request_extra.
331 request_extra goes in BEFORE `body`, so the caller's keys always win; HTTPError propagates.
332 """
333 payload = request_extras(entry)
334 payload.update(body or {})
335 if stream:
336 payload.setdefault("stream", True)
337 req = urllib.request.Request(
338 url, data=json.dumps(payload).encode(), headers=request_headers(entry)
339 )
340 started = time.monotonic()
341 try:
342 raw = urllib.request.urlopen(req, timeout=timeout) # nosec B310 -- endpoint URL from config
343 except urllib.error.HTTPError as e:
344 note_reply(url, payload, e.headers)
345 raise
346 note_reply(url, payload, getattr(raw, "headers", None))
347 return Response(raw, started, deadline)
350def http_error_detail(exc, limit=500):
351 """(status, body) of the HTTPError post_json let through; None for any other."""
352 if not isinstance(exc, urllib.error.HTTPError):
353 return None
354 try:
355 body = exc.read().decode("utf-8", "replace")
356 except Exception: # an unreadable body still has its status
357 body = ""
358 finally:
359 # An unclosed error holds its connection until GC, and Python 3.14 warns then.
360 exc.close()
361 return exc.code, body[:limit]