Coverage for orchestrant/benchmark/depth.py: 95%

1"""How fast does a lane decode once its context is deep? 

2 

3One streamed request after about --context-tokens of filler: the time to the 

4first token, then the decode rate over each --window of generated tokens. 

5 

6 orchestrant-bench depth --backend geniex-cpu --model <id> --output cpu-9b-depth-8k.json 

7""" 

8 

9from __future__ import annotations 

10 

11import argparse 

12import functools 

13import hashlib 

14import json 

15import sys 

16import time 

17 

18from orchestrant.benchmark.answers import MIN_DECODE_WINDOW_S, delta_pieces 

19from orchestrant.benchmark.client import http_error_detail, post_json 

20from orchestrant.benchmark.contract import filler 

21 

22 

23# The tracked traces were measured on exactly this prompt; change it and they stop comparing. 

24FILLER_SEED = 7 

25QUESTION = ( 

26 "Write a detailed technical blog post about the benefits and challenges of " 

27 "running large language models locally." 

28) 

29# This file, and answers.py for what counts as a generated token. 

30TOOL_FILES = ("depth.py", "answers.py") 

31 

32 

33def depth_prompt(context_tokens): 

34 """About `context_tokens` of the contract's fixed filler, then QUESTION.""" 

35 if not context_tokens: 

36 return QUESTION 

37 notes = filler(context_tokens, seed=FILLER_SEED) 

38 return f"Background notes, for reference only:\n\n{notes}\n\n---\n\n{QUESTION}" 

39 

40 

41def windows(stamps, size=256): 

42 """The rate over each `size` generated tokens, cut as the scratch script cut them.""" 

43 out, last = [], len(stamps) - 1 

44 for start in range(1, last, size): 

45 end = min(start + size, last) 

46 if end - start < size // 4: 

47 break 

48 seconds = stamps[end] - stamps[start] 

49 burst = seconds < MIN_DECODE_WINDOW_S 

50 rate = None if burst else round((end - start) / seconds, 2) 

51 out.append({"tokens": f"{start}-{end}", "tok_per_s": rate}) 

52 return out 

53 

54 

55def _error(exc): 

56 """An HTTP error with its body, anything else by type: never just a status.""" 

57 detail = http_error_detail(exc) 

58 if detail: 

59 return f"HTTP {detail[0]}: {detail[1][:300]}" 

60 return f"{type(exc).__name__}: {exc}"[:300] 

61 

62 

63def warm(base_url, model, entry=None, timeout=3600): 

64 """A spacer that waits for the model to load, unlike client.spacer -> {"seconds", "error"}.""" 

65 body = { 

66 "model": model, 

67 "messages": [{"role": "user", "content": "ok"}], 

68 "max_tokens": 1, 

69 } 

70 t0, error = time.monotonic(), None 

71 try: 

72 url = f"{base_url}/v1/chat/completions" 

73 with post_json(url, body, entry=entry, timeout=timeout) as r: 

74 r.json() 

75 except Exception as e: # recorded; the measured request reports its own 

76 error = _error(e) 

77 return {"seconds": round(time.monotonic() - t0, 2), "error": error} 

78 

79 

80def _chunks(lines): 

81 """The parsed chunks of an SSE stream, up to [DONE]; an in-stream server error raises.""" 

82 for line in lines: 

83 if line.startswith("error:"): 

84 raise RuntimeError(f"server error in stream: {line[6:].strip()[:200]}") 

85 if not line.startswith("data:"): 

86 continue 

87 data = line[5:].strip() 

88 if data == "[DONE]": 

89 return 

90 try: 

91 chunk = json.loads(data) 

92 except json.JSONDecodeError: 

93 continue 

94 if not isinstance(chunk, dict): 

95 continue 

96 if chunk.get("error"): 

97 fault = json.dumps(chunk["error"])[:200] 

98 raise RuntimeError(f"server error in stream: {fault}") 

99 yield chunk 

100 

101 

102def measure(base_url, model, context_tokens, max_tokens, *, entry=None, timeout=3600): 

103 """One streamed reply after the filler -> (stamps, state); no generated token is an error.""" 

104 body = { 

105 "model": model, 

106 "messages": [{"role": "user", "content": depth_prompt(context_tokens)}], 

107 "max_tokens": max_tokens, 

108 "stream": True, 

109 "stream_options": {"include_usage": True}, 

110 } 

111 stamps, state = [], {"usage": None, "finish_reason": None, "error": None} 

112 t0 = time.monotonic() 

113 try: 

114 url = f"{base_url}/v1/chat/completions" 

115 with post_json(url, body, entry=entry, stream=True, timeout=timeout) as r: 

116 for chunk in _chunks(r.lines()): 

117 state["usage"] = chunk.get("usage") or state["usage"] 

118 for choice in chunk.get("choices") or []: 

119 finish = choice.get("finish_reason") 

120 state["finish_reason"] = finish or state["finish_reason"] 

121 if any(delta_pieces(choice.get("delta") or {})): 

122 stamps.append(time.monotonic() - t0) 

123 except Exception as e: # a partial trace is still evidence; the error says why 

124 state["error"] = _error(e) 

125 if not stamps and not state["error"]: 

126 # A lane that ignored "stream" answers one JSON body; that is no clean trace. 

