Coverage for src / quber / playground / jobs.py: 29%
126 statements
« prev ^ index » next coverage.py v7.14.0, created at 2026-09-23 22:14 -0400
1"""Job records in Postgres.
3An upload job and a batch run each live in the memory of the task running
4them, as they always have; this module is the durable copy. The running
5task writes every stage transition, every finished batch row, and a
6heartbeat every few seconds. Any task can then answer a status poll or
7serve a run's rows from the tables, and a task that dies leaves a record
8whose heartbeat stops: a job with no heartbeat for ``STALE_SECONDS`` is
9marked failed with that reason the next time anything reads it, instead of
10sitting at its last stage forever.
12The tables are created by ``ensure_tables`` at startup from ``jobs.sql``,
13which is idempotent, so a running database gains them without a migration
14step. Nothing here is imported by the batch runner; the runner takes a store
15object with the same method names, and tests run it without one.
16"""
18from __future__ import annotations
20import json
21import os
22import socket
23import threading
24import time
25from pathlib import Path
26from typing import Any, Callable, Dict, Iterable, List, LiteralString, Optional, cast
28from loguru import logger
30from quber.playground import db
32JOBS_SQL = Path(__file__).with_name("jobs.sql")
34#: Who is running a job: this process, named so a record can say which task
35#: it belonged to. On Fargate the hostname is the task's.
36OWNER = f"{socket.gethostname()}:{os.getpid()}"
38STALE_SECONDS = 90
39HEARTBEAT_SECONDS = 10
40LOG_TAIL = 200
41TERMINAL_STAGES = ("done", "failed")
42STOPPED_REASON = "the task running this job stopped"
44_UPLOAD_COLUMNS: tuple[str, ...] = (
45 "job_id",
46 "filename",
47 "title",
48 "folder",
49 "filing_type",
50 "year",
51 "period",
52 "version",
53 "track",
54 "replace",
55 "replace_id",
56 "replace_key",
57 "doc_key",
58 "content_hash",
59 "doc_id",
60 "source_uri",
61 "stage",
62 "stages",
63 "error",
64 "error_info",
65 "log",
66 "scan_pages",
67 "scan_pages_reason",
68)
69_JSON_COLUMNS = {"stages", "error_info", "log"}
72#: One lock for the table creation. CREATE TABLE IF NOT EXISTS is not safe
73#: when two processes run it at the same instant, which two tasks of the
74#: hosted service starting together do.
75_TABLES_LOCK = 7_290_329
78def ensure_tables() -> None:
79 sql = cast(LiteralString, JOBS_SQL.read_text())
80 with db.connect() as conn, conn.transaction():
81 conn.execute("SELECT pg_advisory_xact_lock(%s)", (_TABLES_LOCK,))
82 conn.execute(sql)
85# ---------------------------------------------------------------- uploads
88def save_upload(job: Dict[str, Any]) -> None:
89 """Write the job as it stands, log tail included, and stamp the heartbeat."""
90 values: List[Any] = []
91 for column in _UPLOAD_COLUMNS:
92 value = job.get(column)
93 if column == "log":
94 value = list(value or [])[-LOG_TAIL:]
95 if column in _JSON_COLUMNS:
96 value = json.dumps(value) if value is not None else None
97 values.append(value)
98 columns: LiteralString = ", ".join(_UPLOAD_COLUMNS) # type: ignore[assignment]
99 placeholders: LiteralString = ", ".join(["%s"] * len(_UPLOAD_COLUMNS)) # type: ignore[assignment]
100 updates: LiteralString = ", ".join(f"{c} = EXCLUDED.{c}" for c in _UPLOAD_COLUMNS[1:]) # type: ignore[assignment]
101 query: LiteralString = (
102 "INSERT INTO ade_playground.upload_jobs (" # noqa: S608 - column names are literals
103 + columns
104 + ", owner, heartbeat, updated_at) VALUES ("
105 + placeholders
106 + ", %s, now(), now()) ON CONFLICT (job_id) DO UPDATE SET "
107 + updates
108 + ", owner = EXCLUDED.owner, heartbeat = now(), updated_at = now()"
109 )
110 with db.connect() as conn:
111 conn.execute(query, (*values, OWNER))
114def _upload_from_row(row: tuple[Any, ...]) -> Dict[str, Any]:
115 job = dict(zip(_UPLOAD_COLUMNS, row, strict=False))
116 for column in _JSON_COLUMNS:
117 value = job.get(column)
118 if isinstance(value, str):
119 job[column] = json.loads(value)
120 job["log"] = job.get("log") or []
121 return job
124def mark_stale_uploads() -> int:
125 """Fail every unfinished job whose heartbeat has stopped. Returns how many."""
126 with db.connect() as conn:
127 rows = conn.execute(
128 """UPDATE ade_playground.upload_jobs
129 SET stage = 'failed', error = %s, updated_at = now()
130 WHERE stage <> ALL(%s)
131 AND (heartbeat IS NULL OR heartbeat < now() - make_interval(secs => %s))
132 RETURNING job_id""",
133 (STOPPED_REASON, list(TERMINAL_STAGES), STALE_SECONDS),
134 ).fetchall()
135 for (job_id,) in rows:
136 logger.warning("upload job {} failed: {}", job_id, STOPPED_REASON)
137 return len(rows)
140def load_upload(job_id: str) -> Optional[Dict[str, Any]]:
141 mark_stale_uploads()
142 columns: LiteralString = ", ".join(_UPLOAD_COLUMNS) # type: ignore[assignment]
143 with db.connect() as conn:
144 row = conn.execute(
145 "SELECT " + columns + " FROM ade_playground.upload_jobs WHERE job_id = %s", # noqa: S608
146 (job_id,),
147 ).fetchone()
148 return _upload_from_row(row) if row else None
151def live_uploads() -> List[Dict[str, Any]]:
152 """Every upload job still running on any task."""
153 mark_stale_uploads()
154 columns: LiteralString = ", ".join(_UPLOAD_COLUMNS) # type: ignore[assignment]
155 with db.connect() as conn:
156 rows = conn.execute(
157 "SELECT " + columns + " FROM ade_playground.upload_jobs WHERE stage <> ALL(%s)", # noqa: S608
158 (list(TERMINAL_STAGES),),
159 ).fetchall()
160 return [_upload_from_row(r) for r in rows]
163def heartbeat_uploads(job_ids: Iterable[str]) -> None:
164 ids = list(job_ids)
165 if not ids:
166 return
167 with db.connect() as conn:
168 conn.execute(
169 "UPDATE ade_playground.upload_jobs SET heartbeat = now() WHERE job_id = ANY(%s)",
170 (ids,),
171 )
174# ---------------------------------------------------------------- batch runs
177class BatchStore:
178 """The batch runner's durable side. Method names are the runner's contract."""
180 def save_run(self, job_id: str, doc_id: int, want: str, questions: List[str]) -> None:
181 with db.connect() as conn:
182 conn.execute(
183 """INSERT INTO ade_playground.batch_runs (job_id, doc_id, want, questions, owner, heartbeat)
184 VALUES (%s, %s, %s, %s, %s, now())
185 ON CONFLICT (job_id) DO NOTHING""",
186 (job_id, doc_id, want, json.dumps(questions), OWNER),
187 )
189 def save_row(
190 self, job_id: str, index: int, question: str, error: Optional[str], response: Optional[Dict[str, Any]]
191 ) -> None:
192 with db.connect() as conn:
193 conn.execute(
194 """INSERT INTO ade_playground.batch_rows (job_id, index, question, error, response)
195 VALUES (%s, %s, %s, %s, %s)
196 ON CONFLICT (job_id, index) DO UPDATE
197 SET error = EXCLUDED.error, response = EXCLUDED.response, finished_at = now()""",
198 (job_id, index, question, error, json.dumps(response) if response is not None else None),
199 )
200 conn.execute(
201 "UPDATE ade_playground.batch_runs SET failed = failed + %s, heartbeat = now(), updated_at = now() WHERE job_id = %s",
202 (1 if error else 0, job_id),
203 )
205 def finish_run(self, job_id: str) -> None:
206 with db.connect() as conn:
207 conn.execute(
208 "UPDATE ade_playground.batch_runs SET running = false, updated_at = now() WHERE job_id = %s",
209 (job_id,),
210 )
212 def mark_stale_runs(self) -> None:
213 """A run whose task stopped: every unanswered question becomes an error
214 row saying so, and the run is closed, so its counts add up."""
215 with db.connect() as conn:
216 stale = conn.execute(
217 """SELECT job_id, questions FROM ade_playground.batch_runs
218 WHERE running AND (heartbeat IS NULL OR heartbeat < now() - make_interval(secs => %s))""",
219 (STALE_SECONDS,),
220 ).fetchall()
221 for job_id, questions in stale:
222 questions = json.loads(questions) if isinstance(questions, str) else questions
223 done = {
224 r[0]
225 for r in conn.execute(
226 "SELECT index FROM ade_playground.batch_rows WHERE job_id = %s", (job_id,)
227 ).fetchall()
228 }
229 missing = [i for i in range(len(questions)) if i not in done]
230 for index in missing:
231 conn.execute(
232 """INSERT INTO ade_playground.batch_rows (job_id, index, question, error)
233 VALUES (%s, %s, %s, %s) ON CONFLICT DO NOTHING""",
234 (job_id, index, questions[index], STOPPED_REASON),
235 )
236 conn.execute(
237 "UPDATE ade_playground.batch_runs SET running = false, failed = failed + %s, updated_at = now() WHERE job_id = %s",
238 (len(missing), job_id),
239 )
240 logger.warning(
241 "batch run {} closed: {} ({} unanswered)", job_id, STOPPED_REASON, len(missing)
242 )
244 def load_run(self, job_id: str) -> Optional[Dict[str, Any]]:
245 """The run and its rows as plain data, or None."""
246 self.mark_stale_runs()
247 with db.connect() as conn:
248 run = conn.execute(
249 "SELECT doc_id, want, questions, failed, running FROM ade_playground.batch_runs WHERE job_id = %s",
250 (job_id,),
251 ).fetchone()
252 if run is None:
253 return None
254 rows = conn.execute(
255 "SELECT index, question, error, response FROM ade_playground.batch_rows WHERE job_id = %s ORDER BY index",
256 (job_id,),
257 ).fetchall()
258 questions = json.loads(run[2]) if isinstance(run[2], str) else run[2]
259 return {
260 "job_id": job_id,
261 "doc_id": run[0],
262 "want": run[1],
263 "questions": questions,
264 "failed": run[3],
265 "running": run[4],
266 "rows": [
267 {
268 "index": r[0],
269 "question": r[1],
270 "error": r[2],
271 "response": json.loads(r[3]) if isinstance(r[3], str) else r[3],
272 }
273 for r in rows
274 ],
275 }
277 def running_for(self, doc_id: int) -> bool:
278 self.mark_stale_runs()
279 with db.connect() as conn:
280 row = conn.execute(
281 "SELECT 1 FROM ade_playground.batch_runs WHERE doc_id = %s AND running LIMIT 1", (doc_id,)
282 ).fetchone()
283 return row is not None
285 def heartbeat(self, job_ids: Iterable[str]) -> None:
286 ids = list(job_ids)
287 if not ids:
288 return
289 with db.connect() as conn:
290 conn.execute(
291 "UPDATE ade_playground.batch_runs SET heartbeat = now() WHERE job_id = ANY(%s)", (ids,)
292 )
295# ---------------------------------------------------------------- heartbeat
298def start_heartbeat(
299 upload_ids: Callable[[], Iterable[str]],
300 run_ids: Callable[[], Iterable[str]],
301 store: BatchStore,
302) -> threading.Thread:
303 """Beat, every ``HEARTBEAT_SECONDS``, for every job this process is still
304 running. The callables name them at each beat, so a job that finished
305 between beats is simply not beaten again."""
307 def beat() -> None:
308 while True:
309 time.sleep(HEARTBEAT_SECONDS)
310 try:
311 heartbeat_uploads(upload_ids())
312 store.heartbeat(run_ids())
313 except Exception as exc:
314 logger.warning("heartbeat failed: {}", exc)
316 thread = threading.Thread(target=beat, name="job-heartbeat", daemon=True)
317 thread.start()
318 return thread