Skip to content

Commit a7762af

Browse files
committed
feat: Refactor task chaining in TPUPerformanceAnalyzer to use chain for improved readability
1 parent 2904720 commit a7762af

1 file changed

Lines changed: 10 additions & 11 deletions

File tree

dags/tpu_observability/tpu_info_streaming_rate.py

Lines changed: 10 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@
2626

2727
from airflow import models
2828
from airflow.decorators import task
29+
from airflow.models.baseoperator import chain
2930
from airflow.utils.trigger_rule import TriggerRule
3031
from airflow.utils.task_group import TaskGroup
3132

@@ -465,16 +466,14 @@ def generate_second_node_pool_name(
465466
setups=create_node_pool,
466467
)
467468

468-
# Airflow uses >> for task chaining, which is pointless for pylint.
469-
# pylint: disable=pointless-statement
470-
(
471-
cluster_info
472-
>> create_node_pool
473-
>> apply_time
474-
>> pod_names
475-
>> wait_for_job_start
476-
>> rate_verification_group
477-
>> cleanup_workload
478-
>> cleanup_node_pool
469+
chain(
470+
cluster_info,
471+
create_node_pool,
472+
apply_time,
473+
pod_names,
474+
wait_for_job_start,
475+
rate_verification_group,
476+
cleanup_workload,
477+
cleanup_node_pool,
479478
)
480479
# pylint: enable=pointless-statement

0 commit comments

Comments
 (0)