fix: Replace task chaining with chain for improved readability (GoogleCloudPlatform#1138)

This PR update the way to chain tasks in tpu_observability, from pylint “pointless-statement” into using from airflow.models.baseoperator import chain. Change 7 files for tpu_observability, including `interruption_validation_dag.py`, `jobset_ttr_rollback.py`, `multi_host_nodepool_rollback_dag.py`, `node_pool_status.py`, `node_pool_ttr_update_label.py`, `tpu_info_format_validation_dags.py` and `update_node_pool_label.py`.

Author: Chrisliao0806

dags/tpu_observability/interruption_validation_dag.py

@@ -8,6 +8,7 @@
 from airflow import models
 from airflow.decorators import task
 from airflow.exceptions import AirflowSkipException
+from airflow.models.baseoperator import chain
 from airflow.utils.task_group import TaskGroup
 
 from dags.common import test_owner
@@ -567,10 +568,10 @@ def fetch_interruption_metric_records_task(
                 log_records,
             )
 
-            (
-                proper_time_range
-                >> [metric_records, log_records]
-                >> check_event_count
+            chain(
+                proper_time_range,
+                [metric_records, log_records],
+                check_event_count,
             )
 
     return dag

dags/tpu_observability/jobset_ttr_rollback.py

@@ -17,6 +17,7 @@
 import datetime
 
 from airflow import models
+from airflow.models.baseoperator import chain
 from airflow.utils.trigger_rule import TriggerRule
 from airflow.utils.task_group import TaskGroup
 
@@ -141,16 +142,13 @@
           setups=create_node_pool,
       )
 
-      # Airflow uses >> for task chaining, which is pointless for pylint.
-      # pylint: disable=pointless-statement
-      (
-          cluster_info
-          >> create_node_pool
-          >> start_workload
-          >> ensure_all_pods_running
-          >> rollback_node_pool
-          >> wait_for_metric_upload
-          >> cleanup_workload
-          >> cleanup_node_pool
+      chain(
+          cluster_info,
+          create_node_pool,
+          start_workload,
+          ensure_all_pods_running,
+          rollback_node_pool,
+          wait_for_metric_upload,
+          cleanup_workload,
+          cleanup_node_pool,
       )
-      # pylint: enable=pointless-statement

dags/tpu_observability/multi_host_nodepool_rollback_dag.py

@@ -20,6 +20,7 @@
 import datetime
 
 from airflow import models
+from airflow.models.baseoperator import chain
 from airflow.utils.task_group import TaskGroup
 from airflow.utils.trigger_rule import TriggerRule
 
@@ -117,15 +118,12 @@
           setups=create_node_pool,
       )
 
-      # Airflow uses >> for task chaining, which is pointless for pylint.
-      # pylint: disable=pointless-statement
-      (
-          node_pool_info
-          >> create_node_pool
-          >> wait_node_pool_available
-          >> rollback_node_pool
-          >> wait_node_pool_unavailable
-          >> wait_node_pool_recovered
-          >> cleanup_node_pool
+      chain(
+          node_pool_info,
+          create_node_pool,
+          wait_node_pool_available,
+          rollback_node_pool,
+          wait_node_pool_unavailable,
+          wait_node_pool_recovered,
+          cleanup_node_pool,
       )
-      # pylint: enable=pointless-statement

dags/tpu_observability/node_pool_status.py

@@ -18,6 +18,7 @@
 
 from airflow import models
 from airflow.decorators import task
+from airflow.models.baseoperator import chain
 from airflow.utils.task_group import TaskGroup
 from airflow.utils.trigger_rule import TriggerRule
 
@@ -172,25 +173,22 @@ def generate_problematic_node_location(
           setups=create_problematic_node_pool_info,
       )
 
-      # Airflow uses >> for task chaining, which is pointless for pylint.
-      # pylint: disable=pointless-statement
-      normal_flow = (
-          node_pool_info
-          >> problematic_node_pool_info
-          >> create_node_pool
-          >> wait_for_provisioning
-          >> wait_for_running
-          >> delete_node
-          >> wait_for_repair
-          >> wait_for_recovered
-          >> delete_node_pool
-          >> wait_for_stopping
-          >> cleanup_node_pool
+      chain(
+          node_pool_info,
+          problematic_node_pool_info,
+          create_node_pool,
+          wait_for_provisioning,
+          wait_for_running,
+          delete_node,
+          wait_for_repair,
+          wait_for_recovered,
+          delete_node_pool,
+          wait_for_stopping,
+          cleanup_node_pool,
       )
 
-      flow_for_error_state = (
-          create_problematic_node_pool_info
-          >> wait_for_error
-          >> cleanup_wrong_node_pool
+      chain(
+          create_problematic_node_pool_info,
+          wait_for_error,
+          cleanup_wrong_node_pool,
       )
-      # pylint: enable=pointless-statement

dags/tpu_observability/node_pool_ttr_update_label.py

@@ -17,6 +17,7 @@
 import datetime
 
 from airflow import models
+from airflow.models.baseoperator import chain
 from airflow.utils.task_group import TaskGroup
 from airflow.utils.trigger_rule import TriggerRule
 
@@ -114,13 +115,13 @@
           setups=create_node_pool,
       )
 
-      _ = (
-          node_pool_info
-          >> create_node_pool
-          >> wait_for_provisioning
-          >> wait_for_running
-          >> update_node_pool_label
-          >> wait_for_recovered
-          >> wait_for_ttr
-          >> cleanup_node_pool
+      chain(
+          node_pool_info,
+          create_node_pool,
+          wait_for_provisioning,
+          wait_for_running,
+          update_node_pool_label,
+          wait_for_recovered,
+          wait_for_ttr,
+          cleanup_node_pool,
       )

dags/tpu_observability/tpu_info_format_validation_dags.py

@@ -28,6 +28,7 @@
 from airflow import models
 from airflow.decorators import task
 from airflow.exceptions import AirflowFailException
+from airflow.models.baseoperator import chain
 from airflow.utils.task_group import TaskGroup
 from airflow.utils.trigger_rule import TriggerRule
 
@@ -350,7 +351,10 @@ def generate_second_node_pool_name(
         tpu_accelerator_type="tpu-v6e-slice",
         tpu_topology="4x4",
         container_name="jax-tpu-worker",
-        image="asia-northeast1-docker.pkg.dev/cienet-cmcs/yuna-docker/tpu-info:v0.5.1",
+        image=(
+            "asia-northeast1-docker.pkg.dev/cienet-cmcs/yuna-docker/"
+            "tpu-info:v0.5.1"
+        ),
         tpu_cores_per_pod=4,
     )
 
@@ -494,32 +498,30 @@ def generate_second_node_pool_name(
             setups=create_node_pool,
         )
 
-      # Airflow uses >> for task chaining, which is pointless for pylint.
-      # pylint: disable=pointless-statement
-      (
-          verify_table_amount_task
-          >> [
+      chain(
+          verify_table_amount_task,
+          [
               validate_tpu_chips_metric,
               validate_runtime_metric,
               validate_tensorcore_metric,
               validate_latency_metric,
-          ]
+          ],
       )
 
       [create_first_node_pool, create_second_node_pool]
-      (cleanup_first_node_pool >> cleanup_second_node_pool)
 
-      (
-          cluster_info
-          >> cluster_info_2
-          >> create_node_pool
-          >> apply_time
-          >> pod_names
-          >> wait_for_job_start
-          >> outputs_of_tpu_info
-          >> output_of_tpu_info
-          >> verification_group
-          >> clean_up_workload
-          >> cleanup_node_pool
+      chain(cleanup_first_node_pool, cleanup_second_node_pool)
+
+      chain(
+          cluster_info,
+          cluster_info_2,
+          create_node_pool,
+          apply_time,
+          pod_names,
+          wait_for_job_start,
+          outputs_of_tpu_info,
+          output_of_tpu_info,
+          verification_group,
+          clean_up_workload,
+          cleanup_node_pool,
       )
-      # pylint: enable=pointless-statement

dags/tpu_observability/update_node_pool_label.py

@@ -19,6 +19,7 @@
 import datetime
 
 from airflow import models
+from airflow.models.baseoperator import chain
 from airflow.utils.task_group import TaskGroup
 from airflow.utils.trigger_rule import TriggerRule
 
@@ -109,15 +110,12 @@
           setups=[create_node_pool],
       )
 
-      # Airflow uses >> for task chaining, which is pointless for pylint.
-      # pylint: disable=pointless-statement
-      (
-          node_pool_info
-          >> create_node_pool
-          >> wait_for_availability
-          >> update_node_pool_label
-          >> wait_for_unavailable
-          >> wait_node_pool_recovered
-          >> cleanup_node_pool
+      chain(
+          node_pool_info,
+          create_node_pool,
+          wait_for_availability,
+          update_node_pool_label,
+          wait_for_unavailable,
+          wait_node_pool_recovered,
+          cleanup_node_pool,
       )
-      # pylint: enable=pointless-statement