Coverage for src/quber/playground/batch.py: 84%

1"""Batch answering: many questions asked of one document as one server-side job.

3Each question is an ordinary answer — retrieval runs for it, the answer step

4fills the same `Answer` a single question produces — so a batch introduces no

5second answer format. What this module adds is the run around them: a job that

6continues if the page is closed, a bounded number of questions in flight, a

7retry with backoff for a provider that returns 529 under load, and one event

8per question completing so progress is a count filling in rather than a log

9being tailed.

11The two steps are two stages with their own bounds. Retrieval for one question

12does not wait for another question's answer: up to `concurrency` retrievals

13run while up to `concurrency` answers run, and at most `concurrency` retrieved

14contexts wait between them. The run then moves at the pace of the slower stage

15rather than the sum of both. Retrieval alone took about as long as the answer

16on the Jev ranker, so serial slots spent half their life on work that held no

17answer-model call.

19The runner holds no reference to the app. The caller hands `start` two

20coroutines, one that retrieves for a question and one that answers it from

21what was retrieved; everything here is scheduling, accounting and delivery,

22which is what makes it testable without a model.

24Runs live in the server's memory, capped at the twenty most recent so a

25long-lived process does not accumulate them. Given a store, the runner also

26records each run and every finished row in the database, and answers for a

27run it does not hold in memory from there: a run started by another task,

28or by a process since replaced. A run known to neither reaches the page as

29an id it kept and a run nobody can answer for. That is never silence — the

30route 404s and the page reports the run as gone, because a batch of forty

31figures quietly becoming thirty-nine is the failure this feature exists to

32prevent.

33"""

35from __future__ import annotations

37import asyncio

38import uuid

39from collections import OrderedDict

40from dataclasses import dataclass, field

41from typing import Any, AsyncIterator, Awaitable, Callable, Dict, List, Optional, Protocol

43from pydantic import BaseModel

45from quber.settings import get_settings

47#: Runs kept in memory, most recent first. Chosen to put eviction out of reach

48#: of a working session rather than to be relied on.

49MAX_RUNS = 20

51#: HTTP statuses worth retrying: overload and transient server failures.

52RETRY_STATUS = {429, 500, 502, 503, 529}

53MAX_ATTEMPTS = 4

54FIRST_RETRY_DELAY = 2.0

56#: The two coroutines the caller supplies. Retrieval: one question in, an

57#: opaque context out, whatever the answer stage needs. Answer: the question

58#: and that context in, the answer contract out.

59RetrieveOne = Callable[[str], Awaitable[Any]]

60AnswerOne = Callable[[str, Any], Awaitable[Dict[str, Any]]]

62#: How often a run served from the store is re-read for new rows.

63STORE_POLL_SECONDS = 2.0

66class RunStore(Protocol):

67 """The durable side of a run: what the runner writes as it goes and what

68 it reads back for a run it does not hold in memory."""

70 def save_run(self, job_id: str, doc_id: int, want: str, questions: List[str]) -> None: ...

71 def save_row(

72 self, job_id: str, index: int, question: str, error: Optional[str], response: Optional[Dict[str, Any]]

73 ) -> None: ...

74 def finish_run(self, job_id: str) -> None: ...

75 def load_run(self, job_id: str) -> Optional[Dict[str, Any]]: ...

76 def running_for(self, doc_id: int) -> bool: ...

79def parse_questions(text: str) -> List[str]:

80 """One question per line, as the batch screen's prompt window states.

82 Lines are trimmed and blank lines dropped. A repeated line is kept — a

83 person who pasted a question twice gets two rows, because a question that

84 silently vanishes from a run is exactly what a batch must never do.

85 """

86 return [line.strip() for line in text.splitlines() if line.strip()]

89class BatchRow(BaseModel):

90 """One finished question. Either `response` (the /api/chat contract,

91 verbatim) or `error` is set, never both."""

93 index: int

94 question: str

95 error: Optional[str] = None

96 response: Optional[Dict[str, Any]] = None

99class RunState(BaseModel):

100 """What `GET /api/batch/{id}` returns: the run's counts and the rows

101 finished so far, in question order."""

102

103 job_id: str

104 doc_id: int

105 want: str

106 total: int

107 done: int

108 failed: int

