Adding a pipe break missing flag that makes the cluster fail quickly

Author: lukebaumann

axlearn/cloud/gcp/pathways_utils.py

@@ -322,6 +322,7 @@ def _build_pathways_head_sidecar_containers(self) -> list[Nested[Any]]:
             f"--gcs_scratch_location={staging_location}",
             # This should be made configurable
             f"--num_elastic_slices={cfg.accelerator.num_replicas}",
+            "--temporary_flags_for_debugging=temporary_flag_for_debugging_pipe_break_on_missing_keepalive=true",
         ]
         cmd_args.extend(xla_flags_from_options(self._xla_options).split())
 
@@ -358,6 +359,7 @@ def _build_pathways_head_sidecar_containers(self) -> list[Nested[Any]]:
                     f"--instance_count={pathways_instance_count}",
                     f"--instance_type={pathways_tpu_version}:{system.topology}",
                     f"--gcs_scratch_location={staging_location}",
+                    "--temporary_flags_for_debugging=temporary_flag_for_debugging_pipe_break_on_missing_keepalive=true",
                 ],
             ),
         ]
@@ -512,6 +514,7 @@ def _build_pathways_worker_container(
             f"--resource_manager_address={pathways_head_address}:"
             + f"{_PATHWAYS_RESOURCE_MANAGER_PORT}",
             f"--gcs_scratch_location={cfg.output_dir}/pathways-staging",
+            "--temporary_flags_for_debugging=temporary_flag_for_debugging_pipe_break_on_missing_keepalive=true",
         ]
         mega_scale_args = xla_flags_from_options(self._mxla_options).split()
         worker_container["args"].extend(mega_scale_args)