Adding an extra slice to the Pathways cluster to swap in when there is a slice failure

Author: lukebaumann

axlearn/cloud/gcp/pathways_utils.py

@@ -356,7 +356,7 @@ def _build_pathways_head_sidecar_containers(self) -> list[Nested[Any]]:
                 args=[
                     f"--server_port={_PATHWAYS_RESOURCE_MANAGER_PORT}",
                     "--node_type=resource_manager",
-                    f"--instance_count={pathways_instance_count}",
+                    f"--instance_count={pathways_instance_count + 1}",
                     f"--instance_type={pathways_tpu_version}:{system.topology}",
                     f"--gcs_scratch_location={staging_location}",
                     "--alsologtostderr",
@@ -626,7 +626,7 @@ def __call__(self) -> Sequence[Nested[Any]]:
             ),
             dict(
                 name=_PATHWAYS_WORKER_REPLICATED_JOB_NAME,
-                replicas=cfg.accelerator.num_replicas,
+                replicas=cfg.accelerator.num_replicas + 1,
                 template=self._build_pathways_worker_job(),
             ),
         ]