backend/queue.py (3728 bytes)
1 """Job dispatch. 2 3 Localhost runs jobs on a small thread pool inside the API process. The public 4 deployment sets SUBPLZ_WEB_QUEUE_BACKEND=redis and runs `worker.py` on separate 5 machines; the API side then only enqueues. Same `enqueue(job_id)` call either way. 6 7 A paid job and a free job go to different Redis queues. The web server's own 8 worker takes both, paid jobs first. A burst worker takes paid jobs only (see 9 infra/burst/README.md), so a free credit never starts a machine that costs money. 10 """ 11 12 from __future__ import annotations 13 14 import logging 15 import threading 16 from abc import ABC, abstractmethod 17 from concurrent.futures import ThreadPoolExecutor 18 19 from .settings import settings 20 21 log = logging.getLogger(__name__) 22 23 PAID_QUEUE = "subplz-jobs-paid" 24 FREE_QUEUE = "subplz-jobs" 25 26 27 class JobQueue(ABC): 28 @abstractmethod 29 def enqueue(self, job_id: str, paid: bool = True) -> None: 30 """Run `job_id`. `paid` is False for a job that a free credit pays for.""" 31 32 @abstractmethod 33 def depth(self) -> int: 34 """Jobs waiting or running. Drives the 'N ahead of you' hint in the UI.""" 35 36 def shutdown(self) -> None: 37 ... 38 39 40 class InProcessQueue(JobQueue): 41 """Thread pool in the API process. Fine for one person on one machine. 42 43 Alignment is CPU-bound and subplz is a subprocess, so the GIL is not the 44 limit here - max_concurrent_jobs is. 45 """ 46 47 def __init__(self, workers: int): 48 self._pool = ThreadPoolExecutor( 49 max_workers=workers, thread_name_prefix="subplz-job" 50 ) 51 self._lock = threading.Lock() 52 self._pending: set[str] = set() 53 54 def enqueue(self, job_id: str, paid: bool = True) -> None: 55 with self._lock: 56 if job_id in self._pending: 57 return 58 self._pending.add(job_id) 59 self._pool.submit(self._run, job_id) 60 61 def _run(self, job_id: str) -> None: 62 # Imported here to avoid a circular import at module load. 63 from .runner import run_job 64 65 try: 66 run_job(job_id) 67 except Exception: 68 log.exception("job %s crashed outside the runner", job_id) 69 finally: 70 with self._lock: 71 self._pending.discard(job_id) 72 73 def depth(self) -> int: 74 with self._lock: 75 return len(self._pending) 76 77 def shutdown(self) -> None: 78 self._pool.shutdown(wait=False, cancel_futures=True) 79 80 81 class RedisQueue(JobQueue): 82 """Public-release backend. Needs `rq` and a reachable Redis.""" 83 84 def __init__(self, url: str): 85 from redis import Redis # lazy: localhost needs neither package 86 from rq import Queue as RQQueue 87 88 self._conn = Redis.from_url(url) 89 self._paid = RQQueue(PAID_QUEUE, connection=self._conn) 90 self._free = RQQueue(FREE_QUEUE, connection=self._conn) 91 92 def enqueue(self, job_id: str, paid: bool = True) -> None: 93 from rq import Retry 94 95 (self._paid if paid else self._free).enqueue( 96 "backend.runner.run_job", 97 job_id, 98 job_timeout=settings.job_timeout_seconds, 99 result_ttl=86400, 100 # A worker can vanish during a job (a spot machine is taken back). 101 # rq then puts the job back on its queue once: two attempts at most. 102 # run_job never raises, so a failed book is not tried again. 103 retry=Retry(max=1), 104 ) 105 106 def depth(self) -> int: 107 return sum(q.count + q.started_job_registry.count for q in (self._paid, self._free)) 108 109 110 def build_queue() -> JobQueue: 111 if settings.queue_backend == "redis": 112 return RedisQueue(settings.redis_url) 113 return InProcessQueue(settings.max_concurrent_jobs) 114 115 116 queue: JobQueue = build_queue()