109 running: bool

110 rows: List[BatchRow]

111

112

113@dataclass

114class _Run:

115 job_id: str

116 doc_id: int

117 want: str

118 questions: List[str]

119 rows: Dict[int, BatchRow] = field(default_factory=dict)

120 failed: int = 0

121 task: Optional[asyncio.Task[None]] = None

122 subscribers: List[asyncio.Queue[Optional[Dict[str, Any]]]] = field(default_factory=list)

123

124 @property

125 def running(self) -> bool:

126 return self.task is not None and not self.task.done()

127

128 def state(self) -> RunState:

129 return RunState(

130 job_id=self.job_id,

131 doc_id=self.doc_id,

132 want=self.want,

133 total=len(self.questions),

134 done=len(self.rows),

135 failed=self.failed,

136 running=self.running,

137 rows=[self.rows[i] for i in sorted(self.rows)],

138 )

139

140 def publish(self, event: Dict[str, Any]) -> None:

141 for q in self.subscribers:

142 q.put_nowait(event)

143

144 def close_subscribers(self) -> None:

145 for q in self.subscribers:

146 q.put_nowait(None)

147

148

149class BatchRunner:

150 """The run registry. One per server process."""

151

152 def __init__(

153 self,

154 concurrency: Optional[int] = None,

155 max_runs: int = MAX_RUNS,

156 store: Optional[RunStore] = None,

157 ) -> None:

158 # Questions in flight at once, from settings unless the caller says.

159 # Each holds an answer-model call, and the provider returns 529 under

160 # load (observed 2026-07-29 during a sweep), so the bound is deliberate.

161 self._concurrency = (

162 concurrency if concurrency is not None else get_settings().playground.question_concurrency

163 )

164 self._max_runs = max_runs

165 self._runs: OrderedDict[str, _Run] = OrderedDict()

166 self._store = store

167

168 def running_ids(self) -> List[str]:

169 """The runs this process is still working on."""

170 return [run.job_id for run in self._runs.values() if run.running]

171

172 def start(

173 self, doc_id: int, questions: List[str], want: str, retrieve_one: RetrieveOne, answer_one: AnswerOne

174 ) -> RunState:

175 """Register a run and schedule it on the running event loop."""

176 job_id = uuid.uuid4().hex[:12]

177 run = _Run(job_id=job_id, doc_id=doc_id, want=want, questions=questions)

178 if self._store is not None:

179 self._store.save_run(job_id, doc_id, want, questions)

180 self._runs[job_id] = run

181 while len(self._runs) > self._max_runs:

182 _, evicted = self._runs.popitem(last=False)

183 if evicted.task is not None and not evicted.task.done():

184 evicted.task.cancel()

185 evicted.close_subscribers()

186 run.task = asyncio.get_running_loop().create_task(self._run(run, retrieve_one, answer_one))

187 return run.state()

188

189 def get(self, job_id: str) -> Optional[RunState]:

190 run = self._runs.get(job_id)

191 if run is not None:

192 return run.state()

193 stored = self._stored(job_id)

194 return self._state_from_stored(stored) if stored else None

195

196 def _stored(self, job_id: str) -> Optional[Dict[str, Any]]:

197 return self._store.load_run(job_id) if self._store is not None else None

198

199 @staticmethod

200 def _state_from_stored(stored: Dict[str, Any]) -> RunState:

201 rows = [BatchRow(**r) for r in stored["rows"]]

202 return RunState(

203 job_id=stored["job_id"],

204 doc_id=stored["doc_id"],

205 want=stored["want"],

206 total=len(stored["questions"]),

207 done=len(rows),

208 failed=stored["failed"],

209 running=stored["running"],

210 rows=rows,

211 )

212

213 def running_for(self, doc_id: int) -> bool:

214 """Whether any run against this document is still going — what lets

215 a document removal or replacement refuse instead of pulling the

216 corpus out from under a run mid-flight."""

217 if any(run.doc_id == doc_id and run.running for run in self._runs.values()):

218 return True

219 return self._store is not None and self._store.running_for(doc_id)

220

221 async def events(self, job_id: str) -> AsyncIterator[Dict[str, Any]]:

