[Router] Drain readiness before SIGTERM shutdown so k8s deregisters the pod first (#39016)

Co-authored-by: Kangyan Zhou <kangyan.zhou@radixark.ai>
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>

experimental/sgl-router/tests/e2e/k8s_integration/conftest.py

@@ -109,40 +109,30 @@ def _wait_for_replacement_pod_ready(
     last_observed = "no pods"
 
     while time.time() < deadline:
-        result = _kubectl(
-            "get",
-            "pods",
-            "-n",
-            namespace,
-            "-l",
-            selector,
-            "-o",
-            "json",
-            check=False,
-        )
-        if getattr(result, "returncode", 0) == 0:
-            pods = json.loads(result.stdout or "{}").get("items", [])
-            names = [pod.get("metadata", {}).get("name", "") for pod in pods]
-            last_observed = ", ".join(filter(None, names)) or "no pods"
+        pods = _pods(selector, namespace, check=False)
+        names = [pod.get("metadata", {}).get("name", "") for pod in pods]
+        # An empty list is "nothing observed" whether the pods are gone or the
+        # kubectl call failed; both read the same in a timeout message.
+        last_observed = ", ".join(filter(None, names)) or "no pods"
 
-            if old_pod not in names:
-                for pod in sorted(
-                    pods, key=lambda item: item.get("metadata", {}).get("name", "")
+        if old_pod not in names:
+            for pod in sorted(
+                pods, key=lambda item: item.get("metadata", {}).get("name", "")
+            ):
+                metadata = pod.get("metadata", {})
+                status = pod.get("status", {})
+                ready = any(
+                    condition.get("type") == "Ready"
+                    and condition.get("status") == "True"
+                    for condition in status.get("conditions", [])
+                )
+                if (
+                    metadata.get("name") != old_pod
+                    and _is_live(pod)
+                    and status.get("phase") == "Running"
+                    and ready
                 ):
-                    metadata = pod.get("metadata", {})
-                    status = pod.get("status", {})
-                    ready = any(
-                        condition.get("type") == "Ready"
-                        and condition.get("status") == "True"
-                        for condition in status.get("conditions", [])
-                    )
-                    if (
-                        metadata.get("name") != old_pod
-                        and not metadata.get("deletionTimestamp")
-                        and status.get("phase") == "Running"
-                        and ready
-                    ):
-                        return metadata["name"]
+                    return metadata["name"]
 
         time.sleep(interval)
 
@@ -172,14 +162,20 @@ def _port_forward_start(
     service: str,
     local_port: int,
     remote_port: int,
+    resource: str = "svc",
 ) -> subprocess.Popen:
-    """Start kubectl port-forward and wait until the port is reachable."""
+    """Start kubectl port-forward and wait until the port is reachable.
+
+    `resource="pod"` binds one specific pod instead of the Service. A draining
+    pod is removed from the Service's ready endpoints, so a test that needs to
+    keep talking to it through the drain must address the pod directly.
+    """
     cmd = [
         "kubectl",
         "--context",
         KUBECTL_CONTEXT,
         "port-forward",
-        f"svc/{service}",
+        f"{resource}/{service}",
         f"{local_port}:{remote_port}",
         "-n",
         namespace,
@@ -215,6 +211,66 @@ def _cleanup_port_forward(name: str, pf: subprocess.Popen) -> None:
         logger.debug("Port-forward %s exited cleanly (rc=%s)", name, rc)
 
 
+def _pod_json(pod: str, namespace: str = NAMESPACE) -> dict:
+    """One pod's full object. The `or "{}"` mirrors
+    `_wait_for_replacement_pod_ready`: kubectl can hand back empty stdout, and a
+    JSONDecodeError there says nothing about what went wrong."""
+    result = _kubectl("get", "pod", pod, "-n", namespace, "-o", "json")
+    return json.loads(result.stdout or "{}")
+
+
+def _pods(
+    selector: str,
+    namespace: str = NAMESPACE,
+    check: bool = True,
+) -> list[dict]:
+    """Pod objects matching `selector`. `check=False` yields `[]` on a failed
+    kubectl instead of raising, for poll loops that expect the API server to be
+    briefly unavailable mid-rollout. The `or "{}"` guards kubectl handing back
+    empty stdout, where a JSONDecodeError would say nothing about what went
+    wrong."""
+    result = _kubectl(
+        "get", "pods", "-n", namespace, "-l", selector, "-o", "json", check=check
+    )
+    if getattr(result, "returncode", 0) != 0:
+        return []
+    return json.loads(result.stdout or "{}").get("items", [])
+
+
+def _is_live(pod: dict) -> bool:
+    """Whether a pod object is not already terminating. One predicate rather
+    than two copies of `deletionTimestamp`, so the replacement-pod poll and
+    `_pod_names` cannot drift apart on what counts as gone."""
+    return not pod.get("metadata", {}).get("deletionTimestamp")
+
+
+def _pod_names(selector: str, namespace: str = NAMESPACE) -> list[str]:
+    """Names of pods matching `selector`, excluding any already terminating."""
+    return [p["metadata"]["name"] for p in _pods(selector, namespace) if _is_live(p)]
+
+
+def _container_restart_count(
+    pod: str,
+    container: str,
+    namespace: str = NAMESPACE,
+) -> int:
+    """`restartCount` for one container — how a test observes that the process
+    exited and kubelet restarted it in place (no new pod, same name)."""
+    statuses = _pod_json(pod, namespace).get("status", {}).get("containerStatuses", [])
+    for status in statuses:
+        if status["name"] == container:
+            return int(status["restartCount"])
+    raise AssertionError(f"container {container!r} not found on pod {pod!r}")
+
+
+def _pod_ready_condition(pod: str, namespace: str = NAMESPACE) -> str:
+    """The pod's `Ready` condition as k8s currently sees it ("True"/"False")."""
+    for cond in _pod_json(pod, namespace).get("status", {}).get("conditions", []):
+        if cond["type"] == "Ready":
+            return cond["status"]
+    return "Unknown"
+
+
 def _poll_until(
     predicate,
     description: str,