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

Author: lukebaumann

axlearn/cloud/gcp/pathways_utils.py

@@ -324,6 +324,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",
             "--vmodule=grpc_host_buffer=3,rpc_helper=3,host_buffer=3,ifrt_backend=3,grpc_service_impl=3",
         ]
         cmd_args.extend(xla_flags_from_options(self._xla_options).split())
@@ -361,6 +362,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",
                 ],
             ),
         ]
@@ -515,6 +517,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)