| 39 | return row["id"] if row else None |
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 | kind_filter = "and kind = any(cast(:kinds as text[]))" if kinds else "" |
45 | kind_filter = "and kind = any(cast(:kinds as text[]))" if kinds else "" |
| 44 | row = await fetch_one(conn, f"""with next as ( |
46 | row = await fetch_one(conn, f"""with next as ( |
| 45 | select id from jobs where status = 'queued' and run_after <= now() {kind_filter} |
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 | update jobs j set status = 'running', locked_by = :w, locked_at = now(), started_at = now(), attempts = attempts + 1 |
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 | return row |
51 | return row |
| 50 | |
52 | |
| 51 | |
53 | |
| 98 | sem.release() |
100 | sem.release() |
| 99 | |
101 | |
| 100 | tasks: set[asyncio.Task[None]] = set() |
102 | tasks: set[asyncio.Task[None]] = set() |
| |
103 | + last_kind: str | None = None |
| 101 | while not stop.is_set(): |
104 | while not stop.is_set(): |
| 102 | await sem.acquire() # hold a slot BEFORE claiming, so concurrency bounds claims too |
105 | await sem.acquire() # hold a slot BEFORE claiming, so concurrency bounds claims too |
| 103 | if stop.is_set(): |
106 | if stop.is_set(): |
| 104 | sem.release() |
107 | sem.release() |
| 105 | break |
108 | break |
| 106 | async with transaction() as conn: |
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 | if job is None: |
113 | if job is None: |
| 109 | sem.release() |
114 | sem.release() |
| 110 | try: |
115 | try: |