fix: Correct code style (GoogleCloudPlatform#1283)

This change focused on resolving code style issues, specifically organizing and correcting import statements across the affected DAG files. No logical changes were introduced.

Co-authored-by: Jim Tseng <jim.tseng@cienet.com>

Author: yuna-tzeng

dags/tpu_observability/configs/common.py

@@ -1,5 +1,7 @@
-from dataclasses import dataclass
+"""Common configurations and constants for TPU observability DAGs."""
+
 import enum
+from dataclasses import dataclass
 
 from dags.common.vm_resource import MachineVersion, TpuVersion
 

dags/tpu_observability/interruption_validation_dag.py

@@ -10,14 +10,13 @@
 from airflow.exceptions import AirflowSkipException
 from airflow.models.baseoperator import chain
 from airflow.utils.task_group import TaskGroup
+from google.cloud import monitoring_v3
 
+from dags import composer_env
 from dags.common import test_owner
 from dags.common.vm_resource import Project
 from dags.multipod.configs.common import Platform
 from dags.tpu_observability.utils import gcp_util, time_util
-from google.cloud import monitoring_v3
-from dags import composer_env
-
 
 _UNKNOWN_RESOURCE_NAME = 'Unknown'
 
@@ -49,8 +48,10 @@ def metric_label(self) -> str:
   def log_filter(self) -> str:
     """Returns the corresponding filter for the interruption reason.
 
-    These filters are in accordance with the definitions in this file from Google3:
-    //depot/google3/java/com/google/cloud/cluster/manager/compute/services/instancemanagerevent/InstanceEventNotificationAction.java
+    These filters are in accordance with the definitions in this file from
+    Google3:
+    //depot/google3/java/com/google/cloud/cluster/manager/compute/services/
+    instancemanagerevent/InstanceEventNotificationAction.java
     """
 
     filters = []
@@ -157,7 +158,8 @@ def fetch_interruption_metric_records(
       f'resource.labels.project_id = "{configs.project_id}" '
       f'metric.type = "{metric_type}" '
       f'resource.type = "{resource_type}" '
-      f'metric.labels.interruption_reason = "{configs.interruption_reason.metric_label()}" '
+      'metric.labels.interruption_reason = '
+      f'"{configs.interruption_reason.metric_label()}" '
   )
 
   # key: resource_name, value: EventRecord
@@ -432,7 +434,9 @@ def validate_interruption_count(
 
     if len(log_timestamps) != len(metric.record_timestamps):
       mismatch_nodes.append(
-          f'mismatch resource name: {resource_name}, metric_count: {len(metric.record_timestamps)}, log_count: {len(log_timestamps)}'
+          f'mismatch resource name: {resource_name}, '
+          f'metric_count: {len(metric.record_timestamps)}, '
+          f'log_count: {len(log_timestamps)}'
       )
 
   if mismatch_nodes:
@@ -527,13 +531,13 @@ def create_interruption_dag(
     for project in Project:
       match project:
         case Project.TPU_PROD_ENV_AUTOMATED | Project.CLOUD_TPU_INFERENCE_TEST:
-          # Production composer lacks permission for these projects; ignore them.
+          # Production composer lacks permission for these projects; ignore them
           continue
         case _:
           with TaskGroup(
               group_id=f'validation_for_{project.value}',
               tooltip=f'Validation pipeline for Project ID: {project.value}',
-          ) as group:
+          ):
             configs = Configs(
                 project_id=project.value,
                 platform=platform,

dags/tpu_observability/jobset_ttr_drain_restart.py

@@ -17,25 +17,26 @@
 import datetime
 
 from airflow import models
+from airflow.decorators import task
 from airflow.models.baseoperator import chain
-from airflow.utils.trigger_rule import TriggerRule
 from airflow.utils.task_group import TaskGroup
-
-from airflow.decorators import task
+from airflow.utils.trigger_rule import TriggerRule
 
 from dags import composer_env
-from dags.tpu_observability.utils import jobset_util as jobset
-from dags.tpu_observability.utils import subprocess_util as subprocess
-from dags.tpu_observability.utils import node_pool_util as node_pool
-from dags.tpu_observability.utils.node_pool_util import Info
-from dags.tpu_observability.utils.node_pool_util import NodeOperationSpec
-from dags.tpu_observability.utils.jobset_util import Workload
+from dags.common.scheduling_helper.scheduling_helper import (
+    SchedulingHelper,
+    get_dag_timeout,
+)
 from dags.tpu_observability.configs.common import (
-    MachineConfigMap,
     GCS_CONFIG_PATH,
     GCS_JOBSET_CONFIG_PATH,
+    MachineConfigMap,
 )
-from dags.common.scheduling_helper.scheduling_helper import SchedulingHelper, get_dag_timeout
+from dags.tpu_observability.utils import jobset_util as jobset
+from dags.tpu_observability.utils import node_pool_util as node_pool
+from dags.tpu_observability.utils import subprocess_util as subprocess
+from dags.tpu_observability.utils.jobset_util import Workload
+from dags.tpu_observability.utils.node_pool_util import Info, NodeOperationSpec
 
 DAG_ID = "jobset_ttr_drain_restart"
 DAGRUN_TIMEOUT = get_dag_timeout(DAG_ID)

dags/tpu_observability/jobset_ttr_kill_process.py

@@ -19,26 +19,29 @@
 
 import datetime
 import logging
-import tempfile
 import os
+import tempfile
 
 from airflow import models
 from airflow.decorators import task
 from airflow.models.baseoperator import chain
-from airflow.utils.trigger_rule import TriggerRule
 from airflow.utils.task_group import TaskGroup
+from airflow.utils.trigger_rule import TriggerRule
 
 from dags import composer_env
-from dags.tpu_observability.utils import jobset_util as jobset
-from dags.tpu_observability.utils import node_pool_util as node_pool
-from dags.tpu_observability.utils import subprocess_util as subprocess
-from dags.tpu_observability.utils.jobset_util import JobSet, Workload
+from dags.common.scheduling_helper.scheduling_helper import (
+    SchedulingHelper,
+    get_dag_timeout,
+)
 from dags.tpu_observability.configs.common import (
-    MachineConfigMap,
     GCS_CONFIG_PATH,
     GCS_JOBSET_CONFIG_PATH,
+    MachineConfigMap,
 )
-from dags.common.scheduling_helper.scheduling_helper import SchedulingHelper, get_dag_timeout
+from dags.tpu_observability.utils import jobset_util as jobset
+from dags.tpu_observability.utils import node_pool_util as node_pool
+from dags.tpu_observability.utils import subprocess_util as subprocess
+from dags.tpu_observability.utils.jobset_util import Workload
 
 DAG_ID = "jobset_ttr_kill_process"
 DAGRUN_TIMEOUT = get_dag_timeout(DAG_ID)

dags/tpu_observability/jobset_ttr_node_pool_resize.py

@@ -18,20 +18,22 @@
 
 from airflow import models
 from airflow.models.baseoperator import chain
-from airflow.utils.trigger_rule import TriggerRule
 from airflow.utils.task_group import TaskGroup
+from airflow.utils.trigger_rule import TriggerRule
 
 from dags import composer_env
-from dags.tpu_observability.utils import jobset_util as jobset
-from dags.tpu_observability.utils import node_pool_util as node_pool
-from dags.tpu_observability.utils.jobset_util import JobSet, Workload
+from dags.common.scheduling_helper.scheduling_helper import (
+    SchedulingHelper,
+    get_dag_timeout,
+)
 from dags.tpu_observability.configs.common import (
-    MachineConfigMap,
     GCS_CONFIG_PATH,
     GCS_JOBSET_CONFIG_PATH,
+    MachineConfigMap,
 )
-from dags.common.scheduling_helper.scheduling_helper import SchedulingHelper, get_dag_timeout
-
+from dags.tpu_observability.utils import jobset_util as jobset
+from dags.tpu_observability.utils import node_pool_util as node_pool
+from dags.tpu_observability.utils.jobset_util import Workload
 
 DAG_ID = "jobset_ttr_node_pool_resize"
 DAGRUN_TIMEOUT = get_dag_timeout(DAG_ID)

dags/tpu_observability/jobset_ttr_node_reboot.py

@@ -12,7 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""A DAG to test JobSet Time-To-Recover (TTR) metric by triggering a node reboot."""
+"""A DAG to test JobSet Time-To-Recover (TTR) metric by triggering a node
+reboot."""
 
 import datetime
 from datetime import timedelta
@@ -22,19 +23,19 @@
 from airflow.utils.trigger_rule import TriggerRule
 
 from dags import composer_env
-from dags.tpu_observability.utils import jobset_util as jobset
-from dags.tpu_observability.utils import node_pool_util as node_pool
-from dags.tpu_observability.utils.jobset_util import Workload
-from dags.tpu_observability.configs.common import (
-    MachineConfigMap,
-    GCS_CONFIG_PATH,
-    GCS_JOBSET_CONFIG_PATH,
-)
 from dags.common.scheduling_helper.scheduling_helper import (
     SchedulingHelper,
     get_dag_timeout,
 )
 from dags.common.task_group_with_timeout import TaskGroupWithTimeout
+from dags.tpu_observability.configs.common import (
+    GCS_CONFIG_PATH,
+    GCS_JOBSET_CONFIG_PATH,
+    MachineConfigMap,
+)
+from dags.tpu_observability.utils import jobset_util as jobset
+from dags.tpu_observability.utils import node_pool_util as node_pool
+from dags.tpu_observability.utils.jobset_util import Workload
 
 DAG_ID = "jobset_ttr_node_reboot"
 DAGRUN_TIMEOUT = get_dag_timeout(DAG_ID)
@@ -131,7 +132,7 @@
 
       reboot_node = jobset.operate_pod.override(task_id="reboot_node")(
           node_pool=cluster_info,
-          operation=jobset.PodOperationSpec.Reboot(),
+          operation=jobset.PodOperationSpec.reboot(),
           pod_name=target_pod,
           namespace="default",
       )

dags/tpu_observability/jobset_ttr_pod_delete.py

@@ -18,19 +18,22 @@
 
 from airflow import models
 from airflow.models.baseoperator import chain
-from airflow.utils.trigger_rule import TriggerRule
 from airflow.utils.task_group import TaskGroup
+from airflow.utils.trigger_rule import TriggerRule
 
 from dags import composer_env
-from dags.tpu_observability.utils import jobset_util as jobset
-from dags.tpu_observability.utils import node_pool_util as node_pool
-from dags.tpu_observability.utils.jobset_util import Workload
+from dags.common.scheduling_helper.scheduling_helper import (
+    SchedulingHelper,
+    get_dag_timeout,
+)
 from dags.tpu_observability.configs.common import (
-    MachineConfigMap,
     GCS_CONFIG_PATH,
     GCS_JOBSET_CONFIG_PATH,
+    MachineConfigMap,
 )
-from dags.common.scheduling_helper.scheduling_helper import SchedulingHelper, get_dag_timeout
+from dags.tpu_observability.utils import jobset_util as jobset
+from dags.tpu_observability.utils import node_pool_util as node_pool
+from dags.tpu_observability.utils.jobset_util import Workload
 
 DAG_ID = "jobset_ttr_pod_delete"
 DAGRUN_TIMEOUT = get_dag_timeout(DAG_ID)

dags/tpu_observability/jobset_ttr_rollback.py

@@ -18,19 +18,22 @@
 
 from airflow import models
 from airflow.models.baseoperator import chain
-from airflow.utils.trigger_rule import TriggerRule
 from airflow.utils.task_group import TaskGroup
+from airflow.utils.trigger_rule import TriggerRule
 
 from dags import composer_env
-from dags.tpu_observability.utils import jobset_util as jobset
-from dags.tpu_observability.utils import node_pool_util as node_pool
-from dags.tpu_observability.utils.jobset_util import Workload
+from dags.common.scheduling_helper.scheduling_helper import (
+    SchedulingHelper,
+    get_dag_timeout,
+)
 from dags.tpu_observability.configs.common import (
-    MachineConfigMap,
     GCS_CONFIG_PATH,
     GCS_JOBSET_CONFIG_PATH,
+    MachineConfigMap,
 )
-from dags.common.scheduling_helper.scheduling_helper import SchedulingHelper, get_dag_timeout
+from dags.tpu_observability.utils import jobset_util as jobset
+from dags.tpu_observability.utils import node_pool_util as node_pool
+from dags.tpu_observability.utils.jobset_util import Workload
 
 DAG_ID = "jobset_rollback_ttr"
 DAGRUN_TIMEOUT = get_dag_timeout(DAG_ID)

dags/tpu_observability/jobset_uptime_validation.py

@@ -23,16 +23,19 @@
 from airflow.utils.trigger_rule import TriggerRule
 
 from dags import composer_env
+from dags.common.scheduling_helper.scheduling_helper import (
+    SchedulingHelper,
+    get_dag_timeout,
+)
 from dags.tpu_observability.configs.common import (
-    MachineConfigMap,
     GCS_CONFIG_PATH,
     GCS_JOBSET_CONFIG_PATH,
+    MachineConfigMap,
 )
 from dags.tpu_observability.utils import jobset_util as jobset
 from dags.tpu_observability.utils import node_pool_util as node_pool
 from dags.tpu_observability.utils.jobset_util import Workload
 from dags.tpu_observability.utils.time_util import TimeUtil
-from dags.common.scheduling_helper.scheduling_helper import SchedulingHelper, get_dag_timeout
 
 DAG_ID = "jobset_uptime_validation"
 DAGRUN_TIMEOUT = get_dag_timeout(DAG_ID)
@@ -149,8 +152,10 @@ def get_current_time() -> TimeUtil:
           task_id="get_current_time"
       )()
 
-      ensure_no_jobset_uptime_data = jobset.ensure_no_jobset_uptime_data.override(
-          task_id="ensure_no_jobset_uptime_data"
+      ensure_no_jobset_uptime_data = (
+          jobset.ensure_no_jobset_uptime_data.override(
+              task_id="ensure_no_jobset_uptime_data"
+          )
       )(
           node_pool=cluster_info,
           jobset_name=jobset_name,

dags/tpu_observability/multi_host_nodepool_rollback_dag.py

@@ -26,10 +26,15 @@
 
 from dags import composer_env
 from dags.common import test_owner
-from dags.tpu_observability.configs.common import MachineConfigMap, GCS_CONFIG_PATH
+from dags.common.scheduling_helper.scheduling_helper import (
+    SchedulingHelper,
+    get_dag_timeout,
+)
+from dags.tpu_observability.configs.common import (
+    GCS_CONFIG_PATH,
+    MachineConfigMap,
+)
 from dags.tpu_observability.utils import node_pool_util as node_pool
-from dags.common.scheduling_helper.scheduling_helper import SchedulingHelper, get_dag_timeout
-
 
 DAG_ID = "multi_host_nodepool_rollback"
 DAGRUN_TIMEOUT = get_dag_timeout(DAG_ID)
@@ -78,7 +83,7 @@
 ) as dag:
   for machine in MachineConfigMap:
     config = machine.value
-    LABELS_TO_UPDATE = (
+    labels_to_update = (
         {"env": "prod"} if composer_env.is_prod_env() else {"env": "dev"}
     )