[iris] Fix VM bootstrap monitor killing healthy workers (#4180)

rjpower · web-flow · commit 079e8cd1f696 · 2026-03-26T12:02:26.000-07:00
- Split the shared 600s deadline in `_run_vm_slice_bootstrap` into
independent phase 1 (cloud READY, 600s) and phase 2 (bootstrap complete,
300s) timeouts
- Add HTTP `/health` probing as the primary bootstrap completion signal
instead of relying solely on GCE serial port text matching ("Bootstrap
complete")
- Serial port polling retained for diagnostics and error detection

## Root cause
A worker VM registered at 17:39:12 and was serving tasks, but the
bootstrap monitor's shared 600s deadline expired at 17:47:33 because the
"Bootstrap complete" line never appeared in GCE serial port output
(likely a flush/buffering issue). This killed a healthy worker and its
running task.

The serial port approach is inherently fragile: the GCE serial port API
uses byte offsets but the poller tracked Python character offsets
(`len(output)`), and the output depends on the metadata_script_runner
flushing before `exit 0`.
diff --git a/lib/iris/src/iris/cluster/providers/gcp/fake.py b/lib/iris/src/iris/cluster/providers/gcp/fake.py
@@ -104,6 +104,9 @@ def __init__(
         self._label_prefix = label_prefix
         self._iris_labels = Labels(label_prefix) if mode == ServiceMode.LOCAL else None
 
+        # Serial port output injection for testing bootstrap monitoring
+        self._serial_port_output: dict[tuple[str, str], str] = {}
+
         # LOCAL mode: track spawned workers per slice for cleanup
         self._local_slices: dict[str, LocalSliceHandle] = {}
 
@@ -321,10 +324,14 @@ def vm_set_metadata(self, name: str, zone: str, metadata: dict[str, str]) -> Non
             raise InfraError(f"VM {name!r} not found in zone {zone!r}")
         vm.metadata.update(metadata)
 
+    def set_serial_port_output(self, name: str, zone: str, output: str) -> None:
+        """Inject serial port output for a VM. Used by tests to simulate GCE serial console."""
+        self._serial_port_output[(name, zone)] = output
+
     def vm_get_serial_port_output(self, name: str, zone: str, start: int = 0) -> str:
         self._check_injected_failure("vm_get_serial_port_output")
-        # DRY_RUN / LOCAL: return empty string (no serial console)
-        return ""
+        full_output = self._serial_port_output.get((name, zone), "")
+        return full_output[start:]
 
     # ========================================================================
     # LOCAL mode: worker spawning
diff --git a/lib/iris/src/iris/cluster/providers/gcp/workers.py b/lib/iris/src/iris/cluster/providers/gcp/workers.py
@@ -640,24 +640,37 @@ def _fetch_bootstrap_logs(project_id: str, handle: GcpSliceHandle) -> None:
         )
 
 
+def _probe_worker_health(address: str, port: int) -> bool:
+    """Probe the worker's HTTP health endpoint. Returns True if healthy."""
+    try:
+        resp = urllib.request.urlopen(f"http://{address}:{port}/health", timeout=5)
+        return resp.status == 200
+    except (urllib.error.URLError, urllib.error.HTTPError, OSError, TimeoutError):
+        return False
+
+
 def _run_vm_slice_bootstrap(
     gcp_service: GcpService,
     handle: GcpVmSliceHandle,
     worker_config: config_pb2.WorkerConfig,
     poll_interval: float = 5.0,
     cloud_ready_timeout: float = 600.0,
+    bootstrap_timeout: float = 300.0,
 ) -> None:
-    """Monitor GCE startup-script bootstrap via serial port output.
+    """Monitor GCE startup-script bootstrap via health probe and serial port.
+
+    Phase 1: Wait for the VM to reach cloud READY with an internal IP.
+    Phase 2: Poll the worker's /health endpoint (primary signal) and serial
+    port output (secondary signal / diagnostics) until bootstrap completes.
 
-    The bootstrap script was baked into VM metadata at creation time, so the
-    VM self-bootstraps on first boot. This method polls serial port output
-    for [iris-init] log lines until the script emits ``Bootstrap complete``
-    or the timeout expires.
+    Each phase has its own independent timeout to prevent a slow phase 1
+    from starving phase 2.
     """
-    deadline = Deadline.from_now(Duration.from_seconds(cloud_ready_timeout))
+    cloud_deadline = Deadline.from_now(Duration.from_seconds(cloud_ready_timeout))
     poll_duration = Duration.from_seconds(poll_interval)
 
-    while not deadline.expired():
+    # Phase 1: wait for cloud READY with an internal IP
+    while not cloud_deadline.expired():
         cloud_status = handle._describe_cloud()
         if cloud_status.state in (CloudSliceState.FAILED, CloudSliceState.DELETING):
             raise InfraError(f"VM slice {handle.slice_id} entered {cloud_status.state} while waiting for cloud READY")
@@ -668,32 +681,45 @@ def _run_vm_slice_bootstrap(
     else:
         raise InfraError(f"VM slice {handle.slice_id} did not reach cloud READY within {cloud_ready_timeout}s")
 
+    worker_address = cloud_status.workers[0].internal_address
+    worker_port = worker_config.port
+
+    # Phase 2: poll health endpoint + serial port with a fresh deadline
+    bootstrap_deadline = Deadline.from_now(Duration.from_seconds(bootstrap_timeout))
     serial_offset = 0
-    bootstrap_complete = False
-    bootstrap_failed = False
 
-    while not deadline.expired():
+    while not bootstrap_deadline.expired():
+        # Primary signal: HTTP health probe
+        if _probe_worker_health(worker_address, worker_port):
+            logger.info("Worker health probe succeeded for VM slice %s", handle.slice_id)
+            break
+
+        # Secondary signal: serial port for diagnostics and error detection
         output = gcp_service.vm_get_serial_port_output(handle._vm_name, handle._zone, start=serial_offset)
         if output:
+            serial_offset += len(output)
+            bootstrap_complete = False
+            bootstrap_errored = False
             for line in output.splitlines():
                 if "[iris-init]" in line:
                     logger.info("[%s serial] %s", handle.slice_id, line.strip())
                 if "Bootstrap complete" in line:
                     bootstrap_complete = True
                 if "[iris-init] ERROR" in line:
-                    bootstrap_failed = True
-
-            serial_offset += len(output)
+                    bootstrap_errored = True
 
-        if bootstrap_complete:
-            break
-        if bootstrap_failed:
-            raise InfraError(f"Startup-script bootstrap failed for VM slice {handle.slice_id} (see serial output above)")
+            if bootstrap_complete:
+                logger.info("Serial port reports bootstrap complete for VM slice %s", handle.slice_id)
+                break
+            if bootstrap_errored:
+                raise InfraError(
+                    f"Startup-script bootstrap failed for VM slice {handle.slice_id} (see serial output above)"
+                )
 
         time.sleep(poll_duration.to_seconds())
     else:
-        raise InfraError(f"VM slice {handle.slice_id} startup-script did not complete within {cloud_ready_timeout}s")
+        raise InfraError(f"VM slice {handle.slice_id} bootstrap did not complete within {bootstrap_timeout}s")
 
-    logger.info("Bootstrap completed for VM slice %s (via startup-script)", handle.slice_id)
+    logger.info("Bootstrap completed for VM slice %s", handle.slice_id)
     with handle._bootstrap_lock:
         handle._bootstrap_state = CloudSliceState.READY
diff --git a/lib/iris/tests/cluster/providers/gcp/test_platform.py b/lib/iris/tests/cluster/providers/gcp/test_platform.py
@@ -17,16 +17,22 @@
 import pytest
 
 from iris.cluster.providers.gcp.fake import InMemoryGcpService
-from iris.cluster.providers.gcp.handles import _build_gce_resource_name
-from iris.cluster.providers.gcp.workers import GcpWorkerProvider, _validate_slice_config
+from iris.cluster.providers.gcp.handles import GcpVmSliceHandle, _build_gce_resource_name
+from iris.cluster.providers.gcp.workers import (
+    GcpWorkerProvider,
+    _run_vm_slice_bootstrap,
+    _validate_slice_config,
+)
 from iris.cluster.providers.manual.provider import ManualControllerProvider, ManualWorkerProvider
 from iris.cluster.providers.types import (
     CloudSliceState,
+    InfraError,
     Labels,
     QuotaExhaustedError,
 )
 from iris.cluster.service_mode import ServiceMode
 from iris.rpc import config_pb2
+from iris.time_utils import Timestamp
 
 # =============================================================================
 # Fixture infrastructure
@@ -633,3 +639,183 @@ def test_gcp_tpu_slice_passes_startup_script_metadata():
     assert "startup-script" in metadata
     assert "[iris-init]" in metadata["startup-script"]
     assert "test-image:latest" in metadata["startup-script"]
+
+
+# =============================================================================
+# Section 6: VM Slice Bootstrap Tests
+#
+# Tests for _run_vm_slice_bootstrap with split timeouts and health probing.
+# =============================================================================
+
+
+def _make_vm_slice_for_bootstrap(
+    gcp_service: InMemoryGcpService,
+    zone: str = "us-central2-b",
+) -> tuple[GcpVmSliceHandle, str]:
+    """Create a VM in InMemoryGcpService and return a handle + vm_name for bootstrap testing."""
+    from iris.cluster.providers.gcp.service import VmCreateRequest
+
+    vm_name = "test-bootstrap-vm"
+    gcp_service.vm_create(
+        VmCreateRequest(
+            name=vm_name,
+            zone=zone,
+            machine_type="n2-standard-4",
+            labels={Labels("iris").iris_slice_id: vm_name},
+        )
+    )
+    handle = GcpVmSliceHandle(
+        _slice_id=vm_name,
+        _vm_name=vm_name,
+        _zone=zone,
+        _project_id="test-project",
+        _gcp_service=gcp_service,
+        _labels={Labels("iris").iris_slice_id: vm_name},
+        _created_at=Timestamp.now(),
+        _label_prefix="iris",
+        _bootstrapping=True,
+    )
+    return handle, vm_name
+
+
+def test_vm_bootstrap_health_probe_succeeds_without_serial_port():
+    """Bootstrap completes when health probe succeeds, even if serial port never shows 'Bootstrap complete'."""
+    gcp_service = InMemoryGcpService(mode=ServiceMode.DRY_RUN, project_id="test-project")
+    handle, _vm_name = _make_vm_slice_for_bootstrap(gcp_service)
+    worker_config = config_pb2.WorkerConfig(port=10001)
+
+    with unittest.mock.patch(
+        "iris.cluster.providers.gcp.workers._probe_worker_health",
+        return_value=True,
+    ):
+        _run_vm_slice_bootstrap(
+            gcp_service,
+            handle,
+            worker_config,
+            poll_interval=0.01,
+            cloud_ready_timeout=5.0,
+            bootstrap_timeout=5.0,
+        )
+
+    assert handle._bootstrap_state == CloudSliceState.READY
+
+
+def test_vm_bootstrap_serial_port_succeeds_without_health_probe():
+    """Bootstrap completes via serial port 'Bootstrap complete' when health probe fails."""
+    gcp_service = InMemoryGcpService(mode=ServiceMode.DRY_RUN, project_id="test-project")
+    handle, vm_name = _make_vm_slice_for_bootstrap(gcp_service)
+    worker_config = config_pb2.WorkerConfig(port=10001)
+
+    gcp_service.set_serial_port_output(
+        vm_name,
+        "us-central2-b",
+        "[iris-init] Starting bootstrap\n[iris-init] Bootstrap complete\n",
+    )
+
+    with unittest.mock.patch(
+        "iris.cluster.providers.gcp.workers._probe_worker_health",
+        return_value=False,
+    ):
+        _run_vm_slice_bootstrap(
+            gcp_service,
+            handle,
+            worker_config,
+            poll_interval=0.01,
+            cloud_ready_timeout=5.0,
+            bootstrap_timeout=5.0,
+        )
+
+    assert handle._bootstrap_state == CloudSliceState.READY
+
+
+def test_vm_bootstrap_serial_port_error_raises():
+    """Bootstrap fails immediately when serial port shows '[iris-init] ERROR'."""
+    gcp_service = InMemoryGcpService(mode=ServiceMode.DRY_RUN, project_id="test-project")
+    handle, vm_name = _make_vm_slice_for_bootstrap(gcp_service)
+    worker_config = config_pb2.WorkerConfig(port=10001)
+
+    gcp_service.set_serial_port_output(
+        vm_name,
+        "us-central2-b",
+        "[iris-init] ERROR: Docker pull failed\n",
+    )
+
+    with unittest.mock.patch(
+        "iris.cluster.providers.gcp.workers._probe_worker_health",
+        return_value=False,
+    ):
+        with pytest.raises(InfraError, match="bootstrap failed"):
+            _run_vm_slice_bootstrap(
+                gcp_service,
+                handle,
+                worker_config,
+                poll_interval=0.01,
+                cloud_ready_timeout=5.0,
+                bootstrap_timeout=5.0,
+            )
+
+
+def test_vm_bootstrap_phase2_has_independent_timeout():
+    """Phase 2 uses its own timeout, not the remainder from phase 1."""
+    gcp_service = InMemoryGcpService(mode=ServiceMode.DRY_RUN, project_id="test-project")
+    handle, _vm_name = _make_vm_slice_for_bootstrap(gcp_service)
+    worker_config = config_pb2.WorkerConfig(port=10001)
+
+    # Health probe never succeeds, serial port never shows complete.
+    # With a very short bootstrap_timeout, this should fail with phase 2 message.
+    with unittest.mock.patch(
+        "iris.cluster.providers.gcp.workers._probe_worker_health",
+        return_value=False,
+    ):
+        with pytest.raises(InfraError, match=r"bootstrap did not complete within 0\.05s"):
+            _run_vm_slice_bootstrap(
+                gcp_service,
+                handle,
+                worker_config,
+                poll_interval=0.01,
+                cloud_ready_timeout=600.0,
+                bootstrap_timeout=0.05,
+            )
+
+
+def test_vm_bootstrap_cloud_not_ready_raises_phase1_timeout():
+    """Phase 1 timeout triggers when VM never reaches READY."""
+    gcp_service = InMemoryGcpService(mode=ServiceMode.DRY_RUN, project_id="test-project")
+
+    # Create a VM but set it to non-READY state
+    from iris.cluster.providers.gcp.service import VmCreateRequest
+
+    vm_name = "test-stuck-vm"
+    gcp_service.vm_create(
+        VmCreateRequest(
+            name=vm_name,
+            zone="us-central2-b",
+            machine_type="n2-standard-4",
+            labels={Labels("iris").iris_slice_id: vm_name},
+        )
+    )
+    # Set VM to STAGING so it never reaches READY
+    gcp_service._vms[(vm_name, "us-central2-b")].status = "STAGING"
+
+    handle = GcpVmSliceHandle(
+        _slice_id=vm_name,
+        _vm_name=vm_name,
+        _zone="us-central2-b",
+        _project_id="test-project",
+        _gcp_service=gcp_service,
+        _labels={Labels("iris").iris_slice_id: vm_name},
+        _created_at=Timestamp.now(),
+        _label_prefix="iris",
+        _bootstrapping=True,
+    )
+    worker_config = config_pb2.WorkerConfig(port=10001)
+
+    with pytest.raises(InfraError, match=r"did not reach cloud READY within 0\.05s"):
+        _run_vm_slice_bootstrap(
+            gcp_service,
+            handle,
+            worker_config,
+            poll_interval=0.01,
+            cloud_ready_timeout=0.05,
+            bootstrap_timeout=300.0,
+        )
