subplz-web/worker.py - Recently Written

worker.py (2879 bytes)
1 """Standalone job worker for the public deployment.
2 
3 On localhost the API process runs jobs itself and you never need this. When you
4 set SUBPLZ_WEB_QUEUE_BACKEND=redis, the API only enqueues, and these processes
5 do the work:
6 
7     # one per GPU box, or several per box if you have the RAM
8     SUBPLZ_WEB_QUEUE_BACKEND=redis \
9     SUBPLZ_WEB_REDIS_URL=redis://redis:6379/0 \
10     SUBPLZ_WEB_DATABASE_URL=postgresql+psycopg://... \
11     SUBPLZ_WEB_STORAGE_BACKEND=s3 SUBPLZ_WEB_S3_BUCKET=... \
12     SUBPLZ_WEB_DEVICE=cuda SUBPLZ_WEB_MODEL=tiny \
13     python worker.py
14 
15 Workers and the API must share the database and the storage bucket. They do not
16 need to share a filesystem: staged uploads are the one thing that is local to
17 whoever received them, which is why the API writes uploads to shared storage
18 before enqueuing when the queue backend is redis.
19 
20 SUBPLZ_WEB_WORKER_QUEUES says which jobs this worker takes, in order: "paid,free"
21 (the default: paid jobs first) or "paid" (a burst worker, see infra/burst).
22 """
23 
24 from __future__ import annotations
25 
26 import logging
27 import os
28 import sys
29 
30 from backend.settings import settings
31 
32 logging.basicConfig(
33     level=logging.INFO,
34     format="%(asctime)s %(levelname)-7s %(name)s: %(message)s",
35 )
36 log = logging.getLogger("subplz.worker")
37 
38 
39 def queue_names() -> list[str]:
40     """The Redis queues of this worker, the first one first."""
41     from backend.queue import FREE_QUEUE, PAID_QUEUE
42 
43     known = {"paid": PAID_QUEUE, "free": FREE_QUEUE}
44     wanted = [n.strip() for n in os.environ.get("SUBPLZ_WEB_WORKER_QUEUES", "paid,free").split(",")]
45     unknown = [n for n in wanted if n not in known]
46     if unknown or not wanted:
47         raise ValueError(f"SUBPLZ_WEB_WORKER_QUEUES: unknown queue {unknown}; use paid and free")
48     return [known[n] for n in wanted]
49 
50 
51 def main() -> int:
52     if settings.queue_backend != "redis":
53         log.error(
54             "queue_backend is %r. worker.py is only for the redis backend - "
55             "with the memory backend the API process runs jobs itself.",
56             settings.queue_backend,
57         )
58         return 2
59 
60     try:
61         from redis import Redis
62         from rq import Queue, Worker
63     except ImportError:
64         log.error("missing dependencies: pip install redis rq")
65         return 2
66 
67     conn = Redis.from_url(settings.redis_url)
68     queues = [Queue(name, connection=conn) for name in queue_names()]
69 
70     log.info(
71         "worker starting | queues=%s | device=%s model=%s | storage=%s",
72         ",".join(q.name for q in queues), settings.device, settings.model,
73         settings.storage_backend,
74     )
75     # burst=False: stay alive and keep taking jobs. rq takes a job from the
76     # first queue that has one, so the order of the queues is their priority.
77     Worker(queues, connection=conn).work(with_scheduler=False)
78     return 0
79 
80 
81 if __name__ == "__main__":
82     sys.exit(main())