worker.py (2050 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 21 from __future__ import annotations 22 23 import logging 24 import sys 25 26 from backend.settings import settings 27 28 logging.basicConfig( 29 level=logging.INFO, 30 format="%(asctime)s %(levelname)-7s %(name)s: %(message)s", 31 ) 32 log = logging.getLogger("subplz.worker") 33 34 35 def main() -> int: 36 if settings.queue_backend != "redis": 37 log.error( 38 "queue_backend is %r. worker.py is only for the redis backend - " 39 "with the memory backend the API process runs jobs itself.", 40 settings.queue_backend, 41 ) 42 return 2 43 44 try: 45 from redis import Redis 46 from rq import Queue, Worker 47 except ImportError: 48 log.error("missing dependencies: pip install redis rq") 49 return 2 50 51 conn = Redis.from_url(settings.redis_url) 52 queue = Queue("subplz-jobs", connection=conn) 53 54 log.info( 55 "worker starting | redis=%s | device=%s model=%s | storage=%s", 56 settings.redis_url, settings.device, settings.model, 57 settings.storage_backend, 58 ) 59 # burst=False: stay alive and keep taking jobs. 60 Worker([queue], connection=conn).work(with_scheduler=False) 61 return 0 62 63 64 if __name__ == "__main__": 65 sys.exit(main())