222 """Yield one event per question completing, then a final `done` event.

223

224 Subscribing to a finished run yields the `done` event immediately, so

225 a page that attaches late still learns the run is over.

226 """

227 run = self._runs.get(job_id)

228 if run is None:

229 async for event in self._stored_events(job_id):

230 yield event

231 return

232 if not run.running:

233 yield self._done_event(run)

234 return

235 queue: asyncio.Queue[Optional[Dict[str, Any]]] = asyncio.Queue()

236 run.subscribers.append(queue)

237 try:

238 while True:

239 event = await queue.get()

240 if event is None:

241 break

242 yield event

243 if event.get("type") == "done":

244 break

245 finally:

246 run.subscribers.remove(queue)

247

248 async def _stored_events(self, job_id: str) -> AsyncIterator[Dict[str, Any]]:

249 """Events for a run another process is working on, read from the

250 store: each row as it lands there, then done once the run closes."""

251 seen: set[int] = set()

252 while True:

253 stored = self._stored(job_id)

254 if stored is None:

255 return

256 total = len(stored["questions"])

257 for row in stored["rows"]:

258 if row["index"] in seen:

259 continue

260 seen.add(row["index"])

261 yield {

262 "type": "row",

263 "index": row["index"],

264 "total": total,

265 "done": len(seen),

266 "failed": stored["failed"],

267 "ok": row["error"] is None,

268 "row": BatchRow(**row).model_dump(mode="json"),

269 }

270 if not stored["running"]:

271 yield {

272 "type": "done",

273 "total": total,

274 "done": len(stored["rows"]),

275 "failed": stored["failed"],

276 }

277 return

278 await asyncio.sleep(STORE_POLL_SECONDS)

279

280 def _done_event(self, run: _Run) -> Dict[str, Any]:

281 return {

282 "type": "done",

283 "total": len(run.questions),

284 "done": len(run.rows),

285 "failed": run.failed,

286 }

287

288 async def _run(self, run: _Run, retrieve_one: RetrieveOne, answer_one: AnswerOne) -> None:

289 n = self._concurrency

290 # At most 2n questions are in progress at once: n retrieving or waiting

291 # with a context in hand, n answering. Retrieval runs ahead exactly as

292 # far as the answer stage has fallen behind, and no further.

293 in_flight = asyncio.Semaphore(2 * n)

294 retrieving = asyncio.Semaphore(n)

295 answering = asyncio.Semaphore(n)

296

297 async def one(index: int, question: str) -> None:

298 async with in_flight:

299 row = BatchRow(index=index, question=question)

300 try:

301 async with retrieving:

302 context = await self._with_retry(retrieve_one, question)

303 async with answering:

304 row.response = await self._with_retry(answer_one, question, context)

305 except Exception as exc:

306 # The row carries its error rather than vanishing; the run

307 # never reports success after losing a question.

308 row.error = str(exc)

309 run.failed += 1

310 run.rows[index] = row

311 if self._store is not None:

312 self._store.save_row(run.job_id, index, question, row.error, row.response)

313 # The event carries the finished row itself: the page appends

314 # it in place rather than refetching the whole run per event.

315 run.publish(

316 {

317 "type": "row",

318 "index": index,

319 "total": len(run.questions),

320 "done": len(run.rows),

321 "failed": run.failed,

322 "ok": row.error is None,

323 "row": row.model_dump(mode="json"),

324 }

325 )

326

327 try:

328 await asyncio.gather(*(one(i, q) for i, q in enumerate(run.questions)))

329 finally:

330 if self._store is not None:

331 self._store.finish_run(run.job_id)

332 run.publish(self._done_event(run))

333 run.close_subscribers()

334

335 async def _with_retry(self, stage: Callable[..., Awaitable[Any]], *args: Any) -> Any:

336 delay = FIRST_RETRY_DELAY

337 for attempt in range(MAX_ATTEMPTS):

338 try:

339 return await stage(*args)

340 except Exception as exc:

341 status = getattr(exc, "status_code", None)

342 if status not in RETRY_STATUS or attempt == MAX_ATTEMPTS - 1:

343 raise

344 await asyncio.sleep(delay)

345 delay *= 2

346 raise RuntimeError("unreachable") # pragma: no cover

Coverage for src / quber / playground / batch.py: 84%

170 statements