feat: add trigger_rule kwarg to scale ops

dags/helper_ops.py

@@ -131,29 +131,29 @@ def wait_op(dag, process):
     )
 
 
-def scale_up_cluster_op(dag, stage, key, initial_size, total_size, queue):
+def scale_up_cluster_op(dag, stage, key, initial_size, total_size, queue, trigger_rule="one_success"):
     return PythonOperator(
         task_id='resize_{}_{}'.format(stage, total_size),
         python_callable=ramp_up_cluster,
         op_args = [key, initial_size, total_size],
         default_args=default_args,
         weight_rule=WeightRule.ABSOLUTE,
         priority_weight=1000,
-        trigger_rule="one_success",
+        trigger_rule=trigger_rule,
         queue=queue,
         dag=dag
     )
 
 
-def scale_down_cluster_op(dag, stage, key, size, queue):
+def scale_down_cluster_op(dag, stage, key, size, queue, trigger_rule="all_success"):
     return PythonOperator(
         task_id='resize_{}_{}'.format(stage, size),
         python_callable=ramp_down_cluster,
         op_args = [key, size],
         default_args=default_args,
         weight_rule=WeightRule.ABSOLUTE,
         priority_weight=1000,
-        trigger_rule="all_done",
+        trigger_rule=trigger_rule,
         queue=queue,
         dag=dag
     )