Fix: Refine the wasted time during retries for jetstream_benchmark_serving (GoogleCloudPlatform#1051)

With 57 TaskGroups across four parallel branches, the retry wait alone wastes ~570 minutes (~9.5 hours)
This change improves the error handling within the ssh_tpu Airflow task to distinguish between critical SSH authentication failures and other command execution errors.

Solution:
- If any host fails with an `AuthenticationException`, the task is failed with a non-retryable `AirflowFailException`.
- If hosts fail due to other exceptions (e.g., `invoke.UnexpectedExit`), the original `fabric.group.GroupException` is re-raised. This allows Airflow's retry mechanism to handle these potentially transient issues as configured for the task.

Author: RUEI4341

xlml/apis/task.py

@@ -109,12 +109,18 @@ def run_queued_resource_test(
           task_test_config,
       )
 
-      queued_resource_op >> tpu.ssh_tpu.override(task_id="setup")(
+      setup_task = tpu.ssh_tpu.override(
+          task_id="setup",
+          # Setup/install retries don’t need a long cooldown.
+          # 30s is enough for network connection problem; longer delays do not make sense.
+          retry_delay=datetime.timedelta(seconds=30),
+      )(
           queued_resource_name,
           task_test_config.setup_script,
           ssh_keys,
           True if task_test_config.test_name.startswith("tf_") else all_workers,
       )
+      queued_resource_op >> setup_task
 
     run_model = tpu.ssh_tpu.override(
         task_id="run_model",

xlml/utils/tpu.py

@@ -27,6 +27,7 @@
 from airflow.utils.task_group import TaskGroup
 from airflow.operators.python import get_current_context
 from airflow.models import Variable
+from airflow.exceptions import AirflowFailException
 from xlml.apis import gcp_config, test_config
 from xlml.utils import ssh, startup_script, composer
 import fabric
@@ -404,6 +405,27 @@ def ssh_tpu(
       gateway='corp-ssh-helper %h %p' if use_external_ips else None,
   )
 
+  def ssh_group_run(cmds: Iterable[str]):
+    try:
+      ssh_group.run(cmds, env=env)
+    except fabric.group.GroupException as e:
+      for connection, result in e.result.items():
+        if isinstance(result, paramiko.ssh_exception.AuthenticationException):
+          logging.error(
+              f'SSH Authentication Failed on {connection.host}: {result}'
+          )
+          raise AirflowFailException(
+              'SSH Authentication failed on one or more hosts. Check logs for details.'
+          ) from e
+      raise
+    except paramiko.ssh_exception.AuthenticationException as e:
+      error_msg = f'SSH Authentication Failed: {e}'
+      logging.error(error_msg)
+      raise AirflowFailException(error_msg) from e
+
+    finally:
+      ssh_group.close()
+
   context = get_current_context()
   if context['task_instance'].try_number > 1:
     # kill TPU process by pid (if any) to avoid `TPU in use` error in retry
@@ -414,10 +436,9 @@ def ssh_tpu(
         f'set -xue; sudo echo "{script}" > {tmp_file}',
         f'bash {tmp_file} {accelerator_type}',
     )
-    ssh_group.run(';'.join(kill_process_cmds))
-
+    ssh_group_run(';'.join(kill_process_cmds))
   # run provided commands
-  ssh_group.run(cmds, env=env)
+  ssh_group_run(cmds)
 
 
 @task