fix: Switch sensors to poke mode for accurate status duration (GoogleCloudPlatform#1131)

This change improves the observability and efficiency of task sensors under folder `dags/tpu_observability` by switching their operational mode from "reschedule" to "poke" to:
1. Enhanced Troubleshooting Visibility: In poke mode, the sensor remains active in a single session. This allows the Airflow UI (Gantt Chart & Task Duration views) to show the exact, continuous time spent waiting for TPU resources.
2. High-Value Resource Prioritization (TPU/Cluster vs. Worker): TPU/Cluster resources are far more precious and limited than Airflow worker slots. Using poke ensures zero-latency task handoff. Holding a worker slot during the poke interval has a negligible impact on cluster capacity unless the system is running a massive number of highly parallelized DAGs.

Author: chiajunglien

dags/tpu_observability/utils/jobset_util.py

@@ -458,7 +458,7 @@ def list_pod_names(node_pool: node_pool_info, namespace: str) -> list[str]:
     return pod_list
 
 
-@task.sensor(poke_interval=30, timeout=900, mode="reschedule")
+@task.sensor(poke_interval=30, timeout=900, mode="poke")
 def wait_for_jobset_started(
     node_pool: node_pool_info,
     pod_name_list: str,
@@ -525,7 +525,7 @@ def wait_for_jobset_started(
   return all(p > threshold_value for p in last_n_data_points)
 
 
-@task.sensor(poke_interval=60, timeout=3600, mode="reschedule")
+@task.sensor(poke_interval=60, timeout=3600, mode="poke")
 def wait_for_jobset_ttr_to_be_found(node_pool: node_pool_info) -> bool:
   """
   Polls the jobset time_between_interruptions metric.
@@ -559,7 +559,7 @@ def wait_for_jobset_ttr_to_be_found(node_pool: node_pool_info) -> bool:
   return len(time_series) > 0
 
 
-@task.sensor(poke_interval=30, timeout=600, mode="reschedule")
+@task.sensor(poke_interval=30, timeout=600, mode="poke")
 def wait_for_jobset_status_occurrence(
     replica_type: str, job_name: str, node_pool: node_pool_info
 ):
@@ -581,7 +581,7 @@ def wait_for_jobset_status_occurrence(
   return ready_replicas > 0
 
 
-@task.sensor(poke_interval=30, timeout=600, mode="reschedule")
+@task.sensor(poke_interval=30, timeout=600, mode="poke")
 def wait_for_all_pods_running(num_pods: int, node_pool: node_pool_info):
   num_running = len(get_running_pods(node_pool=node_pool, namespace="default"))
   return num_running == num_pods

dags/tpu_observability/utils/node_pool_util.py

@@ -413,7 +413,7 @@ def _query_status_metric(node_pool: Info) -> Status:
   return Status.from_str(latest_status)
 
 
-@task.sensor(poke_interval=60, timeout=600, mode="reschedule")
+@task.sensor(poke_interval=60, timeout=600, mode="poke")
 def wait_for_status(
     node_pool: Info,
     status: Status,
@@ -465,7 +465,7 @@ def rollback(node_pool: Info) -> None:
   subprocess.run_exec(command)
 
 
-@task.sensor(poke_interval=30, timeout=1200, mode="reschedule")
+@task.sensor(poke_interval=30, timeout=1200, mode="poke")
 def wait_for_availability(
     node_pool: Info,
     availability: bool,
@@ -541,7 +541,7 @@ def wait_for_availability(
   return availability == state
 
 
-@task.sensor(poke_interval=30, timeout=3600, mode="reschedule")
+@task.sensor(poke_interval=30, timeout=3600, mode="poke")
 def wait_for_ttr(
     node_pool: Info,
     operation_start_time: TimeUtil,