127 finish = state["finish_reason"] 

128 state["error"] = f"no generated token arrived (finish_reason={finish!r})" 

129 return stamps, state 

130 

131 

132def trace_row(model, context_tokens, window, stamps, state, warmup): 

133 """One report row: the scratch script's six fields first, then what it lacked.""" 

134 return { 

135 "model": model, 

136 "context_tokens_requested": context_tokens, 

137 "usage": state["usage"], 

138 "ttft_s": round(stamps[0], 2) if stamps else None, 

139 "tokens": len(stamps), 

140 "windows": windows(stamps, window), 

141 "finish_reason": state["finish_reason"], 

142 "error": state["error"], 

143 "warmup": warmup, 

144 # The raw stamps, so a trace can be re-cut at another window size. 

145 "delta_times_s": [round(s, 3) for s in stamps], 

146 } 

147 

148 

149def summary_line(row): 

150 """The console line: prefill at depth, then the window rates.""" 

151 prompt = (row["usage"] or {}).get("prompt_tokens") 

152 ttft = row["ttft_s"] 

153 prefill = f" ({prompt / ttft:.1f} tok/s prefill)" if prompt and ttft else "" 

154 rates = [w["tok_per_s"] for w in row["windows"]] 

155 line = f" TTFT {ttft} s after {prompt} prompt tokens{prefill}; " 

156 line += f"{row['tokens']} tokens, windows {rates} tok/s" 

157 return line + (f"\n ERROR {row['error']}" if row["error"] else "") 

158 

159 

160def _parser(): 

161 ap = argparse.ArgumentParser( 

162 description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter 

163 ) 

164 ap.add_argument("--backend", default=None) 

165 ap.add_argument("--base-url", default=None) 

166 ap.add_argument("--model", default=None) 

167 ap.add_argument( 

168 "--context-tokens", 

169 type=int, 

170 default=8000, 

171 help="about this many tokens of filler before the question (default " 

172 "8000: 7154-7159 prompt tokens on the Qwen3 tokenizers)", 

173 ) 

174 ap.add_argument("--max-tokens", type=int, default=1024) 

175 ap.add_argument( 

176 "--window", type=int, default=256, help="tokens per rate window (default 256)" 

177 ) 

178 ap.add_argument( 

179 "--rest", 

180 type=float, 

181 default=30, 

182 help="seconds between the warm-up and the measured request, spent " 

183 "reading the host load (default 30)", 

184 ) 

185 ap.add_argument( 

186 "--timeout", 

187 type=int, 

188 default=3600, 

189 help="per-request timeout in seconds (default 3600: Ollama's 9B at 4 " 

190 "threads took 611.6 s to its first token at ~7.2k tokens)", 

191 ) 

192 ap.add_argument("--output", default=None, help="write the report as JSON") 

193 return ap 

194 

195 

196def main(): 

197 from orchestrant.benchmark.client import entry_config, run_start, write_report 

198 from orchestrant.benchmark.openai_api import ( 

199 detect_model_via_api, 

200 resolve_backend, 

201 resolve_backend_entry, 

202 resolve_model, 

203 ) 

204 

205 ap = _parser() 

206 args = ap.parse_args() 

207 if args.window < 1: 

208 ap.error("--window must be at least 1") 

209 base_url, backend_model, _ = resolve_backend(args.backend, args.base_url) 

210 entry = resolve_backend_entry(args.backend, args.base_url) 

211 detect = functools.partial(detect_model_via_api, base_url) 

212 model = resolve_model(args.model, backend_model, entry, detect) 

213 print(f"\n Depth trace: {model} @ {base_url}, ~{args.context_tokens} tokens\n") 

214 warmup = warm(base_url, model, entry, args.timeout) 

215 # The rest is also the host-load window: the load just before the request slows a CPU lane. 

216 rested = time.monotonic() 

217 started = run_start(TOOL_FILES, base_url, seconds=args.rest) 

218 time.sleep(max(0.0, args.rest - (time.monotonic() - rested))) 

219 stamps, state = measure( 

220 base_url, 

221 model, 

222 args.context_tokens, 

223 args.max_tokens, 

224 entry=entry, 

225 timeout=args.timeout, 

226 ) 

227 row = trace_row(model, args.context_tokens, args.window, stamps, state, warmup) 

228 # Written before anything that merely prints (write_report's rule). 

229 if args.output: 

230 prompt = depth_prompt(args.context_tokens).encode() 

231 config = { 

232 "context_tokens": args.context_tokens, 

233 "max_tokens": args.max_tokens, 

234 "window": args.window, 

235 "rest_s": args.rest, 

236 "timeout_s": args.timeout, 

237 "filler_seed": FILLER_SEED, 

238 "prompt_sha256": hashlib.sha256(prompt).hexdigest(), 

239 "backend_entry": entry_config(entry), 

240 } 

241 label = args.backend or model 

242 write_report( 

243 args.output, 

244 "bench_depth", 

245 config, 

246 [{"label": label, **row}], 

247 base_url, 

248 TOOL_FILES, 

249 run_start=started, 

250 ) 

251 print(summary_line(row)) 

252 if args.output: 

253 print(f"\n Report written to {args.output}") 

254 return 1 if row["error"] else 0 

255 

256 

257if __name__ == "__main__": 

258 sys.exit(main())