diff --git a/.github/workflows/dag-check.yml b/.github/workflows/dag-check.yml
index 96fcb76bf..ac3982d0e 100644
--- a/.github/workflows/dag-check.yml
+++ b/.github/workflows/dag-check.yml
@@ -3,7 +3,6 @@ name: DAG Check
 
 on:
   pull_request:
-    branches: [master]
     types: [opened, synchronize, edited]
 
   push:
diff --git a/.github/workflows/pyink-check.yml b/.github/workflows/pyink-check.yml
index ea7a8f250..c5cb1cd83 100644
--- a/.github/workflows/pyink-check.yml
+++ b/.github/workflows/pyink-check.yml
@@ -2,11 +2,12 @@ name: Formatter
 
 on:
   pull_request:
-    branches: [master]
     types: [opened, synchronize, edited]
   push:
     branches: [master]
 
+  workflow_dispatch: {}
+
 jobs:
   format_check:
     runs-on: ubuntu-latest
diff --git a/.github/workflows/pylint-check.yml b/.github/workflows/pylint-check.yml
index 5e23d2812..02f4cb31e 100644
--- a/.github/workflows/pylint-check.yml
+++ b/.github/workflows/pylint-check.yml
@@ -2,12 +2,13 @@ name: Linter
 
 on:
   pull_request:
-    branches: [master]
     types: [opened, synchronize, edited]
 
   push:
     branches: [master]
 
+  workflow_dispatch: {}
+
 jobs:
   linting_check:
     runs-on: ubuntu-latest
diff --git a/.github/workflows/require-checklist.yml b/.github/workflows/require-checklist.yml
index d15d19d99..4da288575 100644
--- a/.github/workflows/require-checklist.yml
+++ b/.github/workflows/require-checklist.yml
@@ -2,10 +2,13 @@ name: Require Checklist
 on:
   pull_request:
     types: [opened, edited, synchronize]
+
+  workflow_dispatch: {}
+
 jobs:
   check_pr_body:
     runs-on: ubuntu-latest
     steps:
       - uses: mheap/require-checklist-action@v2
         with:
-          requireChecklist: false # If this is true and there are no checklists detected, the action will fail
\ No newline at end of file
+          requireChecklist: false # If this is true and there are no checklists detected, the action will fail
diff --git a/.github/workflows/unit-test.yml b/.github/workflows/unit-test.yml
index 4f1723cb2..b771ca7e9 100644
--- a/.github/workflows/unit-test.yml
+++ b/.github/workflows/unit-test.yml
@@ -3,7 +3,6 @@ name: Unit Test
 
 on:
   pull_request:
-    branches: [master]
     types: [opened, synchronize, edited]
 
   push:
diff --git a/dags/tpu_observability/jobset_ttr_kill_process.py b/dags/tpu_observability/jobset_ttr_kill_process.py
index bcef73adb..af8f70def 100644
--- a/dags/tpu_observability/jobset_ttr_kill_process.py
+++ b/dags/tpu_observability/jobset_ttr_kill_process.py
@@ -24,6 +24,7 @@
 
 from airflow import models
 from airflow.decorators import task
+from airflow.exceptions import AirflowFailException
 from airflow.models.baseoperator import chain
 from airflow.utils.trigger_rule import TriggerRule
 from airflow.utils.task_group import TaskGroup
@@ -48,29 +49,41 @@
 
 
 @task
-def kill_tpu_pod_workload(info: node_pool.Info, pod_name: str) -> None:
+def kill_tpu_pod_workloads(info: node_pool.Info, pod_names: list[str]) -> None:
   """
-  Kills the python process on a single pod.
+  Kills the python process on a list of pods.
 
   This task retrieves cluster credentials, then attempts to kill the JAX
-  python process inside the specified pod. It ignores errors if the pod
+  python process inside each specified pod. It ignores errors if the pod
   has already been deleted to ensure pipeline continuity.
   """
   with tempfile.NamedTemporaryFile() as temp_config_file:
     env = os.environ.copy()
     env["KUBECONFIG"] = temp_config_file.name
 
