[SMG] Expand K8s integration tests: cross-namespace, lifecycle, multi-model (#24375)
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
e8d57e7243
commit
8b9ff4a68c
@@ -0,0 +1,79 @@
|
||||
apiVersion: apps/v1
|
||||
kind: Deployment
|
||||
metadata:
|
||||
name: smg-gateway-cluster
|
||||
namespace: smg-test
|
||||
spec:
|
||||
replicas: 1
|
||||
selector:
|
||||
matchLabels:
|
||||
app: smg-gateway-cluster
|
||||
template:
|
||||
metadata:
|
||||
labels:
|
||||
app: smg-gateway-cluster
|
||||
spec:
|
||||
serviceAccountName: smg-gateway-cluster
|
||||
containers:
|
||||
- name: gateway
|
||||
image: smg-gateway:test
|
||||
imagePullPolicy: Never
|
||||
# No --service-discovery-namespace: watch all namespaces (Api::all
|
||||
# in service_discovery.rs:279). Requires the ClusterRole defined in
|
||||
# rbac-cluster-scoped.yaml.
|
||||
#
|
||||
# Selector includes `cross-ns-test=true` (in addition to the usual
|
||||
# app=fake-worker) so this gateway only registers workers deployed
|
||||
# by test_cross_namespace.py, not the app=fake-worker pods other
|
||||
# tests deploy in smg-test. Without this, the exact-count assertion
|
||||
# in the test would race with concurrent test files.
|
||||
args:
|
||||
- "--service-discovery"
|
||||
- "--selector"
|
||||
- "app=fake-worker"
|
||||
- "cross-ns-test=true"
|
||||
- "--service-discovery-port"
|
||||
- "8000"
|
||||
- "--port"
|
||||
- "30002"
|
||||
- "--prometheus-port"
|
||||
- "29002"
|
||||
- "--disable-health-check"
|
||||
- "--worker-startup-timeout-secs"
|
||||
- "30"
|
||||
- "--log-level"
|
||||
- "debug"
|
||||
ports:
|
||||
- containerPort: 30002
|
||||
name: http
|
||||
- containerPort: 29002
|
||||
name: metrics
|
||||
readinessProbe:
|
||||
httpGet:
|
||||
path: /liveness
|
||||
port: 30002
|
||||
initialDelaySeconds: 3
|
||||
periodSeconds: 3
|
||||
livenessProbe:
|
||||
httpGet:
|
||||
path: /liveness
|
||||
port: 30002
|
||||
initialDelaySeconds: 5
|
||||
periodSeconds: 10
|
||||
---
|
||||
apiVersion: v1
|
||||
kind: Service
|
||||
metadata:
|
||||
name: smg-gateway-cluster
|
||||
namespace: smg-test
|
||||
spec:
|
||||
type: NodePort
|
||||
selector:
|
||||
app: smg-gateway-cluster
|
||||
ports:
|
||||
- name: http
|
||||
port: 30002
|
||||
targetPort: 30002
|
||||
- name: metrics
|
||||
port: 29002
|
||||
targetPort: 29002
|
||||
@@ -0,0 +1,148 @@
|
||||
# Two gateways with disjoint --selector values to verify per-model isolation.
|
||||
# Both watch the smg-test namespace, both use the namespace-scoped RBAC from
|
||||
# rbac.yaml, but each only matches workers carrying its own `model=...` label.
|
||||
apiVersion: apps/v1
|
||||
kind: Deployment
|
||||
metadata:
|
||||
name: smg-gateway-llama
|
||||
namespace: smg-test
|
||||
spec:
|
||||
replicas: 1
|
||||
selector:
|
||||
matchLabels:
|
||||
app: smg-gateway-llama
|
||||
template:
|
||||
metadata:
|
||||
labels:
|
||||
app: smg-gateway-llama
|
||||
spec:
|
||||
serviceAccountName: smg-gateway
|
||||
containers:
|
||||
- name: gateway
|
||||
image: smg-gateway:test
|
||||
imagePullPolicy: Never
|
||||
args:
|
||||
- "--service-discovery"
|
||||
- "--selector"
|
||||
- "app=fake-worker"
|
||||
- "model=llama"
|
||||
- "--service-discovery-port"
|
||||
- "8000"
|
||||
- "--service-discovery-namespace"
|
||||
- "smg-test"
|
||||
- "--port"
|
||||
- "30003"
|
||||
- "--prometheus-port"
|
||||
- "29003"
|
||||
- "--disable-health-check"
|
||||
- "--worker-startup-timeout-secs"
|
||||
- "30"
|
||||
- "--log-level"
|
||||
- "debug"
|
||||
ports:
|
||||
- containerPort: 30003
|
||||
name: http
|
||||
- containerPort: 29003
|
||||
name: metrics
|
||||
readinessProbe:
|
||||
httpGet:
|
||||
path: /liveness
|
||||
port: 30003
|
||||
initialDelaySeconds: 3
|
||||
periodSeconds: 3
|
||||
livenessProbe:
|
||||
httpGet:
|
||||
path: /liveness
|
||||
port: 30003
|
||||
initialDelaySeconds: 5
|
||||
periodSeconds: 10
|
||||
---
|
||||
apiVersion: v1
|
||||
kind: Service
|
||||
metadata:
|
||||
name: smg-gateway-llama
|
||||
namespace: smg-test
|
||||
spec:
|
||||
type: NodePort
|
||||
selector:
|
||||
app: smg-gateway-llama
|
||||
ports:
|
||||
- name: http
|
||||
port: 30003
|
||||
targetPort: 30003
|
||||
- name: metrics
|
||||
port: 29003
|
||||
targetPort: 29003
|
||||
---
|
||||
apiVersion: apps/v1
|
||||
kind: Deployment
|
||||
metadata:
|
||||
name: smg-gateway-qwen
|
||||
namespace: smg-test
|
||||
spec:
|
||||
replicas: 1
|
||||
selector:
|
||||
matchLabels:
|
||||
app: smg-gateway-qwen
|
||||
template:
|
||||
metadata:
|
||||
labels:
|
||||
app: smg-gateway-qwen
|
||||
spec:
|
||||
serviceAccountName: smg-gateway
|
||||
containers:
|
||||
- name: gateway
|
||||
image: smg-gateway:test
|
||||
imagePullPolicy: Never
|
||||
args:
|
||||
- "--service-discovery"
|
||||
- "--selector"
|
||||
- "app=fake-worker"
|
||||
- "model=qwen"
|
||||
- "--service-discovery-port"
|
||||
- "8000"
|
||||
- "--service-discovery-namespace"
|
||||
- "smg-test"
|
||||
- "--port"
|
||||
- "30004"
|
||||
- "--prometheus-port"
|
||||
- "29004"
|
||||
- "--disable-health-check"
|
||||
- "--worker-startup-timeout-secs"
|
||||
- "30"
|
||||
- "--log-level"
|
||||
- "debug"
|
||||
ports:
|
||||
- containerPort: 30004
|
||||
name: http
|
||||
- containerPort: 29004
|
||||
name: metrics
|
||||
readinessProbe:
|
||||
httpGet:
|
||||
path: /liveness
|
||||
port: 30004
|
||||
initialDelaySeconds: 3
|
||||
periodSeconds: 3
|
||||
livenessProbe:
|
||||
httpGet:
|
||||
path: /liveness
|
||||
port: 30004
|
||||
initialDelaySeconds: 5
|
||||
periodSeconds: 10
|
||||
---
|
||||
apiVersion: v1
|
||||
kind: Service
|
||||
metadata:
|
||||
name: smg-gateway-qwen
|
||||
namespace: smg-test
|
||||
spec:
|
||||
type: NodePort
|
||||
selector:
|
||||
app: smg-gateway-qwen
|
||||
ports:
|
||||
- name: http
|
||||
port: 30004
|
||||
targetPort: 30004
|
||||
- name: metrics
|
||||
port: 29004
|
||||
targetPort: 29004
|
||||
@@ -0,0 +1,83 @@
|
||||
# Dedicated gateway used by the gateway-restart test in test_lifecycle.py.
|
||||
# Runs on its own ports (30005/29005) so killing this gateway's pod does
|
||||
# not invalidate the session-scoped port-forward against the default
|
||||
# smg-gateway used by test_reconciliation.py.
|
||||
#
|
||||
# This gateway's selector requires `app=fake-worker,lifecycle=restart`, so
|
||||
# pods from test_reconciliation.py (which carry only `app=fake-worker`) are
|
||||
# ignored here. The default `smg-gateway` does still match this test's
|
||||
# `lifecycle=restart` workers (its selector is the broader subset), but the
|
||||
# restart test never queries port 30000 — it only reads from this gateway,
|
||||
# so worker counts stay deterministic.
|
||||
apiVersion: apps/v1
|
||||
kind: Deployment
|
||||
metadata:
|
||||
name: smg-gateway-restart
|
||||
namespace: smg-test
|
||||
spec:
|
||||
replicas: 1
|
||||
selector:
|
||||
matchLabels:
|
||||
app: smg-gateway-restart
|
||||
template:
|
||||
metadata:
|
||||
labels:
|
||||
app: smg-gateway-restart
|
||||
spec:
|
||||
serviceAccountName: smg-gateway
|
||||
containers:
|
||||
- name: gateway
|
||||
image: smg-gateway:test
|
||||
imagePullPolicy: Never
|
||||
args:
|
||||
- "--service-discovery"
|
||||
- "--selector"
|
||||
- "app=fake-worker"
|
||||
- "lifecycle=restart"
|
||||
- "--service-discovery-port"
|
||||
- "8000"
|
||||
- "--service-discovery-namespace"
|
||||
- "smg-test"
|
||||
- "--port"
|
||||
- "30005"
|
||||
- "--prometheus-port"
|
||||
- "29005"
|
||||
- "--disable-health-check"
|
||||
- "--worker-startup-timeout-secs"
|
||||
- "30"
|
||||
- "--log-level"
|
||||
- "debug"
|
||||
ports:
|
||||
- containerPort: 30005
|
||||
name: http
|
||||
- containerPort: 29005
|
||||
name: metrics
|
||||
readinessProbe:
|
||||
httpGet:
|
||||
path: /liveness
|
||||
port: 30005
|
||||
initialDelaySeconds: 3
|
||||
periodSeconds: 3
|
||||
livenessProbe:
|
||||
httpGet:
|
||||
path: /liveness
|
||||
port: 30005
|
||||
initialDelaySeconds: 5
|
||||
periodSeconds: 10
|
||||
---
|
||||
apiVersion: v1
|
||||
kind: Service
|
||||
metadata:
|
||||
name: smg-gateway-restart
|
||||
namespace: smg-test
|
||||
spec:
|
||||
type: NodePort
|
||||
selector:
|
||||
app: smg-gateway-restart
|
||||
ports:
|
||||
- name: http
|
||||
port: 30005
|
||||
targetPort: 30005
|
||||
- name: metrics
|
||||
port: 29005
|
||||
targetPort: 29005
|
||||
@@ -0,0 +1,30 @@
|
||||
# Cluster-wide RBAC for the multi-namespace discovery test.
|
||||
# Distinct ServiceAccount/ClusterRole names so it does not collide with the
|
||||
# namespace-scoped Role in rbac.yaml that the default gateway uses.
|
||||
apiVersion: v1
|
||||
kind: ServiceAccount
|
||||
metadata:
|
||||
name: smg-gateway-cluster
|
||||
namespace: smg-test
|
||||
---
|
||||
apiVersion: rbac.authorization.k8s.io/v1
|
||||
kind: ClusterRole
|
||||
metadata:
|
||||
name: smg-gateway-cluster
|
||||
rules:
|
||||
- apiGroups: [""]
|
||||
resources: ["pods"]
|
||||
verbs: ["get", "list", "watch"]
|
||||
---
|
||||
apiVersion: rbac.authorization.k8s.io/v1
|
||||
kind: ClusterRoleBinding
|
||||
metadata:
|
||||
name: smg-gateway-cluster
|
||||
roleRef:
|
||||
apiGroup: rbac.authorization.k8s.io
|
||||
kind: ClusterRole
|
||||
name: smg-gateway-cluster
|
||||
subjects:
|
||||
- kind: ServiceAccount
|
||||
name: smg-gateway-cluster
|
||||
namespace: smg-test
|
||||
@@ -0,0 +1,266 @@
|
||||
"""Cross-namespace service discovery integration test.
|
||||
|
||||
When the gateway is started without --service-discovery-namespace, the K8s
|
||||
API client falls through to Api::all (sgl-model-gateway/src/service_discovery.rs:279),
|
||||
watching pods in every namespace. That path requires a ClusterRole rather
|
||||
than the namespace-scoped Role used by the default gateway.
|
||||
|
||||
Validates: a single gateway with cluster-wide RBAC discovers workers running
|
||||
in two distinct namespaces, with each worker registered exactly once.
|
||||
|
||||
Run with:
|
||||
cd e2e_test/k8s_integration
|
||||
pytest test_cross_namespace.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
|
||||
FAKE_WORKER_SCRIPT,
|
||||
KUBECTL_CONTEXT,
|
||||
NAMESPACE,
|
||||
_apply_from_stdin,
|
||||
_cleanup_port_forward,
|
||||
_get_workers,
|
||||
_kubectl,
|
||||
_kubectl_json,
|
||||
_poll_until,
|
||||
_port_forward_start,
|
||||
_wait_for_deployment_ready,
|
||||
_wait_for_pod_ready,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
MANIFESTS_DIR = Path(__file__).parent / "manifests"
|
||||
|
||||
CLUSTER_GATEWAY_HTTP_PORT = 30002
|
||||
EXTRA_NAMESPACE = "smg-test-extra"
|
||||
|
||||
|
||||
def _deploy_worker_pod(name: str, namespace: str):
|
||||
"""Deploy a fake-worker pod in the given namespace.
|
||||
|
||||
Assumes the `fake-worker-script` ConfigMap already exists in that namespace.
|
||||
The `cross-ns-test=true` label is paired with the gateway's selector in
|
||||
gateway-cluster-scoped.yaml so this test owns its worker fleet exclusively
|
||||
and the exact-count assertion isn't racy across files.
|
||||
"""
|
||||
pod_manifest = {
|
||||
"apiVersion": "v1",
|
||||
"kind": "Pod",
|
||||
"metadata": {
|
||||
"name": name,
|
||||
"namespace": namespace,
|
||||
"labels": {"app": "fake-worker", "cross-ns-test": "true"},
|
||||
},
|
||||
"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=False,
|
||||
)
|
||||
if proc.returncode != 0:
|
||||
# kubectl's actual error (admission webhook, missing configmap, schema
|
||||
# validation) lives in stderr; surface it instead of bubbling up an
|
||||
# opaque "non-zero exit status 1".
|
||||
raise RuntimeError(
|
||||
f"Failed to deploy pod {name} in namespace {namespace} "
|
||||
f"(rc={proc.returncode}): stderr={proc.stderr.strip()!r}"
|
||||
)
|
||||
logger.info("Deployed worker %s in namespace %s", name, namespace)
|
||||
|
||||
|
||||
def _get_pod_ip(name: str, namespace: str) -> str:
|
||||
"""Return the current podIP for a pod (must be running)."""
|
||||
pod = _kubectl_json("get", "pod", name, "-n", namespace)
|
||||
ip = pod.get("status", {}).get("podIP")
|
||||
if not ip:
|
||||
raise RuntimeError(
|
||||
f"Pod {namespace}/{name} has no podIP yet: {pod.get('status')}"
|
||||
)
|
||||
return ip
|
||||
|
||||
|
||||
def _safe_delete_pod(name: str, namespace: 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 in ns %s: %s", name, namespace, e)
|
||||
|
||||
|
||||
def _ensure_namespace(name: str):
|
||||
"""Create a namespace if it doesn't already exist (apply is idempotent)."""
|
||||
manifest = {
|
||||
"apiVersion": "v1",
|
||||
"kind": "Namespace",
|
||||
"metadata": {"name": name},
|
||||
}
|
||||
_apply_from_stdin(json.dumps(manifest))
|
||||
|
||||
|
||||
def _ensure_fake_worker_configmap(namespace: str):
|
||||
"""Create the fake-worker-script ConfigMap in the given namespace.
|
||||
|
||||
Mirrors the configmap shipped to the default namespace by the deploy_base
|
||||
fixture, so workers in other namespaces can boot the same script.
|
||||
"""
|
||||
cm_yaml = _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_yaml.stdout)
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def cluster_scoped_gateway(deploy_base):
|
||||
"""Deploy the cluster-scoped RBAC + gateway, plus a second namespace.
|
||||
|
||||
Cleanup runs in `finally:` so a port-forward failure does not leak
|
||||
deployments / RBAC objects across test runs.
|
||||
"""
|
||||
rbac_manifest = MANIFESTS_DIR / "rbac-cluster-scoped.yaml"
|
||||
gateway_manifest = MANIFESTS_DIR / "gateway-cluster-scoped.yaml"
|
||||
|
||||
_kubectl("apply", "-f", str(rbac_manifest))
|
||||
_ensure_namespace(EXTRA_NAMESPACE)
|
||||
_ensure_fake_worker_configmap(EXTRA_NAMESPACE)
|
||||
_kubectl("apply", "-f", str(gateway_manifest))
|
||||
|
||||
pf: subprocess.Popen | None = None
|
||||
try:
|
||||
_wait_for_deployment_ready("smg-gateway-cluster")
|
||||
pf = _port_forward_start(
|
||||
NAMESPACE,
|
||||
"smg-gateway-cluster",
|
||||
CLUSTER_GATEWAY_HTTP_PORT,
|
||||
CLUSTER_GATEWAY_HTTP_PORT,
|
||||
)
|
||||
yield f"http://127.0.0.1:{CLUSTER_GATEWAY_HTTP_PORT}"
|
||||
finally:
|
||||
if pf is not None:
|
||||
_cleanup_port_forward("cluster_gateway", pf)
|
||||
_kubectl(
|
||||
"delete", "-f", str(gateway_manifest), "--ignore-not-found", check=False
|
||||
)
|
||||
_kubectl("delete", "-f", str(rbac_manifest), "--ignore-not-found", check=False)
|
||||
# Drop the extra namespace last so any worker pods left behind go with
|
||||
# it. The bounded --timeout prevents pytest from hanging at session
|
||||
# end if a finalizer in the namespace is stuck (kindnet, kube-system
|
||||
# controllers); the cleanup error gets logged but pytest continues.
|
||||
_kubectl(
|
||||
"delete",
|
||||
"namespace",
|
||||
EXTRA_NAMESPACE,
|
||||
"--ignore-not-found",
|
||||
"--wait=true",
|
||||
"--timeout=60s",
|
||||
check=False,
|
||||
)
|
||||
|
||||
|
||||
class TestClusterWideDiscovery:
|
||||
"""Verify that a gateway with ClusterRole + no namespace filter sees pods
|
||||
in every namespace."""
|
||||
|
||||
def test_workers_in_two_namespaces_are_both_discovered(
|
||||
self, cluster_scoped_gateway
|
||||
):
|
||||
gateway_url = cluster_scoped_gateway
|
||||
worker_a = "cross-ns-worker-a"
|
||||
worker_b = "cross-ns-worker-b"
|
||||
|
||||
try:
|
||||
_deploy_worker_pod(worker_a, NAMESPACE)
|
||||
_deploy_worker_pod(worker_b, EXTRA_NAMESPACE)
|
||||
_wait_for_pod_ready(worker_a, namespace=NAMESPACE)
|
||||
_wait_for_pod_ready(worker_b, namespace=EXTRA_NAMESPACE)
|
||||
|
||||
ip_a = _get_pod_ip(worker_a, NAMESPACE)
|
||||
ip_b = _get_pod_ip(worker_b, EXTRA_NAMESPACE)
|
||||
logger.info(
|
||||
"worker_a (%s/%s) ip=%s, worker_b (%s/%s) ip=%s",
|
||||
NAMESPACE,
|
||||
worker_a,
|
||||
ip_a,
|
||||
EXTRA_NAMESPACE,
|
||||
worker_b,
|
||||
ip_b,
|
||||
)
|
||||
|
||||
_poll_until(
|
||||
lambda: _get_workers(gateway_url)["total"] >= 2,
|
||||
"both workers discovered across namespaces",
|
||||
timeout=30,
|
||||
interval=3,
|
||||
)
|
||||
|
||||
workers = _get_workers(gateway_url)
|
||||
urls = sorted(w["url"] for w in workers.get("workers", []))
|
||||
logger.info("Discovered worker URLs: %s", urls)
|
||||
|
||||
# Each pod should appear once. With two pods deployed and two
|
||||
# discovered, the URLs must be distinct (no duplicate registration
|
||||
# from cluster-wide watcher firing twice).
|
||||
assert len(urls) == len(set(urls)), f"Duplicate worker URLs: {urls}"
|
||||
assert workers["total"] == 2, (
|
||||
f"Expected exactly 2 workers across namespaces, "
|
||||
f"got {workers['total']}: {urls}"
|
||||
)
|
||||
|
||||
# Per-pod IP membership: proves the gateway actually listed both
|
||||
# namespaces, not "two pods from the same namespace by accident".
|
||||
# A regression that quietly hardcoded a namespace filter would
|
||||
# still produce total=2 if labels happened to match elsewhere,
|
||||
# but only one of these IPs would surface.
|
||||
assert any(
|
||||
ip_a in u for u in urls
|
||||
), f"worker_a IP {ip_a} (ns {NAMESPACE}) not in {urls}"
|
||||
assert any(
|
||||
ip_b in u for u in urls
|
||||
), f"worker_b IP {ip_b} (ns {EXTRA_NAMESPACE}) not in {urls}"
|
||||
finally:
|
||||
_safe_delete_pod(worker_a, NAMESPACE)
|
||||
_safe_delete_pod(worker_b, EXTRA_NAMESPACE)
|
||||
@@ -0,0 +1,554 @@
|
||||
"""Worker lifecycle integration tests.
|
||||
|
||||
Covers three scenarios that the existing reconciliation/PD tests don't:
|
||||
|
||||
1. Gateway pod restart with persistent workers — verifies the K8s watcher
|
||||
re-discovers existing pods after the gateway restarts, with no duplicate
|
||||
registrations.
|
||||
2. Pod IP change (same pod name, new IP) — verifies the gateway's worker
|
||||
registry tracks the new IP after a pod is force-deleted and recreated,
|
||||
not the stale one.
|
||||
3. Graceful drain — verifies the gateway deregisters a worker as soon as
|
||||
K8s sets `metadata.deletionTimestamp` (handled by handle_pod_deletion in
|
||||
sgl-model-gateway/src/service_discovery.rs:533), instead of waiting for
|
||||
the pod to fully terminate. This is what keeps the registry fresh during
|
||||
long terminationGracePeriodSeconds windows / preStop hooks.
|
||||
|
||||
Run with:
|
||||
cd e2e_test/k8s_integration
|
||||
pytest test_lifecycle.py -v -s
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
import subprocess
|
||||
import time
|
||||
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,
|
||||
_kubectl_json,
|
||||
_poll_until,
|
||||
_port_forward_start,
|
||||
_wait_for_deployment_ready,
|
||||
_wait_for_pod_ready,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
MANIFESTS_DIR = Path(__file__).parent / "manifests"
|
||||
|
||||
RESTART_GATEWAY_HTTP_PORT = 30005
|
||||
|
||||
|
||||
def _deploy_worker_pod(
|
||||
name: str,
|
||||
extra_labels: dict[str, str] | None = None,
|
||||
grace_period_secs: int | None = None,
|
||||
prestop_sleep_secs: int | None = None,
|
||||
):
|
||||
"""Deploy a fake-worker pod with optional extra labels and grace period.
|
||||
|
||||
`prestop_sleep_secs` adds an exec preStop hook (`sleep N`) so the pod
|
||||
stays in the Terminating state for at least N seconds after a graceful
|
||||
delete — long enough for the test to observe the watcher's
|
||||
deletionTimestamp event firing before SIGKILL.
|
||||
"""
|
||||
labels = {"app": "fake-worker"}
|
||||
if extra_labels:
|
||||
labels.update(extra_labels)
|
||||
|
||||
container: dict = {
|
||||
"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"}],
|
||||
}
|
||||
if prestop_sleep_secs is not None:
|
||||
container["lifecycle"] = {
|
||||
"preStop": {"exec": {"command": ["sleep", str(prestop_sleep_secs)]}}
|
||||
}
|
||||
|
||||
spec: dict = {
|
||||
"containers": [container],
|
||||
"volumes": [{"name": "app", "configMap": {"name": "fake-worker-script"}}],
|
||||
}
|
||||
if grace_period_secs is not None:
|
||||
spec["terminationGracePeriodSeconds"] = grace_period_secs
|
||||
|
||||
pod_manifest = {
|
||||
"apiVersion": "v1",
|
||||
"kind": "Pod",
|
||||
"metadata": {
|
||||
"name": name,
|
||||
"namespace": NAMESPACE,
|
||||
"labels": labels,
|
||||
},
|
||||
"spec": spec,
|
||||
}
|
||||
proc = subprocess.run(
|
||||
["kubectl", "--context", KUBECTL_CONTEXT, "apply", "-f", "-"],
|
||||
input=json.dumps(pod_manifest),
|
||||
capture_output=True,
|
||||
text=True,
|
||||
check=False,
|
||||
)
|
||||
if proc.returncode != 0:
|
||||
# Surface kubectl's actual error (webhook denial, schema, etc.)
|
||||
# instead of an opaque "non-zero exit status 1" CalledProcessError.
|
||||
raise RuntimeError(
|
||||
f"Failed to deploy pod {name} (rc={proc.returncode}): "
|
||||
f"stderr={proc.stderr.strip()!r}"
|
||||
)
|
||||
logger.info(
|
||||
"Deployed pod %s (labels=%s, grace=%s, prestop=%s)",
|
||||
name,
|
||||
labels,
|
||||
grace_period_secs,
|
||||
prestop_sleep_secs,
|
||||
)
|
||||
|
||||
|
||||
def _safe_force_delete(name: str):
|
||||
"""Force-delete a pod, swallowing errors (used in cleanup)."""
|
||||
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)
|
||||
|
||||
|
||||
def _wait_for_pod_gone(name: str, timeout: int = 60):
|
||||
"""Wait until a pod no longer exists.
|
||||
|
||||
Uses `kubectl get -o name --ignore-not-found`: empty stdout means the pod
|
||||
is gone (no need to substring-match "NotFound" against stderr, which is
|
||||
locale- and version-fragile). Any non-zero rc is a real cluster error
|
||||
(apiserver unreachable, RBAC drift) and propagates immediately via
|
||||
`_poll_until` (it only retries transient network errors, not RuntimeError).
|
||||
"""
|
||||
|
||||
def _gone() -> bool:
|
||||
result = _kubectl(
|
||||
"get",
|
||||
"pod",
|
||||
name,
|
||||
"-n",
|
||||
NAMESPACE,
|
||||
"-o",
|
||||
"name",
|
||||
"--ignore-not-found",
|
||||
check=False,
|
||||
)
|
||||
if result.returncode != 0:
|
||||
raise RuntimeError(
|
||||
f"kubectl get pod {name} failed unexpectedly "
|
||||
f"(rc={result.returncode}): {result.stderr.strip()}"
|
||||
)
|
||||
return not result.stdout.strip()
|
||||
|
||||
_poll_until(_gone, f"pod {name} gone", timeout=timeout, interval=2)
|
||||
|
||||
|
||||
def _get_pod_ip(name: str) -> str:
|
||||
"""Return the current podIP for a pod (must be running)."""
|
||||
pod = _kubectl_json("get", "pod", name, "-n", NAMESPACE)
|
||||
ip = pod.get("status", {}).get("podIP")
|
||||
if not ip:
|
||||
raise RuntimeError(f"Pod {name} has no podIP yet: {pod.get('status')}")
|
||||
return ip
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def restart_gateway(deploy_base):
|
||||
"""Deploy a dedicated gateway for the restart test on its own ports.
|
||||
|
||||
Function-scoped so each invocation starts from a clean Deployment — that
|
||||
keeps the per-test pod-restart count predictable. The selector is
|
||||
`app=fake-worker,lifecycle=restart`, distinct from the default gateway,
|
||||
so the test owns the worker fleet exclusively.
|
||||
|
||||
Yields a single-element list `[pf]` so the test can swap in a fresh
|
||||
port-forward after killing the gateway pod (the original `pf` exits with
|
||||
the pod). Fixture teardown cleans up whichever handle is current.
|
||||
"""
|
||||
manifest = MANIFESTS_DIR / "gateway-restart.yaml"
|
||||
_kubectl("apply", "-f", str(manifest))
|
||||
pf_holder: list[subprocess.Popen | None] = [None]
|
||||
try:
|
||||
_wait_for_deployment_ready("smg-gateway-restart")
|
||||
pf_holder[0] = _port_forward_start(
|
||||
NAMESPACE,
|
||||
"smg-gateway-restart",
|
||||
RESTART_GATEWAY_HTTP_PORT,
|
||||
RESTART_GATEWAY_HTTP_PORT,
|
||||
)
|
||||
yield f"http://127.0.0.1:{RESTART_GATEWAY_HTTP_PORT}", pf_holder
|
||||
finally:
|
||||
if pf_holder[0] is not None:
|
||||
_cleanup_port_forward("restart_gateway", pf_holder[0])
|
||||
_kubectl("delete", "-f", str(manifest), "--ignore-not-found", check=False)
|
||||
|
||||
|
||||
class TestGatewayRestart:
|
||||
"""Killing the gateway pod must not lose worker state — the new pod
|
||||
re-lists existing workers via the K8s watcher and registers them once."""
|
||||
|
||||
def test_workers_re_discovered_without_duplicates_after_restart(
|
||||
self, restart_gateway
|
||||
):
|
||||
gateway_url, pf_holder = restart_gateway
|
||||
worker_names = ["restart-worker-a", "restart-worker-b", "restart-worker-c"]
|
||||
|
||||
try:
|
||||
for name in worker_names:
|
||||
_deploy_worker_pod(name, extra_labels={"lifecycle": "restart"})
|
||||
for name in worker_names:
|
||||
_wait_for_pod_ready(name)
|
||||
|
||||
_poll_until(
|
||||
lambda: _get_worker_count(gateway_url) >= len(worker_names),
|
||||
f"all {len(worker_names)} workers discovered (pre-restart)",
|
||||
timeout=30,
|
||||
interval=3,
|
||||
)
|
||||
|
||||
urls_before = sorted(
|
||||
w["url"] for w in _get_workers(gateway_url).get("workers", [])
|
||||
)
|
||||
logger.info("Workers before restart: %s", urls_before)
|
||||
assert len(urls_before) == len(worker_names)
|
||||
|
||||
# Kill the gateway pod. The Deployment's ReplicaSet restarts it.
|
||||
# Guard against an empty list — `_wait_for_deployment_ready` above
|
||||
# gates on the rollout, but a controller race could in principle
|
||||
# leave the selector momentarily empty, and an IndexError here
|
||||
# would mask the real problem.
|
||||
res = _kubectl_json(
|
||||
"get",
|
||||
"pod",
|
||||
"-n",
|
||||
NAMESPACE,
|
||||
"-l",
|
||||
"app=smg-gateway-restart",
|
||||
)
|
||||
assert res.get(
|
||||
"items"
|
||||
), "No pods found for selector app=smg-gateway-restart"
|
||||
old_pod = res["items"][0]["metadata"]["name"]
|
||||
_kubectl(
|
||||
"delete",
|
||||
"pod",
|
||||
old_pod,
|
||||
"-n",
|
||||
NAMESPACE,
|
||||
"--force",
|
||||
"--grace-period=0",
|
||||
)
|
||||
|
||||
# Tear down the fixture's port-forward NOW, before the new pod is
|
||||
# ready. Otherwise the stale kubectl process can briefly hold port
|
||||
# 30005 while attempting to reconnect, and our fresh port-forward
|
||||
# below races with it for "address already in use". After cleanup
|
||||
# we clear the holder so the fixture teardown doesn't double-free.
|
||||
if pf_holder[0] is not None:
|
||||
_cleanup_port_forward("restart_gateway_pre_kill", pf_holder[0])
|
||||
pf_holder[0] = None
|
||||
|
||||
_wait_for_deployment_ready("smg-gateway-restart")
|
||||
|
||||
# Open a fresh port-forward to the new gateway pod and stash it
|
||||
# in the holder so the fixture cleans it up even if the assertion
|
||||
# block raises before the explicit teardown.
|
||||
pf_holder[0] = _port_forward_start(
|
||||
NAMESPACE,
|
||||
"smg-gateway-restart",
|
||||
RESTART_GATEWAY_HTTP_PORT,
|
||||
RESTART_GATEWAY_HTTP_PORT,
|
||||
)
|
||||
|
||||
_poll_until(
|
||||
lambda: _get_worker_count(gateway_url) >= len(worker_names),
|
||||
f"all {len(worker_names)} workers re-discovered (post-restart)",
|
||||
timeout=60,
|
||||
interval=3,
|
||||
)
|
||||
|
||||
urls_after = sorted(
|
||||
w["url"] for w in _get_workers(gateway_url).get("workers", [])
|
||||
)
|
||||
logger.info("Workers after restart: %s", urls_after)
|
||||
|
||||
# No duplicates: each pod should appear exactly once.
|
||||
assert len(urls_after) == len(
|
||||
set(urls_after)
|
||||
), f"Duplicate worker registrations after gateway restart: {urls_after}"
|
||||
# Set equality: the same workers come back, neither dropped
|
||||
# nor duplicated.
|
||||
assert set(urls_after) == set(urls_before), (
|
||||
f"Worker set diverged across gateway restart. "
|
||||
f"before={urls_before}, after={urls_after}"
|
||||
)
|
||||
|
||||
# Liveness probe of the new watcher: deploy one more worker AFTER
|
||||
# the relist has converged. The relist can't pick this up — only
|
||||
# an active watch loop can. A regression where the gateway does
|
||||
# the initial list correctly but exits the watch loop afterward
|
||||
# would silently pass the assertions above; this catches it.
|
||||
extra_name = "restart-worker-d"
|
||||
try:
|
||||
_deploy_worker_pod(extra_name, extra_labels={"lifecycle": "restart"})
|
||||
_wait_for_pod_ready(extra_name)
|
||||
_poll_until(
|
||||
lambda: _get_worker_count(gateway_url) >= len(worker_names) + 1,
|
||||
f"post-restart watcher picked up {extra_name}",
|
||||
timeout=30,
|
||||
interval=2,
|
||||
)
|
||||
finally:
|
||||
_safe_force_delete(extra_name)
|
||||
finally:
|
||||
for name in worker_names:
|
||||
_safe_force_delete(name)
|
||||
|
||||
|
||||
class TestPodIpChange:
|
||||
"""When a pod is deleted and recreated with the same name, the gateway's
|
||||
worker registry must drop the old URL and pick up the new one. The test
|
||||
skips on the rare CNI-IP-reuse case (since both URLs would be identical
|
||||
and the assertions can't distinguish), so it does not exercise the
|
||||
"old IP recycled by a different pod" scenario — covering that would need
|
||||
a static-IP harness or controlled IP exhaustion, neither of which is
|
||||
cheap on kind."""
|
||||
|
||||
def test_recreated_pod_yields_current_ip_in_worker_registry(
|
||||
self, gateway_port_forward
|
||||
):
|
||||
gateway_url, _ = gateway_port_forward
|
||||
pod_name = "ip-change-worker"
|
||||
|
||||
try:
|
||||
_deploy_worker_pod(pod_name)
|
||||
_wait_for_pod_ready(pod_name)
|
||||
|
||||
_poll_until(
|
||||
lambda: _get_worker_count(gateway_url) >= 1,
|
||||
"initial worker discovered",
|
||||
timeout=30,
|
||||
interval=3,
|
||||
)
|
||||
|
||||
ip_before = _get_pod_ip(pod_name)
|
||||
urls_before = {w["url"] for w in _get_workers(gateway_url)["workers"]}
|
||||
assert any(
|
||||
ip_before in url for url in urls_before
|
||||
), f"Expected initial worker URL containing {ip_before}, got {urls_before}"
|
||||
logger.info("Pod IP before: %s, urls: %s", ip_before, urls_before)
|
||||
|
||||
# Force-delete and wait until the registry no longer references
|
||||
# this pod's IP. Asserting on "URL contains ip_before" rather than
|
||||
# "total count == 0" keeps the test robust if a parallel/earlier
|
||||
# test left an unrelated worker in the same gateway's registry.
|
||||
_kubectl(
|
||||
"delete",
|
||||
"pod",
|
||||
pod_name,
|
||||
"-n",
|
||||
NAMESPACE,
|
||||
"--force",
|
||||
"--grace-period=0",
|
||||
)
|
||||
_wait_for_pod_gone(pod_name)
|
||||
_poll_until(
|
||||
lambda: not any(
|
||||
ip_before in w["url"]
|
||||
for w in _get_workers(gateway_url).get("workers", [])
|
||||
),
|
||||
f"stale worker for IP {ip_before} removed",
|
||||
timeout=RECONCILIATION_WAIT_SECS,
|
||||
interval=5,
|
||||
)
|
||||
|
||||
# Recreate with same pod name. CNI typically assigns a new IP
|
||||
# since the previous one isn't immediately recycled, but the
|
||||
# assertion below works either way.
|
||||
_deploy_worker_pod(pod_name)
|
||||
_wait_for_pod_ready(pod_name)
|
||||
|
||||
ip_after = _get_pod_ip(pod_name)
|
||||
# If the CNI happened to recycle the same IP, skip — both the
|
||||
# "no URL contains ip_before" wait and the "URL contains ip_after"
|
||||
# assertion describe the same string, so the test would silently
|
||||
# pass without exercising the IP-change path it's meant to cover.
|
||||
if ip_after == ip_before:
|
||||
pytest.skip(
|
||||
f"CNI reused pod IP {ip_before}; cannot validate IP-change "
|
||||
f"path without two distinct IPs"
|
||||
)
|
||||
|
||||
_poll_until(
|
||||
lambda: any(
|
||||
ip_after in w["url"]
|
||||
for w in _get_workers(gateway_url).get("workers", [])
|
||||
),
|
||||
f"recreated worker (IP {ip_after}) discovered",
|
||||
timeout=30,
|
||||
interval=3,
|
||||
)
|
||||
|
||||
urls_after = {w["url"] for w in _get_workers(gateway_url)["workers"]}
|
||||
logger.info("Pod IP after: %s, urls: %s", ip_after, urls_after)
|
||||
|
||||
# The registry must reflect the current pod's IP, and the old IP
|
||||
# must not linger. Tested by URL membership rather than count to
|
||||
# tolerate unrelated workers from other tests.
|
||||
matching_after = [u for u in urls_after if ip_after in u]
|
||||
assert len(matching_after) == 1, (
|
||||
f"Expected exactly one worker URL containing current IP "
|
||||
f"{ip_after}, got {matching_after} (all urls: {urls_after})"
|
||||
)
|
||||
assert not any(
|
||||
ip_before in u for u in urls_after
|
||||
), f"Stale URL with old IP {ip_before} still in registry: {urls_after}"
|
||||
finally:
|
||||
_safe_force_delete(pod_name)
|
||||
|
||||
|
||||
class TestGracefulDrain:
|
||||
"""A graceful `kubectl delete pod` sets metadata.deletionTimestamp; the
|
||||
watcher's `applied_objects()` stream emits the resulting MODIFIED event,
|
||||
which the if-branch at service_discovery.rs:349 routes to
|
||||
`handle_pod_deletion` (defined at service_discovery.rs:533). That removes
|
||||
the worker immediately rather than waiting out the grace period or the
|
||||
eventual SIGKILL → DELETED event (which `applied_objects()` filters out
|
||||
anyway). This test pins that behavior so a future regression doesn't keep
|
||||
a Terminating pod in the registry while kubelet is already running its
|
||||
preStop hook."""
|
||||
|
||||
def test_deregistration_fires_during_grace_period(self, gateway_port_forward):
|
||||
gateway_url, _ = gateway_port_forward
|
||||
pod_name = "graceful-drain-worker"
|
||||
# Grace period and preStop sleep need to be long enough that the
|
||||
# watcher reliably observes the deletion *while the pod is still
|
||||
# Terminating*. 60s grace + 40s sleep gives ~40s of overlap before
|
||||
# SIGKILL — far more than the watcher's typical event latency
|
||||
# (sub-second under normal conditions).
|
||||
grace_secs = 60
|
||||
prestop_sleep = 40
|
||||
|
||||
try:
|
||||
_deploy_worker_pod(
|
||||
pod_name,
|
||||
grace_period_secs=grace_secs,
|
||||
prestop_sleep_secs=prestop_sleep,
|
||||
)
|
||||
_wait_for_pod_ready(pod_name)
|
||||
pod_ip = _get_pod_ip(pod_name)
|
||||
_poll_until(
|
||||
lambda: any(
|
||||
pod_ip in w["url"]
|
||||
for w in _get_workers(gateway_url).get("workers", [])
|
||||
),
|
||||
f"worker {pod_name} (ip {pod_ip}) registered before drain",
|
||||
timeout=30,
|
||||
interval=3,
|
||||
)
|
||||
|
||||
# Graceful delete: no --force, no --grace-period=0. K8s sets
|
||||
# deletionTimestamp and starts the preStop hook + grace timer.
|
||||
delete_started = time.time()
|
||||
_kubectl(
|
||||
"delete",
|
||||
"pod",
|
||||
pod_name,
|
||||
"-n",
|
||||
NAMESPACE,
|
||||
f"--grace-period={grace_secs}",
|
||||
"--wait=false",
|
||||
)
|
||||
|
||||
# Assertion is URL-membership rather than total count: a stale
|
||||
# worker leaked from a prior test would skew a count comparison
|
||||
# without affecting whether *this* pod's IP got removed. The
|
||||
# meaningful timing guarantee (`elapsed < grace_secs`) is below.
|
||||
_poll_until(
|
||||
lambda: not any(
|
||||
pod_ip in w["url"]
|
||||
for w in _get_workers(gateway_url).get("workers", [])
|
||||
),
|
||||
f"worker for ip {pod_ip} deregistered after graceful delete",
|
||||
timeout=RECONCILIATION_WAIT_SECS,
|
||||
interval=1,
|
||||
)
|
||||
elapsed = time.time() - delete_started
|
||||
logger.info(
|
||||
"Worker deregistered %.1fs after delete (grace_period=%ds, "
|
||||
"prestop_sleep=%ds)",
|
||||
elapsed,
|
||||
grace_secs,
|
||||
prestop_sleep,
|
||||
)
|
||||
assert elapsed < grace_secs, (
|
||||
f"Deregistration took {elapsed:.1f}s — that's >= grace_period "
|
||||
f"({grace_secs}s). The gateway should have acted on the "
|
||||
f"watcher's deletionTimestamp event well before the grace "
|
||||
f"period elapsed (and well before the post-grace SIGKILL "
|
||||
f"would emit a Deleted event that applied_objects() filters)."
|
||||
)
|
||||
|
||||
# The pod should still exist (Terminating) at this point: it
|
||||
# has roughly grace_secs - elapsed seconds left. This guards
|
||||
# against a lucky-pass where the test runs so slowly that the
|
||||
# pod is gone by the time we check, which would no longer prove
|
||||
# the deletionTimestamp path. Use `--ignore-not-found` so a
|
||||
# genuinely-missing pod yields rc=0 with empty stdout; any other
|
||||
# non-zero rc is a real cluster error (apiserver blip, RBAC
|
||||
# drift) and must not masquerade as "pod is gone".
|
||||
get_result = _kubectl(
|
||||
"get",
|
||||
"pod",
|
||||
pod_name,
|
||||
"-n",
|
||||
NAMESPACE,
|
||||
"-o",
|
||||
"name",
|
||||
"--ignore-not-found",
|
||||
check=False,
|
||||
)
|
||||
if get_result.returncode != 0:
|
||||
pytest.fail(
|
||||
f"kubectl get pod {pod_name} failed unexpectedly "
|
||||
f"(rc={get_result.returncode}): {get_result.stderr.strip()}"
|
||||
)
|
||||
still_exists = bool(get_result.stdout.strip())
|
||||
assert still_exists, (
|
||||
f"Pod {pod_name} is already gone — the test ran too slowly "
|
||||
f"to distinguish the deletionTimestamp event from pod-gone "
|
||||
f"reconciliation. Increase grace_period or prestop_sleep."
|
||||
)
|
||||
finally:
|
||||
_safe_force_delete(pod_name)
|
||||
@@ -0,0 +1,204 @@
|
||||
"""Multi-model selector isolation integration test.
|
||||
|
||||
Two gateways watch the same namespace with disjoint --selector values
|
||||
(model=llama vs model=qwen). Each must register only the workers carrying
|
||||
its own label and ignore the other gateway's pool — proving that
|
||||
PodInfo::should_include (sgl-model-gateway/src/service_discovery.rs:99)
|
||||
honors the configured label selector when running multiple gateways
|
||||
side-by-side (the typical multi-tenant deployment pattern).
|
||||
|
||||
Run with:
|
||||
cd e2e_test/k8s_integration
|
||||
pytest test_multi_model.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,
|
||||
_cleanup_port_forward,
|
||||
_get_workers,
|
||||
_kubectl,
|
||||
_poll_until,
|
||||
_port_forward_start,
|
||||
_wait_for_deployment_ready,
|
||||
_wait_for_pod_ready,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
MANIFESTS_DIR = Path(__file__).parent / "manifests"
|
||||
|
||||
LLAMA_GATEWAY_HTTP_PORT = 30003
|
||||
QWEN_GATEWAY_HTTP_PORT = 30004
|
||||
|
||||
|
||||
def _deploy_model_worker(name: str, model: str):
|
||||
"""Deploy a fake-worker with both `app=fake-worker` and `model=<...>` labels.
|
||||
|
||||
Including `app=fake-worker` matches the existing baseline label used by
|
||||
the default gateway's selector — but only the *model-specific* gateway
|
||||
has the second `model=...` constraint, so worker-to-gateway mapping is
|
||||
determined by the model label alone.
|
||||
"""
|
||||
pod_manifest = {
|
||||
"apiVersion": "v1",
|
||||
"kind": "Pod",
|
||||
"metadata": {
|
||||
"name": name,
|
||||
"namespace": NAMESPACE,
|
||||
"labels": {"app": "fake-worker", "model": model},
|
||||
},
|
||||
"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=False,
|
||||
)
|
||||
if proc.returncode != 0:
|
||||
# Surface kubectl's actual error (webhook denial, schema, missing
|
||||
# configmap, etc.) instead of an opaque CalledProcessError.
|
||||
raise RuntimeError(
|
||||
f"Failed to deploy worker {name} (rc={proc.returncode}): "
|
||||
f"stderr={proc.stderr.strip()!r}"
|
||||
)
|
||||
logger.info("Deployed worker %s with model=%s", name, model)
|
||||
|
||||
|
||||
def _safe_force_delete(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 multimodel_gateways(deploy_base):
|
||||
"""Deploy llama + qwen gateways and start port-forwards to both.
|
||||
|
||||
Module-scoped: every test in this file shares the deployment to keep the
|
||||
suite fast (gateway startup + readiness is the slowest step). Cleanup
|
||||
runs in `finally:` so a port-forward failure cannot leak Deployments.
|
||||
"""
|
||||
manifest = MANIFESTS_DIR / "gateway-multimodel.yaml"
|
||||
_kubectl("apply", "-f", str(manifest))
|
||||
|
||||
pf_llama: subprocess.Popen | None = None
|
||||
pf_qwen: subprocess.Popen | None = None
|
||||
try:
|
||||
_wait_for_deployment_ready("smg-gateway-llama")
|
||||
_wait_for_deployment_ready("smg-gateway-qwen")
|
||||
pf_llama = _port_forward_start(
|
||||
NAMESPACE,
|
||||
"smg-gateway-llama",
|
||||
LLAMA_GATEWAY_HTTP_PORT,
|
||||
LLAMA_GATEWAY_HTTP_PORT,
|
||||
)
|
||||
pf_qwen = _port_forward_start(
|
||||
NAMESPACE,
|
||||
"smg-gateway-qwen",
|
||||
QWEN_GATEWAY_HTTP_PORT,
|
||||
QWEN_GATEWAY_HTTP_PORT,
|
||||
)
|
||||
yield (
|
||||
f"http://127.0.0.1:{LLAMA_GATEWAY_HTTP_PORT}",
|
||||
f"http://127.0.0.1:{QWEN_GATEWAY_HTTP_PORT}",
|
||||
)
|
||||
finally:
|
||||
if pf_llama is not None:
|
||||
_cleanup_port_forward("llama_gateway", pf_llama)
|
||||
if pf_qwen is not None:
|
||||
_cleanup_port_forward("qwen_gateway", pf_qwen)
|
||||
_kubectl("delete", "-f", str(manifest), "--ignore-not-found", check=False)
|
||||
|
||||
|
||||
class TestMultiModelSelectorIsolation:
|
||||
"""Each gateway sees only the worker pool that matches its selector."""
|
||||
|
||||
def test_each_gateway_sees_only_its_model_pool(self, multimodel_gateways):
|
||||
llama_url, qwen_url = multimodel_gateways
|
||||
llama_workers = ["model-llama-a", "model-llama-b"]
|
||||
qwen_workers = ["model-qwen-a", "model-qwen-b"]
|
||||
|
||||
try:
|
||||
for name in llama_workers:
|
||||
_deploy_model_worker(name, model="llama")
|
||||
for name in qwen_workers:
|
||||
_deploy_model_worker(name, model="qwen")
|
||||
for name in llama_workers + qwen_workers:
|
||||
_wait_for_pod_ready(name)
|
||||
|
||||
_poll_until(
|
||||
lambda: _get_workers(llama_url)["total"] >= len(llama_workers),
|
||||
f"llama gateway sees {len(llama_workers)} workers",
|
||||
timeout=30,
|
||||
interval=3,
|
||||
)
|
||||
_poll_until(
|
||||
lambda: _get_workers(qwen_url)["total"] >= len(qwen_workers),
|
||||
f"qwen gateway sees {len(qwen_workers)} workers",
|
||||
timeout=30,
|
||||
interval=3,
|
||||
)
|
||||
|
||||
llama_view = _get_workers(llama_url)
|
||||
qwen_view = _get_workers(qwen_url)
|
||||
llama_urls = sorted(w["url"] for w in llama_view.get("workers", []))
|
||||
qwen_urls = sorted(w["url"] for w in qwen_view.get("workers", []))
|
||||
logger.info("Llama gateway workers: %s", llama_urls)
|
||||
logger.info("Qwen gateway workers: %s", qwen_urls)
|
||||
|
||||
assert llama_view["total"] == len(llama_workers), (
|
||||
f"Llama gateway should see exactly {len(llama_workers)} workers, "
|
||||
f"got {llama_view['total']}: {llama_urls}"
|
||||
)
|
||||
assert qwen_view["total"] == len(qwen_workers), (
|
||||
f"Qwen gateway should see exactly {len(qwen_workers)} workers, "
|
||||
f"got {qwen_view['total']}: {qwen_urls}"
|
||||
)
|
||||
|
||||
# No URL should appear in both views — that would mean a
|
||||
# selector mismatch leaked a worker into the wrong gateway.
|
||||
cross_talk = set(llama_urls) & set(qwen_urls)
|
||||
assert (
|
||||
not cross_talk
|
||||
), f"Workers leaked across model selectors: {cross_talk}"
|
||||
finally:
|
||||
for name in llama_workers + qwen_workers:
|
||||
_safe_force_delete(name)
|
||||
Reference in New Issue
Block a user