Scale gpu workers matching the number of chunkflow tasks

seung-lab · Jun 15, 2024 · f5d0862 · f5d0862
1 parent 623e393
commit f5d0862
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/dags/cluster_dag.py b/dags/cluster_dag.py
@@ -59,7 +59,10 @@ def estimate_optimal_number_of_workers(cluster, cluster_info):
         tasks = [check_queue(f"{cluster}_{layer}") for layer in layers]
         num_workers = sum(tasks)
     else:
-        num_tasks = check_queue(cluster)
+        if cluster == "gpu":
+            num_tasks = check_queue(queue="chunkflow")
+        else:
+            num_tasks = check_queue(queue=cluster)
         num_workers = estimate_worker_instances(num_tasks, cluster_info)
 
     return num_workers