-    cmd = " && ".join([
-        jobset.Command.get_credentials_command(info),
-        f"kubectl exec {pod_name} -n default -- pkill -9 -f python",
-    ])
-
-    try:
-      subprocess.run_exec(cmd, env=env)
-    except subprocess.ProcessKilledException:
-      logging.info("Process was terminated with SIGKILL")
-    except Exception as e:
-      raise e
+    failed_pods = []
+    for pod_name in pod_names:
+      cmd = " && ".join([
+          jobset.Command.get_credentials_command(info),
+          f"kubectl exec {pod_name} -n default -- pkill -9 -f python",
+      ])
+
+      try:
+        subprocess.run_exec(cmd, env=env)
+        logging.info("Execution succeeded for pod: %s", pod_name)
+      except subprocess.ProcessKilledException:
+        logging.info(f"Process in pod {pod_name} was terminated with SIGKILL")
+        logging.info("Execution failed for pod: %s", pod_name)
+      except Exception as e:
+        logging.error("Execution failed for pod: %s. Error: %s", pod_name, e)
+        failed_pods.append((pod_name, e))
+
+    if failed_pods:
+      failed_pod_names = [name for name, _ in failed_pods]
+      logging.error("The following pods failed execution: %s", failed_pod_names)
+      raise AirflowFailException(
+          f"Task failed because execution failed on pods: {failed_pod_names}"
+      )
 
 
 # Keyword arguments are generated dynamically at runtime (pylint does not
@@ -156,10 +169,11 @@ def kill_tpu_pod_workload(info: node_pool.Info, pod_name: str) -> None:
           workload_type=Workload.JAX_TPU_BENCHMARK,
       )
 
