From f5d0862b2ea3bc5ca6a8d63f24cdcf32152015da Mon Sep 17 00:00:00 2001 From: Ran Lu Date: Fri, 14 Jun 2024 20:21:03 -0400 Subject: [PATCH] Scale gpu workers matching the number of chunkflow tasks --- dags/cluster_dag.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/dags/cluster_dag.py b/dags/cluster_dag.py index 287b6285..1152a894 100644 --- 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