136 lines
3.5 KiB
Python
136 lines
3.5 KiB
Python
"""A queued run must be claimed once, even if two workers look at it.
|
|
|
|
The worker selected queued jobs and then set them to running in a second
|
|
statement. Two workers — an API restart overlapping the previous process, or a
|
|
second replica — could both select the same row and both start tiled GPU
|
|
inference on it, producing duplicate analysis runs and doubling the GPU load.
|
|
|
|
The AOI worker beside it already claims with ``FOR UPDATE SKIP LOCKED``. This
|
|
uses a conditional update, which is the same guarantee expressed in one
|
|
statement: exactly one caller sees a row count of 1.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from uuid import uuid4
|
|
|
|
from app.models import Job
|
|
from app.services.analysis_job_worker import AnalysisJobWorker
|
|
|
|
|
|
class _Update:
|
|
"""Mimics a conditional UPDATE: the first caller wins, the rest see zero."""
|
|
|
|
def __init__(self, store: dict, job_id):
|
|
self.store = store
|
|
self.job_id = job_id
|
|
|
|
def update(self, values, **_kwargs) -> int:
|
|
if self.store.get(self.job_id) != "queued":
|
|
return 0
|
|
self.store[self.job_id] = "running"
|
|
return 1
|
|
|
|
|
|
class _Query:
|
|
def __init__(self, session, model):
|
|
self.session = session
|
|
self.model = model
|
|
self.job_id = None
|
|
|
|
def filter(self, *criteria):
|
|
for criterion in criteria:
|
|
right = getattr(criterion, "right", None)
|
|
value = getattr(right, "value", None)
|
|
if isinstance(value, type(uuid4())):
|
|
self.job_id = value
|
|
return self
|
|
|
|
def update(self, values, **kwargs) -> int:
|
|
return _Update(self.session.statuses, self.job_id).update(values, **kwargs)
|
|
|
|
def order_by(self, *_args):
|
|
return self
|
|
|
|
def limit(self, _count):
|
|
return self
|
|
|
|
def all(self):
|
|
return list(self.session.rows)
|
|
|
|
|
|
class _Session:
|
|
def __init__(self, rows: list[Job]):
|
|
self.rows = rows
|
|
self.statuses = {row.id: row.status for row in rows}
|
|
self.committed = 0
|
|
|
|
def query(self, model):
|
|
return _Query(self, model)
|
|
|
|
def get(self, _model, item_id):
|
|
return next((row for row in self.rows if row.id == item_id), None)
|
|
|
|
def add(self, _item):
|
|
return None
|
|
|
|
def commit(self):
|
|
self.committed += 1
|
|
|
|
def rollback(self):
|
|
return None
|
|
|
|
def close(self):
|
|
return None
|
|
|
|
|
|
def _job() -> Job:
|
|
return Job(
|
|
id=uuid4(),
|
|
job_type="detection.run",
|
|
status="queued",
|
|
project_id=uuid4(),
|
|
parameters_json={},
|
|
)
|
|
|
|
|
|
def test_the_first_claim_wins() -> None:
|
|
job = _job()
|
|
session = _Session([job])
|
|
|
|
assert AnalysisJobWorker.claim(session, job) is True
|
|
assert job.status == "running"
|
|
|
|
|
|
def test_a_second_claim_on_the_same_job_is_refused() -> None:
|
|
job = _job()
|
|
session = _Session([job])
|
|
|
|
assert AnalysisJobWorker.claim(session, job) is True
|
|
assert AnalysisJobWorker.claim(session, job) is False
|
|
|
|
|
|
def test_a_job_that_is_no_longer_queued_cannot_be_claimed() -> None:
|
|
job = _job()
|
|
session = _Session([job])
|
|
session.statuses[job.id] = "success"
|
|
|
|
assert AnalysisJobWorker.claim(session, job) is False
|
|
|
|
|
|
def test_an_unclaimable_job_is_skipped_rather_than_run(monkeypatch) -> None:
|
|
job = _job()
|
|
session = _Session([job])
|
|
session.statuses[job.id] = "running"
|
|
dispatched: list[Job] = []
|
|
monkeypatch.setattr(
|
|
AnalysisJobWorker,
|
|
"_dispatch",
|
|
staticmethod(lambda _db, item: dispatched.append(item)),
|
|
)
|
|
|
|
processed = AnalysisJobWorker.run_once(db=session)
|
|
|
|
assert processed == 0
|
|
assert dispatched == []
|