added new rq-gpu-worker (that uses gpu queue, old gpu queue renamed to api)

bigabig · bigabig · commit 2df90b30dbd5 · 2025-09-05T13:38:14.000Z
diff --git a/backend/.vscode/launch.json b/backend/.vscode/launch.json
@@ -27,6 +27,7 @@
       "env": {
         "PYTHONPATH": "${workspaceFolder}/src",
         "RQ_WORKERS_CPU": "1",
+        "RQ_WORKERS_API": "1",
         "RQ_WORKERS_GPU": "1",
         "CUDA_VISIBLE_DEVICES": "1"
       }
diff --git a/backend/src/modules/concept_over_time_analysis/cota_refinement_job.py b/backend/src/modules/concept_over_time_analysis/cota_refinement_job.py
@@ -19,7 +19,7 @@
     job_type=JobType.COTA_REFINEMENT,
     input_type=COTARefinementJobInput,
     output_type=None,
-    device="gpu",
+    device="api",
     result_ttl=JobResultTTL.NINETY_DAYS,
 )
 def cota_refinement(payload: COTARefinementJobInput, job: Job) -> None:
diff --git a/backend/src/modules/doc_processing/audio/transcription_job.py b/backend/src/modules/doc_processing/audio/transcription_job.py
@@ -41,7 +41,7 @@ class TranscriptionJobOutput(JobOutputBase):
     job_type=JobType.AUDIO_TRANSCRIPTION,
     input_type=TranscriptionJobInput,
     output_type=TranscriptionJobOutput,
-    device="gpu",
+    device="api",
 )
 def handle_transcription_job(
     payload: TranscriptionJobInput, job: Job
diff --git a/backend/src/modules/doc_processing/html/sentence_embedding_job.py b/backend/src/modules/doc_processing/html/sentence_embedding_job.py
@@ -22,7 +22,7 @@ class TextSentenceEmbeddingJobInput(SdocProcessingJobInput):
 @register_job(
     job_type=JobType.TEXT_SENTENCE_EMBEDDING,
     input_type=TextSentenceEmbeddingJobInput,
-    device="gpu",
+    device="api",
 )
 def handle_text_sentence_embedding_job(
     payload: TextSentenceEmbeddingJobInput, job: Job
diff --git a/backend/src/modules/doc_processing/image/image_caption_job.py b/backend/src/modules/doc_processing/image/image_caption_job.py
@@ -30,7 +30,7 @@ class ImageCaptionJobOutput(JobOutputBase):
     job_type=JobType.IMAGE_CAPTION,
     input_type=ImageCaptionJobInput,
     output_type=ImageCaptionJobOutput,
-    device="gpu",
+    device="api",
 )
 def handle_image_caption_job(
     payload: ImageCaptionJobInput, job: Job
diff --git a/backend/src/modules/doc_processing/image/image_embedding_job.py b/backend/src/modules/doc_processing/image/image_embedding_job.py
@@ -20,7 +20,7 @@ class ImageEmbeddingJobInput(SdocProcessingJobInput):
 
 
 @register_job(
-    job_type=JobType.IMAGE_EMBEDDING, input_type=ImageEmbeddingJobInput, device="gpu"
+    job_type=JobType.IMAGE_EMBEDDING, input_type=ImageEmbeddingJobInput, device="api"
 )
 def handle_image_embedding_job(payload: ImageEmbeddingJobInput, job: Job) -> None:
     # embed the image
diff --git a/backend/src/modules/doc_processing/image/object_detection_job.py b/backend/src/modules/doc_processing/image/object_detection_job.py
@@ -28,7 +28,7 @@ class ObjectDetectionJobInput(SdocProcessingJobInput):
 @register_job(
     job_type=JobType.IMAGE_OBJECT_DETECTION,
     input_type=ObjectDetectionJobInput,
-    device="gpu",
+    device="api",
 )
 def handle_object_detection_job(payload: ObjectDetectionJobInput, job: Job) -> None:
     # Run object detection with ray
diff --git a/backend/src/modules/llm_assistant/llm_job.py b/backend/src/modules/llm_assistant/llm_job.py
@@ -14,7 +14,7 @@
     output_type=LLMJobOutput,
     generate_endpoints=EndpointGeneration.ALL,
     router=router,
-    device="gpu",
+    device="api",
     result_ttl=JobResultTTL.NINETY_DAYS,
     timeout=-1,  # infinite/no timeout
 )
diff --git a/backend/src/modules/ml/ml_job.py b/backend/src/modules/ml/ml_job.py
@@ -27,7 +27,7 @@
     job_type=JobType.ML,
     input_type=MLJobInput,
     generate_endpoints=EndpointGeneration.ALL,
-    device="gpu",
+    device="api",
     result_ttl=JobResultTTL.NINETY_DAYS,
 )
 def ml_job(payload: MLJobInput, job: Job) -> None:
