[SMG][CI] Add K8s integration tests + wire into pr-test-rust (#24278)
Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
Kangyan-Zhou
co-authored by
Claude Opus 4.7
Claude Opus 4.7
parent
e5c58eb9d6
commit
c545a5b1c0
@@ -349,8 +349,86 @@ jobs:
|
||||
cache-from: type=gha
|
||||
cache-to: type=gha,mode=max
|
||||
|
||||
k8s-integration:
|
||||
# Runs SMG against a kind cluster with fake worker pods to exercise the
|
||||
# K8s service discovery / reconciliation path. No GPU required (workers
|
||||
# are python:3.12-slim mocks); the h100 matrix runners are unsuitable
|
||||
# because they're containers without a Docker daemon.
|
||||
if: |
|
||||
github.event_name != 'pull_request' ||
|
||||
(github.event.action != 'labeled' && contains(github.event.pull_request.labels.*.name, 'run-ci')) ||
|
||||
(github.event.action == 'labeled' && github.event.label.name == 'run-ci')
|
||||
runs-on: ubuntu-22.04
|
||||
timeout-minutes: 30
|
||||
steps:
|
||||
- name: Checkout code
|
||||
uses: actions/checkout@v4
|
||||
|
||||
- name: Set up Python
|
||||
uses: actions/setup-python@v5
|
||||
with:
|
||||
python-version: "3.12"
|
||||
|
||||
- name: Install kind and kubectl
|
||||
run: |
|
||||
curl -fsSLo /tmp/kind https://kind.sigs.k8s.io/dl/v0.24.0/kind-linux-amd64
|
||||
chmod +x /tmp/kind && sudo mv /tmp/kind /usr/local/bin/kind
|
||||
KUBECTL_VERSION=$(curl -fsSL https://dl.k8s.io/release/stable.txt)
|
||||
curl -fsSLo /tmp/kubectl "https://dl.k8s.io/release/${KUBECTL_VERSION}/bin/linux/amd64/kubectl"
|
||||
chmod +x /tmp/kubectl && sudo mv /tmp/kubectl /usr/local/bin/kubectl
|
||||
kind --version
|
||||
kubectl version --client
|
||||
|
||||
- name: Set up Docker Buildx
|
||||
uses: docker/setup-buildx-action@v3
|
||||
|
||||
- name: Build smg-gateway:test image
|
||||
uses: docker/build-push-action@v5
|
||||
with:
|
||||
context: sgl-model-gateway
|
||||
file: sgl-model-gateway/e2e_test/k8s_integration/Dockerfile.gateway
|
||||
tags: smg-gateway:test
|
||||
load: true
|
||||
cache-from: type=gha,scope=k8s-integration
|
||||
cache-to: type=gha,scope=k8s-integration,mode=max
|
||||
|
||||
- name: Install Python test dependencies
|
||||
run: |
|
||||
python3 -m pip install --upgrade pip
|
||||
python3 -m pip install pytest httpx
|
||||
|
||||
- name: Set up kind cluster and deploy
|
||||
env:
|
||||
SKIP_DOCKER_BUILD: "1"
|
||||
run: |
|
||||
cd sgl-model-gateway
|
||||
bash e2e_test/k8s_integration/setup.sh
|
||||
|
||||
- name: Run K8s integration tests
|
||||
run: |
|
||||
cd sgl-model-gateway
|
||||
# confcutdir avoids loading the parent e2e_test/conftest.py, which
|
||||
# pulls in heavy infra deps (requests, sglang_router) that this job
|
||||
# intentionally doesn't install.
|
||||
pytest e2e_test/k8s_integration/ \
|
||||
--confcutdir=e2e_test/k8s_integration \
|
||||
-v -s -o log_cli=true --log-cli-level=INFO
|
||||
|
||||
- name: Dump cluster state on failure
|
||||
if: failure()
|
||||
run: |
|
||||
kubectl --context kind-smg-test get all -A || true
|
||||
kubectl --context kind-smg-test -n smg-test describe pods || true
|
||||
kubectl --context kind-smg-test -n smg-test logs deploy/smg-gateway --tail=200 || true
|
||||
|
||||
- name: Tear down kind cluster
|
||||
if: always()
|
||||
run: |
|
||||
cd sgl-model-gateway
|
||||
bash e2e_test/k8s_integration/setup.sh teardown || true
|
||||
|
||||
finish:
|
||||
needs: [build-wheel, python-unit-tests, unit-tests, gateway-e2e, docker-build-test]
|
||||
needs: [build-wheel, python-unit-tests, unit-tests, gateway-e2e, docker-build-test, k8s-integration]
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- name: Finish
|
||||
|
||||
@@ -125,8 +125,13 @@ _SRC = _ROOT / "bindings" / "python"
|
||||
if str(_E2E_TEST) not in sys.path:
|
||||
sys.path.insert(0, str(_E2E_TEST))
|
||||
|
||||
# Add bindings/python to path if the wheel is not installed (for local development)
|
||||
_wheel_installed = find_spec("sglang_router.sglang_router_rs") is not None
|
||||
# Add bindings/python to path if the wheel is not installed (for local development).
|
||||
# find_spec raises ModuleNotFoundError when the parent package itself is absent,
|
||||
# which is the case in CI jobs that don't install the sglang_router wheel.
|
||||
try:
|
||||
_wheel_installed = find_spec("sglang_router.sglang_router_rs") is not None
|
||||
except ModuleNotFoundError:
|
||||
_wheel_installed = False
|
||||
|
||||
if not _wheel_installed and str(_SRC) not in sys.path:
|
||||
sys.path.insert(0, str(_SRC))
|
||||
|
||||
@@ -0,0 +1,31 @@
|
||||
# syntax=docker/dockerfile:1.6
|
||||
# Lightweight Dockerfile for integration testing.
|
||||
# Builds the smg binary directly (no Python/maturin/wheel overhead).
|
||||
# Uses the "ci" cargo profile (thin LTO, 16 codegen units) for fast builds.
|
||||
#
|
||||
# The repo's docker/gateway.Dockerfile builds a Python wheel via maturin for
|
||||
# production. This Dockerfile builds just the Rust binary in ~5 min on a
|
||||
# warm cache.
|
||||
|
||||
FROM rust:1.90-bookworm AS builder
|
||||
|
||||
RUN apt-get update && apt-get install -y \
|
||||
libssl-dev pkg-config protobuf-compiler cmake \
|
||||
&& rm -rf /var/lib/apt/lists/*
|
||||
|
||||
WORKDIR /build
|
||||
COPY . .
|
||||
|
||||
RUN --mount=type=cache,target=/usr/local/cargo/registry \
|
||||
--mount=type=cache,target=/usr/local/cargo/git \
|
||||
--mount=type=cache,target=/build/target \
|
||||
cargo build --profile ci --bin smg --features vendored-openssl \
|
||||
&& cp target/ci/smg /usr/local/bin/smg
|
||||
|
||||
FROM debian:bookworm-slim
|
||||
|
||||
RUN apt-get update && apt-get install -y ca-certificates && rm -rf /var/lib/apt/lists/*
|
||||
|
||||
COPY --from=builder /usr/local/bin/smg /usr/local/bin/smg
|
||||
|
||||
ENTRYPOINT ["smg"]
|
||||
@@ -0,0 +1,343 @@
|
||||
"""Pytest configuration for K8s integration tests.
|
||||
|
||||
These tests require:
|
||||
- A kind cluster named 'smg-test'
|
||||
- The smg-gateway:test image loaded into kind
|
||||
- kubectl configured to use the kind-smg-test context
|
||||
|
||||
Setup: ./e2e_test/k8s_integration/setup.sh
|
||||
Teardown: ./e2e_test/k8s_integration/setup.sh teardown
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
import socket
|
||||
import subprocess
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Callable
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
NAMESPACE = "smg-test"
|
||||
MANIFESTS_DIR = Path(__file__).parent / "manifests"
|
||||
FAKE_WORKER_SCRIPT = Path(__file__).parent / "fake_worker.py"
|
||||
KUBECTL_CONTEXT = "kind-smg-test"
|
||||
|
||||
# Reconciliation interval matches ServiceDiscoveryConfig.check_interval in
|
||||
# sgl-model-gateway/src/service_discovery.rs (currently 60s).
|
||||
RECONCILIATION_INTERVAL_SECS = 60
|
||||
RECONCILIATION_WAIT_SECS = RECONCILIATION_INTERVAL_SECS + 30
|
||||
|
||||
# Errors safe to retry while polling: connection-level failures only.
|
||||
# httpx.HTTPStatusError (4xx/5xx) is intentionally NOT included — a gateway
|
||||
# returning 5xx is the kind of regression these tests should surface, not
|
||||
# silently swallow as "transient".
|
||||
_TRANSIENT_ERRORS = (
|
||||
httpx.TransportError,
|
||||
httpx.TimeoutException,
|
||||
ConnectionError,
|
||||
OSError,
|
||||
)
|
||||
|
||||
|
||||
def pytest_configure(config):
|
||||
config.addinivalue_line(
|
||||
"markers",
|
||||
"slow: marks tests that wait for multiple reconciliation cycles "
|
||||
"(deselect with '-m \"not slow\"')",
|
||||
)
|
||||
|
||||
|
||||
def _kubectl(
|
||||
*args: str, check: bool = True, capture: bool = True
|
||||
) -> subprocess.CompletedProcess:
|
||||
cmd = ["kubectl", "--context", KUBECTL_CONTEXT, *args]
|
||||
logger.debug("Running: %s", " ".join(cmd))
|
||||
return subprocess.run(cmd, capture_output=capture, text=True, check=check)
|
||||
|
||||
|
||||
def _kubectl_json(*args: str) -> dict:
|
||||
result = _kubectl(*args, "-o", "json")
|
||||
try:
|
||||
return json.loads(result.stdout)
|
||||
except json.JSONDecodeError as e:
|
||||
raise RuntimeError(
|
||||
f"Failed to parse kubectl JSON output for args={args!r}. "
|
||||
f"stdout={result.stdout!r}, stderr={result.stderr!r}"
|
||||
) from e
|
||||
|
||||
|
||||
def _wait_for_pod_ready(name: str, namespace: str = NAMESPACE, timeout: int = 120):
|
||||
"""Wait until a pod is Ready."""
|
||||
logger.info("Waiting for pod %s to be ready (timeout=%ds)", name, timeout)
|
||||
_kubectl(
|
||||
"wait",
|
||||
"--for=condition=Ready",
|
||||
f"pod/{name}",
|
||||
"-n",
|
||||
namespace,
|
||||
f"--timeout={timeout}s",
|
||||
)
|
||||
|
||||
|
||||
def _wait_for_deployment_ready(
|
||||
name: str, namespace: str = NAMESPACE, timeout: int = 180
|
||||
):
|
||||
"""Wait until a deployment has all replicas available."""
|
||||
logger.info("Waiting for deployment %s to be ready (timeout=%ds)", name, timeout)
|
||||
_kubectl(
|
||||
"rollout",
|
||||
"status",
|
||||
f"deployment/{name}",
|
||||
"-n",
|
||||
namespace,
|
||||
f"--timeout={timeout}s",
|
||||
)
|
||||
|
||||
|
||||
def _get_gateway_url() -> str:
|
||||
"""Return the gateway URL, assuming port-forward is active on localhost:30000."""
|
||||
return "http://127.0.0.1:30000"
|
||||
|
||||
|
||||
def _get_metrics_url() -> str:
|
||||
"""Return the metrics URL, assuming port-forward is active on localhost:29000."""
|
||||
return "http://127.0.0.1:29000"
|
||||
|
||||
|
||||
def _wait_for_port(port: int, proc: subprocess.Popen, timeout: int = 15):
|
||||
"""Poll until a TCP connection to localhost:port succeeds."""
|
||||
deadline = time.time() + timeout
|
||||
while time.time() < deadline:
|
||||
if proc.poll() is not None:
|
||||
stderr = proc.stderr.read().decode() if proc.stderr else ""
|
||||
raise RuntimeError(f"port-forward process exited early: {stderr}")
|
||||
try:
|
||||
with socket.create_connection(("127.0.0.1", port), timeout=1):
|
||||
return
|
||||
except OSError:
|
||||
time.sleep(0.5)
|
||||
raise TimeoutError(f"Port {port} not ready after {timeout}s")
|
||||
|
||||
|
||||
def _get_workers(gateway_url: str) -> dict:
|
||||
"""GET /workers from the gateway and validate the response shape."""
|
||||
resp = httpx.get(f"{gateway_url}/workers", timeout=10)
|
||||
resp.raise_for_status()
|
||||
data = resp.json()
|
||||
if not isinstance(data, dict) or "total" not in data:
|
||||
raise ValueError(
|
||||
f"/workers returned unexpected structure (missing 'total' key): "
|
||||
f"{json.dumps(data)[:200]}"
|
||||
)
|
||||
return data
|
||||
|
||||
|
||||
def _get_worker_count(gateway_url: str) -> int:
|
||||
"""Return total worker count from /workers."""
|
||||
return _get_workers(gateway_url)["total"]
|
||||
|
||||
|
||||
def _poll_until(
|
||||
predicate: Callable[[], bool],
|
||||
description: str,
|
||||
timeout: int,
|
||||
interval: float = 5,
|
||||
) -> bool:
|
||||
"""Poll a predicate until it returns True, or raise TimeoutError.
|
||||
|
||||
Only transient network errors (see _TRANSIENT_ERRORS) are retried.
|
||||
Programming errors (KeyError, TypeError, etc.) and HTTP status errors
|
||||
(httpx.HTTPStatusError on 4xx/5xx) propagate immediately so real bugs
|
||||
aren't masked as "still polling".
|
||||
"""
|
||||
deadline = time.time() + timeout
|
||||
last_error: Exception | None = None
|
||||
attempts = 0
|
||||
while time.time() < deadline:
|
||||
try:
|
||||
attempts += 1
|
||||
if predicate():
|
||||
logger.info(
|
||||
"Condition met: %s (after %d attempts)", description, attempts
|
||||
)
|
||||
return True
|
||||
except _TRANSIENT_ERRORS as e:
|
||||
last_error = e
|
||||
logger.debug("Transient error on attempt %d: %s", attempts, e)
|
||||
time.sleep(interval)
|
||||
msg = f"Timeout waiting for: {description} (after {timeout}s, {attempts} attempts)"
|
||||
if last_error:
|
||||
msg += f" — last error: {last_error}"
|
||||
raise TimeoutError(msg)
|
||||
|
||||
|
||||
def _port_forward_start(
|
||||
namespace: str, service: str, local_port: int, remote_port: int
|
||||
) -> subprocess.Popen:
|
||||
"""Start kubectl port-forward and verify the port is reachable."""
|
||||
cmd = [
|
||||
"kubectl",
|
||||
"--context",
|
||||
KUBECTL_CONTEXT,
|
||||
"port-forward",
|
||||
f"svc/{service}",
|
||||
f"{local_port}:{remote_port}",
|
||||
"-n",
|
||||
namespace,
|
||||
]
|
||||
logger.info("Starting port-forward: %s", " ".join(cmd))
|
||||
proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
|
||||
_wait_for_port(local_port, proc)
|
||||
return proc
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def k8s_cluster():
|
||||
"""Ensure the kind cluster exists and is reachable."""
|
||||
result = subprocess.run(
|
||||
["kind", "get", "clusters"],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
check=True,
|
||||
)
|
||||
if "smg-test" not in result.stdout:
|
||||
pytest.skip("kind cluster 'smg-test' not found — run setup first")
|
||||
|
||||
# Verify kubectl connectivity
|
||||
_kubectl("cluster-info")
|
||||
return True
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def deploy_base(k8s_cluster):
|
||||
"""Ensure namespace, RBAC, configmap, and gateway are deployed.
|
||||
|
||||
Does NOT tear down at session end — use setup.sh teardown for that.
|
||||
This allows running pytest multiple times without full re-setup.
|
||||
"""
|
||||
# Create namespace (apply is idempotent — succeeds if already exists)
|
||||
_kubectl("apply", "-f", str(MANIFESTS_DIR / "namespace.yaml"))
|
||||
|
||||
# Create/update the fake-worker script as a ConfigMap
|
||||
cm_result = _kubectl(
|
||||
"create",
|
||||
"configmap",
|
||||
"fake-worker-script",
|
||||
f"--from-file=fake_worker.py={FAKE_WORKER_SCRIPT}",
|
||||
"-n",
|
||||
NAMESPACE,
|
||||
"--dry-run=client",
|
||||
"-o",
|
||||
"yaml",
|
||||
)
|
||||
_apply_from_stdin(cm_result.stdout)
|
||||
|
||||
# Apply RBAC
|
||||
_kubectl("apply", "-f", str(MANIFESTS_DIR / "rbac.yaml"))
|
||||
|
||||
# Apply gateway deployment
|
||||
_kubectl("apply", "-f", str(MANIFESTS_DIR / "gateway.yaml"))
|
||||
|
||||
# Wait for gateway to be ready
|
||||
_wait_for_deployment_ready("smg-gateway")
|
||||
|
||||
# Clean up any residual test pods from previous runs
|
||||
result = _kubectl(
|
||||
"get",
|
||||
"pods",
|
||||
"-n",
|
||||
NAMESPACE,
|
||||
"-l",
|
||||
"app=fake-worker",
|
||||
"-o",
|
||||
"jsonpath={.items[*].metadata.name}",
|
||||
check=False,
|
||||
)
|
||||
if result.returncode != 0:
|
||||
# Listing failed during fixture setup — surface it instead of silently
|
||||
# leaving stale workers around, which would skew worker-count assertions
|
||||
# in subsequent tests.
|
||||
logger.warning(
|
||||
"Failed to list residual fake-worker pods (rc=%d): %s",
|
||||
result.returncode,
|
||||
result.stderr.strip(),
|
||||
)
|
||||
elif result.stdout.strip():
|
||||
for pod_name in result.stdout.strip().split():
|
||||
_kubectl(
|
||||
"delete",
|
||||
"pod",
|
||||
pod_name,
|
||||
"-n",
|
||||
NAMESPACE,
|
||||
"--force",
|
||||
"--grace-period=0",
|
||||
"--ignore-not-found",
|
||||
)
|
||||
# Wait a bit for cleanup
|
||||
time.sleep(5)
|
||||
|
||||
yield
|
||||
|
||||
|
||||
def _apply_from_stdin(yaml_content: str):
|
||||
"""Apply a YAML manifest from stdin."""
|
||||
proc = subprocess.run(
|
||||
["kubectl", "--context", KUBECTL_CONTEXT, "apply", "-f", "-"],
|
||||
input=yaml_content,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
check=True,
|
||||
)
|
||||
return proc
|
||||
|
||||
|
||||
def _cleanup_port_forward(name: str, pf: subprocess.Popen):
|
||||
"""Terminate a port-forward process; always log final exit state."""
|
||||
try:
|
||||
pf.terminate()
|
||||
pf.wait(timeout=10)
|
||||
except subprocess.TimeoutExpired:
|
||||
logger.warning(
|
||||
"Port-forward %s did not exit on SIGTERM after 10s; killing", name
|
||||
)
|
||||
pf.kill()
|
||||
try:
|
||||
pf.wait(timeout=5)
|
||||
except subprocess.TimeoutExpired:
|
||||
logger.warning("Port-forward %s still running after SIGKILL", name)
|
||||
except Exception as e:
|
||||
logger.warning("Error cleaning up %s port-forward: %s", name, e)
|
||||
|
||||
rc = pf.returncode
|
||||
stderr = pf.stderr.read().decode() if pf.stderr else ""
|
||||
# rc == -15 (-SIGTERM) is the clean shutdown case; anything else is worth
|
||||
# surfacing — including rc is None, which means terminate failed silently
|
||||
# and the process may still be alive.
|
||||
if rc != -15:
|
||||
suffix = f": {stderr.strip()}" if stderr.strip() else ""
|
||||
logger.warning("Port-forward %s exited rc=%s%s", name, rc, suffix)
|
||||
else:
|
||||
logger.debug("Port-forward %s exited cleanly (rc=%s)", name, rc)
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def gateway_port_forward(deploy_base):
|
||||
"""Set up port-forwarding to the gateway service."""
|
||||
pf_http = _port_forward_start(NAMESPACE, "smg-gateway", 30000, 30000)
|
||||
try:
|
||||
pf_metrics = _port_forward_start(NAMESPACE, "smg-gateway", 29000, 29000)
|
||||
except Exception:
|
||||
pf_http.terminate()
|
||||
pf_http.wait()
|
||||
raise
|
||||
yield _get_gateway_url(), _get_metrics_url()
|
||||
_cleanup_port_forward("http", pf_http)
|
||||
_cleanup_port_forward("metrics", pf_metrics)
|
||||
@@ -0,0 +1,81 @@
|
||||
"""Minimal fake worker that mimics an SGLang worker for integration testing.
|
||||
|
||||
Responds to:
|
||||
GET /health -> 200 OK
|
||||
GET /v1/models -> {"data": [{"id": "fake-model", "owned_by": "sglang"}]}
|
||||
GET /server_info, /get_server_info -> {"model_path": ..., "version": ..., "tp_size": ..., "dp_size": ...}
|
||||
GET /model_info, /get_model_info -> {"model_path": ..., "is_generation": true}
|
||||
"""
|
||||
|
||||
import json
|
||||
from http.server import BaseHTTPRequestHandler, HTTPServer
|
||||
|
||||
PORT = 8000
|
||||
|
||||
|
||||
class FakeWorkerHandler(BaseHTTPRequestHandler):
|
||||
def do_GET(self):
|
||||
if self.path == "/health":
|
||||
self.send_response(200)
|
||||
self.send_header("Content-Type", "text/plain")
|
||||
self.end_headers()
|
||||
self.wfile.write(b"OK")
|
||||
|
||||
elif self.path == "/v1/models":
|
||||
body = json.dumps(
|
||||
{
|
||||
"object": "list",
|
||||
"data": [
|
||||
{
|
||||
"id": "fake-model",
|
||||
"object": "model",
|
||||
"created": 0,
|
||||
"owned_by": "sglang",
|
||||
}
|
||||
],
|
||||
}
|
||||
)
|
||||
self.send_response(200)
|
||||
self.send_header("Content-Type", "application/json")
|
||||
self.end_headers()
|
||||
self.wfile.write(body.encode())
|
||||
|
||||
elif self.path in ("/server_info", "/get_server_info"):
|
||||
body = json.dumps(
|
||||
{
|
||||
"model_path": "fake-model",
|
||||
"version": "0.0.0-test",
|
||||
"tp_size": 1,
|
||||
"dp_size": 1,
|
||||
}
|
||||
)
|
||||
self.send_response(200)
|
||||
self.send_header("Content-Type", "application/json")
|
||||
self.end_headers()
|
||||
self.wfile.write(body.encode())
|
||||
|
||||
elif self.path in ("/model_info", "/get_model_info"):
|
||||
body = json.dumps(
|
||||
{
|
||||
"model_path": "fake-model",
|
||||
"is_generation": True,
|
||||
}
|
||||
)
|
||||
self.send_response(200)
|
||||
self.send_header("Content-Type", "application/json")
|
||||
self.end_headers()
|
||||
self.wfile.write(body.encode())
|
||||
|
||||
else:
|
||||
self.send_response(404)
|
||||
self.end_headers()
|
||||
|
||||
def log_message(self, format, *args):
|
||||
# Suppress per-request logs to keep test output clean
|
||||
pass
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
server = HTTPServer(("0.0.0.0", PORT), FakeWorkerHandler)
|
||||
print(f"Fake worker listening on port {PORT}", flush=True)
|
||||
server.serve_forever()
|
||||
@@ -0,0 +1,74 @@
|
||||
apiVersion: apps/v1
|
||||
kind: Deployment
|
||||
metadata:
|
||||
name: smg-gateway-pd
|
||||
namespace: smg-test
|
||||
spec:
|
||||
replicas: 1
|
||||
selector:
|
||||
matchLabels:
|
||||
app: smg-gateway-pd
|
||||
template:
|
||||
metadata:
|
||||
labels:
|
||||
app: smg-gateway-pd
|
||||
spec:
|
||||
serviceAccountName: smg-gateway
|
||||
containers:
|
||||
- name: gateway
|
||||
image: smg-gateway:test
|
||||
imagePullPolicy: Never
|
||||
args:
|
||||
- "--service-discovery"
|
||||
- "--pd-disaggregation"
|
||||
- "--prefill-selector"
|
||||
- "role=prefill"
|
||||
- "--decode-selector"
|
||||
- "role=decode"
|
||||
- "--service-discovery-port"
|
||||
- "8000"
|
||||
- "--service-discovery-namespace"
|
||||
- "smg-test"
|
||||
- "--port"
|
||||
- "30001"
|
||||
- "--prometheus-port"
|
||||
- "29001"
|
||||
- "--disable-health-check"
|
||||
- "--worker-startup-timeout-secs"
|
||||
- "30"
|
||||
- "--log-level"
|
||||
- "debug"
|
||||
ports:
|
||||
- containerPort: 30001
|
||||
name: http
|
||||
- containerPort: 29001
|
||||
name: metrics
|
||||
readinessProbe:
|
||||
httpGet:
|
||||
path: /liveness
|
||||
port: 30001
|
||||
initialDelaySeconds: 3
|
||||
periodSeconds: 3
|
||||
livenessProbe:
|
||||
httpGet:
|
||||
path: /liveness
|
||||
port: 30001
|
||||
initialDelaySeconds: 5
|
||||
periodSeconds: 10
|
||||
---
|
||||
apiVersion: v1
|
||||
kind: Service
|
||||
metadata:
|
||||
name: smg-gateway-pd
|
||||
namespace: smg-test
|
||||
spec:
|
||||
type: NodePort
|
||||
selector:
|
||||
app: smg-gateway-pd
|
||||
ports:
|
||||
- name: http
|
||||
port: 30001
|
||||
targetPort: 30001
|
||||
- name: metrics
|
||||
port: 29001
|
||||
targetPort: 29001
|
||||
@@ -0,0 +1,74 @@
|
||||
apiVersion: apps/v1
|
||||
kind: Deployment
|
||||
metadata:
|
||||
name: smg-gateway
|
||||
namespace: smg-test
|
||||
spec:
|
||||
replicas: 1
|
||||
selector:
|
||||
matchLabels:
|
||||
app: smg-gateway
|
||||
template:
|
||||
metadata:
|
||||
labels:
|
||||
app: smg-gateway
|
||||
spec:
|
||||
serviceAccountName: smg-gateway
|
||||
containers:
|
||||
- name: gateway
|
||||
image: smg-gateway:test
|
||||
imagePullPolicy: Never
|
||||
args:
|
||||
- "--service-discovery"
|
||||
- "--selector"
|
||||
- "app=fake-worker"
|
||||
- "--service-discovery-port"
|
||||
- "8000"
|
||||
- "--service-discovery-namespace"
|
||||
- "smg-test"
|
||||
- "--port"
|
||||
- "30000"
|
||||
- "--prometheus-port"
|
||||
- "29000"
|
||||
- "--disable-health-check"
|
||||
- "--worker-startup-timeout-secs"
|
||||
- "30"
|
||||
- "--log-level"
|
||||
- "debug"
|
||||
ports:
|
||||
- containerPort: 30000
|
||||
name: http
|
||||
- containerPort: 29000
|
||||
name: metrics
|
||||
# Use /liveness instead of /readiness for the K8s readiness probe
|
||||
# because /readiness returns 503 when no healthy workers are registered,
|
||||
# and with service discovery the gateway starts with 0 workers.
|
||||
readinessProbe:
|
||||
httpGet:
|
||||
path: /liveness
|
||||
port: 30000
|
||||
initialDelaySeconds: 3
|
||||
periodSeconds: 3
|
||||
livenessProbe:
|
||||
httpGet:
|
||||
path: /liveness
|
||||
port: 30000
|
||||
initialDelaySeconds: 5
|
||||
periodSeconds: 10
|
||||
---
|
||||
apiVersion: v1
|
||||
kind: Service
|
||||
metadata:
|
||||
name: smg-gateway
|
||||
namespace: smg-test
|
||||
spec:
|
||||
type: NodePort
|
||||
selector:
|
||||
app: smg-gateway
|
||||
ports:
|
||||
- name: http
|
||||
port: 30000
|
||||
targetPort: 30000
|
||||
- name: metrics
|
||||
port: 29000
|
||||
targetPort: 29000
|
||||
@@ -0,0 +1,4 @@
|
||||
apiVersion: v1
|
||||
kind: Namespace
|
||||
metadata:
|
||||
name: smg-test
|
||||
@@ -0,0 +1,29 @@
|
||||
apiVersion: v1
|
||||
kind: ServiceAccount
|
||||
metadata:
|
||||
name: smg-gateway
|
||||
namespace: smg-test
|
||||
---
|
||||
apiVersion: rbac.authorization.k8s.io/v1
|
||||
kind: Role
|
||||
metadata:
|
||||
name: smg-gateway
|
||||
namespace: smg-test
|
||||
rules:
|
||||
- apiGroups: [""]
|
||||
resources: ["pods"]
|
||||
verbs: ["get", "list", "watch"]
|
||||
---
|
||||
apiVersion: rbac.authorization.k8s.io/v1
|
||||
kind: RoleBinding
|
||||
metadata:
|
||||
name: smg-gateway
|
||||
namespace: smg-test
|
||||
roleRef:
|
||||
apiGroup: rbac.authorization.k8s.io
|
||||
kind: Role
|
||||
name: smg-gateway
|
||||
subjects:
|
||||
- kind: ServiceAccount
|
||||
name: smg-gateway
|
||||
namespace: smg-test
|
||||
+108
@@ -0,0 +1,108 @@
|
||||
#!/usr/bin/env bash
|
||||
# Setup script for K8s integration tests.
|
||||
#
|
||||
# Prerequisites:
|
||||
# - Docker running
|
||||
# - kind, kubectl installed
|
||||
#
|
||||
# Usage:
|
||||
# ./e2e_test/k8s_integration/setup.sh # full setup
|
||||
# ./e2e_test/k8s_integration/setup.sh teardown # cleanup
|
||||
|
||||
set -euo pipefail
|
||||
|
||||
SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
|
||||
REPO_ROOT="$(cd "$SCRIPT_DIR/../.." && pwd)"
|
||||
CLUSTER_NAME="smg-test"
|
||||
NAMESPACE="smg-test"
|
||||
CONTEXT="kind-${CLUSTER_NAME}"
|
||||
MANIFESTS_DIR="${SCRIPT_DIR}/manifests"
|
||||
|
||||
log() { echo "==> $*"; }
|
||||
|
||||
teardown() {
|
||||
log "Tearing down..."
|
||||
if kind get clusters 2>/dev/null | grep -q "^${CLUSTER_NAME}$"; then
|
||||
kind delete cluster --name "$CLUSTER_NAME"
|
||||
else
|
||||
log "Cluster '${CLUSTER_NAME}' not found, nothing to tear down."
|
||||
fi
|
||||
log "Done."
|
||||
}
|
||||
|
||||
if [[ "${1:-}" == "teardown" ]]; then
|
||||
teardown
|
||||
exit 0
|
||||
fi
|
||||
|
||||
# Step 1: Create kind cluster (skip if exists)
|
||||
if kind get clusters 2>/dev/null | grep -q "^${CLUSTER_NAME}$"; then
|
||||
log "Kind cluster '${CLUSTER_NAME}' already exists"
|
||||
else
|
||||
log "Creating kind cluster '${CLUSTER_NAME}'..."
|
||||
kind create cluster --name "$CLUSTER_NAME"
|
||||
fi
|
||||
|
||||
kubectl config use-context "$CONTEXT"
|
||||
|
||||
# Step 2: Build the gateway Docker image.
|
||||
# Uses a lightweight test Dockerfile that builds just the Rust binary with
|
||||
# the "ci" cargo profile (~5 min), instead of the repo's
|
||||
# docker/gateway.Dockerfile which builds a full Python wheel via maturin.
|
||||
#
|
||||
# CI sets SKIP_DOCKER_BUILD=1 after pre-building smg-gateway:test via
|
||||
# docker/build-push-action with GHA cache, so we don't rebuild here.
|
||||
cd "$REPO_ROOT"
|
||||
if [[ "${SKIP_DOCKER_BUILD:-}" == "1" ]]; then
|
||||
log "SKIP_DOCKER_BUILD=1 — skipping docker build, expecting smg-gateway:test to exist"
|
||||
if ! docker image inspect smg-gateway:test >/dev/null 2>&1; then
|
||||
log "ERROR: smg-gateway:test not found locally; cannot continue"
|
||||
exit 1
|
||||
fi
|
||||
else
|
||||
log "Building gateway Docker image (this may take 5-10 minutes on first run)..."
|
||||
docker build -f e2e_test/k8s_integration/Dockerfile.gateway -t smg-gateway:test .
|
||||
fi
|
||||
|
||||
# Step 3: Load the image into kind
|
||||
log "Loading smg-gateway:test image into kind..."
|
||||
kind load docker-image smg-gateway:test --name "$CLUSTER_NAME"
|
||||
|
||||
# Step 4: Ensure python:3.12-slim is available inside kind (for fake workers).
|
||||
# Pull it locally if not present, then try loading into kind.
|
||||
# If kind load fails (common with multi-arch images), fall back to pulling
|
||||
# directly inside the kind node.
|
||||
log "Ensuring python:3.12-slim is available in kind..."
|
||||
if ! docker image inspect python:3.12-slim >/dev/null 2>&1; then
|
||||
log "Pulling python:3.12-slim..."
|
||||
docker pull python:3.12-slim
|
||||
fi
|
||||
if ! kind load docker-image python:3.12-slim --name "$CLUSTER_NAME" 2>/dev/null; then
|
||||
log "kind load failed (multi-arch image), pulling inside kind node..."
|
||||
docker exec "${CLUSTER_NAME}-control-plane" crictl pull docker.io/library/python:3.12-slim
|
||||
fi
|
||||
|
||||
# Step 5: Apply base manifests
|
||||
log "Applying namespace and RBAC..."
|
||||
kubectl --context "$CONTEXT" apply -f "${MANIFESTS_DIR}/namespace.yaml"
|
||||
kubectl --context "$CONTEXT" apply -f "${MANIFESTS_DIR}/rbac.yaml"
|
||||
|
||||
# Step 6: Create the fake-worker ConfigMap
|
||||
log "Creating fake-worker ConfigMap..."
|
||||
kubectl --context "$CONTEXT" -n "$NAMESPACE" create configmap fake-worker-script \
|
||||
--from-file="fake_worker.py=${SCRIPT_DIR}/fake_worker.py" \
|
||||
--dry-run=client -o yaml | kubectl --context "$CONTEXT" apply -f -
|
||||
|
||||
# Step 7: Apply the gateway deployment
|
||||
log "Deploying SMG gateway..."
|
||||
kubectl --context "$CONTEXT" apply -f "${MANIFESTS_DIR}/gateway.yaml"
|
||||
|
||||
log "Waiting for gateway to be ready..."
|
||||
kubectl --context "$CONTEXT" -n "$NAMESPACE" rollout status deployment/smg-gateway --timeout=180s
|
||||
|
||||
log ""
|
||||
log "Setup complete! Run the integration tests with:"
|
||||
log " pytest e2e_test/k8s_integration/ -v -s"
|
||||
log ""
|
||||
log "To tear down:"
|
||||
log " ./e2e_test/k8s_integration/setup.sh teardown"
|
||||
@@ -0,0 +1,316 @@
|
||||
"""Integration test for PD mode pod type change during hostNetwork rollout.
|
||||
|
||||
Scenario: With hostNetwork, the pod IP = node IP. During a rolling update,
|
||||
an old prefill pod is deleted and a new decode pod comes up on the same node
|
||||
with the same IP but a new UID. The gateway must:
|
||||
|
||||
1. Remove the stale prefill worker (via watcher delete event or reconciliation)
|
||||
2. Register the new decode worker
|
||||
3. End up with the correct worker_type=decode, not the old prefill
|
||||
|
||||
This covers the UID-based eviction path in handle_pod_event (same name,
|
||||
different UID) and the reconciliation diff (stale uid-A, missing uid-B).
|
||||
|
||||
Run with:
|
||||
cd e2e_test/k8s_integration
|
||||
source .venv/bin/activate
|
||||
pytest test_pd_type_change.py -v -s
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
import subprocess
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
from conftest import ( # pytest's rootdir adds the test dir to sys.path
|
||||
KUBECTL_CONTEXT,
|
||||
NAMESPACE,
|
||||
RECONCILIATION_WAIT_SECS,
|
||||
_cleanup_port_forward,
|
||||
_get_worker_count,
|
||||
_get_workers,
|
||||
_kubectl,
|
||||
_poll_until,
|
||||
_wait_for_deployment_ready,
|
||||
_wait_for_pod_ready,
|
||||
_wait_for_port,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
MANIFESTS_DIR = Path(__file__).parent / "manifests"
|
||||
|
||||
PD_GATEWAY_HTTP_PORT = 30001
|
||||
|
||||
|
||||
def _get_workers_by_type(gateway_url: str) -> dict[str, list[dict]]:
|
||||
"""Return workers grouped by worker_type."""
|
||||
data = _get_workers(gateway_url)
|
||||
result: dict[str, list[dict]] = {}
|
||||
for w in data.get("workers", []):
|
||||
wtype = w.get("worker_type", "unknown")
|
||||
result.setdefault(wtype, []).append(w)
|
||||
return result
|
||||
|
||||
|
||||
def _deploy_pd_worker(name: str, role: str):
|
||||
"""Deploy a fake worker pod with a role label for PD mode."""
|
||||
pod_manifest = {
|
||||
"apiVersion": "v1",
|
||||
"kind": "Pod",
|
||||
"metadata": {
|
||||
"name": name,
|
||||
"namespace": NAMESPACE,
|
||||
"labels": {"role": role},
|
||||
},
|
||||
"spec": {
|
||||
"containers": [
|
||||
{
|
||||
"name": "worker",
|
||||
"image": "python:3.12-slim",
|
||||
"imagePullPolicy": "IfNotPresent",
|
||||
"command": ["python3", "/app/fake_worker.py"],
|
||||
"ports": [{"containerPort": 8000}],
|
||||
"readinessProbe": {
|
||||
"httpGet": {"path": "/health", "port": 8000},
|
||||
"initialDelaySeconds": 2,
|
||||
"periodSeconds": 3,
|
||||
},
|
||||
"volumeMounts": [{"name": "app", "mountPath": "/app"}],
|
||||
}
|
||||
],
|
||||
"volumes": [{"name": "app", "configMap": {"name": "fake-worker-script"}}],
|
||||
},
|
||||
}
|
||||
subprocess.run(
|
||||
["kubectl", "--context", KUBECTL_CONTEXT, "apply", "-f", "-"],
|
||||
input=json.dumps(pod_manifest),
|
||||
capture_output=True,
|
||||
text=True,
|
||||
check=True,
|
||||
)
|
||||
logger.info("Deployed PD pod %s with role=%s", name, role)
|
||||
|
||||
|
||||
def _safe_delete_pod(name: str):
|
||||
try:
|
||||
_kubectl(
|
||||
"delete",
|
||||
"pod",
|
||||
name,
|
||||
"-n",
|
||||
NAMESPACE,
|
||||
"--ignore-not-found",
|
||||
"--force",
|
||||
"--grace-period=0",
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning("Cleanup failed for pod %s: %s", name, e)
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def pd_gateway():
|
||||
"""Deploy the PD-mode gateway and set up port-forwarding.
|
||||
|
||||
Cleanup runs in `finally:` so a failure in port-forward setup does not
|
||||
leak the kubectl process or leave the gateway-pd Deployment behind.
|
||||
"""
|
||||
manifest = MANIFESTS_DIR / "gateway-pd.yaml"
|
||||
_kubectl("apply", "-f", str(manifest))
|
||||
pf: subprocess.Popen | None = None
|
||||
try:
|
||||
_wait_for_deployment_ready("smg-gateway-pd")
|
||||
|
||||
cmd = [
|
||||
"kubectl",
|
||||
"--context",
|
||||
KUBECTL_CONTEXT,
|
||||
"port-forward",
|
||||
"svc/smg-gateway-pd",
|
||||
f"{PD_GATEWAY_HTTP_PORT}:{PD_GATEWAY_HTTP_PORT}",
|
||||
"-n",
|
||||
NAMESPACE,
|
||||
]
|
||||
pf = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
|
||||
_wait_for_port(PD_GATEWAY_HTTP_PORT, pf)
|
||||
|
||||
yield f"http://127.0.0.1:{PD_GATEWAY_HTTP_PORT}"
|
||||
finally:
|
||||
if pf is not None:
|
||||
_cleanup_port_forward("pd_gateway", pf)
|
||||
result = _kubectl(
|
||||
"delete",
|
||||
"-f",
|
||||
str(manifest),
|
||||
"--ignore-not-found",
|
||||
check=False,
|
||||
)
|
||||
if result.returncode != 0:
|
||||
logger.warning(
|
||||
"Teardown delete of gateway-pd failed (rc=%d): %s",
|
||||
result.returncode,
|
||||
result.stderr.strip(),
|
||||
)
|
||||
|
||||
|
||||
class TestPDRolloutTypeChange:
|
||||
"""Test that the gateway correctly transitions worker type when a pod
|
||||
is deleted and recreated with a different role (prefill -> decode).
|
||||
|
||||
This simulates a hostNetwork rolling update where a node changes role.
|
||||
The new pod has the same name and IP but a different UID and labels.
|
||||
"""
|
||||
|
||||
def test_prefill_discovered_as_prefill(self, pd_gateway):
|
||||
"""Baseline: a prefill pod is correctly discovered as prefill type."""
|
||||
pod_name = "test-pd-baseline"
|
||||
|
||||
try:
|
||||
_deploy_pd_worker(pod_name, role="prefill")
|
||||
_wait_for_pod_ready(pod_name)
|
||||
|
||||
_poll_until(
|
||||
lambda: _get_worker_count(pd_gateway) >= 1,
|
||||
"prefill worker discovered",
|
||||
timeout=30,
|
||||
interval=3,
|
||||
)
|
||||
|
||||
by_type = _get_workers_by_type(pd_gateway)
|
||||
logger.info("Workers by type: %s", json.dumps(by_type, indent=2))
|
||||
assert (
|
||||
"prefill" in by_type
|
||||
), f"Expected prefill, got: {list(by_type.keys())}"
|
||||
|
||||
finally:
|
||||
_safe_delete_pod(pod_name)
|
||||
_poll_until(
|
||||
lambda: _get_worker_count(pd_gateway) == 0,
|
||||
"cleanup: worker count back to 0",
|
||||
timeout=RECONCILIATION_WAIT_SECS,
|
||||
interval=5,
|
||||
)
|
||||
|
||||
def test_delete_prefill_recreate_as_decode(self, pd_gateway):
|
||||
"""Delete a prefill pod and recreate with the same name as decode.
|
||||
|
||||
This is the realistic rollout path: the old pod is deleted (new UID),
|
||||
and a new pod with a different role comes up. The gateway should
|
||||
transition the worker from prefill to decode.
|
||||
"""
|
||||
pod_name = "test-pd-rollout"
|
||||
|
||||
try:
|
||||
# Step 1: Deploy as prefill
|
||||
_deploy_pd_worker(pod_name, role="prefill")
|
||||
_wait_for_pod_ready(pod_name)
|
||||
|
||||
_poll_until(
|
||||
lambda: _get_worker_count(pd_gateway) >= 1,
|
||||
"prefill worker discovered",
|
||||
timeout=30,
|
||||
interval=3,
|
||||
)
|
||||
|
||||
by_type = _get_workers_by_type(pd_gateway)
|
||||
logger.info("Before rollout: %s", json.dumps(by_type, indent=2))
|
||||
assert "prefill" in by_type
|
||||
|
||||
# Capture the prefill worker's URL for later comparison
|
||||
prefill_url = by_type["prefill"][0]["url"]
|
||||
logger.info("Prefill worker URL: %s", prefill_url)
|
||||
|
||||
# Step 2: Delete the prefill pod (simulates rollout termination)
|
||||
_kubectl(
|
||||
"delete",
|
||||
"pod",
|
||||
pod_name,
|
||||
"-n",
|
||||
NAMESPACE,
|
||||
"--force",
|
||||
"--grace-period=0",
|
||||
)
|
||||
|
||||
# Step 3: Wait for the gateway to remove the stale prefill worker
|
||||
_poll_until(
|
||||
lambda: _get_worker_count(pd_gateway) == 0,
|
||||
"prefill worker removed",
|
||||
timeout=RECONCILIATION_WAIT_SECS,
|
||||
interval=5,
|
||||
)
|
||||
|
||||
# Step 4: Recreate with same name as decode (new UID!)
|
||||
_deploy_pd_worker(pod_name, role="decode")
|
||||
_wait_for_pod_ready(pod_name)
|
||||
|
||||
# Step 5: Verify the gateway discovers it as decode
|
||||
_poll_until(
|
||||
lambda: _get_worker_count(pd_gateway) >= 1,
|
||||
"decode worker discovered after rollout",
|
||||
timeout=30,
|
||||
interval=3,
|
||||
)
|
||||
|
||||
by_type = _get_workers_by_type(pd_gateway)
|
||||
logger.info("After rollout: %s", json.dumps(by_type, indent=2))
|
||||
|
||||
assert (
|
||||
"decode" in by_type
|
||||
), f"Expected decode worker after rollout, got: {list(by_type.keys())}"
|
||||
assert (
|
||||
"prefill" not in by_type
|
||||
), "Stale prefill worker persists after rollout"
|
||||
|
||||
finally:
|
||||
_safe_delete_pod(pod_name)
|
||||
|
||||
def test_simultaneous_prefill_and_decode(self, pd_gateway):
|
||||
"""Both prefill and decode pods exist at the same time.
|
||||
|
||||
During a rolling update there may be a brief overlap where both
|
||||
old and new pods are running. The gateway should track both.
|
||||
"""
|
||||
prefill_pod = "test-pd-both-p"
|
||||
decode_pod = "test-pd-both-d"
|
||||
|
||||
try:
|
||||
_deploy_pd_worker(prefill_pod, role="prefill")
|
||||
_deploy_pd_worker(decode_pod, role="decode")
|
||||
_wait_for_pod_ready(prefill_pod)
|
||||
_wait_for_pod_ready(decode_pod)
|
||||
|
||||
_poll_until(
|
||||
lambda: _get_worker_count(pd_gateway) >= 2,
|
||||
"both prefill and decode workers discovered",
|
||||
timeout=30,
|
||||
interval=3,
|
||||
)
|
||||
|
||||
by_type = _get_workers_by_type(pd_gateway)
|
||||
logger.info("Both pods: %s", json.dumps(by_type, indent=2))
|
||||
|
||||
assert "prefill" in by_type, f"Missing prefill, got: {list(by_type.keys())}"
|
||||
assert "decode" in by_type, f"Missing decode, got: {list(by_type.keys())}"
|
||||
|
||||
# Now remove prefill, only decode should remain
|
||||
_safe_delete_pod(prefill_pod)
|
||||
|
||||
_poll_until(
|
||||
lambda: _get_worker_count(pd_gateway) == 1,
|
||||
"only decode worker remains",
|
||||
timeout=RECONCILIATION_WAIT_SECS,
|
||||
interval=5,
|
||||
)
|
||||
|
||||
by_type = _get_workers_by_type(pd_gateway)
|
||||
logger.info("After prefill removed: %s", json.dumps(by_type, indent=2))
|
||||
|
||||
assert "decode" in by_type, "Decode worker should still exist"
|
||||
assert "prefill" not in by_type, "Prefill worker should be gone"
|
||||
|
||||
finally:
|
||||
_safe_delete_pod(prefill_pod)
|
||||
_safe_delete_pod(decode_pod)
|
||||
@@ -0,0 +1,506 @@
|
||||
"""Integration tests for K8s service discovery and reconciliation.
|
||||
|
||||
Tests verify that:
|
||||
1. The K8s watcher correctly discovers new pods
|
||||
2. Stale workers are removed after pod deletion (watcher or reconciliation)
|
||||
3. All pods are eventually discovered and tracked consistently
|
||||
4. Prometheus discovery metrics are emitted correctly
|
||||
5. Reconciliation does not cause instability over multiple cycles
|
||||
|
||||
These tests require a kind cluster with the gateway deployed. The
|
||||
reconciliation interval is 60s (see ServiceDiscoveryConfig.check_interval
|
||||
in sgl-model-gateway/src/service_discovery.rs), so tests exercising
|
||||
reconciliation must wait ~90s for a tick to fire.
|
||||
|
||||
Run with:
|
||||
cd e2e_test/k8s_integration
|
||||
source .venv/bin/activate
|
||||
pytest test_reconciliation.py -v -s
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
import subprocess
|
||||
import time
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from conftest import ( # pytest's rootdir adds the test dir to sys.path
|
||||
KUBECTL_CONTEXT,
|
||||
NAMESPACE,
|
||||
RECONCILIATION_WAIT_SECS,
|
||||
_get_worker_count,
|
||||
_get_workers,
|
||||
_kubectl,
|
||||
_poll_until,
|
||||
_wait_for_pod_ready,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _get_metrics(metrics_url: str) -> str:
|
||||
"""GET /metrics from the gateway (Prometheus text format)."""
|
||||
resp = httpx.get(f"{metrics_url}/metrics", timeout=10)
|
||||
resp.raise_for_status()
|
||||
return resp.text
|
||||
|
||||
|
||||
def _get_worker_urls(gateway_url: str) -> set[str]:
|
||||
"""Return the set of worker URLs currently registered in the gateway."""
|
||||
data = _get_workers(gateway_url)
|
||||
return {w["url"] for w in data.get("workers", [])}
|
||||
|
||||
|
||||
def _parse_metric_value(
|
||||
metrics_text: str, metric_name: str, labels: dict | None = None
|
||||
) -> float | None:
|
||||
"""Parse a specific metric value from Prometheus text format.
|
||||
|
||||
Uses exact metric name matching (line must start with the metric name)
|
||||
and logs diagnostics when the metric is not found.
|
||||
"""
|
||||
matching_lines = []
|
||||
for line in metrics_text.splitlines():
|
||||
if line.startswith("#"):
|
||||
continue
|
||||
# Exact metric name match: name must be followed by '{' or ' '
|
||||
if not line.startswith(metric_name):
|
||||
continue
|
||||
rest = line[len(metric_name) :]
|
||||
if rest and rest[0] not in ("{", " "):
|
||||
continue
|
||||
matching_lines.append(line)
|
||||
|
||||
if not matching_lines:
|
||||
logger.debug("Metric %s not found in output", metric_name)
|
||||
return None
|
||||
|
||||
for line in matching_lines:
|
||||
if labels:
|
||||
if not all(f'{k}="{v}"' in line for k, v in labels.items()):
|
||||
continue
|
||||
parts = line.split()
|
||||
if len(parts) >= 2:
|
||||
try:
|
||||
return float(parts[-1])
|
||||
except ValueError:
|
||||
logger.warning("Could not parse float from metric line: %s", line)
|
||||
continue
|
||||
|
||||
logger.debug(
|
||||
"Metric %s found but no line matched labels %s. Lines: %s",
|
||||
metric_name,
|
||||
labels,
|
||||
matching_lines,
|
||||
)
|
||||
return None
|
||||
|
||||
|
||||
def _deploy_worker_pod(name: str, labels: dict[str, str] | None = None):
|
||||
"""Deploy a single fake worker pod via kubectl apply from stdin."""
|
||||
pod_labels = {"app": "fake-worker"}
|
||||
if labels:
|
||||
pod_labels.update(labels)
|
||||
|
||||
pod_manifest = {
|
||||
"apiVersion": "v1",
|
||||
"kind": "Pod",
|
||||
"metadata": {
|
||||
"name": name,
|
||||
"namespace": NAMESPACE,
|
||||
"labels": pod_labels,
|
||||
},
|
||||
"spec": {
|
||||
"containers": [
|
||||
{
|
||||
"name": "worker",
|
||||
"image": "python:3.12-slim",
|
||||
"imagePullPolicy": "IfNotPresent",
|
||||
"command": ["python3", "/app/fake_worker.py"],
|
||||
"ports": [{"containerPort": 8000}],
|
||||
"readinessProbe": {
|
||||
"httpGet": {"path": "/health", "port": 8000},
|
||||
"initialDelaySeconds": 2,
|
||||
"periodSeconds": 3,
|
||||
},
|
||||
"volumeMounts": [
|
||||
{
|
||||
"name": "app",
|
||||
"mountPath": "/app",
|
||||
}
|
||||
],
|
||||
}
|
||||
],
|
||||
"volumes": [
|
||||
{
|
||||
"name": "app",
|
||||
"configMap": {"name": "fake-worker-script"},
|
||||
}
|
||||
],
|
||||
},
|
||||
}
|
||||
proc = subprocess.run(
|
||||
["kubectl", "--context", KUBECTL_CONTEXT, "apply", "-f", "-"],
|
||||
input=json.dumps(pod_manifest),
|
||||
capture_output=True,
|
||||
text=True,
|
||||
check=True,
|
||||
)
|
||||
logger.info("Deployed pod %s: %s", name, proc.stdout.strip())
|
||||
|
||||
|
||||
def _delete_worker_pod(name: str, force: bool = False):
|
||||
"""Delete a fake worker pod."""
|
||||
args = ["delete", "pod", name, "-n", NAMESPACE, "--ignore-not-found"]
|
||||
if force:
|
||||
args.extend(["--grace-period=0", "--force"])
|
||||
_kubectl(*args)
|
||||
logger.info("Deleted pod %s (force=%s)", name, force)
|
||||
|
||||
|
||||
def _safe_delete_worker_pod(name: str):
|
||||
"""Delete a worker pod in a cleanup context, logging errors instead of raising."""
|
||||
try:
|
||||
_delete_worker_pod(name, force=True)
|
||||
except Exception as e:
|
||||
logger.warning("Cleanup failed for pod %s: %s", name, e)
|
||||
|
||||
|
||||
def _wait_for_pod_gone(name: str, timeout: int = 60):
|
||||
"""Wait until a pod no longer exists in K8s.
|
||||
|
||||
Raises TimeoutError if the pod still exists after timeout, or RuntimeError
|
||||
if kubectl returns an unexpected error (e.g., apiserver unreachable, RBAC
|
||||
drift). The latter would otherwise surface as a misleading "still exists"
|
||||
timeout.
|
||||
"""
|
||||
deadline = time.time() + timeout
|
||||
while time.time() < deadline:
|
||||
result = _kubectl(
|
||||
"get",
|
||||
"pod",
|
||||
name,
|
||||
"-n",
|
||||
NAMESPACE,
|
||||
check=False,
|
||||
)
|
||||
if result.returncode == 0:
|
||||
time.sleep(2)
|
||||
continue
|
||||
stderr = result.stderr.strip()
|
||||
if "NotFound" in stderr or "not found" in stderr.lower():
|
||||
logger.info("Pod %s is gone", name)
|
||||
return
|
||||
# Anything else is a real cluster-level error — fail loudly so the
|
||||
# caller sees the actual problem instead of a generic timeout.
|
||||
raise RuntimeError(
|
||||
f"kubectl get pod {name} failed unexpectedly (rc={result.returncode}): {stderr}"
|
||||
)
|
||||
raise TimeoutError(f"Pod {name} still exists after {timeout}s")
|
||||
|
||||
|
||||
class TestWatcherDiscovery:
|
||||
"""Tests that the K8s watcher correctly discovers pods on creation."""
|
||||
|
||||
def test_watcher_discovers_new_pod(self, gateway_port_forward):
|
||||
"""Deploy a new worker pod and verify the watcher picks it up quickly."""
|
||||
gateway_url, metrics_url = gateway_port_forward
|
||||
pod_name = "test-watcher-discovery"
|
||||
|
||||
try:
|
||||
initial_count = _get_worker_count(gateway_url)
|
||||
logger.info("Initial worker count: %d", initial_count)
|
||||
|
||||
_deploy_worker_pod(pod_name)
|
||||
_wait_for_pod_ready(pod_name)
|
||||
|
||||
# The watcher should pick up the pod within seconds
|
||||
_poll_until(
|
||||
lambda: _get_worker_count(gateway_url) > initial_count,
|
||||
f"worker count > {initial_count}",
|
||||
timeout=30,
|
||||
interval=3,
|
||||
)
|
||||
|
||||
workers = _get_workers(gateway_url)
|
||||
logger.info(
|
||||
"Workers after pod creation: %s",
|
||||
json.dumps(workers, indent=2),
|
||||
)
|
||||
assert workers["total"] > initial_count
|
||||
|
||||
finally:
|
||||
_safe_delete_worker_pod(pod_name)
|
||||
|
||||
|
||||
class TestReconciliationStaleWorkerRemoval:
|
||||
"""Test that stale workers are removed after pod deletion.
|
||||
|
||||
The watcher DELETE event typically handles this immediately.
|
||||
If the watcher misses it (e.g., during restart or backoff),
|
||||
reconciliation catches it within ~60s.
|
||||
"""
|
||||
|
||||
def test_stale_worker_removed_after_pod_deletion(self, gateway_port_forward):
|
||||
"""Deploy a worker, verify discovery, delete the pod, and verify
|
||||
the worker is removed (by either watcher or reconciliation)."""
|
||||
gateway_url, metrics_url = gateway_port_forward
|
||||
pod_name = "test-stale-removal"
|
||||
|
||||
try:
|
||||
_deploy_worker_pod(pod_name)
|
||||
_wait_for_pod_ready(pod_name)
|
||||
|
||||
_poll_until(
|
||||
lambda: _get_worker_count(gateway_url) >= 1,
|
||||
"at least 1 worker discovered",
|
||||
timeout=30,
|
||||
interval=3,
|
||||
)
|
||||
|
||||
count_with_pod = _get_worker_count(gateway_url)
|
||||
logger.info("Worker count with test pod: %d", count_with_pod)
|
||||
|
||||
# Force-delete the pod (instant removal from K8s API)
|
||||
_delete_worker_pod(pod_name, force=True)
|
||||
_wait_for_pod_gone(pod_name)
|
||||
|
||||
# Wait for the gateway to remove the stale worker.
|
||||
# The watcher DELETE event may handle this immediately.
|
||||
# If it doesn't, reconciliation will catch it within ~60s.
|
||||
_poll_until(
|
||||
lambda: _get_worker_count(gateway_url) < count_with_pod,
|
||||
f"worker count < {count_with_pod} (stale worker removed)",
|
||||
timeout=RECONCILIATION_WAIT_SECS,
|
||||
interval=5,
|
||||
)
|
||||
|
||||
final_count = _get_worker_count(gateway_url)
|
||||
logger.info("Worker count after stale removal: %d", final_count)
|
||||
assert final_count < count_with_pod
|
||||
|
||||
finally:
|
||||
_safe_delete_worker_pod(pod_name)
|
||||
|
||||
|
||||
class TestReconciliationMissedPodDiscovery:
|
||||
"""Verify reconciliation coexists with watcher discovery without interference.
|
||||
|
||||
Note: this test cannot force the watcher to miss events, so it does NOT
|
||||
prove that reconciliation discovers missed pods in isolation. It validates
|
||||
that reconciliation maintains consistency when pods are already discovered
|
||||
by the watcher, and that no pods are lost.
|
||||
"""
|
||||
|
||||
def test_all_workers_eventually_discovered(self, gateway_port_forward):
|
||||
"""Deploy multiple worker pods and verify they are all discovered."""
|
||||
gateway_url, metrics_url = gateway_port_forward
|
||||
pod_names = ["test-reconcile-a", "test-reconcile-b"]
|
||||
|
||||
try:
|
||||
for name in pod_names:
|
||||
_deploy_worker_pod(name)
|
||||
|
||||
for name in pod_names:
|
||||
_wait_for_pod_ready(name)
|
||||
|
||||
# Wait for watcher (or reconciliation) to discover all pods
|
||||
_poll_until(
|
||||
lambda: _get_worker_count(gateway_url) >= len(pod_names),
|
||||
f"at least {len(pod_names)} workers discovered",
|
||||
timeout=RECONCILIATION_WAIT_SECS,
|
||||
interval=5,
|
||||
)
|
||||
|
||||
workers = _get_workers(gateway_url)
|
||||
logger.info(
|
||||
"Workers after discovery: %s",
|
||||
json.dumps(workers, indent=2),
|
||||
)
|
||||
assert workers["total"] >= len(pod_names)
|
||||
|
||||
finally:
|
||||
for name in pod_names:
|
||||
_safe_delete_worker_pod(name)
|
||||
for name in pod_names:
|
||||
try:
|
||||
_wait_for_pod_gone(name, timeout=30)
|
||||
except TimeoutError:
|
||||
logger.warning("Pod %s still present after cleanup", name)
|
||||
|
||||
|
||||
class TestReconciliationMetrics:
|
||||
"""Test that the gateway emits expected Prometheus discovery metrics."""
|
||||
|
||||
def test_discovery_metrics_populated(self, gateway_port_forward):
|
||||
"""After pods are discovered, verify registration and gauge metrics."""
|
||||
gateway_url, metrics_url = gateway_port_forward
|
||||
pod_name = "test-metrics"
|
||||
|
||||
try:
|
||||
_deploy_worker_pod(pod_name)
|
||||
_wait_for_pod_ready(pod_name)
|
||||
|
||||
# Wait for the watcher to discover the pod
|
||||
_poll_until(
|
||||
lambda: _get_worker_count(gateway_url) >= 1,
|
||||
"at least 1 worker",
|
||||
timeout=30,
|
||||
interval=3,
|
||||
)
|
||||
|
||||
# Poll for the registration metric instead of a fixed sleep
|
||||
def _registration_metric_exists():
|
||||
text = _get_metrics(metrics_url)
|
||||
val = _parse_metric_value(
|
||||
text,
|
||||
"smg_discovery_registrations_total",
|
||||
{"source": "kubernetes", "result": "success"},
|
||||
)
|
||||
return val is not None and val >= 1
|
||||
|
||||
_poll_until(
|
||||
_registration_metric_exists,
|
||||
"registration success metric >= 1",
|
||||
timeout=30,
|
||||
interval=3,
|
||||
)
|
||||
|
||||
metrics_text = _get_metrics(metrics_url)
|
||||
|
||||
reg_value = _parse_metric_value(
|
||||
metrics_text,
|
||||
"smg_discovery_registrations_total",
|
||||
{"source": "kubernetes", "result": "success"},
|
||||
)
|
||||
logger.info("Registration success metric: %s", reg_value)
|
||||
assert (
|
||||
reg_value is not None and reg_value >= 1
|
||||
), f"Expected at least 1 registration, got {reg_value}"
|
||||
|
||||
gauge_value = _parse_metric_value(
|
||||
metrics_text,
|
||||
"smg_discovery_workers_discovered",
|
||||
{"source": "kubernetes"},
|
||||
)
|
||||
logger.info("Workers discovered gauge: %s", gauge_value)
|
||||
assert (
|
||||
gauge_value is not None and gauge_value >= 1
|
||||
), f"Expected workers_discovered >= 1, got {gauge_value}"
|
||||
|
||||
finally:
|
||||
_safe_delete_worker_pod(pod_name)
|
||||
|
||||
def test_deregistration_metric_after_pod_deletion(self, gateway_port_forward):
|
||||
"""Deploy a pod, delete it, and verify a deregistration metric fires.
|
||||
|
||||
Either 'pod_deleted' (from the watcher) or 'reconciled' (from
|
||||
periodic reconciliation) should increment. Which one fires depends
|
||||
on whether the watcher sees the deletion event first.
|
||||
"""
|
||||
gateway_url, metrics_url = gateway_port_forward
|
||||
pod_name = "test-dereg-metric"
|
||||
|
||||
try:
|
||||
_deploy_worker_pod(pod_name)
|
||||
_wait_for_pod_ready(pod_name)
|
||||
|
||||
_poll_until(
|
||||
lambda: _get_worker_count(gateway_url) >= 1,
|
||||
"at least 1 worker",
|
||||
timeout=30,
|
||||
interval=3,
|
||||
)
|
||||
|
||||
count_before = _get_worker_count(gateway_url)
|
||||
|
||||
_delete_worker_pod(pod_name, force=True)
|
||||
_wait_for_pod_gone(pod_name)
|
||||
|
||||
_poll_until(
|
||||
lambda: _get_worker_count(gateway_url) < count_before,
|
||||
"worker count decreased after pod deletion",
|
||||
timeout=RECONCILIATION_WAIT_SECS,
|
||||
interval=5,
|
||||
)
|
||||
|
||||
metrics_text = _get_metrics(metrics_url)
|
||||
|
||||
pod_deleted = _parse_metric_value(
|
||||
metrics_text,
|
||||
"smg_discovery_deregistrations_total",
|
||||
{"source": "kubernetes", "reason": "pod_deleted"},
|
||||
)
|
||||
reconciled = _parse_metric_value(
|
||||
metrics_text,
|
||||
"smg_discovery_deregistrations_total",
|
||||
{"source": "kubernetes", "reason": "reconciled"},
|
||||
)
|
||||
|
||||
logger.info(
|
||||
"Deregistration metrics — pod_deleted: %s, reconciled: %s",
|
||||
pod_deleted,
|
||||
reconciled,
|
||||
)
|
||||
|
||||
total_dereg = (pod_deleted or 0) + (reconciled or 0)
|
||||
assert total_dereg >= 1, (
|
||||
f"Expected at least 1 deregistration, got pod_deleted={pod_deleted}, "
|
||||
f"reconciled={reconciled}"
|
||||
)
|
||||
|
||||
finally:
|
||||
_safe_delete_worker_pod(pod_name)
|
||||
|
||||
|
||||
class TestReconciliationConsistency:
|
||||
"""Test that reconciliation maintains consistency over multiple cycles."""
|
||||
|
||||
@pytest.mark.slow
|
||||
def test_repeated_reconciliation_is_stable(self, gateway_port_forward):
|
||||
"""Deploy pods, wait for 2+ reconciliation cycles, and verify worker
|
||||
count stays stable (no duplicate additions or spurious removals)."""
|
||||
gateway_url, metrics_url = gateway_port_forward
|
||||
pod_names = ["test-stable-a", "test-stable-b"]
|
||||
|
||||
try:
|
||||
for name in pod_names:
|
||||
_deploy_worker_pod(name)
|
||||
for name in pod_names:
|
||||
_wait_for_pod_ready(name)
|
||||
|
||||
# Wait for initial discovery
|
||||
_poll_until(
|
||||
lambda: _get_worker_count(gateway_url) >= len(pod_names),
|
||||
f"at least {len(pod_names)} workers",
|
||||
timeout=30,
|
||||
interval=3,
|
||||
)
|
||||
|
||||
stable_count = _get_worker_count(gateway_url)
|
||||
logger.info("Stable worker count: %d", stable_count)
|
||||
|
||||
# Wait for 2+ reconciliation cycles: 2*60s interval + 30s margin = 150s total
|
||||
wait_time = RECONCILIATION_WAIT_SECS + 60
|
||||
logger.info("Waiting %ds for 2+ reconciliation cycles...", wait_time)
|
||||
|
||||
# Sample count periodically to verify stability
|
||||
end_time = time.time() + wait_time
|
||||
samples = []
|
||||
while time.time() < end_time:
|
||||
count = _get_worker_count(gateway_url)
|
||||
samples.append(count)
|
||||
time.sleep(15)
|
||||
|
||||
logger.info("Worker count samples over time: %s", samples)
|
||||
|
||||
assert all(
|
||||
s == stable_count for s in samples
|
||||
), f"Worker count fluctuated: {samples} (expected stable at {stable_count})"
|
||||
|
||||
finally:
|
||||
for name in pod_names:
|
||||
_safe_delete_worker_pod(name)
|
||||
Reference in New Issue
Block a user