| 39 |
39 |
return row["id"] if row else None |
| 40 |
40 |
|
| 41 |
41 |
|
| 42 |
|
−async def claim_next(conn: AsyncConnection, worker: str, kinds: list[str] | None = None) -> dict[str, Any] | None: |
|
42 |
+async def claim_next(conn: AsyncConnection, worker: str, kinds: list[str] | None = None, *, sticky_kind: str | None = None) -> dict[str, Any] | None: |
|
43 |
+ """Claim the next job. `sticky_kind` keeps a worker on the kind it just ran while any remain: LLM jobs of one kind share a loaded |
|
44 |
+ model, and alternating kinds (extraction ↔ embeddings) makes the inference server swap models every job.""" |
| 43 |
45 |
kind_filter = "and kind = any(cast(:kinds as text[]))" if kinds else "" |
| 44 |
46 |
row = await fetch_one(conn, f"""with next as ( |
| 45 |
47 |
select id from jobs where status = 'queued' and run_after <= now() {kind_filter} |
| 46 |
|
− order by priority, run_after limit 1 for update skip locked) |
|
48 |
+ order by case when kind = cast(:sticky as text) then 0 else 1 end, priority, run_after limit 1 for update skip locked) |
| 47 |
49 |
update jobs j set status = 'running', locked_by = :w, locked_at = now(), started_at = now(), attempts = attempts + 1 |
| 48 |
|
− from next where j.id = next.id returning j.*""", w=worker, kinds=kinds or []) |
|
50 |
+ from next where j.id = next.id returning j.*""", w=worker, kinds=kinds or [], sticky=sticky_kind or "") |
| 49 |
51 |
return row |
| 50 |
52 |
|
| 51 |
53 |
|
| 98 |
100 |
sem.release() |
| 99 |
101 |
|
| 100 |
102 |
tasks: set[asyncio.Task[None]] = set() |
|
103 |
+ last_kind: str | None = None |
| 101 |
104 |
while not stop.is_set(): |
| 102 |
105 |
await sem.acquire() # hold a slot BEFORE claiming, so concurrency bounds claims too |
| 103 |
106 |
if stop.is_set(): |
| 104 |
107 |
sem.release() |
| 105 |
108 |
break |
| 106 |
109 |
async with transaction() as conn: |
| 107 |
|
− job = await claim_next(conn, worker, kinds) |
|
110 |
+ job = await claim_next(conn, worker, kinds, sticky_kind=last_kind) |
|
111 |
+ if job is not None: |
|
112 |
+ last_kind = job["kind"] |
| 108 |
113 |
if job is None: |
| 109 |
114 |
sem.release() |
| 110 |
115 |
try: |