diff --git a/backend/src/modules/perspectives/perspectives_job.py b/backend/src/modules/perspectives/perspectives_job.py
@@ -10,7 +10,7 @@
 @register_job(
     job_type=JobType.PERSPECTIVES,
     input_type=PerspectivesJobInput,
-    device="gpu",
+    device="api",
     result_ttl=JobResultTTL.NINETY_DAYS,
 )
 def perspectives_job(payload: PerspectivesJobInput, job: Job) -> None:
diff --git a/backend/src/rq_worker_entrypoint.sh b/backend/src/rq_worker_entrypoint.sh
@@ -3,4 +3,5 @@
 set -e
 source .venv/bin/activate
 
-python src/worker.py work
+DEVICE=${1:-cpu}
+python src/worker.py work "$DEVICE"
diff --git a/backend/src/systems/job_system/job_register_decorator.py b/backend/src/systems/job_system/job_register_decorator.py
@@ -21,7 +21,7 @@ def register_job(
     input_type: type[InputT],
     output_type: type[OutputT] | None = None,
     priority: JobPriority = JobPriority.DEFAULT,
-    device: Literal["gpu", "cpu"] = "cpu",
+    device: Literal["gpu", "cpu", "api"] = "cpu",
     generate_endpoints: EndpointGeneration = EndpointGeneration.NONE,
     router: APIRouter | None = None,
     result_ttl: JobResultTTL = JobResultTTL.DEFAULT,
diff --git a/backend/src/systems/job_system/job_service.py b/backend/src/systems/job_system/job_service.py
@@ -37,7 +37,7 @@ class RegisteredJob(TypedDict):
     output_type: type[JobOutputBase] | None
     generate_endpoints: EndpointGeneration
     priority: JobPriority
-    device: Literal["gpu", "cpu"]
+    device: Literal["gpu", "cpu", "api"]
     router: APIRouter | None
     result_ttl: JobResultTTL  # how long to keep successful jobs and their results (defaults to 500 seconds)
     # failure_ttl  # how long to keep failed jobs (defaults to 1 year)
@@ -75,7 +75,7 @@ def __new__(cls, *args, **kwargs):
         # Define priority queues and their registries (every queue has its own 5 registries)
         cls.queues: Dict[tuple[str, JobPriority], rq.Queue] = {}
         cls.registries: dict[tuple[str, JobPriority], dict[str, BaseRegistry]] = {}
-        for device in ["cpu", "gpu"]:
+        for device in ["cpu", "gpu", "api"]:
             for priority in [JobPriority.HIGH, JobPriority.DEFAULT, JobPriority.LOW]:
                 qk = (device, priority)  # queue key
                 qn = f"{device}-{priority.value}"  # queue name
@@ -98,7 +98,7 @@ def register_job(
         input_type: type[InputT],
         output_type: type[OutputT] | None,
         priority: JobPriority,
-        device: Literal["gpu", "cpu"],
+        device: Literal["gpu", "cpu", "api"],
         generate_endpoints: EndpointGeneration,
         router: APIRouter | None,
         result_ttl: JobResultTTL,
diff --git a/backend/src/worker.py b/backend/src/worker.py
@@ -43,7 +43,7 @@ def do_healthcheck():
         sys.exit(1)
 
 
-def do_work():
+def do_work(device: str):
     # import all expensive stuff before forking, so that imports are only done once
     import_by_suffix("_repo.py")
     import_by_suffix("_service.py")
@@ -56,16 +56,26 @@ def do_work():
 
     ctx = mp.get_context("fork")
 
-    cpu = ctx.Process(
-        target=create_pool, args=("cpu", int(environ.get("RQ_WORKERS_CPU", "8")))
-    )
-    gpu = ctx.Process(
-        target=create_pool, args=("gpu", int(environ.get("RQ_WORKERS_GPU", "16")))
-    )
-    cpu.start()
-    gpu.start()
-    cpu.join()
-    gpu.join()
+    if device == "cpu":
+        cpu = ctx.Process(
+            target=create_pool, args=("cpu", int(environ.get("RQ_WORKERS_CPU", "8")))
+        )
+        api = ctx.Process(
+            target=create_pool, args=("api", int(environ.get("RQ_WORKERS_API", "16")))
+        )
+        cpu.start()
+        api.start()
+        cpu.join()
+        api.join()
+    elif device == "gpu":
+        gpu = ctx.Process(
+            target=create_pool, args=("gpu", int(environ.get("RQ_WORKERS_GPU", "1")))
+        )
+        gpu.start()
+        gpu.join()
+    else:
+        print("Usage: worker.py healthcheck or worker.py work [cpu|gpu]")
+        sys.exit(1)
 
 
 def create_pool(queue_name: str, num_workers: int):
@@ -90,11 +100,15 @@ def create_pool(queue_name: str, num_workers: int):
     if len(sys.argv) == 2:
         if sys.argv[1] == "healthcheck":
             do_healthcheck()
-        elif sys.argv[1] == "work":
-            do_work()
         else:
-            print("Usage: worker.py [healthcheck|work]")
+            print("Usage: worker.py healthcheck or worker.py work [cpu|gpu]")
+            sys.exit(1)
+    elif len(sys.argv) == 3:
+        if sys.argv[1] == "work":
+            do_work(sys.argv[2])
+        else:
+            print("Usage: worker.py healthcheck or worker.py work [cpu|gpu]")
             sys.exit(1)
     else:
-        print("Usage: worker.py [healthcheck|work]")
+        print("Usage: worker.py healthcheck or worker.py work [cpu|gpu]")
         sys.exit(1)
diff --git a/bin/setup-folders.sh b/bin/setup-folders.sh
@@ -15,7 +15,8 @@ mkdir -p ../docker/backend_repo
 
 # cache directories
 mkdir -p ../docker/api_cache
-mkdir -p ../docker/rq_cache
+mkdir -p ../docker/rq_cpu_cache
+mkdir -p ../docker/rq_gpu_cache
 mkdir -p ../docker/ray_cache
 chmod 777 ../docker/ray_cache
 mkdir -p ../docker/vllm_cache
diff --git a/docker/.env.example b/docker/.env.example
@@ -98,7 +98,9 @@ POSTGRES_MAX_CONNECTIONS=200
 # LOG LEVEL ('critical', 'error', 'warning', 'info', 'debug', 'trace')
 RQ_LOG_LEVEL=info
 RQ_WORKERS_CPU=8
-RQ_WORKERS_GPU=16
+RQ_WORKERS_API=16
+RQ_WORKERS_GPU=1
+RQ_DEVICE_IDS=1
 
 
 ### RQ Dashboard settings
diff --git a/docker/.gitignore b/docker/.gitignore
@@ -3,7 +3,8 @@
 /api_cache/*
 /backend_repo/*
 /ray_cache/*
-/rq_cache/*
+/rq_cpu_cache/*
+/rq_gpu_cache/*
 /vllm_cache/*
 /vllm_config/*
 
diff --git a/docker/compose.override.yml b/docker/compose.override.yml
@@ -58,7 +58,11 @@ services:
     networks:
       - dats_network
 
-  rq-worker:
+  rq-cpu-worker:
+    volumes:
+      - ../backend/src:/dats_code/src
+
+  rq-gpu-worker:
     volumes:
       - ../backend/src:/dats_code/src
 
diff --git a/docker/compose.yml b/docker/compose.yml
@@ -34,6 +34,11 @@ x-dats-common: &dats-common
     IS_STABLE: ${IS_STABLE?is-stable-required!}
     NUMBA_CACHE_DIR: /cache/numba_cache
     MPLCONFIGDIR: /cache/matplotlib
+    HUGGINGFACE_HUB_CACHE: /cache/hf_hub_cache
+    HF_HOME: /cache/hf_cache
+    TORCH_HOME: /cache/torch_cache
+    TORCHINDUCTOR_CACHE_DIR: /cache/torch_inductor_cache
+    TRITON_CACHE_DIR: /cache/triton_cache
   depends_on: &dats-common-depends-on
     postgres:
       condition: service_healthy
@@ -144,22 +149,47 @@ services:
     networks:
       - dats_network
 
-  rq-worker:
+  rq-cpu-worker:
     <<: *dats-common
     healthcheck:
       test: "/dats_code/.venv/bin/python /dats_code/src/worker.py healthcheck"
       interval: 30s
       timeout: 10s
       retries: 5
-    command: /dats_code/src/rq_worker_entrypoint.sh
+    command: /dats_code/src/rq_worker_entrypoint.sh cpu
     environment:
       <<: *dats-common-env
       RQ_WORKERS_CPU: ${RQ_WORKERS_CPU?rq-workers-cpu-required!}
+      RQ_WORKERS_API: ${RQ_WORKERS_API?rq-workers-api-required!}
+      LOG_LEVEL: ${RQ_LOG_LEVEL?rq-log-level-required!}
+    volumes:
+      - ./backend_repo:/tmp/dats
+      - ./rq_cpu_cache:/cache
+    profiles:
+      - rq
+
+  rq-gpu-worker:
+    <<: *dats-common
+    healthcheck:
+      test: "/dats_code/.venv/bin/python /dats_code/src/worker.py healthcheck"
+      interval: 30s
+      timeout: 10s
+      retries: 5
+    command: /dats_code/src/rq_worker_entrypoint.sh gpu
+    environment:
+      <<: *dats-common-env
       RQ_WORKERS_GPU: ${RQ_WORKERS_GPU?rq-workers-gpu-required!}
       LOG_LEVEL: ${RQ_LOG_LEVEL?rq-log-level-required!}
     volumes:
       - ./backend_repo:/tmp/dats
-      - ./rq_cache:/cache
+      - ./rq_gpu_cache:/cache
+    deploy:
+      resources:
+        reservations:
+          devices:
+            - driver: nvidia
+              device_ids: ["${RQ_DEVICE_IDS?rq-device-ids-required!}"]
+              capabilities: [gpu]
     profiles:
       - rq
 
@@ -170,7 +200,10 @@ services:
       RQ_DASHBOARD_USERNAME: ${RQ_DASHBOARD_USERNAME?rq-dashboard-username-required!}
       RQ_DASHBOARD_PASSWORD: ${RQ_DASHBOARD_PASSWORD?rq-dashboard-password-required!}
     depends_on:
-      rq-worker:
+      rq-cpu-worker:
+        condition: service_healthy
+        restart: true
+      rq-gpu-worker:
         condition: service_healthy
         restart: true
     ports:
@@ -198,7 +231,10 @@ services:
       - ./api_cache:/cache
     depends_on:
       <<: *dats-common-depends-on
-      rq-worker:
+      rq-cpu-worker:
+        condition: service_healthy
+        restart: true
+      rq-gpu-worker:
         condition: service_healthy
         restart: true
     ports:

Original file line number	Diff line number	Diff line change
`@@ -27,6 +27,7 @@`
`27`	`27`	`"env": {`
`28`	`28`	`"PYTHONPATH": "${workspaceFolder}/src",`
`29`	`29`	`"RQ_WORKERS_CPU": "1",`
	`30`	`+ "RQ_WORKERS_API": "1",`
`30`	`31`	`"RQ_WORKERS_GPU": "1",`
`31`	`32`	`"CUDA_VISIBLE_DEVICES": "1"`
`32`	`33`	`}`
Original file line number	Diff line number	Diff line change
`@@ -19,7 +19,7 @@`
`19`	`19`	`job_type=JobType.COTA_REFINEMENT,`
`20`	`20`	`input_type=COTARefinementJobInput,`
`21`	`21`	`output_type=None,`
`22`		`- device="gpu",`
	`22`	`+ device="api",`
`23`	`23`	`result_ttl=JobResultTTL.NINETY_DAYS,`
`24`	`24`	`)`
`25`	`25`	`def cota_refinement(payload: COTARefinementJobInput, job: Job) -> None:`
Original file line number	Diff line number	Diff line change
`@@ -41,7 +41,7 @@ class TranscriptionJobOutput(JobOutputBase):`
`41`	`41`	`job_type=JobType.AUDIO_TRANSCRIPTION,`
`42`	`42`	`input_type=TranscriptionJobInput,`
`43`	`43`	`output_type=TranscriptionJobOutput,`
`44`		`- device="gpu",`
	`44`	`+ device="api",`
`45`	`45`	`)`
`46`	`46`	`def handle_transcription_job(`
`47`	`47`	`payload: TranscriptionJobInput, job: Job`
Original file line number	Diff line number	Diff line change
`@@ -22,7 +22,7 @@ class TextSentenceEmbeddingJobInput(SdocProcessingJobInput):`
`22`	`22`	`@register_job(`
`23`	`23`	`job_type=JobType.TEXT_SENTENCE_EMBEDDING,`
`24`	`24`	`input_type=TextSentenceEmbeddingJobInput,`
`25`		`- device="gpu",`
	`25`	`+ device="api",`
`26`	`26`	`)`
`27`	`27`	`def handle_text_sentence_embedding_job(`
`28`	`28`	`payload: TextSentenceEmbeddingJobInput, job: Job`
Original file line number	Diff line number	Diff line change
`@@ -30,7 +30,7 @@ class ImageCaptionJobOutput(JobOutputBase):`
`30`	`30`	`job_type=JobType.IMAGE_CAPTION,`
`31`	`31`	`input_type=ImageCaptionJobInput,`
`32`	`32`	`output_type=ImageCaptionJobOutput,`
`33`		`- device="gpu",`
	`33`	`+ device="api",`
`34`	`34`	`)`
`35`	`35`	`def handle_image_caption_job(`
`36`	`36`	`payload: ImageCaptionJobInput, job: Job`
Original file line number	Diff line number	Diff line change
`@@ -20,7 +20,7 @@ class ImageEmbeddingJobInput(SdocProcessingJobInput):`
`20`	`20`
`21`	`21`
`22`	`22`	`@register_job(`
`23`		`- job_type=JobType.IMAGE_EMBEDDING, input_type=ImageEmbeddingJobInput, device="gpu"`
	`23`	`+ job_type=JobType.IMAGE_EMBEDDING, input_type=ImageEmbeddingJobInput, device="api"`
`24`	`24`	`)`
`25`	`25`	`def handle_image_embedding_job(payload: ImageEmbeddingJobInput, job: Job) -> None:`
`26`	`26`	`# embed the image`
Original file line number	Diff line number	Diff line change
`@@ -28,7 +28,7 @@ class ObjectDetectionJobInput(SdocProcessingJobInput):`
`28`	`28`	`@register_job(`
`29`	`29`	`job_type=JobType.IMAGE_OBJECT_DETECTION,`
`30`	`30`	`input_type=ObjectDetectionJobInput,`
`31`		`- device="gpu",`
	`31`	`+ device="api",`
`32`	`32`	`)`
`33`	`33`	`def handle_object_detection_job(payload: ObjectDetectionJobInput, job: Job) -> None:`
`34`	`34`	`# Run object detection with ray`
Original file line number	Diff line number	Diff line change
`@@ -14,7 +14,7 @@`
`14`	`14`	`output_type=LLMJobOutput,`
`15`	`15`	`generate_endpoints=EndpointGeneration.ALL,`
`16`	`16`	`router=router,`
`17`		`- device="gpu",`
	`17`	`+ device="api",`
`18`	`18`	`result_ttl=JobResultTTL.NINETY_DAYS,`
`19`	`19`	`timeout=-1, # infinite/no timeout`
`20`	`20`	`)`
Original file line number	Diff line number	Diff line change
`@@ -27,7 +27,7 @@`
`27`	`27`	`job_type=JobType.ML,`
`28`	`28`	`input_type=MLJobInput,`
`29`	`29`	`generate_endpoints=EndpointGeneration.ALL,`
`30`		`- device="gpu",`
	`30`	`+ device="api",`
`31`	`31`	`result_ttl=JobResultTTL.NINETY_DAYS,`
`32`	`32`	`)`
`33`	`33`	`def ml_job(payload: MLJobInput, job: Job) -> None:`
Original file line number	Diff line number	Diff line change
`@@ -10,7 +10,7 @@`
`10`	`10`	`@register_job(`
`11`	`11`	`job_type=JobType.PERSPECTIVES,`
`12`	`12`	`input_type=PerspectivesJobInput,`
`13`		`- device="gpu",`
	`13`	`+ device="api",`
`14`	`14`	`result_ttl=JobResultTTL.NINETY_DAYS,`
`15`	`15`	`)`
`16`	`16`	`def perspectives_job(payload: PerspectivesJobInput, job: Job) -> None:`