-      kill_tasks = (
-          kill_tpu_pod_workload.override(task_id="kill_tpu_pod_workload")
-          .partial(info=cluster_info)
-          .expand(pod_name=startup.running_pods)
+      kill_tasks = kill_tpu_pod_workloads.override(
+          task_id="kill_tpu_pod_workloads"
+      )(
+          info=cluster_info,
+          pod_names=startup.running_pods,
       )
 
       wait_for_metric_upload = jobset.wait_for_jobset_ttr_to_be_found.override(
diff --git a/dags/tpu_observability/tpu_info_format_validation_dags.py b/dags/tpu_observability/tpu_info_format_validation_dags.py
index 5cdf4d2b4..67fe91a77 100644
--- a/dags/tpu_observability/tpu_info_format_validation_dags.py
+++ b/dags/tpu_observability/tpu_info_format_validation_dags.py
@@ -23,6 +23,7 @@
 import re
 import subprocess
 import tempfile
+import logging
 
 from airflow import models
 from airflow.decorators import task
@@ -43,10 +44,10 @@
 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 import tpu_info_util as tpu_info
+from dags.tpu_observability.utils.tpu_info_util import parse_tpu_info_output
 from dags.tpu_observability.utils.jobset_util import Workload
 from dags.common.scheduling_helper.scheduling_helper import SchedulingHelper, get_dag_timeout
 
-
 DAG_ID = "tpu_info_format_validation_dag"
 DAGRUN_TIMEOUT = get_dag_timeout(DAG_ID)
 SCHEDULE = SchedulingHelper.arrange_schedule_time(DAG_ID)
@@ -81,6 +82,52 @@ def get_tpu_info_from_pod(info: node_pool.Info, pod_name: str) -> str:
 
 
 @task
+def validate_tpu_info_format(
+    info: node_pool.Info,
+    tpu_config: TpuConfig,
+    pod_names: list[str],
+) -> None:
+  """Executes tpu-info command and runs all validations sequentially across all pods."""
+
+  failed_pods = []
+  for pod_name in pod_names:
+    logging.info(
+        "Executing tpu-info and performing format validation for pod: %s",
+        pod_name,
+    )
+
+    try:
+      with tempfile.NamedTemporaryFile() as temp_config_file:
+        env = os.environ.copy()
+        env["KUBECONFIG"] = temp_config_file.name
+
+        cmd = " && ".join([
+            jobset.Command.get_credentials_command(info),
+            f"kubectl exec {pod_name} -n default -- tpu-info",
+        ])
+
+        raw_output = subprocess.run_exec(cmd, env=env)
+
+      tpu_info_output = parse_tpu_info_output(raw_output)
+
+      verify_table_amount(tpu_info_output)
+      validate_chips_table(tpu_info_output, tpu_config)
+      validate_runtime_table(tpu_info_output)
+      validate_tensorcore_table(tpu_info_output)
+      validate_latency_table(tpu_info_output)
+      logging.info("Validation succeeded for pod: %s", pod_name)
+    except Exception as e:
+      logging.error("Validation failed for pod: %s. Error: %s", pod_name, e)
+      failed_pods.append((pod_name, e))
+
+  if failed_pods:
+    failed_pod_names = [name for name, _ in failed_pods]
+    logging.error("The following pods failed validation: %s", failed_pod_names)
+    raise AirflowFailException(
+        f"Task failed because validation failed on pods: {failed_pod_names}"
+    )
+
+
 def verify_table_amount(tpu_info_output: list[tpu_info.Table]):
   """
   Verifies if all expected tables are present.
@@ -103,7 +150,6 @@ def verify_table_amount(tpu_info_output: list[tpu_info.Table]):
     )
 
 
-@task
 def validate_chips_table(
     tpu_info_output: list[tpu_info.Table],
     tpu_config: TpuConfig,
@@ -156,7 +202,6 @@ def validate_chips_table(
     )
 
 
-@task
 def validate_runtime_table(tpu_info_output: list[tpu_info.Table]):
   """
   Validates the row count and content of table 'TPU Runtime Utilization'
@@ -211,7 +256,6 @@ def validate_runtime_table(tpu_info_output: list[tpu_info.Table]):
     )
 
 
-@task
 def validate_tensorcore_table(tpu_info_output: list[tpu_info.Table]):
   """
   Validates the row count and content of table 'TensorCore Utilization'
@@ -251,7 +295,6 @@ def validate_tensorcore_table(tpu_info_output: list[tpu_info.Table]):
     )
 
 
-@task
 def validate_latency_table(tpu_info_output: list[tpu_info.Table]):
   """
   Validates the row count and content of table 'TPU Buffer Transfer Latency'
@@ -407,57 +450,14 @@ def generate_second_node_pool_name(
           workload_type=Workload.JAX_TPU_BENCHMARK,
       )
 
-      outputs_of_tpu_info = (
-          get_tpu_info_from_pod.override(task_id="get_tpu_info")
-          .partial(info=cluster_info)
-          .expand(pod_name=startup.running_pods)
-      )
-
-      output_of_tpu_info = (
-          tpu_info.parse_tpu_info_output.override(
-              task_id="get_each_metric_table"
-          )
-          .partial()
-          .expand(output=outputs_of_tpu_info)
+      validate_format = validate_tpu_info_format.override(
+          task_id="validate_tpu_info_format"
+      )(
+          info=cluster_info,
+          tpu_config=config,
+          pod_names=startup.running_pods,
       )
 
-      # Keyword arguments are generated dynamically at runtime (pylint does not
-      # know this signature).
-      with TaskGroup(  # pylint: disable=unexpected-keyword-arg
-          group_id="verification_group"
-      ) as verification_group:
-        verify_table_amount_task = (
-            verify_table_amount.override(task_id="verify_table_amount_task")
-            .partial()
-            .expand(tpu_info_output=output_of_tpu_info)
-        )
-
-        validate_tpu_chips_metric = (
-            validate_chips_table.override(task_id="validate_tpu_chips_metric")
-            .partial(tpu_config=config)
-            .expand(tpu_info_output=output_of_tpu_info)
-        )
-
-        validate_runtime_metric = (
-            validate_runtime_table.override(task_id="validate_runtime_metric")
-            .partial()
-            .expand(tpu_info_output=output_of_tpu_info)
-        )
-
-        validate_tensorcore_metric = (
-            validate_tensorcore_table.override(
-                task_id="validate_tensorcore_metric"
-            )
-            .partial()
-            .expand(tpu_info_output=output_of_tpu_info)
-        )
-
-        validate_latency_metric = (
-            validate_latency_table.override(task_id="validate_latency_metric")
-            .partial()
-            .expand(tpu_info_output=output_of_tpu_info)
-        )
-
       clean_up_workload = jobset.end_workload.override(
           task_id="clean_up_workload", trigger_rule=TriggerRule.ALL_DONE
       )(
@@ -488,16 +488,6 @@ def generate_second_node_pool_name(
             setups=create_node_pool,
         )
 
-      chain(
-          verify_table_amount_task,
-          [
-              validate_tpu_chips_metric,
-              validate_runtime_metric,
-              validate_tensorcore_metric,
-              validate_latency_metric,
-          ],
-      )
-
       chain(create_first_node_pool, create_second_node_pool)
 
       chain(cleanup_first_node_pool, cleanup_second_node_pool)
@@ -509,9 +499,7 @@ def generate_second_node_pool_name(
           cluster_info_2,
           create_node_pool,
           *startup.tasks,
-          outputs_of_tpu_info,
-          output_of_tpu_info,
-          verification_group,
+          validate_format,
           clean_up_workload,
           cleanup_node_pool,
       )
diff --git a/dags/tpu_observability/tpu_info_metrics_dag.py b/dags/tpu_observability/tpu_info_metrics_dag.py
index 10fef9784..c957c437a 100644
--- a/dags/tpu_observability/tpu_info_metrics_dag.py
+++ b/dags/tpu_observability/tpu_info_metrics_dag.py
@@ -25,7 +25,7 @@
 
 from airflow import models
 from airflow.decorators import task
-from airflow.exceptions import AirflowException
+from airflow.exceptions import AirflowException, AirflowFailException
 from airflow.models.baseoperator import chain
 from airflow.utils.task_group import TaskGroup
 from airflow.utils.trigger_rule import TriggerRule
@@ -46,6 +46,7 @@
 from dags.tpu_observability.utils.node_pool_util import Info
 from dags.tpu_observability.utils.time_util import TimeUtil
 from dags.tpu_observability.utils.jobset_util import Workload
+from dags.tpu_observability.utils.tpu_info_util import parse_tpu_info_output
 from dags.common.scheduling_helper.scheduling_helper import SchedulingHelper, get_dag_timeout
 
 
@@ -114,72 +115,94 @@ def compare_metric_values(
 
 
 @task
-def get_tpu_info_metric_from_pod(
-    node_pool: node_pool.Info,
-    pod_name: str,
-    jobset_config: jobset,
-    metric_name: str,
-) -> str:
-  """Executes the 'tpu-info' command in the specified pod and returns its output."""
-  with tempfile.TemporaryDirectory() as tmpdir:
-    kube_dir = tmpdir + "/kubeconfig"
-    env = os.environ.copy()
-    env["KUBECONFIG"] = kube_dir
-
-    cmd = " && ".join([
-        jobset.Command.get_credentials_command(node_pool),
-        (
-            f"kubectl --kubeconfig={kube_dir} "
-            f"exec {pod_name} -n {jobset_config.namespace} "
-            f"-- tpu-info --metric {metric_name}"
-        ),
-    ])
-
-    return subprocess.run_exec(cmd=cmd, env=env)
-
-
-@task
-def run_metric_verification(
+def verify_metric_for_all_pods(
     node_pool: Info,
+    jobset_config: jobset.JobSet,
     job_apply_time: TimeUtil,
     metric_strategy: BaseMetricStrategy,
-    comparison_data: tuple[str, list[tpu_info.Table]],
-):
-  """A generic task that uses a strategy object to verify a metric."""
-  pod_name, tpu_info_output = comparison_data
-  metric_name = metric_strategy.metric_name
-  logging.info("Verifying metric '%s' for pod: %s...", metric_name, pod_name)
-
-  start_time = job_apply_time
-  end_time = job_apply_time + datetime.timedelta(minutes=10)
-
-  time_series_data = metric_strategy.list_or_query_metric(
-      project_id=node_pool.project_id,
-      cluster_name=node_pool.cluster_name,
-      pod_name=pod_name,
-      start_time=start_time,
-      end_time=end_time,
-  )
+    pod_names: list[str],
+) -> list[bool]:
+  """Runs metric verification across all pods sequentially."""
 
-  monitoring_values = metric_strategy.parse_from_monitoring(time_series_data)
-  cmd_values = metric_strategy.parse_from_tpu_info(tpu_info_output)
+  results = []
+  failed_pods = []
+  for pod_name in pod_names:
+    logging.info(
+        "Fetching tpu-info metric '%s' for pod: %s...",
+        metric_strategy.tpu_info_metric_name,
+        pod_name,
+    )
 
-  tolerance_for_metric = metric_strategy.tolerance_percent
-  logging.info(
-      "Using a tolerance of %.2f%% for metric '%s' comparison.",
-      tolerance_for_metric,
-      metric_strategy.dag_id_suffix,
-  )
+    try:
+      with tempfile.TemporaryDirectory() as tmpdir:
+        kube_dir = tmpdir + "/kubeconfig"
+        env = os.environ.copy()
+        env["KUBECONFIG"] = kube_dir
+
+        cmd = " && ".join([
+            jobset.Command.get_credentials_command(node_pool),
+            (
+                f"kubectl --kubeconfig={kube_dir} "
+                f"exec {pod_name} -n {jobset_config.namespace} "
+                f"-- tpu-info --metric {metric_strategy.tpu_info_metric_name}"
+            ),
+        ])
+
+        output = subprocess.run_exec(cmd=cmd, env=env)
+
+      tpu_info_output = parse_tpu_info_output(output)
+
+      logging.info(
+          "Verifying metric '%s' for pod: %s...",
+          metric_strategy.metric_name,
+          pod_name,
+      )
 
-  compare_metric_values(
-      cmd_values,
-      monitoring_values,
-      pod_name,
-      metric_display_name=metric_strategy.dag_id_suffix,
-      tolerance_percent=tolerance_for_metric,
-  )
+      start_time = job_apply_time
+      end_time = job_apply_time + datetime.timedelta(minutes=10)
+
+      time_series_data = metric_strategy.list_or_query_metric(
+          project_id=node_pool.project_id,
+          cluster_name=node_pool.cluster_name,
+          pod_name=pod_name,
+          start_time=start_time,
+          end_time=end_time,
+      )
+
+      monitoring_values = metric_strategy.parse_from_monitoring(
+          time_series_data
+      )
+      cmd_values = metric_strategy.parse_from_tpu_info(tpu_info_output)
+
+      tolerance_for_metric = metric_strategy.tolerance_percent
+      logging.info(
+          "Using a tolerance of %.2f%% for metric '%s' comparison on pod %s.",
+          tolerance_for_metric,
+          metric_strategy.dag_id_suffix,
+          pod_name,
+      )
 
-  return True
+      compare_metric_values(
+          cmd_values,
+          monitoring_values,
+          pod_name,
+          metric_display_name=metric_strategy.dag_id_suffix,
+          tolerance_percent=tolerance_for_metric,
+      )
+      results.append(True)
+      logging.info("Validation succeeded for pod: %s", pod_name)
+    except Exception as e:
+      logging.error("Validation failed for pod: %s. Error: %s", pod_name, e)
+      failed_pods.append((pod_name, e))
+
+  if failed_pods:
+    failed_pod_names = [name for name, _ in failed_pods]
+    logging.error("The following pods failed validation: %s", failed_pod_names)
+    raise AirflowFailException(
+        f"Task failed because validation failed on pods: {failed_pod_names}"
+    )
+
+  return results
 
 
 @task
@@ -334,48 +357,20 @@ def generate_second_node_pool_name(
       )
 
       verification_results = {}
-      all_verification_groups = []
+      all_verification_tasks = []
 
       for strategy in ALL_METRIC_STRATEGIES:
-        group_id = f"verify_{strategy.dag_id_suffix}"
-
-        with TaskGroup(group_id=group_id) as verification_group:
-          tpu_info_metric_outputs = (
-              get_tpu_info_metric_from_pod.override(
-                  task_id="get_tpu_info_metric_table"
-              )
-              .partial(
-                  node_pool=cluster_info,
-                  jobset_config=jobset_config,
-                  metric_name=strategy.tpu_info_metric_name,
-              )
-              .expand(pod_name=startup.running_pods)
-          )
-
-          tpu_info_metric_output = (
-              tpu_info.parse_tpu_info_output.override(
-                  task_id="get_each_metric_table"
-              )
-              .partial()
-              .expand(output=tpu_info_metric_outputs)
-          )
-
-          verify_metric = (
-              run_metric_verification.override(task_id="run_verification")
-              .partial(
-                  node_pool=cluster_info,
-                  job_apply_time=startup.jobset_start_time,
-                  metric_strategy=strategy,
-              )
-              .expand(
-                  comparison_data=startup.running_pods.zip(
-                      tpu_info_metric_output
-                  )
-              )
-          )
-
-        all_verification_groups.append(verification_group)
+        verify_metric = verify_metric_for_all_pods.override(
+            task_id=f"verify_{strategy.dag_id_suffix}"
+        )(
+            node_pool=cluster_info,
+            jobset_config=jobset_config,
+            job_apply_time=startup.jobset_start_time,
+            metric_strategy=strategy,
+            pod_names=startup.running_pods,
+        )
 
+        all_verification_tasks.append(verify_metric)
         verification_results[strategy.dag_id_suffix] = verify_metric
 
       summary = summarize_results.override(
@@ -418,7 +413,7 @@ def generate_second_node_pool_name(
           cluster_info_2,
           create_node_pool,
           *startup.tasks,
-          all_verification_groups,
+          all_verification_tasks,
           summary,
           clean_up_workload,
           cleanup_node_pool,
diff --git a/dags/tpu_observability/tpu_sdk_monitoring_validation_dag.py b/dags/tpu_observability/tpu_sdk_monitoring_validation_dag.py
index 644d1971d..4dd5ae521 100644
--- a/dags/tpu_observability/tpu_sdk_monitoring_validation_dag.py
+++ b/dags/tpu_observability/tpu_sdk_monitoring_validation_dag.py
@@ -16,12 +16,14 @@
 list_supported_metrics() are functional inside TPU worker pods."""
 
 import datetime
+import logging
 
 from airflow import models
+from airflow.decorators import task
+from airflow.exceptions import AirflowFailException
 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
@@ -42,7 +44,7 @@
 
 
 @task
-def validate_monitoring_sdk(info: node_pool.Info, pod_name: str) -> None:
+def validate_monitoring_sdk(info: node_pool.Info, pod_names: list[str]) -> None:
   """Validates the tpumonitoring SDK functions inside TPU worker pods.
 
   This task executes both help() and list_supported_metrics() via the SDK
@@ -50,7 +52,7 @@ def validate_monitoring_sdk(info: node_pool.Info, pod_name: str) -> None:
 
   Args:
     info: Cluster info for gcloud credentials.
-    pod_name: Pod name provided by dynamic task mapping.
+    pod_names: List of pod names.
   """
   # A dict of script to its expected result patterns.
   validate_spec: dict[sdk.TpuMonitoringScript, list[str]] = {
@@ -74,14 +76,29 @@ def validate_monitoring_sdk(info: node_pool.Info, pod_name: str) -> None:
       ],
   }
 
-  for script, patterns in validate_spec.items():
-    output = sdk.execute_sdk_command(info, pod_name, script)
-    for pattern in patterns:
-      if pattern not in output:
-        raise AssertionError(
-            f"Validation failed for 'tpumonitoring.{script.name.lower()}()': "
-            f"Missing '{pattern}'."
-        )
+  failed_pods = []
+  for pod_name in pod_names:
+    logging.info("Validating tpumonitoring SDK for pod: %s", pod_name)
+    try:
+      for script, patterns in validate_spec.items():
+        output = sdk.execute_sdk_command(info, pod_name, script)
+        for pattern in patterns:
+          if pattern not in output:
+            raise AssertionError(
+                f"Validation failed for 'tpumonitoring.{script.name.lower()}()' inside pod '{pod_name}': "
+                f"Missing '{pattern}'."
+            )
+      logging.info("Validation succeeded for pod: %s", pod_name)
+    except Exception as e:
+      logging.error("Validation failed for pod: %s. Error: %s", pod_name, e)
+      failed_pods.append((pod_name, e))
+
+  if failed_pods:
+    failed_pod_names = [name for name, _ in failed_pods]
+    logging.error("The following pods failed validation: %s", failed_pod_names)
+    raise AirflowFailException(
+        f"Task failed because validation failed on pods: {failed_pod_names}"
+    )
 
 
 with models.DAG(
@@ -162,10 +179,9 @@ def validate_monitoring_sdk(info: node_pool.Info, pod_name: str) -> None:
         workload_type=Workload.JAX_TPU_BENCHMARK,
     )
 
-    sdk_validation = (
-        validate_monitoring_sdk.override(task_id="sdk_validation")
-        .partial(info=cluster_info)
-        .expand(pod_name=startup.running_pods)
+    sdk_validation = validate_monitoring_sdk.override(task_id="sdk_validation")(
+        info=cluster_info,
+        pod_names=startup.running_pods,
     )
 
     cleanup_workload = jobset.end_workload.override(
diff --git a/dags/tpu_observability/utils/tpu_info_util.py b/dags/tpu_observability/utils/tpu_info_util.py
index c408463d4..89f4c5a70 100644
--- a/dags/tpu_observability/utils/tpu_info_util.py
+++ b/dags/tpu_observability/utils/tpu_info_util.py
@@ -5,7 +5,6 @@
 from enum import IntEnum
 import re
 
-from airflow.decorators import task
 
 # A type alias for a parsed row, mapping column headers to their values.
 _TableRow = dict[str, str]
@@ -65,7 +64,6 @@ class TableLineIndex(IntEnum):
     self.body = parsed_body
 
 
-@task
 def parse_tpu_info_output(output: str) -> list[Table]:
   """Splits a multi-table string from tpu-info into a structured TpuInfo object.
 
diff --git a/scripts/code-style.sh b/scripts/code-style.sh
index 36cfa13e9..ea50709ac 100755
--- a/scripts/code-style.sh
+++ b/scripts/code-style.sh
@@ -19,14 +19,37 @@ set -e
 
 FOLDERS_TO_FORMAT=("dags" "xlml")
 
-for folder in "${FOLDERS_TO_FORMAT[@]}"
-do
-  pyink "$folder" --pyink-indentation=2 --pyink-use-majority-quotes --line-length=80 --check --diff
-done
-
-for folder in "${FOLDERS_TO_FORMAT[@]}"
-do
-  pylint "./$folder" --fail-under=9.6
-done
+HEAD_SHA="$(git rev-parse HEAD)"
+BASE_BRANCH="dev"
+
+if ! git rev-parse --verify "$BASE_BRANCH" >/dev/null 2>&1; then
+  git fetch origin "$BASE_BRANCH":"$BASE_BRANCH" || {
+    echo "[code-style] base branch '$BASE_BRANCH' not found, skip diff-based check."
+    exit 0
+  }
+fi
+
+CHANGED_PY_FILES="$(
+  git diff --name-only --diff-filter=ACM "${BASE_BRANCH}" "${HEAD_SHA}" \
+    | grep '\.py$' \
+    | while read -r f; do
+        for folder in "${FOLDERS_TO_FORMAT[@]}"; do
+          if [[ "$f" == "$folder/"* ]]; then
+            echo "$f"
+            break
+          fi
+        done
+      done \
+    | sort -u
+)"
+
+if [[ -z "${CHANGED_PY_FILES}" ]]; then
+  echo "[pre-push hook] no changed files detected between ${HEAD_SHA} and ${BASE_BRANCH}"
+  exit 1
+fi
+
+pyink ${CHANGED_PY_FILES} --pyink-indentation=2 --pyink-use-majority-quotes --line-length=80 --check --diff
+
+pylint ${CHANGED_PY_FILES} --fail-under=9.6 --disable=E1123
 
 echo "Successfully clean up all codes."