Increase reserved TPU cloud startup timeout

Author: dlwh

lib/iris/src/iris/cluster/providers/gcp/workers.py

@@ -76,6 +76,7 @@ def _run():
 DEFAULT_BOOT_DISK_SIZE_GB = 50
 # pd-ssd provides ~6000 IOPS vs ~38 on pd-standard, critical for controller DB
 DEFAULT_BOOT_DISK_TYPE = "pd-ssd"
+RESERVED_TPU_CLOUD_READY_TIMEOUT = 7200.0
 
 
 def _gcp_instance_metadata(
@@ -692,8 +693,12 @@ def _run_tpu_bootstrap(
     Phase 2: Poll worker health endpoints until all respond healthy.
     On timeout: query Cloud Logging for [iris-init] entries for diagnostics.
     """
+    queued_resource_ready_timeout = (
+        RESERVED_TPU_CLOUD_READY_TIMEOUT if handle.is_queued_resource else cloud_ready_timeout
+    )
+
     # Single deadline covers Phase 0 (queued resource wait) + Phase 1 (cloud READY).
-    cloud_deadline = Deadline.from_now(Duration.from_seconds(cloud_ready_timeout))
+    cloud_deadline = Deadline.from_now(Duration.from_seconds(queued_resource_ready_timeout))
 
     # Phase 0: If this is a queued resource (reserved TPU), wait for ACTIVE
     # before polling the TPU VM state. The queued resource may sit in QUEUED
@@ -712,7 +717,7 @@ def _run_tpu_bootstrap(
             time.sleep(queued_resource_poll_interval)
         else:
             raise InfraError(
-                f"Queued resource {handle.slice_id} did not become ACTIVE " f"within {cloud_ready_timeout}s"
+                f"Queued resource {handle.slice_id} did not become ACTIVE " f"within {queued_resource_ready_timeout}s"
             )
 
     while not cloud_deadline.expired():
@@ -731,7 +736,7 @@ def _run_tpu_bootstrap(
             )
         time.sleep(poll_interval)
     else:
-        raise InfraError(f"Slice {handle.slice_id} did not reach cloud READY within {cloud_ready_timeout}s")
+        raise InfraError(f"Slice {handle.slice_id} did not reach cloud READY within {queued_resource_ready_timeout}s")
 
     workers = cloud_status.workers
     worker_addrs = [(w.worker_id, w.internal_address) for w in workers]

lib/iris/tests/cluster/providers/gcp/test_platform.py

@@ -18,10 +18,12 @@
 
 from iris.cluster.providers.gcp.controller import GcpControllerProvider
 from iris.cluster.providers.gcp.fake import InMemoryGcpService
-from iris.cluster.providers.gcp.handles import GcpVmSliceHandle, _build_gce_resource_name
+from iris.cluster.providers.gcp.handles import GcpSliceHandle, GcpVmSliceHandle, _build_gce_resource_name
 from iris.cluster.providers.remote_exec import DirectSshRemoteExec, GceRemoteExec, GcloudRemoteExec
 from iris.cluster.providers.gcp.workers import (
     GcpWorkerProvider,
+    RESERVED_TPU_CLOUD_READY_TIMEOUT,
+    _run_tpu_bootstrap,
     _run_vm_slice_bootstrap,
     _validate_slice_config,
 )
@@ -929,6 +931,49 @@ def test_gcp_tpu_slice_os_login_prefers_external_ip_for_direct_ssh():
     assert status.workers[0]._remote_exec.host == "34.1.2.3"
 
 
+# =============================================================================
+# Section 6: TPU Slice Bootstrap Tests
+# =============================================================================
+
+
+class _ImmediateDeadline:
+    def expired(self) -> bool:
+        return True
+
+
+def test_reserved_tpu_bootstrap_uses_extended_cloud_timeout():
+    """Reserved TPU bootstrap uses the longer queued-resource timeout."""
+    gcp_service = InMemoryGcpService(mode=ServiceMode.DRY_RUN, project_id="test-project")
+    handle = GcpSliceHandle(
+        _slice_id="test-reserved-tpu",
+        _zone="us-central2-b",
+        _project_id="test-project",
+        _labels={},
+        _created_at=Timestamp.now(),
+        _label_prefix="iris",
+        _accelerator_variant="v4-32",
+        _gcp_service=gcp_service,
+        _ssh_config=config_pb2.SshConfig(),
+        _bootstrapping=True,
+        _is_queued_resource=True,
+    )
+    worker_config = config_pb2.WorkerConfig(port=10001)
+    seen_deadlines = []
+
+    def _fake_deadline_from_now(duration):
+        seen_deadlines.append(duration.to_seconds())
+        return _ImmediateDeadline()
+
+    with unittest.mock.patch(
+        "iris.cluster.providers.gcp.workers.Deadline.from_now",
+        side_effect=_fake_deadline_from_now,
+    ):
+        with pytest.raises(InfraError, match=rf"within {RESERVED_TPU_CLOUD_READY_TIMEOUT}s"):
+            _run_tpu_bootstrap(gcp_service, "test-project", handle, worker_config)
+
+    assert seen_deadlines == [RESERVED_TPU_CLOUD_READY_TIMEOUT]
+
+
 # =============================================================================
 # Section 6: VM Slice Bootstrap Tests
 #