Merge pull request #99 from seung-lab/split_synaptor_queue

Split synaptor queue

Split synaptor queue into multiple queues, one for each worker groups. This prevent wrong types of workers still other workers' tasks and causing the dag to go out of sync.

dags/cluster_dag.py

@@ -61,6 +61,8 @@ def estimate_optimal_number_of_workers(cluster, cluster_info):
     else:
         if cluster == "gpu":
             num_tasks = check_queue(queue="chunkflow")
+        elif cluster.startswith("synaptor"):
+            num_tasks = check_queue(queue=f"{cluster}-tasks")
         else:
             num_tasks = check_queue(queue=cluster)
         num_workers = estimate_worker_instances(num_tasks, cluster_info)

dags/synaptor_dags.py

@@ -6,7 +6,7 @@
 from airflow import DAG
 from airflow.models import Variable, BaseOperator
 
-from helper_ops import scale_up_cluster_op, scale_down_cluster_op, collect_metrics_op
+from helper_ops import placeholder_op, scale_up_cluster_op, scale_down_cluster_op, collect_metrics_op
 from param_default import synaptor_param_default, default_synaptor_image
 from synaptor_ops import manager_op, drain_op, self_destruct_op
 from synaptor_ops import synaptor_op, wait_op, generate_op, nglink_op
@@ -92,10 +92,10 @@ def __init__(self, name):
 
 def fill_dag(dag: DAG, tasklist: list[Task], collect_metrics: bool = True) -> DAG:
     """Fills a synaptor DAG from a list of Tasks."""
-    drain = drain_op(dag)
+    drain_tasks = [drain_op(dag, task_queue_name=f"synaptor-{t}-tasks") for t in ["cpu", "gpu", "seggraph"]]
     init_cloudvols = manager_op(dag, "init_cloudvols", image=SYNAPTOR_IMAGE)
 
-    drain >> init_cloudvols
+    start_nfs_server >> drain_tasks >> init_cloudvols
 
     curr_operator = init_cloudvols
     if WORKFLOW_PARAMS.get("workspacetype", "File") == "Database":
@@ -105,7 +105,7 @@ def fill_dag(dag: DAG, tasklist: list[Task], collect_metrics: bool = True) -> DA
 
     if collect_metrics:
         metrics = collect_metrics_op(dag)
-        metrics >> drain
+        metrics >> drain_tasks
 
     curr_cluster = ""
     for task in tasklist:
@@ -167,6 +167,7 @@ def change_cluster_if_required(
             i,
             image=SYNAPTOR_IMAGE,
             op_queue_name=next_task.cluster_key,
+            task_queue_name=f"{next_task.cluster_key}-tasks",
             tag=new_tag,
             use_gpus=next_task.use_gpus
         )
@@ -233,12 +234,17 @@ def add_task(
     elif task.cluster_key == "manager":
         generate = manager_op(dag, task.name, image=SYNAPTOR_IMAGE)
     else:
-        generate = generate_op(dag, task.name, image=SYNAPTOR_IMAGE, tag=tag)
+        generate = generate_op(dag, task.name, image=SYNAPTOR_IMAGE, tag=tag, task_queue_name=f"{task.cluster_key}-tasks")
 
     if tag:
-        wait = wait_op(dag, f"{task.name}_{tag}")
+        wait_task_name = f"{task.name}_{tag}"
     else:
-        wait = wait_op(dag, task.name)
+        wait_task_name = task.name
+
+    if task.cluster_key == "manager":
+        wait = placeholder_op(dag, wait_task_name)
+    else:
+        wait = wait_op(dag, wait_task_name, task_queue_name=f"{task.cluster_key}-tasks")
 
     prev_operator >> generate >> wait
 

dags/synaptor_ops.py

@@ -193,7 +193,7 @@ def drain_op(
     """Drains leftover messages from the RabbitMQ."""
 
     return PythonOperator(
-        task_id="drain_messages",
+        task_id=f"drain_messages_{task_queue_name}",
         python_callable=drain_messages,
         priority_weight=100_000,
         op_args=(airflow_broker_url, task_queue_name),
@@ -317,12 +317,12 @@ def synaptor_op(
     )
 
 
-def wait_op(dag: DAG, taskname: str) -> PythonOperator:
+def wait_op(dag: DAG, taskname: str, task_queue_name: Optional[str] = TASK_QUEUE_NAME) -> PythonOperator:
     """Waits for a task to finish."""
     return PythonOperator(
         task_id=f"wait_for_queue_{taskname}",
         python_callable=check_queue,
-        op_args=(TASK_QUEUE_NAME,),
+        op_args=(task_queue_name,),
         priority_weight=100_000,
         weight_rule=WeightRule.ABSOLUTE,
         on_success_callback=task_done_alert,

slackbot/cancel_run_commands.py

@@ -64,6 +64,9 @@ def cancel_run(msg):
     drain_messages(broker_url, "custom-gpu")
     drain_messages(broker_url, "chunkflow")
     drain_messages(broker_url, "synaptor")
+    drain_messages(broker_url, "synaptor-cpu-tasks")
+    drain_messages(broker_url, "synaptor-gpu-tasks")
+    drain_messages(broker_url, "synaptor-seggraph-tasks")
     drain_messages(broker_url, "deepem-gpu")
     time.sleep(10)