79 lines
3.0 KiB
Python
79 lines
3.0 KiB
Python
"""Persistent extraction progress and per-batch worker ownership."""
|
|
import json
|
|
import time
|
|
import uuid
|
|
|
|
from db.connection import execute, query
|
|
|
|
LEASE_SECONDS = 300
|
|
|
|
|
|
def claim(batch_id, total):
|
|
"""Claim a batch atomically, allowing retry after failure or expired lease."""
|
|
owner = str(uuid.uuid4())
|
|
now = time.time()
|
|
changed = execute(
|
|
"""INSERT INTO extraction_jobs (batch_id, owner, state, total, updated_at)
|
|
VALUES (%s, %s, 'running', %s, %s)
|
|
ON CONFLICT(batch_id) DO UPDATE SET owner=excluded.owner, state='running',
|
|
total=excluded.total, done=0, failed=0, classes_json='{}',
|
|
result_json=NULL, error=NULL, updated_at=excluded.updated_at
|
|
WHERE extraction_jobs.state != 'running' OR extraction_jobs.updated_at < %s""",
|
|
(batch_id, owner, total, now, now - LEASE_SECONDS),
|
|
)
|
|
return owner if changed else None
|
|
|
|
|
|
def heartbeat(batch_id, owner):
|
|
return execute(
|
|
"UPDATE extraction_jobs SET updated_at=%s WHERE batch_id=%s AND owner=%s AND state='running'",
|
|
(time.time(), batch_id, owner),
|
|
)
|
|
|
|
|
|
def progress(batch_id, owner, completed, result):
|
|
rows = query(
|
|
"SELECT failed, classes_json FROM extraction_jobs WHERE batch_id=%s AND owner=%s AND state='running'",
|
|
(batch_id, owner),
|
|
)
|
|
if not rows:
|
|
return
|
|
failed = rows[0]["failed"]
|
|
classes = json.loads(rows[0]["classes_json"])
|
|
if not result.get("ok"):
|
|
failed += 1
|
|
else:
|
|
doc_class = (result.get("extraction") or {}).get("doc_class", "unknown")
|
|
classes[doc_class] = classes.get(doc_class, 0) + 1
|
|
execute(
|
|
"""UPDATE extraction_jobs SET done=%s, failed=%s, classes_json=%s, updated_at=%s
|
|
WHERE batch_id=%s AND owner=%s AND state='running'""",
|
|
(completed, failed, json.dumps(classes), time.time(), batch_id, owner),
|
|
)
|
|
|
|
|
|
def finish(batch_id, owner, result=None, error=None):
|
|
state = "done" if result and result.get("ok") and not error else "failed"
|
|
return execute(
|
|
"""UPDATE extraction_jobs SET state=%s, result_json=%s, error=%s, updated_at=%s
|
|
WHERE batch_id=%s AND owner=%s AND state='running'""",
|
|
(state, json.dumps(result, ensure_ascii=False) if result is not None else None,
|
|
error or ("extraction failed" if state == "failed" else None), time.time(), batch_id, owner),
|
|
)
|
|
|
|
|
|
def get(batch_id):
|
|
"""Mark an abandoned lease interrupted; callers may explicitly retry it."""
|
|
execute(
|
|
"""UPDATE extraction_jobs SET state='interrupted', error='worker stopped; retry the batch'
|
|
WHERE batch_id=%s AND state='running' AND updated_at < %s""",
|
|
(batch_id, time.time() - LEASE_SECONDS),
|
|
)
|
|
rows = query("SELECT * FROM extraction_jobs WHERE batch_id=%s", (batch_id,))
|
|
if not rows:
|
|
return None
|
|
job = rows[0]
|
|
job["classes"] = json.loads(job.pop("classes_json"))
|
|
result_json = job.pop("result_json")
|
|
job["result"] = json.loads(result_json) if result_json else None
|
|
return job |