diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 5abd08ba..d3750dcd 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -39,6 +39,14 @@ repos: files: ^rust/.*\.rs$ pass_filenames: false + # Go formatting gate (needs a local toolchain). Only the operator is Go. + # `go vet` and the tests are heavier, so they run in CI. + - id: gofmt + name: gofmt (go) + entry: scripts/check-gofmt.sh + language: script + files: ^deploy/operator/.*\.go$ + # Refuse a commit whose git email is a machine-generated local hostname # (git's fallback when user.email is unset), so an internal build-host # name can't leak into permanent public history. Skipped in CI. diff --git a/README.md b/README.md index 51a97b55..081baa9f 100644 --- a/README.md +++ b/README.md @@ -49,7 +49,8 @@ Around them: - **Multi-engine** — run vLLM, SGLang, or ATOM behind one common serving interface. - **OpenAI- and Anthropic-compatible API** — `/v1/chat/completions`, `/v1/completions`, and `/v1/messages` (Anthropic Messages, translated in-process). -- **Self-registering fleet** — workers register into etcd and heartbeat, so the router works from a live view and never routes to a worker that is gone; run any number of stateless server replicas. +- **Self-registering fleet** — workers register into etcd (or their own Pod annotation on Kubernetes) and heartbeat, so the router works from a live view and never routes to a worker that is gone; run any number of stateless server replicas. +- **Scale without dropping requests** — a worker joins when it is ready and leaves by draining: it announces `DRAINING`, finishes the generations it already accepted, and only then deregisters. Measured on MI355X: a worker stops receiving new work **under a second** after `SIGTERM` while its in-flight 4000-token generations all complete, and adding or removing instances under continuous traffic costs **zero failed requests**. See [Scaling a fleet](https://rocm.docs.amd.com/projects/infera/en/latest/features/scaling.html). - **Kubernetes-native** — an operator reconciles an `InferaDeployment` CRD (aggregated / PD / multi-node), with an optional Gateway API (GAIE) endpoint picker. ## Architecture @@ -146,6 +147,24 @@ kubectl apply -f examples/k8s-deployments/single-node-aggregated.yaml Ready-to-fill deployment templates (single-node, prefill/decode, multi-node TP, GAIE) and their placeholders are in [`examples/k8s-deployments/`](examples/k8s-deployments/README.md). +**Scaling.** There is no scaling controller — workers self-register when ready and deregister when +they drain, and the router routes to whatever is registered at that instant. Scaling is therefore +`kubectl scale` on the CR, or starting and stopping workers; nothing has to be told about it. + +The two directions cost very different things, and it shapes everything built on top: + +| | measured | +|---|---| +| scale up: `docker run` → serving | **140 s**, almost all of it weight loading | +| scale down: `SIGTERM` → stops receiving | **< 1 s** | +| scale down: in-flight generations | run to completion, bounded by `--drain-timeout` | +| adding + removing under traffic | **260 requests, 0 failures** | + +Because a cold start is minutes and the control loop is seconds, **a burst shorter than a cold +start cannot be answered by adding workers** — keep headroom, or steer traffic to instances that +are already running. Infera does not ship an autoscaler; [Scaling a +fleet](https://rocm.docs.amd.com/projects/infera/en/latest/features/scaling.html) documents what is in place for one, and what is not. + ## Engine images Prebuilt images are published to the `rocm/infera` repository on diff --git a/deploy/operator/cmd/main.go b/deploy/operator/cmd/main.go index b0da1976..f370c031 100644 --- a/deploy/operator/cmd/main.go +++ b/deploy/operator/cmd/main.go @@ -9,9 +9,9 @@ import ( "flag" "os" - clientgoscheme "k8s.io/client-go/kubernetes/scheme" "k8s.io/apimachinery/pkg/runtime" utilruntime "k8s.io/apimachinery/pkg/util/runtime" + clientgoscheme "k8s.io/client-go/kubernetes/scheme" ctrl "sigs.k8s.io/controller-runtime" "sigs.k8s.io/controller-runtime/pkg/healthz" "sigs.k8s.io/controller-runtime/pkg/log/zap" diff --git a/deploy/operator/internal/controller/apply_idempotence_test.go b/deploy/operator/internal/controller/apply_idempotence_test.go new file mode 100644 index 00000000..3910a535 --- /dev/null +++ b/deploy/operator/internal/controller/apply_idempotence_test.go @@ -0,0 +1,243 @@ +/* +Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. + +SPDX-License-Identifier: MIT +*/ + +package controller + +import ( + "context" + "testing" + + "k8s.io/apimachinery/pkg/apis/meta/v1/unstructured" + "k8s.io/apimachinery/pkg/types" + "sigs.k8s.io/controller-runtime/pkg/client/fake" + + inferav1alpha1 "github.com/amd/infera/deploy/operator/api/v1alpha1" +) + +// The operator writes three fields of a LeaderWorkerSet; the API server fills +// in the rest from the CRD's defaults -- startupPolicy, rolloutStrategy, +// leaderWorkerTemplate.restartPolicy and more, nine of them on LWS v1. +// +// Replacing the whole .spec strips every one of those on each pass, the API +// server puts them back, and the next pass strips them again. That was a +// wasted write every resync; it becomes a hot loop now that the reconciler +// watches LeaderWorkerSet, because the write it just made enqueues the +// request that makes the next one. +// +// So: reconciling an object that is already in the desired state must not +// write to it. +func TestApplyingTheSameLwsTwiceDoesNotWriteAgain(t *testing.T) { + s := testScheme(t) + if err := inferav1alpha1.AddToScheme(s); err != nil { + t.Fatalf("scheme: %v", err) + } + c := fake.NewClientBuilder().WithScheme(s).Build() + r := &InferaDeploymentReconciler{Client: c, Scheme: s} + ctx := context.Background() + + idep := &inferav1alpha1.InferaDeployment{} + idep.Name = "qwen" + idep.Namespace = "default" + idep.UID = "uid-1" + + desired := func() *unstructured.Unstructured { + u := &unstructured.Unstructured{} + u.SetGroupVersionKind(lwsGVK()) + u.SetName("qwen-worker") + u.SetNamespace("default") + _ = unstructured.SetNestedField(u.Object, int64(2), "spec", "replicas") + _ = unstructured.SetNestedField(u.Object, int64(2), + "spec", "leaderWorkerTemplate", "size") + return u + } + + if err := r.applyUnstructured(ctx, idep, desired()); err != nil { + t.Fatalf("first apply: %v", err) + } + + // Stand in for the API server defaulting the fields the operator omits. + live := &unstructured.Unstructured{} + live.SetGroupVersionKind(lwsGVK()) + if err := c.Get(ctx, types.NamespacedName{Name: "qwen-worker", Namespace: "default"}, live); err != nil { + t.Fatalf("get after create: %v", err) + } + _ = unstructured.SetNestedField(live.Object, "LeaderCreated", "spec", "startupPolicy") + _ = unstructured.SetNestedField(live.Object, "RollingUpdate", "spec", "rolloutStrategy", "type") + _ = unstructured.SetNestedField(live.Object, "RecreateGroupOnPodRestart", + "spec", "leaderWorkerTemplate", "restartPolicy") + if err := c.Update(ctx, live); err != nil { + t.Fatalf("apply defaults: %v", err) + } + before := live.GetResourceVersion() + + if err := r.applyUnstructured(ctx, idep, desired()); err != nil { + t.Fatalf("second apply: %v", err) + } + + after := &unstructured.Unstructured{} + after.SetGroupVersionKind(lwsGVK()) + if err := c.Get(ctx, types.NamespacedName{Name: "qwen-worker", Namespace: "default"}, after); err != nil { + t.Fatalf("get after reconcile: %v", err) + } + + if got := after.GetResourceVersion(); got != before { + t.Errorf("reconcile rewrote an unchanged object (resourceVersion %s -> %s); "+ + "with the LWS watch registered this is a write loop", before, got) + } + for _, f := range [][]string{ + {"spec", "startupPolicy"}, + {"spec", "rolloutStrategy", "type"}, + {"spec", "leaderWorkerTemplate", "restartPolicy"}, + } { + if v, ok, _ := unstructured.NestedString(after.Object, f...); !ok || v == "" { + t.Errorf("%v was stripped; the API server will re-default it and the "+ + "next pass strips it again", f) + } + } +} + +// Merging must not turn into "never update": a genuine spec change still has +// to reach the child, or scaling through the CR would silently do nothing. +func TestApplyStillPushesAChangedField(t *testing.T) { + s := testScheme(t) + c := fake.NewClientBuilder().WithScheme(s).Build() + r := &InferaDeploymentReconciler{Client: c, Scheme: s} + ctx := context.Background() + + idep := &inferav1alpha1.InferaDeployment{} + idep.Name = "qwen" + idep.Namespace = "default" + idep.UID = "uid-1" + + build := func(replicas int64) *unstructured.Unstructured { + u := &unstructured.Unstructured{} + u.SetGroupVersionKind(lwsGVK()) + u.SetName("qwen-worker") + u.SetNamespace("default") + _ = unstructured.SetNestedField(u.Object, replicas, "spec", "replicas") + return u + } + + if err := r.applyUnstructured(ctx, idep, build(2)); err != nil { + t.Fatalf("create: %v", err) + } + if err := r.applyUnstructured(ctx, idep, build(5)); err != nil { + t.Fatalf("scale: %v", err) + } + + got := &unstructured.Unstructured{} + got.SetGroupVersionKind(lwsGVK()) + if err := c.Get(ctx, types.NamespacedName{Name: "qwen-worker", Namespace: "default"}, got); err != nil { + t.Fatalf("get: %v", err) + } + if v, _, _ := unstructured.NestedInt64(got.Object, "spec", "replicas"); v != 5 { + t.Fatalf("replicas = %d, want 5 -- scaling through the CR did not land", v) + } +} + +// A field the builder emits only when the CR asks for it -- HTTPRoute's +// hostnames -- has to disappear from the child when it disappears from the CR. +// Merging alone cannot do that: an absent field looks the same as one the +// operator does not manage, so the old value would survive and the route would +// keep matching a host the user deleted. +func TestRemovingAConditionalFieldClearsItOnTheChild(t *testing.T) { + s := testScheme(t) + c := fake.NewClientBuilder().WithScheme(s).Build() + r := &InferaDeploymentReconciler{Client: c, Scheme: s} + ctx := context.Background() + + idep := &inferav1alpha1.InferaDeployment{} + idep.Name = "qwen" + idep.Namespace = "default" + idep.UID = "uid-1" + + route := func(hostnames []any) *unstructured.Unstructured { + u := &unstructured.Unstructured{} + u.SetAPIVersion(httpRouteAPIVersion) + u.SetKind(httpRouteKind) + u.SetName("qwen-route") + u.SetNamespace("default") + spec := map[string]any{"rules": []any{}} + if len(hostnames) > 0 { + spec["hostnames"] = hostnames + } + _ = unstructured.SetNestedMap(u.Object, spec, "spec") + return u + } + + if err := r.applyUnstructured(ctx, idep, route([]any{"a.example.com"})); err != nil { + t.Fatalf("create with hostnames: %v", err) + } + if err := r.applyUnstructured(ctx, idep, route(nil)); err != nil { + t.Fatalf("reapply without hostnames: %v", err) + } + + got := &unstructured.Unstructured{} + got.SetAPIVersion(httpRouteAPIVersion) + got.SetKind(httpRouteKind) + if err := c.Get(ctx, types.NamespacedName{Name: "qwen-route", Namespace: "default"}, got); err != nil { + t.Fatalf("get: %v", err) + } + if v, ok, _ := unstructured.NestedSlice(got.Object, "spec", "hostnames"); ok { + t.Fatalf("hostnames still %v after removal from the CR; the route keeps "+ + "matching a host the user deleted", v) + } +} + +// Pruning owned fields must not start pruning the server's defaults again -- +// that is the write loop this merge exists to stop. +func TestPruningLeavesServerDefaultsAlone(t *testing.T) { + s := testScheme(t) + c := fake.NewClientBuilder().WithScheme(s).Build() + r := &InferaDeploymentReconciler{Client: c, Scheme: s} + ctx := context.Background() + + idep := &inferav1alpha1.InferaDeployment{} + idep.Name = "qwen" + idep.Namespace = "default" + idep.UID = "uid-1" + + desired := func() *unstructured.Unstructured { + u := &unstructured.Unstructured{} + u.SetGroupVersionKind(lwsGVK()) + u.SetName("qwen-worker") + u.SetNamespace("default") + _ = unstructured.SetNestedField(u.Object, int64(2), "spec", "replicas") + _ = unstructured.SetNestedField(u.Object, int64(2), "spec", "leaderWorkerTemplate", "size") + return u + } + + if err := r.applyUnstructured(ctx, idep, desired()); err != nil { + t.Fatalf("create: %v", err) + } + + live := &unstructured.Unstructured{} + live.SetGroupVersionKind(lwsGVK()) + if err := c.Get(ctx, types.NamespacedName{Name: "qwen-worker", Namespace: "default"}, live); err != nil { + t.Fatalf("get: %v", err) + } + _ = unstructured.SetNestedField(live.Object, "LeaderCreated", "spec", "startupPolicy") + if err := c.Update(ctx, live); err != nil { + t.Fatalf("apply defaults: %v", err) + } + before := live.GetResourceVersion() + + if err := r.applyUnstructured(ctx, idep, desired()); err != nil { + t.Fatalf("second apply: %v", err) + } + + after := &unstructured.Unstructured{} + after.SetGroupVersionKind(lwsGVK()) + if err := c.Get(ctx, types.NamespacedName{Name: "qwen-worker", Namespace: "default"}, after); err != nil { + t.Fatalf("get: %v", err) + } + if v, ok, _ := unstructured.NestedString(after.Object, "spec", "startupPolicy"); !ok || v == "" { + t.Error("a server default the operator does not set was pruned") + } + if got := after.GetResourceVersion(); got != before { + t.Errorf("reconcile rewrote an unchanged object (%s -> %s)", before, got) + } +} diff --git a/deploy/operator/internal/controller/builders.go b/deploy/operator/internal/controller/builders.go index d88384a3..c28c3376 100644 --- a/deploy/operator/internal/controller/builders.go +++ b/deploy/operator/internal/controller/builders.go @@ -7,13 +7,18 @@ package controller import ( "fmt" + "math" + "strconv" + "strings" appsv1 "k8s.io/api/apps/v1" corev1 "k8s.io/api/core/v1" rbacv1 "k8s.io/api/rbac/v1" + "k8s.io/apimachinery/pkg/api/meta" "k8s.io/apimachinery/pkg/api/resource" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" "k8s.io/apimachinery/pkg/apis/meta/v1/unstructured" + "k8s.io/apimachinery/pkg/runtime/schema" "k8s.io/apimachinery/pkg/util/intstr" inferav1alpha1 "github.com/amd/infera/deploy/operator/api/v1alpha1" @@ -30,19 +35,166 @@ const ( lwsKind = "LeaderWorkerSet" // Graceful rolling-upgrade tuning for GPU worker pods. - workerPreStopDrainSeconds = 15 // preStop sleep: let the router drop us before SIGTERM - workerTerminationGraceSeconds int64 = 120 // must exceed preStop + the worker --drain-timeout + workerPreStopDrainSeconds = 15 // preStop sleep: let the router drop us before SIGTERM + workerDefaultDrainTimeoutSeconds = 30 // matches the worker's --drain-timeout default + // Teardown after the drain finishes: stopping the KV plane and + // engine.stop(), which SIGTERMs the engine's process group and waits up + // to 30s before escalating to SIGKILL. + workerTeardownHeadroomSeconds = 50 + // Floor, so short drain timeouts still leave room for a slow engine exit. + workerTerminationGraceSeconds int64 = 120 + + // The worker reads this as the default for --drain-timeout, so it sets the + // drain just as effectively as the flag does. + drainTimeoutEnvVar = "INFERA_DRAIN_TIMEOUT" +) + +// drainSeconds parses a worker --drain-timeout value. The worker takes a +// float; round up so a fractional value never shortens the budget. +// Ceiling on a parsed drain timeout, and therefore on the grace period derived +// from it. An hour is far past any real generation; beyond that the value is +// more likely a typo than an intention, and it is written into +// terminationGracePeriodSeconds, where too large means a stuck Pod that only +// `--force` can delete. +const maxDrainTimeoutSeconds = 3600 + +func drainSeconds(v string) (int, bool) { + f, err := strconv.ParseFloat(v, 64) + // NaN fails every comparison, so `f <= 0` does not catch it, and neither it + // nor an infinity survives the conversion below: Go leaves out-of-range + // float-to-int implementation-defined, and on amd64 both land on minInt64, + // which the floor then quietly turns back into the default budget. Refusing + // them means the fallback is at least a deliberate one. + if err != nil || math.IsNaN(f) || math.IsInf(f, 0) || f <= 0 { + return 0, false + } + if f > maxDrainTimeoutSeconds { + return maxDrainTimeoutSeconds, true + } + return int(math.Ceil(f)), true +} + +// graceSecondsFor sizes terminationGracePeriodSeconds so the kubelet cannot +// SIGKILL a worker in the middle of shutting down. +// +// The budget is preStop + the worker's --drain-timeout + teardown. That last +// term is not small: engine.stop() alone waits up to 30s for the engine's +// process group before escalating. Leaving the grace at a fixed 120s was fine +// for the default 30s drain, but --drain-timeout lives in free-form args that +// nothing here parsed -- so raising it for long generations (the exact reason +// anyone raises it) silently pushed shutdown past the grace and turned a +// graceful drain back into a kill. +// +// The drain can be set two ways and both have to be read. The worker takes +// $INFERA_DRAIN_TIMEOUT as the flag's *default*, so an env var raises the drain +// exactly as effectively as the flag does -- and parsing only the flag left the +// same silent overrun through a different door. +// +// Sources are given in increasing priority, and precedence is resolved per +// source rather than globally. It has to be: on the extraPodSpec path the +// template is passed through verbatim, so a --drain-timeout in ServiceSpec.Args +// is never rendered into the container and does not affect the drain at all. +// Letting that inert flag outrank the variable the container really reads sizes +// the budget for a drain that never happens, while the real one runs long and +// is killed partway through -- the exact failure this function exists to stop. +func graceSecondsFor(sources ...drainSource) int64 { + drain := workerDefaultDrainTimeoutSeconds + for _, s := range sources { + if d, ok := s.drainTimeout(); ok { + drain = d + } + } + need := int64(workerPreStopDrainSeconds + drain + workerTeardownHeadroomSeconds) + if need < workerTerminationGraceSeconds { + return workerTerminationGraceSeconds + } + return need +} + +// drainSource is one place a drain timeout can be configured: a set of args and +// env vars that travel together, either both from ServiceSpec or both from the +// container itself. +type drainSource struct { + args []string + env []corev1.EnvVar +} + +// drainTimeout resolves this source alone, reporting whether it set anything. +// Env first so an explicit flag overrides it, matching argparse: the variable +// supplies the default, the flag replaces it. +func (s drainSource) drainTimeout() (int, bool) { + out, found := 0, false + for _, e := range s.env { + if e.Name != drainTimeoutEnvVar { + continue + } + // A valueFrom reference is resolved by the kubelet, not here, so its + // value is unknowable at build time and the budget falls back to the + // flag or the default. Worth knowing if a drain is ever cut short + // despite a ConfigMap saying otherwise. + if d, ok := drainSeconds(e.Value); ok { + out, found = d, true + } + } + for i, a := range s.args { + v := "" + if a == "--drain-timeout" && i+1 < len(s.args) { + v = s.args[i+1] + } else if strings.HasPrefix(a, "--drain-timeout=") { + v = strings.TrimPrefix(a, "--drain-timeout=") + } + if v == "" { + continue + } + if d, ok := drainSeconds(v); ok { + out, found = d, true + } + } + return out, found +} + +// Identity labels on every workload this operator builds. They are the only +// link back from a Deployment/LeaderWorkerSet to the CR and service that +// produced it, so the watch handlers that map a workload event to the objects +// interested in it read these rather than re-deriving the name. +const ( + labelKeyDeployment = "infera.amd.com/deployment" + labelKeyService = "infera.amd.com/service" ) // labelsFor returns the selector/identity labels for a service's workload. func labelsFor(idepName, svcName string) map[string]string { return map[string]string{ "app.kubernetes.io/managed-by": "infera-operator", - "infera.amd.com/deployment": idepName, - "infera.amd.com/service": svcName, + labelKeyDeployment: idepName, + labelKeyService: svcName, } } +// lwsInstalled reports whether the LeaderWorkerSet CRD is served by the API. +// +// It gates registering a watch on LWS: controller-runtime builds an informer +// for every watched type at startup, and one for a kind the API server does not +// serve fails the manager outright. LWS is an optional dependency here -- only +// multi-node services use it -- so a single-node cluster without the CRD must +// still be able to run the operator. +// +// The check runs once, at setup. Installing the CRD afterwards therefore needs +// an operator restart to pick up the watch; until then multi-node status still +// refreshes on the reconciler's periodic resync, just not immediately. +func lwsInstalled(mapper meta.RESTMapper) bool { + _, err := mapper.RESTMapping( + schema.GroupKind{Group: lwsGVK().Group, Kind: lwsGVK().Kind}, lwsGVK().Version) + return err == nil +} + +// lwsObject returns an empty LeaderWorkerSet for use as a watch target. +func lwsObject() *unstructured.Unstructured { + u := &unstructured.Unstructured{} + u.SetGroupVersionKind(lwsGVK()) + return u +} + // podLabelsFor returns the operator's selector labels merged with any // caller-supplied ServiceSpec.PodLabels (e.g. an external orchestrator's // workload-id label used by its pod syncer). Operator selector labels always @@ -216,11 +368,25 @@ var mainContainerNames = map[string]struct{}{"main": {}, "infera": {}} // generations, plus a /health readiness probe for single-node workers (skipped // for multi-node LWS groups whose follower ranks > 0 do not serve /health). // Existing values are preserved; the grace is only raised, never lowered. -func injectWorkerRolloutDefaults(spec *corev1.PodSpec, idx int, port int32, addReadiness bool) { +func injectWorkerRolloutDefaults( + spec *corev1.PodSpec, idx int, port int32, addReadiness bool, + args []string, env []corev1.EnvVar, +) { if idx < 0 || idx >= len(spec.Containers) { return } c := &spec.Containers[idx] + // The drain can arrive two ways: via ServiceSpec.Args/Env on the rendered + // path, or written straight into the container by an extraPodSpec template, + // which is passed through verbatim. Reading only the first would miss + // exactly the deployments most likely to have tuned it. They stay separate + // sources, listed in increasing priority, because the container's is what + // the process actually reads -- see graceSecondsFor. + fromService := drainSource{args: args, env: env} + fromContainer := drainSource{ + args: append(append([]string{}, c.Command...), c.Args...), + env: c.Env, + } if addReadiness && c.ReadinessProbe == nil { // SGLang's /health runs a tiny prefill self-check that often takes // >1s, so a 1s probe timeout (the k8s default) flaps the pod between @@ -246,8 +412,9 @@ func injectWorkerRolloutDefaults(spec *corev1.PodSpec, idx int, port int32, addR }, } } - if spec.TerminationGracePeriodSeconds == nil || *spec.TerminationGracePeriodSeconds < workerTerminationGraceSeconds { - grace := workerTerminationGraceSeconds + if want := graceSecondsFor(fromService, fromContainer); spec.TerminationGracePeriodSeconds == nil || + *spec.TerminationGracePeriodSeconds < want { + grace := want spec.TerminationGracePeriodSeconds = &grace } } @@ -296,7 +463,7 @@ func podTemplateFromExtra(idep *inferav1alpha1.InferaDeployment, svcName string, // Graceful rolling-upgrade defaults for worker pods rendered by an external // template: inject readiness/preStop/grace the template omitted. if svc.ComponentType == inferav1alpha1.ComponentTypeWorker { - injectWorkerRolloutDefaults(&spec, idx, port, svc.NumberOfNodes <= 1 && !svc.SkipReadinessProbe) + injectWorkerRolloutDefaults(&spec, idx, port, svc.NumberOfNodes <= 1 && !svc.SkipReadinessProbe, svc.Args, svc.Env) } return corev1.PodTemplateSpec{ ObjectMeta: metav1.ObjectMeta{Labels: podLabelsFor(idep.Name, svcName, svc)}, @@ -355,7 +522,7 @@ func podTemplate(idep *inferav1alpha1.InferaDeployment, svcName string, svc infe // readiness is skipped for multi-node LWS groups (follower ranks have no // /health). The server (CPU-only) keeps the default fast shutdown. if svc.ComponentType == inferav1alpha1.ComponentTypeWorker { - injectWorkerRolloutDefaults(&podSpec, 0, port, svc.NumberOfNodes <= 1 && !svc.SkipReadinessProbe) + injectWorkerRolloutDefaults(&podSpec, 0, port, svc.NumberOfNodes <= 1 && !svc.SkipReadinessProbe, svc.Args, svc.Env) } tmpl := corev1.PodTemplateSpec{ ObjectMeta: metav1.ObjectMeta{Labels: podLabelsFor(idep.Name, svcName, svc)}, diff --git a/deploy/operator/internal/controller/builders_test.go b/deploy/operator/internal/controller/builders_test.go new file mode 100644 index 00000000..daf0d36d --- /dev/null +++ b/deploy/operator/internal/controller/builders_test.go @@ -0,0 +1,250 @@ +/* +Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. + +SPDX-License-Identifier: MIT +*/ + +package controller + +import ( + "testing" + + corev1 "k8s.io/api/core/v1" +) + +// The grace period is the only thing standing between a graceful drain and a +// SIGKILL halfway through one. It has to cover preStop, the worker's own +// --drain-timeout, and the teardown that follows -- of which engine.stop() +// alone can take 30s waiting on the engine's process group. +// +// The failure this guards against is quiet: raising --drain-timeout is exactly +// what an operator does when generations are long, and until the grace was +// derived from it that made shutdown *less* graceful, not more. +func TestGraceSecondsFor(t *testing.T) { + cases := []struct { + name string + args []string + want int64 + }{ + {"no args uses the floor", nil, 120}, + {"default drain stays at the floor", []string{"--drain-timeout", "30"}, 120}, + { + "a long drain raises the grace above the floor", + []string{"--model-path", "/m", "--drain-timeout", "120"}, + 185, // 15 preStop + 120 drain + 50 teardown + }, + {"equals form is parsed too", []string{"--drain-timeout=120"}, 185}, + { + "fractional values round up rather than shortening the budget", + []string{"--drain-timeout", "60.5"}, + 126, // 15 + 61 + 50 + }, + {"a short drain does not lower the floor", []string{"--drain-timeout", "1"}, 120}, + {"garbage falls back to the default", []string{"--drain-timeout", "abc"}, 120}, + {"a trailing flag with no value is ignored", []string{"--drain-timeout"}, 120}, + {"non-positive is ignored", []string{"--drain-timeout", "0"}, 120}, + {"the last occurrence wins", []string{"--drain-timeout", "5", "--drain-timeout", "200"}, 265}, + } + for _, c := range cases { + t.Run(c.name, func(t *testing.T) { + if got := graceSecondsFor(drainSource{args: c.args}); got != c.want { + t.Fatalf("graceSecondsFor(%v) = %d, want %d", c.args, got, c.want) + } + }) + } +} + +// The budget must actually hold, not merely be larger than the old constant. +func TestGraceCoversTheWholeShutdown(t *testing.T) { + for _, drain := range []int{30, 60, 120, 300} { + args := []string{"--drain-timeout", itoa(drain)} + grace := graceSecondsFor(drainSource{args: args}) + need := int64(workerPreStopDrainSeconds + drain + workerTeardownHeadroomSeconds) + if grace < need { + t.Fatalf("drain=%d: grace %d < required %d -- kubelet would SIGKILL mid-drain", + drain, grace, need) + } + } +} + +func itoa(i int) string { + if i == 0 { + return "0" + } + var b []byte + for i > 0 { + b = append([]byte{byte('0' + i%10)}, b...) + i /= 10 + } + return string(b) +} + +// extraPodSpec templates are passed through verbatim, so --drain-timeout may +// live on the container rather than in ServiceSpec.Args. Reading only the +// latter would miss precisely the deployments that tuned it. +func TestGraceReadsDrainTimeoutFromTheContainerToo(t *testing.T) { + spec := &corev1.PodSpec{Containers: []corev1.Container{{ + Name: "main", + Command: []string{"python3", "-m", "infera.engine.sglang"}, + Args: []string{"--model-path", "/m", "--drain-timeout", "240"}, + }}} + injectWorkerRolloutDefaults(spec, 0, 8080, false, nil, nil) + if spec.TerminationGracePeriodSeconds == nil { + t.Fatal("grace not set") + } + want := int64(workerPreStopDrainSeconds + 240 + workerTeardownHeadroomSeconds) + if *spec.TerminationGracePeriodSeconds != want { + t.Fatalf("grace = %d, want %d", *spec.TerminationGracePeriodSeconds, want) + } +} + +// The worker takes $INFERA_DRAIN_TIMEOUT as the default for --drain-timeout, so +// setting it raises the drain exactly as the flag does. Sizing the grace from +// the flag alone left the same silent overrun through a different door: the +// worker would drain for its full timeout and be SIGKILLed partway through. +func TestGraceReadsDrainTimeoutFromTheEnvironment(t *testing.T) { + env := []corev1.EnvVar{ + {Name: "HF_HOME", Value: "/models"}, + {Name: drainTimeoutEnvVar, Value: "300"}, + } + want := int64(workerPreStopDrainSeconds + 300 + workerTeardownHeadroomSeconds) + if got := graceSecondsFor(drainSource{env: env}); got != want { + t.Fatalf("env-set drain: grace = %d, want %d", got, want) + } +} + +// argparse reads the variable as the flag's default, so an explicit flag wins. +// Sizing the budget off the larger of the two would be safe but wrong, and +// wrong here means a pod that lingers minutes longer than its config says. +func TestGraceFlagOverridesTheEnvironment(t *testing.T) { + env := []corev1.EnvVar{{Name: drainTimeoutEnvVar, Value: "300"}} + args := []string{"--drain-timeout", "60"} + want := int64(workerPreStopDrainSeconds + 60 + workerTeardownHeadroomSeconds) + if got := graceSecondsFor(drainSource{args: args, env: env}); got != want { + t.Fatalf("flag with env set: grace = %d, want the flag's %d", got, want) + } +} + +func TestGraceIgnoresUnreadableEnvValues(t *testing.T) { + // valueFrom resolves in the kubelet; nothing is readable here, so the + // budget has to fall back rather than treat the empty value as zero. + from := []corev1.EnvVar{{ + Name: drainTimeoutEnvVar, + ValueFrom: &corev1.EnvVarSource{ + ConfigMapKeyRef: &corev1.ConfigMapKeySelector{Key: "drain"}, + }, + }} + if got := graceSecondsFor(drainSource{env: from}); got != workerTerminationGraceSeconds { + t.Fatalf("valueFrom: grace = %d, want the floor %d", got, workerTerminationGraceSeconds) + } + for _, v := range []string{"", "abc", "0", "-5"} { + env := []corev1.EnvVar{{Name: drainTimeoutEnvVar, Value: v}} + if got := graceSecondsFor(drainSource{env: env}); got != workerTerminationGraceSeconds { + t.Fatalf("env %q: grace = %d, want the floor %d", v, got, workerTerminationGraceSeconds) + } + } +} + +// An extraPodSpec template is passed through verbatim, so the variable may sit +// on the container rather than in ServiceSpec.Env -- the same asymmetry the +// flag has, and the deployments most likely to have tuned the drain. +func TestGraceReadsDrainEnvFromTheContainerToo(t *testing.T) { + spec := &corev1.PodSpec{Containers: []corev1.Container{{ + Name: "main", + Env: []corev1.EnvVar{{Name: drainTimeoutEnvVar, Value: "240"}}, + }}} + injectWorkerRolloutDefaults(spec, 0, 8080, false, nil, nil) + if spec.TerminationGracePeriodSeconds == nil { + t.Fatal("grace not set") + } + want := int64(workerPreStopDrainSeconds + 240 + workerTeardownHeadroomSeconds) + if *spec.TerminationGracePeriodSeconds != want { + t.Fatalf("grace = %d, want %d", *spec.TerminationGracePeriodSeconds, want) + } +} + +// On the extraPodSpec path the template is passed through verbatim, so +// ServiceSpec.Args is never rendered into the container -- a --drain-timeout +// sitting there is inert. It must not outrank the variable the container will +// actually read, or the budget is sized for a drain that never happens while +// the real one runs long and gets SIGKILLed partway through. Precedence is by +// source: what the process sees wins, and only within a source does a flag +// beat a variable. +func TestAnInertServiceSpecFlagDoesNotOutrankTheContainer(t *testing.T) { + spec := &corev1.PodSpec{Containers: []corev1.Container{{ + Name: "main", + Env: []corev1.EnvVar{{Name: drainTimeoutEnvVar, Value: "600"}}, + }}} + injectWorkerRolloutDefaults(spec, 0, 8080, false, []string{"--drain-timeout", "30"}, nil) + + want := int64(workerPreStopDrainSeconds + 600 + workerTeardownHeadroomSeconds) + if got := *spec.TerminationGracePeriodSeconds; got != want { + t.Fatalf("grace = %d, want %d -- the container drains for 600s, so %d "+ + "leaves the kubelet killing it partway through", got, want, got) + } +} + +// The same precedence, the other way round: a flag the container really runs +// with beats a variable from ServiceSpec. +func TestTheContainerFlagBeatsAServiceSpecVariable(t *testing.T) { + spec := &corev1.PodSpec{Containers: []corev1.Container{{ + Name: "main", + Args: []string{"--drain-timeout=300"}, + }}} + env := []corev1.EnvVar{{Name: drainTimeoutEnvVar, Value: "45"}} + injectWorkerRolloutDefaults(spec, 0, 8080, false, nil, env) + + want := int64(workerPreStopDrainSeconds + 300 + workerTeardownHeadroomSeconds) + if got := *spec.TerminationGracePeriodSeconds; got != want { + t.Fatalf("grace = %d, want %d", got, want) + } +} + +// A drain timeout arrives as free-form text from args or an env var, so the +// parse has to survive whatever is there. Two directions matter. +// +// Below: Go leaves float-to-int conversion implementation-defined when the +// value does not fit, and on amd64 `inf` and `NaN` both land on minInt64. The +// floor then hides it, so a worker configured with an unusable value silently +// gets the default budget instead of anything signalling a mistake. Python's +// argparse accepts `inf` as a float, so this is reachable. +// +// Above: nothing bounded the result, so a typo like 86400 renders a Pod that +// takes a day to delete, and 9e18 overflows into a nonsensical grace period. +func TestDrainTimeoutRejectsValuesItCannotUse(t *testing.T) { + for _, v := range []string{"inf", "+Inf", "-Inf", "NaN", "abc", "", "0", "-5"} { + if got, ok := drainSeconds(v); ok { + t.Errorf("drainSeconds(%q) = %d, accepted; an unusable value must be refused "+ + "so the budget falls back to the default", v, got) + } + } +} + +func TestDrainTimeoutIsCappedAtSomethingSurvivable(t *testing.T) { + // Finite but implausible: clamped rather than refused, since the intent is + // legible even when the number is not. 9e18 also overflows an int, which is + // what made an unbounded path dangerous rather than merely silly. + for _, v := range []string{"86400", "1e30", "9e18"} { + got, ok := drainSeconds(v) + if !ok { + t.Fatalf("drainSeconds(%q): a finite positive value should parse", v) + } + if got != maxDrainTimeoutSeconds { + t.Errorf("drainSeconds(%q) = %d, want it clamped to %d: an unbounded grace "+ + "period leaves a stuck Pod deletable only with --force", + v, got, maxDrainTimeoutSeconds) + } + } +} + +func TestDrainTimeoutStillAcceptsOrdinaryValues(t *testing.T) { + for _, c := range []struct { + in string + want int + }{{"30", 30}, {"0.5", 1}, {"120.4", 121}, {"300", 300}} { + got, ok := drainSeconds(c.in) + if !ok || got != c.want { + t.Errorf("drainSeconds(%q) = %d,%v; want %d,true", c.in, got, ok, c.want) + } + } +} diff --git a/deploy/operator/internal/controller/discovery_backend_test.go b/deploy/operator/internal/controller/discovery_backend_test.go new file mode 100644 index 00000000..f44ff27e --- /dev/null +++ b/deploy/operator/internal/controller/discovery_backend_test.go @@ -0,0 +1,102 @@ +/* +Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. + +SPDX-License-Identifier: MIT +*/ + +package controller + +import ( + "context" + "errors" + "strings" + "testing" + + appsv1 "k8s.io/api/apps/v1" + "k8s.io/apimachinery/pkg/types" + ctrl "sigs.k8s.io/controller-runtime" + "sigs.k8s.io/controller-runtime/pkg/client/fake" + "sigs.k8s.io/controller-runtime/pkg/reconcile" + + inferav1alpha1 "github.com/amd/infera/deploy/operator/api/v1alpha1" +) + +// An operator-managed deployment is by definition running in Kubernetes, and +// there the orchestrator is what knows a worker is going away: a condemned Pod +// carries deletionTimestamp before the process is even signalled, and the +// registry drops it from routing then. +// +// Pointing such a deployment at an external etcd throws that away. The router +// stops watching Pods, so nothing reads the deletionTimestamp, and the only +// remaining signal is the worker's record disappearing on SIGTERM -- which +// arrives once the preStop delay the operator itself injects has elapsed. The +// combination keeps that delay and loses the early notice it exists to give, +// so for its whole duration the router keeps handing new work to a Pod that is +// already condemned. Refusing is better than rendering a deployment whose +// drain is worse than either backend alone. +func TestAnOperatorDeploymentRefusesTheExternalEtcdBackend(t *testing.T) { + s := scaleScheme(t) + idep := idepWith(1) + idep.Spec.DiscoveryBackend = "etcd" + idep.Spec.EtcdEndpoint = "etcd:2379" + cl := fake.NewClientBuilder().WithScheme(s). + WithObjects(idep).WithStatusSubresource(idep).Build() + + r := &InferaDeploymentReconciler{Client: cl, Scheme: s} + res, err := r.Reconcile(context.Background(), ctrl.Request{ + NamespacedName: types.NamespacedName{Name: "qwen", Namespace: "ns"}, + }) + if err == nil { + t.Fatal("reconcile accepted discoveryBackend=etcd; it must be refused") + } + if !strings.Contains(err.Error(), "discoveryBackend") { + t.Fatalf("error should name the field so the cause is obvious, got: %v", err) + } + + // No amount of retrying changes a spec field, and controller-runtime + // re-queues a plain error with exponential backoff forever -- two error + // logs and a status write on every attempt, and reconcile_errors_total + // climbing until someone edits the CR. A terminal error is recorded once + // and dropped. + if !errors.Is(err, reconcile.TerminalError(nil)) { + t.Errorf("error must be terminal, or the request is re-queued forever: %v", err) + } + if res.Requeue || res.RequeueAfter != 0 { //nolint:staticcheck // Requeue kept for clarity + t.Errorf("refusal must not ask to be retried, got %+v", res) + } + + // Refusing must not leave a half-built deployment behind. + dep := &appsv1.Deployment{} + key := types.NamespacedName{Name: "qwen-decode", Namespace: "ns"} + if err := cl.Get(context.Background(), key, dep); err == nil { + t.Fatal("a child workload was created for a configuration that was refused") + } + + // The reason belongs on the object, not only in the operator's log. + got := &inferav1alpha1.InferaDeployment{} + if err := cl.Get(context.Background(), types.NamespacedName{Name: "qwen", Namespace: "ns"}, got); err != nil { + t.Fatalf("get idep: %v", err) + } + if got.Status.State != inferav1alpha1.StateFailed { + t.Errorf("status.state = %q, want %q so `kubectl get idep` shows it", + got.Status.State, inferav1alpha1.StateFailed) + } +} + +// The default and an explicit "kubernetes" both reconcile normally. +func TestTheKubernetesBackendIsAccepted(t *testing.T) { + for _, backend := range []string{"", "kubernetes"} { + s := scaleScheme(t) + idep := idepWith(1) + idep.Spec.DiscoveryBackend = backend + cl := fake.NewClientBuilder().WithScheme(s). + WithObjects(idep).WithStatusSubresource(idep).Build() + reconcileOnce(t, cl, s) + + dep := &appsv1.Deployment{} + key := types.NamespacedName{Name: "qwen-decode", Namespace: "ns"} + if err := cl.Get(context.Background(), key, dep); err != nil { + t.Fatalf("discoveryBackend=%q: child Deployment not created: %v", backend, err) + } + } +} diff --git a/deploy/operator/internal/controller/gaie.go b/deploy/operator/internal/controller/gaie.go index 6cf3c2bd..adfe54db 100644 --- a/deploy/operator/internal/controller/gaie.go +++ b/deploy/operator/internal/controller/gaie.go @@ -237,7 +237,7 @@ func buildInferencePool(idep *inferav1alpha1.InferaDeployment) *unstructured.Uns "selector": map[string]any{ "matchLabels": map[string]any{ "infera.amd.com/deployment": idep.Name, - gaieFrontendLabel: "true", + gaieFrontendLabel: "true", }, }, "endpointPickerRef": map[string]any{ diff --git a/deploy/operator/internal/controller/inferadeployment_controller.go b/deploy/operator/internal/controller/inferadeployment_controller.go index 167c3570..0e5f361e 100644 --- a/deploy/operator/internal/controller/inferadeployment_controller.go +++ b/deploy/operator/internal/controller/inferadeployment_controller.go @@ -7,6 +7,7 @@ package controller import ( "context" + "fmt" "sort" "time" @@ -19,6 +20,7 @@ import ( "sigs.k8s.io/controller-runtime/pkg/client" "sigs.k8s.io/controller-runtime/pkg/controller/controllerutil" "sigs.k8s.io/controller-runtime/pkg/log" + "sigs.k8s.io/controller-runtime/pkg/reconcile" inferav1alpha1 "github.com/amd/infera/deploy/operator/api/v1alpha1" ) @@ -67,6 +69,43 @@ func (r *InferaDeploymentReconciler) Reconcile(ctx context.Context, req ctrl.Req return ctrl.Result{}, nil } + // Anything this operator builds runs in Kubernetes, and there the + // orchestrator is what knows a worker is leaving: a condemned Pod carries + // deletionTimestamp before the process is signalled, and the registry drops + // it from routing at that moment. + // + // Pointing such a deployment at an external etcd discards that. The server + // stops watching Pods, so nothing reads the deletionTimestamp, and the only + // remaining signal is the worker's record disappearing when it deregisters + // on SIGTERM -- which it receives only once the preStop delay injected + // below has elapsed. The + // combination keeps that delay while losing the early notice it exists to + // provide, so for its whole duration the router keeps handing new work to a + // Pod already on its way out. Refusing beats rendering a deployment whose + // drain is worse than either backend on its own. + if !useK8sDiscovery(idep) { + err := fmt.Errorf( + "spec.discoveryBackend=%q is not supported by the operator: an in-cluster "+ + "deployment must use the default \"kubernetes\" backend, which learns of a "+ + "departing worker from its Pod's deletionTimestamp. External etcd is for "+ + "deployments outside Kubernetes", + idep.Spec.DiscoveryBackend, + ) + lg.Error(err, "refusing to reconcile") + idep.Status.ObservedGeneration = idep.Generation + idep.Status.State = inferav1alpha1.StateFailed + if uerr := r.Status().Update(ctx, idep); uerr != nil { + lg.Error(uerr, "status update failed") + } + // Wrapped as terminal so it is recorded once and dropped. A plain error + // is re-queued with exponential backoff and retried forever, and no + // amount of retrying edits a spec field -- it would only produce two + // error logs and a status write per attempt, with + // reconcile_errors_total climbing until someone changes the CR. Editing + // the CR re-triggers reconciliation on its own. + return ctrl.Result{}, reconcile.TerminalError(err) + } + // 0. Kubernetes-native discovery RBAC: a namespaced ServiceAccount + Role so // workers can patch their own Pod annotation and the server can list/watch // this deployment's worker Pods (no external etcd). @@ -143,7 +182,7 @@ func (r *InferaDeploymentReconciler) Reconcile(ctx context.Context, req ctrl.Req } else { idep.Status.GAIE = nil } - idep.Status.State = rollupState(status) + idep.Status.State = rollupState(status, idep.Spec.Services) if err := r.Status().Update(ctx, idep); err != nil { lg.Error(err, "status update failed") return ctrl.Result{RequeueAfter: 5 * time.Second}, nil @@ -171,27 +210,107 @@ func (r *InferaDeploymentReconciler) applyUnstructured(ctx context.Context, idep existing.SetNamespace(desired.GetNamespace()) _, err := controllerutil.CreateOrUpdate(ctx, r.Client, existing, func() error { spec, _, _ := unstructured.NestedMap(desired.Object, "spec") - _ = unstructured.SetNestedMap(existing.Object, spec, "spec") + current, _, _ := unstructured.NestedMap(existing.Object, "spec") + if current == nil { + current = map[string]any{} + } + merged := mergeSpec(current, spec, ownedSpecFields(desired.GetKind())) + _ = unstructured.SetNestedMap(existing.Object, merged, "spec") existing.SetLabels(desired.GetLabels()) return controllerutil.SetControllerReference(idep, existing, r.Scheme) }) return err } +// ownedSpecFields lists the top-level spec fields the operator owns outright +// for a kind: it decides their entire contents, so one absent from the desired +// object has been removed and must be cleared rather than kept. +// +// The distinction matters for fields the builders emit conditionally. +// HTTPRoute.spec.hostnames is only written when the CR lists any, so without +// this a user who deletes their hostnames would keep matching them forever -- +// the merge below would see nothing to overlay and leave the old value in +// place. Fields not listed here belong to someone else, almost always the API +// server's defaulting, and are left untouched. +func ownedSpecFields(kind string) map[string]bool { + switch kind { + case httpRouteKind: + return map[string]bool{"parentRefs": true, "rules": true, "hostnames": true} + case inferencePoolKind: + return map[string]bool{"targetPorts": true, "selector": true, "endpointPickerRef": true} + case lwsKind: + return map[string]bool{"replicas": true, "leaderWorkerTemplate": true} + } + return nil +} + +// mergeSpec overlays the fields the operator sets onto what is already there, +// leaving anything it does not mention alone. +// +// Replacing .spec wholesale would be simpler, but these objects are only +// partly ours: buildLeaderWorkerSet writes three fields and the API server +// defaults the other nine from the CRD. Overwriting the whole map strips those +// defaults on every pass, the API server restores them, and the next pass +// strips them again -- so CreateOrUpdate sees a difference every single time +// and issues a write. Harmless-but-wasteful once per resync; a write loop now +// that the reconciler watches LeaderWorkerSet, since each write enqueues the +// reconcile that produces the next one. +// +// Nested maps merge; anything else replaces. Lists are owned outright -- a +// container list merged element-wise would be neither what was asked for nor +// what was there. +// +// `owned` names the top-level fields the operator decides entirely. One of +// those missing from `from` has been removed rather than left unmanaged, so it +// is deleted; that is what lets a conditionally-emitted field like +// HTTPRoute's hostnames be taken away again. Nested levels are not pruned: +// below the top level the desired object and the server's defaults are +// interleaved, with no way to tell them apart. +func mergeSpec(into, from map[string]any, owned map[string]bool) map[string]any { + for k := range owned { + if _, still := from[k]; !still { + delete(into, k) + } + } + for k, v := range from { + sub, isMap := v.(map[string]any) + if !isMap { + into[k] = v + continue + } + existing, ok := into[k].(map[string]any) + if !ok { + into[k] = v + continue + } + into[k] = mergeSpec(existing, sub, nil) + } + return into +} + func (r *InferaDeploymentReconciler) deploymentStatus(ctx context.Context, idep *inferav1alpha1.InferaDeployment, name string, svc inferav1alpha1.ServiceSpec) inferav1alpha1.ServiceStatus { - st := inferav1alpha1.ServiceStatus{Kind: "Deployment", Replicas: replicasOf(svc)} + // Replicas is what the workload reports, not what the spec asked for. + // Echoing the desired value back makes status useless for exactly the + // reader that needs it most -- an autoscaler computing + // `desired = ceil(current * metric/target)` cannot tell a scale-up has not + // landed if `current` is the number it just asked for. + st := inferav1alpha1.ServiceStatus{Kind: "Deployment"} dep := &appsv1.Deployment{} if err := r.Get(ctx, client.ObjectKey{Name: idep.Name + "-" + name, Namespace: idep.Namespace}, dep); err == nil { + st.Replicas = dep.Status.Replicas st.ReadyReplicas = dep.Status.ReadyReplicas } return st } func (r *InferaDeploymentReconciler) lwsStatus(ctx context.Context, idep *inferav1alpha1.InferaDeployment, name string, svc inferav1alpha1.ServiceSpec) inferav1alpha1.ServiceStatus { - st := inferav1alpha1.ServiceStatus{Kind: "LeaderWorkerSet", Replicas: replicasOf(svc)} + st := inferav1alpha1.ServiceStatus{Kind: "LeaderWorkerSet"} u := &unstructured.Unstructured{} u.SetGroupVersionKind(lwsGVK()) if err := r.Get(ctx, client.ObjectKey{Name: idep.Name + "-" + name, Namespace: idep.Namespace}, u); err == nil { + if v, ok, _ := unstructured.NestedInt64(u.Object, "status", "replicas"); ok { + st.Replicas = int32(v) + } if v, ok, _ := unstructured.NestedInt64(u.Object, "status", "readyReplicas"); ok { st.ReadyReplicas = int32(v) } @@ -234,20 +353,39 @@ func copySpec(existing, desired client.Object) { } } -func rollupState(svcs map[string]inferav1alpha1.ServiceStatus) inferav1alpha1.DeploymentState { +// rollupState answers whether every service has the capacity the spec asked +// for, so it compares against the spec rather than against the status. +// +// ServiceStatus.Replicas is what the workload reports, which is the right +// thing for a reader watching a scale-up land but the wrong side of this +// comparison: `ReadyReplicas < Replicas` only sees Pods that exist and are +// not ready. A replica that was never created -- unschedulable, out of quota, +// no GPU -- is absent from both numbers, so they agree and the deployment +// calls itself ready on a fraction of its fleet. +func rollupState( + svcs map[string]inferav1alpha1.ServiceStatus, + specs map[string]inferav1alpha1.ServiceSpec, +) inferav1alpha1.DeploymentState { if len(svcs) == 0 { return inferav1alpha1.StatePending } - allReady := true - for _, s := range svcs { - if s.ReadyReplicas < s.Replicas || s.Replicas == 0 { - allReady = false + for name, s := range svcs { + spec, ok := specs[name] + if !ok { + // Reported but no longer in the spec: on its way out, and not a + // reason to hold the deployment back. + continue + } + want := replicasOf(spec) + if want == 0 { + // Deliberately scaled to zero; nothing to wait for. + continue + } + if s.ReadyReplicas < want { + return inferav1alpha1.StatePending } } - if allReady { - return inferav1alpha1.StateReady - } - return inferav1alpha1.StatePending + return inferav1alpha1.StateReady } func sortedKeys(m map[string]inferav1alpha1.ServiceSpec) []string { @@ -261,10 +399,16 @@ func sortedKeys(m map[string]inferav1alpha1.ServiceSpec) []string { // SetupWithManager registers the controller. func (r *InferaDeploymentReconciler) SetupWithManager(mgr ctrl.Manager) error { - return ctrl.NewControllerManagedBy(mgr). + b := ctrl.NewControllerManagedBy(mgr). For(&inferav1alpha1.InferaDeployment{}). Owns(&appsv1.Deployment{}). Owns(&appsv1.StatefulSet{}). - Owns(&corev1.Service{}). - Complete(r) + Owns(&corev1.Service{}) + // Multi-node services are LeaderWorkerSets, so their status only reaches + // InferaDeployment.status on a resync unless we watch them. Guarded because + // the CRD is optional -- see lwsInstalled. + if lwsInstalled(mgr.GetRESTMapper()) { + b = b.Owns(lwsObject()) + } + return b.Complete(r) } diff --git a/deploy/operator/internal/controller/nats.go b/deploy/operator/internal/controller/nats.go index 545f4f20..d234cfe4 100644 --- a/deploy/operator/internal/controller/nats.go +++ b/deploy/operator/internal/controller/nats.go @@ -19,8 +19,8 @@ import ( func natsLabels(idepName string) map[string]string { return map[string]string{ "app.kubernetes.io/managed-by": "infera-operator", - "infera.amd.com/deployment": idepName, - "infera.amd.com/component": "nats", + "infera.amd.com/deployment": idepName, + "infera.amd.com/component": "nats", } } diff --git a/deploy/operator/internal/controller/scale_paths_test.go b/deploy/operator/internal/controller/scale_paths_test.go new file mode 100644 index 00000000..1d8a48d7 --- /dev/null +++ b/deploy/operator/internal/controller/scale_paths_test.go @@ -0,0 +1,316 @@ +/* +Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. + +SPDX-License-Identifier: MIT +*/ + +package controller + +import ( + "context" + "testing" + + appsv1 "k8s.io/api/apps/v1" + rbacv1 "k8s.io/api/rbac/v1" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + "k8s.io/apimachinery/pkg/apis/meta/v1/unstructured" + "k8s.io/apimachinery/pkg/runtime" + "k8s.io/apimachinery/pkg/types" + ctrl "sigs.k8s.io/controller-runtime" + "sigs.k8s.io/controller-runtime/pkg/client" + "sigs.k8s.io/controller-runtime/pkg/client/fake" + + inferav1alpha1 "github.com/amd/infera/deploy/operator/api/v1alpha1" +) + +// Which of the three ways to write a replica count actually reaches the pods. +// These are easy to conflate and they behave differently on purpose, so the +// distinction is pinned here rather than left to a README. + +func scaleScheme(t *testing.T) *runtime.Scheme { + t.Helper() + s := testScheme(t) + if err := rbacv1.AddToScheme(s); err != nil { + t.Fatalf("add rbac scheme: %v", err) + } + return s +} + +func idepWith(replicas int32) *inferav1alpha1.InferaDeployment { + r := replicas + return &inferav1alpha1.InferaDeployment{ + ObjectMeta: metav1.ObjectMeta{Name: "qwen", Namespace: "ns"}, + Spec: inferav1alpha1.InferaDeploymentSpec{ + Image: "infera:test", + Services: map[string]inferav1alpha1.ServiceSpec{ + "decode": { + ComponentType: inferav1alpha1.ComponentTypeWorker, + Replicas: &r, + NumberOfNodes: 1, + }, + }, + }, + } +} + +func reconcileOnce(t *testing.T, cl client.Client, s *runtime.Scheme) { + t.Helper() + r := &InferaDeploymentReconciler{Client: cl, Scheme: s} + _, err := r.Reconcile(context.Background(), ctrl.Request{ + NamespacedName: types.NamespacedName{Name: "qwen", Namespace: "ns"}, + }) + if err != nil { + t.Fatalf("reconcile: %v", err) + } +} + +func childReplicas(t *testing.T, cl client.Client) int32 { + t.Helper() + dep := &appsv1.Deployment{} + key := types.NamespacedName{Name: "qwen-decode", Namespace: "ns"} + if err := cl.Get(context.Background(), key, dep); err != nil { + t.Fatalf("get child Deployment: %v", err) + } + if dep.Spec.Replicas == nil { + t.Fatal("child Deployment has no replicas set") + } + return *dep.Spec.Replicas +} + +// Editing the CR is the normal path and it works: the CR is the desired state, +// so a change to it is what reconciliation exists to propagate. +func TestEditingTheCRScales(t *testing.T) { + s := scaleScheme(t) + idep := idepWith(2) + cl := fake.NewClientBuilder().WithScheme(s).WithObjects(idep). + WithStatusSubresource(idep).Build() + + reconcileOnce(t, cl, s) + if got := childReplicas(t, cl); got != 2 { + t.Fatalf("initial: child has %d replicas, want 2", got) + } + + // The user edits the CR. + live := &inferav1alpha1.InferaDeployment{} + key := types.NamespacedName{Name: "qwen", Namespace: "ns"} + if err := cl.Get(context.Background(), key, live); err != nil { + t.Fatalf("get idep: %v", err) + } + five := int32(5) + svc := live.Spec.Services["decode"] + svc.Replicas = &five + live.Spec.Services["decode"] = svc + if err := cl.Update(context.Background(), live); err != nil { + t.Fatalf("update idep: %v", err) + } + + reconcileOnce(t, cl, s) + if got := childReplicas(t, cl); got != 5 { + t.Fatalf("after editing the CR: child has %d replicas, want 5", got) + } +} + +// Editing the *child* is the path that does not survive, and that is the +// intended behaviour of any operator: the child is derived state, so the next +// pass restores it from the CR. The write succeeds and nothing reports an +// error, which is why pointing an autoscaler at the generated Deployment looks +// like it works right up until the next reconcile. +func TestEditingTheChildIsReverted(t *testing.T) { + s := scaleScheme(t) + idep := idepWith(2) + cl := fake.NewClientBuilder().WithScheme(s).WithObjects(idep). + WithStatusSubresource(idep).Build() + + reconcileOnce(t, cl, s) + + // Something scales the generated Deployment directly. + dep := &appsv1.Deployment{} + key := types.NamespacedName{Name: "qwen-decode", Namespace: "ns"} + if err := cl.Get(context.Background(), key, dep); err != nil { + t.Fatalf("get child: %v", err) + } + three := int32(3) + dep.Spec.Replicas = &three + if err := cl.Update(context.Background(), dep); err != nil { + t.Fatalf("update child: %v", err) + } + + reconcileOnce(t, cl, s) + if got := childReplicas(t, cl); got != 2 { + t.Fatalf("child edit survived reconciliation: %d, want it reverted to the CR's 2", got) + } +} + +// Scaling down through the CR has to land on pods that drain, not pods that +// get cut. The two features are built separately -- the replica count comes +// from the CR, the graceful shutdown from what the operator injects into the +// pod template -- so this checks they meet: a Deployment produced by a normal +// reconcile carries the preStop delay and a grace period long enough to cover +// the whole shutdown. +// +// Without preStop the router keeps assigning work for the entire termination; +// without the grace covering preStop + drain + teardown the kubelet SIGKILLs +// mid-drain. Either one silently turns a graceful scale-down back into a kill. +func TestScalingDownThroughTheCRLandsOnDrainablePods(t *testing.T) { + s := scaleScheme(t) + idep := idepWith(3) + // A long drain, the case where a fixed grace period used to fall short. + svc := idep.Spec.Services["decode"] + svc.Args = []string{"--drain-timeout", "300"} + idep.Spec.Services["decode"] = svc + + cl := fake.NewClientBuilder().WithScheme(s).WithObjects(idep). + WithStatusSubresource(idep).Build() + reconcileOnce(t, cl, s) + + dep := &appsv1.Deployment{} + key := types.NamespacedName{Name: "qwen-decode", Namespace: "ns"} + if err := cl.Get(context.Background(), key, dep); err != nil { + t.Fatalf("get child: %v", err) + } + spec := dep.Spec.Template.Spec + + if len(spec.Containers) == 0 { + t.Fatal("no containers in the pod template") + } + lc := spec.Containers[0].Lifecycle + if lc == nil || lc.PreStop == nil || lc.PreStop.Exec == nil { + t.Fatal("no preStop hook: the router would keep routing to a condemned pod") + } + + if spec.TerminationGracePeriodSeconds == nil { + t.Fatal("no terminationGracePeriodSeconds: the kubelet default is 30s, far under a 300s drain") + } + want := int64(workerPreStopDrainSeconds + 300 + workerTeardownHeadroomSeconds) + if got := *spec.TerminationGracePeriodSeconds; got != want { + t.Fatalf("grace = %ds, want %ds (preStop %d + drain 300 + teardown %d)", + got, want, workerPreStopDrainSeconds, workerTeardownHeadroomSeconds) + } + + // And the scale-down itself still works on that same object. + live := &inferav1alpha1.InferaDeployment{} + if err := cl.Get(context.Background(), types.NamespacedName{Name: "qwen", Namespace: "ns"}, live); err != nil { + t.Fatalf("get idep: %v", err) + } + one := int32(1) + svc = live.Spec.Services["decode"] + svc.Replicas = &one + live.Spec.Services["decode"] = svc + if err := cl.Update(context.Background(), live); err != nil { + t.Fatalf("update idep: %v", err) + } + reconcileOnce(t, cl, s) + if got := childReplicas(t, cl); got != 1 { + t.Fatalf("scale down through the CR: child has %d replicas, want 1", got) + } +} + +// Multi-node workers are torn down a whole group at a time, so the same +// guarantees have to hold on the LWS path -- where the pod template travels +// through a different builder. +func TestMultiNodePodsAlsoDrain(t *testing.T) { + s := scaleScheme(t) + s.AddKnownTypeWithName(lwsGVK(), &unstructured.Unstructured{}) + + idep := idepWith(2) + svc := idep.Spec.Services["decode"] + svc.NumberOfNodes = 3 + svc.Args = []string{"--drain-timeout", "180"} + idep.Spec.Services["decode"] = svc + + cl := fake.NewClientBuilder().WithScheme(s).WithObjects(idep). + WithStatusSubresource(idep).Build() + reconcileOnce(t, cl, s) + + u := &unstructured.Unstructured{} + u.SetGroupVersionKind(lwsGVK()) + if err := cl.Get(context.Background(), + types.NamespacedName{Name: "qwen-decode", Namespace: "ns"}, u); err != nil { + t.Fatalf("get child LWS: %v", err) + } + + grace, found, err := unstructured.NestedInt64(u.Object, + "spec", "leaderWorkerTemplate", "workerTemplate", "spec", "terminationGracePeriodSeconds") + if err != nil || !found { + t.Fatalf("LWS pod template has no terminationGracePeriodSeconds (found=%v, err=%v)", found, err) + } + want := int64(workerPreStopDrainSeconds + 180 + workerTeardownHeadroomSeconds) + if grace != want { + t.Fatalf("LWS grace = %ds, want %ds", grace, want) + } + + containers, found, err := unstructured.NestedSlice(u.Object, + "spec", "leaderWorkerTemplate", "workerTemplate", "spec", "containers") + if err != nil || !found || len(containers) == 0 { + t.Fatalf("LWS pod template has no containers (found=%v, err=%v)", found, err) + } + c, _ := containers[0].(map[string]any) + if _, ok := c["lifecycle"]; !ok { + t.Fatal("LWS container has no lifecycle/preStop: a condemned group keeps receiving work") + } +} + +// A LeaderWorkerSet carries a real scale subresource of its own, so an HPA can +// be pointed straight at the generated LWS and the write will succeed. It still +// does not work, for the same reason it does not work on the generated +// Deployment: reconciliation assigns the whole child spec every pass, replicas +// included. The scale write lands, and the next reconcile overwrites it. +// +// This is worth pinning because the LWS case looks different from the outside +// -- `kubectl get lws` shows a scale subresource, HPA reports success, nothing +// errors -- and the only symptom is a replica count that keeps snapping back. +func TestEditingTheChildLWSIsAlsoReverted(t *testing.T) { + s := scaleScheme(t) + s.AddKnownTypeWithName(lwsGVK(), &unstructured.Unstructured{}) + s.AddKnownTypeWithName(lwsGVK().GroupVersion().WithKind(lwsKind+"List"), + &unstructured.UnstructuredList{}) + + idep := idepWith(2) + svc := idep.Spec.Services["decode"] + svc.NumberOfNodes = 3 // multi-node -> LeaderWorkerSet instead of Deployment + idep.Spec.Services["decode"] = svc + + cl := fake.NewClientBuilder().WithScheme(s).WithObjects(idep). + WithStatusSubresource(idep).Build() + reconcileOnce(t, cl, s) + + get := func() *unstructured.Unstructured { + u := &unstructured.Unstructured{} + u.SetGroupVersionKind(lwsGVK()) + key := types.NamespacedName{Name: "qwen-decode", Namespace: "ns"} + if err := cl.Get(context.Background(), key, u); err != nil { + t.Fatalf("get child LWS: %v", err) + } + return u + } + replicas := func(u *unstructured.Unstructured) int64 { + v, found, err := unstructured.NestedInt64(u.Object, "spec", "replicas") + if err != nil || !found { + t.Fatalf("LWS has no spec.replicas (found=%v, err=%v)", found, err) + } + return v + } + + lws := get() + if got := replicas(lws); got != 2 { + t.Fatalf("initial: LWS has %d groups, want 2", got) + } + + // An HPA scales the LWS directly -- exactly what its scale subresource + // invites, and exactly what does not survive. + if err := unstructured.SetNestedField(lws.Object, int64(6), "spec", "replicas"); err != nil { + t.Fatalf("set replicas: %v", err) + } + if err := cl.Update(context.Background(), lws); err != nil { + t.Fatalf("update child LWS: %v", err) + } + if got := replicas(get()); got != 6 { + t.Fatalf("precondition: the scale write itself must land, got %d", got) + } + + reconcileOnce(t, cl, s) + + if got := replicas(get()); got != 2 { + t.Fatalf("LWS scale survived reconciliation: %d groups, want it reverted to 2", got) + } +} diff --git a/deploy/operator/internal/controller/status_rollup_test.go b/deploy/operator/internal/controller/status_rollup_test.go new file mode 100644 index 00000000..826971c0 --- /dev/null +++ b/deploy/operator/internal/controller/status_rollup_test.go @@ -0,0 +1,89 @@ +/* +Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. + +SPDX-License-Identifier: MIT +*/ + +package controller + +import ( + "testing" + + inferav1alpha1 "github.com/amd/infera/deploy/operator/api/v1alpha1" +) + +// ServiceStatus.Replicas reports what the workload has, not what was asked +// for -- an autoscaler cannot tell a scale-up has not landed if `current` is +// the number it just requested. That makes it the wrong side of the readiness +// comparison: `ready < observed` only catches Pods that exist and are not +// ready, and says nothing about Pods that were never created at all. +// +// Which is the case that matters. A worker Pod that cannot be scheduled -- no +// GPU, quota exhausted, a node taint -- never reaches the ReplicaSet's +// status.replicas, so ready equals observed and the whole deployment reports +// itself ready on a fraction of its capacity. `.status.state` is a +// printcolumn and the natural readiness gate for anything orchestrating on +// top, so this is what decides whether traffic is sent. +func TestReadyNeedsTheReplicaCountThatWasAskedFor(t *testing.T) { + three := int32(3) + svcs := map[string]inferav1alpha1.ServiceSpec{ + "decode": {Replicas: &three}, + } + + cases := []struct { + name string + observed inferav1alpha1.ServiceStatus + want inferav1alpha1.DeploymentState + }{ + { + name: "every requested replica is up", + observed: inferav1alpha1.ServiceStatus{Replicas: 3, ReadyReplicas: 3}, + want: inferav1alpha1.StateReady, + }, + { + name: "a replica could not be scheduled, so it is not in the workload at all", + observed: inferav1alpha1.ServiceStatus{Replicas: 2, ReadyReplicas: 2}, + want: inferav1alpha1.StatePending, + }, + { + name: "all present, one still starting", + observed: inferav1alpha1.ServiceStatus{Replicas: 3, ReadyReplicas: 2}, + want: inferav1alpha1.StatePending, + }, + { + name: "nothing up yet", + observed: inferav1alpha1.ServiceStatus{Replicas: 0, ReadyReplicas: 0}, + want: inferav1alpha1.StatePending, + }, + } + + for _, c := range cases { + t.Run(c.name, func(t *testing.T) { + got := rollupState(map[string]inferav1alpha1.ServiceStatus{"decode": c.observed}, svcs) + if got != c.want { + t.Fatalf("state = %q, want %q (spec asked for %d, workload has %d/%d)", + got, c.want, three, c.observed.ReadyReplicas, c.observed.Replicas) + } + }) + } +} + +// A service the spec no longer mentions must not hold the deployment back, +// and one with no status yet must not read as satisfied. +func TestRollupHandlesServicesMissingFromEitherSide(t *testing.T) { + one := int32(1) + specs := map[string]inferav1alpha1.ServiceSpec{"decode": {Replicas: &one}} + + if got := rollupState(map[string]inferav1alpha1.ServiceStatus{}, specs); got != inferav1alpha1.StatePending { + t.Fatalf("no status reported yet: state = %q, want pending", got) + } + + // Status carries a service the spec dropped; the live one is satisfied. + svcs := map[string]inferav1alpha1.ServiceStatus{ + "decode": {Replicas: 1, ReadyReplicas: 1}, + "stale": {Replicas: 0, ReadyReplicas: 0}, + } + if got := rollupState(svcs, specs); got != inferav1alpha1.StateReady { + t.Fatalf("a service no longer in the spec blocked readiness: state = %q", got) + } +} diff --git a/deploy/operator/internal/controller/watches_test.go b/deploy/operator/internal/controller/watches_test.go new file mode 100644 index 00000000..4725d912 --- /dev/null +++ b/deploy/operator/internal/controller/watches_test.go @@ -0,0 +1,51 @@ +/* +Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. + +SPDX-License-Identifier: MIT +*/ + +package controller + +import ( + "testing" + + appsv1 "k8s.io/api/apps/v1" + corev1 "k8s.io/api/core/v1" + "k8s.io/apimachinery/pkg/api/meta" + "k8s.io/apimachinery/pkg/runtime" + "k8s.io/apimachinery/pkg/runtime/schema" + + inferav1alpha1 "github.com/amd/infera/deploy/operator/api/v1alpha1" +) + +func testScheme(t *testing.T) *runtime.Scheme { + t.Helper() + s := runtime.NewScheme() + if err := inferav1alpha1.AddToScheme(s); err != nil { + t.Fatalf("add infera scheme: %v", err) + } + if err := appsv1.AddToScheme(s); err != nil { + t.Fatalf("add apps scheme: %v", err) + } + if err := corev1.AddToScheme(s); err != nil { + t.Fatalf("add core scheme: %v", err) + } + return s +} + +// The LWS watch is registered only when the CRD is served. controller-runtime +// builds an informer per watched type at startup and one for an unserved kind +// fails the manager, so a single-node cluster without LWS installed must not +// have the operator refuse to start. +func TestLwsInstalled(t *testing.T) { + empty := meta.NewDefaultRESTMapper(nil) + if lwsInstalled(empty) { + t.Fatal("no LWS CRD: reported installed, the manager would fail to start") + } + + withLWS := meta.NewDefaultRESTMapper([]schema.GroupVersion{lwsGVK().GroupVersion()}) + withLWS.Add(lwsGVK(), meta.RESTScopeNamespace) + if !lwsInstalled(withLWS) { + t.Fatal("LWS CRD present: reported missing, multi-node status would lag a resync") + } +} diff --git a/infera/common/discovery_k8s.py b/infera/common/discovery_k8s.py index 5c3bfe8d..69704df8 100644 --- a/infera/common/discovery_k8s.py +++ b/infera/common/discovery_k8s.py @@ -34,6 +34,7 @@ CanaryVerifier, WorkerInfo, WorkerPool, + WorkerStatus, ) logger = logging.getLogger(__name__) @@ -94,8 +95,13 @@ async def _relist(self) -> str | None: resp = await self._http.get(f"/api/v1/namespaces/{self._namespace}/pods", params=params) resp.raise_for_status() body = resp.json() + seen: set[str] = set() for pod in body.get("items", []) or []: + name = ((pod.get("metadata") or {}).get("name")) or "" + if name: + seen.add(name) self._handle_pod(pod, deleted=False) + self._reconcile_absent(seen) rv = (body.get("metadata") or {}).get("resourceVersion") logger.info( "k8s (re)list: %d worker(s) for selector %r in ns %s (rv=%s)", @@ -211,6 +217,26 @@ def _pod_running(pod: dict) -> bool: phase = ((pod.get("status") or {}).get("phase")) or "" return phase == "Running" + @staticmethod + def _pod_terminating(pod: dict) -> bool: + """True once the API server has stamped the Pod for deletion. + + A terminating Pod keeps ``phase: Running`` until its containers exit, so + a liveness check alone cannot see it. That gap is not academic: the + operator injects a ``preStop sleep`` before SIGTERM, and for the whole + of that delay the Pod is condemned, still Running, and — without this + check — still a routing candidate. The router would keep assigning new + work right up to the moment the process is killed, so the sleep that + exists to make shutdown graceful instead buys more requests that are + guaranteed to be cut. + + Reading ``deletionTimestamp`` closes that: the worker leaves the pool + the instant deletion is requested, in-flight work finishes on the + connections it already has, and the preStop delay becomes what it was + meant to be — drain time. + """ + return bool((pod.get("metadata") or {}).get("deletionTimestamp")) + def _handle_pod(self, pod: dict, *, deleted: bool) -> None: meta = pod.get("metadata") or {} pod_name = meta.get("name") or "" @@ -219,13 +245,24 @@ def _handle_pod(self, pod: dict, *, deleted: bool) -> None: annotations = meta.get("annotations") or {} raw = annotations.get(WORKER_INFO_ANNOTATION) - # Removal: explicit DELETE, pod no longer Running, or annotation gone. + # Not a routing candidate any more: explicit DELETE, annotation cleared + # (the worker deregistered, which is how its drain begins), or no longer + # Running. if deleted or raw is None or not self._pod_running(pod): worker_id = self._pod_to_worker.pop(pod_name, None) if worker_id is not None: self._remove(worker_id) return + # Condemned but still serving. Checked separately from Running because a + # terminating Pod stays Running until its containers exit; it leaves + # routing now and the record survives until it actually goes. + if self._pod_terminating(pod): + worker_id = self._pod_to_worker.get(pod_name) + if worker_id is not None: + self._mark_draining(worker_id) + return + try: info = worker_info_from_json(json.loads(raw)) except Exception as exc: @@ -273,17 +310,73 @@ def _handle_pod(self, pod: dict, *, deleted: bool) -> None: except Exception: logger.exception("on_worker_added callback failed") + def _reconcile_absent(self, seen: set[str]) -> None: + """Drop workers whose Pod is missing from a full list. + + A list is a complete snapshot, so anything still tracked that it does + not mention has been deleted -- and the event saying so was lost. That + is a routine occurrence rather than an edge case: the re-list exists + because etcd compaction expires the watch's resourceVersion every few + minutes, and any Pod deleted inside a reconnect window produces no + event anyone observes. + + It matters more now that a draining worker keeps its record. Removal + used to happen the instant a Pod was condemned, which bounded how long + a stale entry could survive; waiting for a later event instead makes a + missed one permanent. A phantom is filtered out of routing by its + DRAINING status, so nothing is dispatched to it -- but it is reported + by `/v1/workers` forever, and it keeps its model's tokenizer canary + pinned, which would reject a genuinely different worker later. + """ + for pod_name in [p for p in self._pod_to_worker if p not in seen]: + worker_id = self._pod_to_worker.pop(pod_name) + logger.info("k8s: pod %s absent from list; dropping worker %s", pod_name, worker_id) + self._remove(worker_id) + + def _mark_draining(self, worker_id: str) -> None: + """Take a condemned worker out of routing without dropping its record. + + ``list_active`` filters DRAINING, so this stops new work reaching the + worker just as removal would -- but ``list_all`` still shows it, and + that difference is what makes a condemned Pod visible as such on + ``/v1/workers`` instead of looking like one that crashed. The window is + the preStop delay: from the deletion being requested to the process + being signalled. + + The record does not linger: the worker clears its own annotation on + SIGTERM, before draining, which lands here as "annotation gone" and + removes it for real. + + Callbacks fire here rather than at that later removal because routing + is what they act on -- the KV subscriber and the policy's block + accounting must stop treating a departing worker as a target now, not + when its Pod object finally disappears. ``_remove`` therefore skips + them if it already announced, so each worker is announced exactly once. + """ + existing = self._pool.get(worker_id) + if existing is None or existing.status is WorkerStatus.DRAINING: + return # never registered, or already announced + existing.status = WorkerStatus.DRAINING + logger.info("k8s: worker %s draining (pod terminating, record kept)", worker_id) + self._notify_removed(worker_id) + def _remove(self, worker_id: str) -> None: existing = self._pool.get(worker_id) if existing is None: return + announced = existing.status is WorkerStatus.DRAINING self._pool.remove(worker_id) remaining = [w for w in self._pool.list_all() if w.model_name == existing.model_name] if not remaining: self._canary.forget(existing.model_name) logger.info("k8s: worker %s removed (pod deleted / not ready)", worker_id) - if self._on_removed is not None: - try: - self._on_removed(worker_id) - except Exception: - logger.exception("on_worker_removed callback failed") + if not announced: + self._notify_removed(worker_id) + + def _notify_removed(self, worker_id: str) -> None: + if self._on_removed is None: + return + try: + self._on_removed(worker_id) + except Exception: + logger.exception("on_worker_removed callback failed") diff --git a/infera/common/engine_metrics.py b/infera/common/engine_metrics.py new file mode 100644 index 00000000..712a021a --- /dev/null +++ b/infera/common/engine_metrics.py @@ -0,0 +1,162 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# SPDX-License-Identifier: MIT +############################################################################### +"""What each engine calls the metrics we need, in one place. + +Every engine exposes the same three facts — requests running, requests queued, +KV cache in use — under a different name, and the names drift between releases. +Anything that reads them (graceful drain, an autoscaler) needs the same mapping, +so it lives here rather than being spelled out at each call site where one of +them would quietly rot. + +Provenance, because it is uneven and matters: + +* **vLLM** — verified against a running engine (vLLM 0.1.dev19253, Qwen3-8B on + MI355X). Note ``kv_cache_usage_perc`` was ``gpu_cache_usage_perc`` in older + builds; the alias list below covers both. +* **SGLang** — verified against a running engine (SGLang 0.5.15, Qwen3-8B on + MI355X). Note sglang serves ``/metrics`` only with ``--enable-metrics``; the + worker entrypoint injects it. Treat a lookup failure as "unknown", never as + "zero"; the difference decides whether a drain waits or gives up. +* **ATOM** — unknown. Deliberately absent rather than guessed: a wrong name + reads as an idle engine, and an idle engine is exactly the answer that makes a + drain cut live requests. +""" + +from __future__ import annotations + +import logging +import re + +from infera.common.worker_pool import EngineType + +logger = logging.getLogger(__name__) + +#: metric key -> per-engine exposition name(s). A missing engine means "we do +#: not know", which callers must distinguish from "the value is zero". Several +#: names per entry means the engine renamed the series between releases and both +#: spellings are in the wild. +_ALIASES: dict[str, dict[EngineType, tuple[str, ...]]] = { + "kv_cache_usage": { + EngineType.VLLM: ("vllm:kv_cache_usage_perc", "vllm:gpu_cache_usage_perc"), + }, +} + +_NAMES: dict[str, dict[EngineType, str]] = { + "requests_running": { + EngineType.VLLM: "vllm:num_requests_running", + EngineType.SGLANG: "sglang:num_running_reqs", + }, + "requests_waiting": { + EngineType.VLLM: "vllm:num_requests_waiting", + EngineType.SGLANG: "sglang:num_queue_reqs", + }, + # vLLM renamed this: older builds expose gpu_cache_usage_perc, current ones + # kv_cache_usage_perc. Both are listed and callers sum whichever is present, + # because pinning one silently returns "no KV in use" on the other. + "kv_cache_usage": { + EngineType.VLLM: "vllm:kv_cache_usage_perc", + EngineType.SGLANG: "sglang:token_usage", + }, +} + +#: Extra per-engine gauges that also represent unfinished work, counted only +#: when draining. These are the PD handoff queues: a prefill worker can show no +#: running and no queued requests while KV transfers are still outstanding, and +#: killing it there strands the decode workers waiting on that KV -- the failure +#: every PD system in the field documents and none of them prevents. +#: Verified present on SGLang 0.5.15 (`--enable-metrics`). +_DRAIN_EXTRA: dict[EngineType, tuple[str, ...]] = { + EngineType.SGLANG: ( + "sglang:num_prefill_bootstrap_queue_reqs", + "sglang:num_prefill_inflight_queue_reqs", + "sglang:num_decode_prealloc_queue_reqs", + "sglang:num_decode_transfer_queue_reqs", + ), +} + + +def metric_name(key: str, engine: EngineType) -> str | None: + """Primary exposition name for ``key`` on ``engine``, or None if unknown.""" + return _NAMES[key].get(engine) + + +def metric_names(key: str, engine: EngineType) -> tuple[str, ...]: + """Every spelling of ``key`` on ``engine``, newest first.""" + alias = _ALIASES.get(key, {}).get(engine) + if alias: + return alias + name = _NAMES[key].get(engine) + return (name,) if name else () + + +def parse_metric(text: str, name: str) -> float | None: + """Sum every label set of a gauge in Prometheus text exposition. + + Engines label these per rank -- SGLang emits + ``sglang:num_running_reqs{tp_rank="0",...}`` and one series per rank -- so + reading only the first match would let a busy rank hide behind an idle one. + Summing is safe for the question a drain asks, because the sum is zero + exactly when every rank is zero. + + Returns None when the series is absent, which is not the same as 0.0: a + caller draining in-flight work must not read "metric missing" as "idle". + """ + total = 0.0 + found = False + for m in re.finditer( + rf"^{re.escape(name)}(?:\{{[^}}]*\}})?\s+([0-9.eE+-]+)\s*$", text, re.MULTILINE + ): + try: + total += float(m.group(1)) + except ValueError: + continue + found = True + return total if found else None + + +def inflight_from_metrics(text: str, engine: EngineType) -> float | None: + """Requests the engine is running plus those it has queued. + + Queued requests count: a request the engine has accepted but not started is + still work the client is waiting on, and killing the process loses it just + as surely as one mid-generation. So do the PD handoff queues, where the + request may be finished locally while its KV is still in transit. + + None means the engine's in-flight count could not be determined. + """ + total = 0.0 + seen = False + missing: list[str] = [] + for key in ("requests_running", "requests_waiting"): + name = metric_name(key, engine) + if name is None: + continue + value = parse_metric(text, name) + if value is None: + missing.append(name) + continue + total += value + seen = True + if not seen: + return None + if missing: + # Partial readings are still worth acting on -- refusing them outright + # would mean one renamed series stops the drain waiting for anything at + # all, which is the worse failure. But it must not pass silently: the + # absent series contributes zero, so a drain can finish while the work + # it describes is still outstanding. These names do drift; the vLLM KV + # gauge was renamed under exactly this module. + logger.warning( + "drain: %s not found on the %s metrics page; its work counts as zero, " + "so in-flight requests may be cut. Check the exposition names.", + ", ".join(missing), + engine.value, + ) + # Absent PD queues are genuinely zero here rather than unknown: the engine + # published a metrics page and simply is not running disaggregated. + for name in _DRAIN_EXTRA.get(engine, ()): + total += parse_metric(text, name) or 0.0 + return total diff --git a/infera/common/nats_request.py b/infera/common/nats_request.py index 9e7cad40..0894d951 100644 --- a/infera/common/nats_request.py +++ b/infera/common/nats_request.py @@ -71,6 +71,10 @@ # stream's per-worker consumer pending count is a live backlog gauge. REQUEST_STREAM = "INFERA_REQUESTS" +# How often a draining worker re-checks its JetStream backlog. Short, because +# the messages are already accepted and every poll is one round-trip. +_QUEUED_POLL_INTERVAL_S = 0.2 + # Reply framing headers. HDR_TYPE = "rs-type" HDR_STATUS = "rs-status" @@ -440,6 +444,7 @@ def __init__( self._url = url self._nc = None self._sub = None + self._js = None # set only under the throttle (JetStream-backed path) self._cancel_sub = None self._http: httpx.AsyncClient | None = None # flag (entry-point CLI) > env > built-in default. @@ -470,6 +475,7 @@ async def start(self) -> None: from nats.js.api import AckPolicy, ConsumerConfig js = self._nc.jetstream() + self._js = js await _ensure_request_stream(js) self._sub = await js.subscribe( subject, @@ -502,6 +508,18 @@ async def stop(self, *, drain: bool = False, drain_timeout: float = 0.0) -> None generations finish for up to ``drain_timeout`` seconds before cancelling any leftovers, so a worker being rolled does not sever active streams. With ``drain=False`` (default) in-flight tasks are cancelled at once.""" + deadline = time.monotonic() + drain_timeout if (drain and drain_timeout > 0) else None + + # 0. Under the throttle, requests land in a WorkQueue stream and are + # pulled from it -- so one can be accepted, queued, and invisible to + # `_inflight`, which only tracks what has been *delivered* here. + # `unsubscribe()` discards remaining messages ("remaining messages will + # be discarded", nats-py), so closing the door first would strand work + # the router already handed us: the client waits out the full idle + # timeout (900s by default) for a reply nobody will ever send. + if deadline is not None: + await self._await_queued(deadline) + # 1. Stop accepting NEW requests immediately so nothing new lands while # we drain (unsubscribe the request subject first). if self._sub is not None: @@ -511,15 +529,16 @@ async def stop(self, *, drain: bool = False, drain_timeout: float = 0.0) -> None pass self._sub = None # 2. Optionally let in-flight requests finish (bounded by drain_timeout). - if drain and drain_timeout > 0: + if deadline is not None: inflight = [t for t in self._inflight.values() if not t.done()] if inflight: + remaining = max(0.0, deadline - time.monotonic()) logger.info( "draining %d in-flight NATS request(s), up to %.0fs", len(inflight), - drain_timeout, + remaining, ) - _done, pending = await asyncio.wait(inflight, timeout=drain_timeout) + _done, pending = await asyncio.wait(inflight, timeout=remaining) if pending: logger.warning( "drain timeout; cancelling %d unfinished request(s)", len(pending) @@ -529,7 +548,18 @@ async def stop(self, *, drain: bool = False, drain_timeout: float = 0.0) -> None if not task.done(): task.cancel() self._inflight.clear() - # 4. Drop the cancel listener and the connection. + # 4. Delete this worker's durable consumer. `unsubscribe()` only tears + # down the local subscription -- a durable survives on the server by + # definition, and its name is derived from worker_id, which a rebuilt + # Pod never reuses (the IP changes). Left behind, every rollout adds an + # orphan holding WorkQueue quota that nothing will ever consume. + if self._js is not None: + try: + await self._js.delete_consumer(REQUEST_STREAM, request_durable(self._worker_id)) + except Exception as exc: # noqa: BLE001 - shutdown must continue + logger.debug("could not delete request consumer: %s", exc) + self._js = None + # 5. Drop the cancel listener and the connection. if self._cancel_sub is not None: try: await self._cancel_sub.unsubscribe() @@ -546,6 +576,41 @@ async def stop(self, *, drain: bool = False, drain_timeout: float = 0.0) -> None pass self._nc = None + async def _await_queued(self, deadline: float) -> None: + """Let JetStream hand over everything already queued for this worker. + + Only ``num_pending`` (accepted, not yet delivered) is waited on: + ``num_ack_pending`` is work already delivered, which is exactly what + ``_inflight`` tracks and step 2 waits for. Counting both would double + the wait for the same requests. + + Never raises, and gives up rather than hanging when the consumer cannot + be read -- a shutdown that stalls on a broker hiccup is worse than one + that drops a queued request, and the caller's deadline is shared with + the in-flight wait that follows. + """ + if self._js is None: + return + durable = request_durable(self._worker_id) + while True: + try: + info = await self._js.consumer_info(REQUEST_STREAM, durable) + except Exception as exc: # noqa: BLE001 - shutdown must continue + logger.debug("drain: cannot read consumer backlog (%s); not waiting", exc) + return + queued = int(getattr(info, "num_pending", 0) or 0) + if queued <= 0: + return + if time.monotonic() >= deadline: + logger.warning( + "drain timeout with %d request(s) still queued in JetStream; " + "they will not be served", + queued, + ) + return + logger.info("drain: waiting for %d queued request(s) to be delivered", queued) + await asyncio.sleep(min(_QUEUED_POLL_INTERVAL_S, max(0.0, deadline - time.monotonic()))) + async def _reply( self, inbox: str, rtype: str, data: bytes = b"", status: int | None = None ) -> None: diff --git a/infera/common/registration.py b/infera/common/registration.py index 26059c05..631f0b1e 100644 --- a/infera/common/registration.py +++ b/infera/common/registration.py @@ -25,6 +25,13 @@ def build_worker_payload(config: EngineConfig) -> dict: The same dict is PUT to etcd (RegistrationClient) or stored in the worker Pod annotation (K8sRegistrationClient), so the server-side parse (discovery.worker_info_from_json) is transport-agnostic. + + Identity only, and every field is fixed for the life of the process. No + status is written: the record's own presence is the etcd backend's answer to + "is this worker available", and under Kubernetes the Pod's deletionTimestamp + answers it earlier than the worker could. Keeping the payload stateless is + also what makes the heartbeat safe -- it rebuilds this from config, so any + state written here would be erased by the next refresh. """ worker_id = f"{config.host}:{config.port}" payload: dict = { @@ -98,17 +105,39 @@ async def register(self, config: EngineConfig) -> str: ) return worker_id - async def deregister(self) -> None: + async def deregister(self) -> bool: + """Revoke the lease, reporting whether the record is actually gone. + + The caller drains after this, and this is what stops new work arriving, + so a failure here is not cosmetic: the record survives on an unrenewed + lease for up to its TTL -- long enough to cover the whole drain -- and + the router keeps dispatching to a worker on its way out. Never raises, + because the teardown that follows still has to run. + """ + ok = True if self._lease_id is not None: try: - await self._http.post("/v3/lease/revoke", json={"ID": self._lease_id}) + r = await self._http.post("/v3/lease/revoke", json={"ID": self._lease_id}) + # httpx does not raise on 4xx/5xx, and etcd answering "lease not + # found" or a proxy returning 503 is exactly the failure this + # reports -- without this the refusal would be logged as a + # successful revoke. + r.raise_for_status() logger.info( "deregistered worker %s (lease %d revoked)", self._worker_id, self._lease_id, ) - except Exception as exc: - logger.warning("lease revoke failed: %s", exc) + except Exception as exc: # noqa: BLE001 - shutdown must continue + ok = False + logger.error( + "lease revoke failed for worker %s (%s); the record survives until " + "its %ds lease expires, so the router may keep sending work here " + "for the whole drain", + self._worker_id, + exc, + self._lease_ttl, + ) self._lease_id = None self._worker_id = None self._key = None @@ -116,6 +145,7 @@ async def deregister(self) -> None: await self._http.aclose() except Exception: pass + return ok async def heartbeat_loop(self, interval: float | None = None) -> None: """Refresh the etcd lease until cancelled. diff --git a/infera/common/registration_k8s.py b/infera/common/registration_k8s.py index 747abf83..72ca429d 100644 --- a/infera/common/registration_k8s.py +++ b/infera/common/registration_k8s.py @@ -36,7 +36,21 @@ class K8sRegistrationClient: - """Worker-side self-registration by patching its own Pod annotation.""" + """Worker-side self-registration by patching its own Pod annotation. + + The annotation carries identity and nothing else -- who this worker is, + where to reach it, and what it can serve. All of that is fixed for the + life of the process, which is what makes the refresh below safe to run at + any time, including throughout a shutdown. + + State is deliberately absent. Kubernetes stamps a condemned Pod with + ``deletionTimestamp`` before the worker is even signalled, and + ``KubernetesRegistry`` reads it, so the orchestrator already answers "is + this worker going away" earlier and more authoritatively than the worker + could. Writing the same fact into the annotation as well would make it a + race between two writers of one truth -- and the refresh below, which + rebuilds the payload from config, would be the one to win it. + """ def __init__( self, @@ -80,19 +94,36 @@ async def register(self, config: EngineConfig) -> str: ) return worker_id - async def deregister(self) -> None: - # Best-effort: clear the annotation so a terminating-but-lingering Pod - # stops being routed before its DELETE event lands. + async def deregister(self) -> bool: + """Clear the annotation, reporting whether the record is actually gone. + + This is what takes the worker out of routing, and the caller drains + afterwards -- so a failure means draining while still being dispatched + to. On a Pod that is being deleted the registry has already dropped it + from its deletionTimestamp and this is only cleanup; on every other way + a process gets SIGTERM (a probe restart, a node shutdown, a manual + kill) there is no such signal and this patch is the only one. Never + raises, because the teardown that follows still has to run. + """ + ok = True try: await self._patch_annotation(None) logger.info("deregistered worker %s (annotation cleared)", self._worker_id) - except Exception as exc: - logger.warning("k8s deregister failed (pod likely terminating): %s", exc) + except Exception as exc: # noqa: BLE001 - shutdown must continue + ok = False + logger.error( + "could not clear the worker annotation for %s (%s); if this Pod is not " + "being deleted, the router has no other signal and may keep sending " + "work here for the whole drain", + self._worker_id, + exc, + ) self._worker_id = None try: await self._http.aclose() except Exception: pass + return ok async def heartbeat_loop(self, interval: float | None = None) -> None: """Periodically re-assert the annotation (self-heal); never expires it.""" diff --git a/infera/engine/atom/__main__.py b/infera/engine/atom/__main__.py index 1be94a06..a0c6f123 100644 --- a/infera/engine/atom/__main__.py +++ b/infera/engine/atom/__main__.py @@ -156,7 +156,11 @@ async def main() -> None: except asyncio.CancelledError: pass - await reg_client.deregister() + if not await reg_client.deregister(): + # deregister() already logged why. No drain step here, so nothing waits + # on the record being gone -- but a silent branch would be the wrong + # thing to inherit if one is ever added. + logger.warning("stopping anyway") await engine.stop() diff --git a/infera/engine/drain.py b/infera/engine/drain.py new file mode 100644 index 00000000..7fc52537 --- /dev/null +++ b/infera/engine/drain.py @@ -0,0 +1,142 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# SPDX-License-Identifier: MIT +############################################################################### +"""Let in-flight generations finish before the engine is stopped. + +On the NATS transport infera owns the request path, so it knows exactly what is +in flight and ``NatsRequestServer.stop(drain=True)`` waits for it. On HTTP the +router talks straight to the engine's own server: infera never sees the request, +cannot count it, and so — until this — did not wait for it. Shutdown went +``deregister()`` then ``engine.stop()``, cutting every active generation. + +The way out is to ask the engine, which does know. It publishes its running and +queued request counts on ``/metrics``; poll until both reach zero or the timeout +expires. That is a poll rather than a signal, so it is bounded by +``poll_interval`` rather than exact — acceptable, because the alternative is not +draining at all. + +Two behaviours are deliberate: + +* **Deregister first, then drain.** Ordering is the whole point. Draining while + still a routing candidate just means more work arrives, and the count this + polls never reaches zero. Removing the record is what takes the worker out of + the candidate list, so it happens before waiting for the work already in hand. + The cost is that a worker finishing its in-flight requests is indistinguishable + from one that crashed; the alternative is a drain that cannot converge. +* **An unreadable metric does not block shutdown.** If the engine's in-flight + count cannot be determined — an unknown engine, a renamed series, a dead HTTP + server — this logs loudly and returns rather than hanging until the timeout. + A rolling update that stalls on a parse failure is a worse outcome than one + that cuts a request, and a silent full-timeout wait would look identical to a + genuinely busy worker. +""" + +from __future__ import annotations + +import asyncio +import logging +import time + +import httpx + +from infera.common.engine_metrics import inflight_from_metrics +from infera.common.worker_pool import EngineType + +logger = logging.getLogger(__name__) + +_POLL_INTERVAL_S = 0.5 + +#: How long the engine must report zero before we believe it. +#: +#: These gauges are refreshed on the engine's own schedule, not per request. +#: Measured on SGLang 0.5.15: ``num_running_reqs`` stayed at 12 for 5-15s after +#: the last HTTP response completed. The lag is safe in the direction that +#: matters (stale-high just makes the drain wait), but it is dangerous at the +#: start: a request accepted moments before SIGTERM may not be in the gauge yet, +#: so a single zero reading can mean "idle" or "not counted yet". Requiring the +#: zero to persist past one refresh cycle tells those apart. Costs a few seconds +#: on every shutdown; cheap against cutting a live generation. +_SETTLE_S = 6.0 + + +async def drain_engine_inflight( + *, + host: str, + port: int, + engine: EngineType, + timeout: float, + poll_interval: float = _POLL_INTERVAL_S, + settle: float = _SETTLE_S, +) -> bool: + """Wait until the engine reports no in-flight work, bounded by ``timeout``. + + Returns True if it drained, False if it timed out or could not be measured. + Never raises: this runs on the shutdown path, where an exception would skip + the engine teardown that follows. + """ + if timeout <= 0: + return False + + # The engine binds the advertised port, but 0.0.0.0 is not a destination. + probe_host = "127.0.0.1" if host in ("0.0.0.0", "", "::") else host + url = f"http://{probe_host}:{port}/metrics" + deadline = time.monotonic() + timeout + peak = 0.0 + zero_since: float | None = None + + try: + async with httpx.AsyncClient(timeout=5.0) as client: + while True: + try: + resp = await client.get(url) + inflight = ( + inflight_from_metrics(resp.text, engine) + if resp.status_code == 200 + else None + ) + except httpx.HTTPError as exc: + logger.info("drain: engine metrics unreachable (%s); not waiting", exc) + return False + + if inflight is None: + logger.warning( + "drain: cannot read in-flight count for %s from %s -- shutting down " + "WITHOUT draining. In-flight generations will be cut.", + engine.value, + url, + ) + return False + + peak = max(peak, inflight) + now = time.monotonic() + if inflight <= 0: + if zero_since is None: + zero_since = now + elif now - zero_since >= settle: + if peak > 0: + logger.info( + "drain: engine idle for %.0fs, %.0f request(s) completed", + settle, + peak, + ) + return True + else: + # A late gauge refresh revealed work we had not seen; the + # settle window has to start over. + zero_since = None + + if time.monotonic() >= deadline: + logger.warning( + "drain: timeout after %.0fs with %.0f request(s) still in flight; " + "they will be cut", + timeout, + inflight, + ) + return False + + await asyncio.sleep(min(poll_interval, max(0.0, deadline - time.monotonic()))) + except Exception as exc: # noqa: BLE001 - shutdown must continue regardless + logger.warning("drain: aborted (%s: %s); not waiting", type(exc).__name__, exc) + return False diff --git a/infera/engine/sglang/__main__.py b/infera/engine/sglang/__main__.py index 75f6ae69..c334e3fb 100644 --- a/infera/engine/sglang/__main__.py +++ b/infera/engine/sglang/__main__.py @@ -30,6 +30,7 @@ ) from infera.common.registration import RegistrationClient from infera.engine.base import watch_engine_death +from infera.engine.drain import drain_engine_inflight from infera.engine.sglang.args import SglangWorkerArgs, parse_sglang_args from infera.engine.sglang.kv_wiring import ( SglangKvWiring, @@ -408,15 +409,47 @@ async def _run_after_start(args: SglangWorkerArgs, engine: SglangEngine, config) except asyncio.CancelledError: pass + # Stop the heartbeat before touching the record: it re-asserts registration + # from config, so a refresh landing after deregistration would put the + # worker straight back into the pool. hb_task.cancel() try: await hb_task except asyncio.CancelledError: pass - await reg_client.deregister() - if nats_req_server is not None: - await nats_req_server.stop(drain=True, drain_timeout=args.drain_timeout) + async def _drain() -> None: + if nats_req_server is not None: + await nats_req_server.stop(drain=True, drain_timeout=args.drain_timeout) + else: + # HTTP transport: the router talks straight to the engine, so infera + # never saw these requests and has to ask the engine what is still + # in flight. + await drain_engine_inflight( + host=config.host, + port=config.port, + engine=config.engine, + timeout=args.drain_timeout, + ) + + # Deregister before draining, on every backend: removing the record is what + # stops new work arriving, and waiting on in-flight work while still being + # dispatched to just races arrivals. + # + # On Kubernetes the registry does drop a Pod on its deletionTimestamp, well + # before this process is signalled -- but only when the Pod is being + # deleted. A liveness-probe restart, a node graceful shutdown or a manual + # kill all deliver SIGTERM with the Pod object untouched, and on those paths + # the annotation is still there and still parsed, so this worker stays + # routable until it clears it. Draining first would hand it new work for the + # whole drain window. + # + # The cost is that the worker is gone from /v1/workers while it finishes, + # rather than visibly draining. + if not await reg_client.deregister(): + # deregister() already logged why, including whether it matters here. + logger.warning("draining anyway") + await _drain() if kv_relay is not None: await kv_relay.stop() diff --git a/infera/engine/sglang/worker.py b/infera/engine/sglang/worker.py index 2fb3f533..613b73dd 100644 --- a/infera/engine/sglang/worker.py +++ b/infera/engine/sglang/worker.py @@ -80,6 +80,14 @@ def __init__( async def start(self) -> EngineConfig: argv = list(self.sglang_argv) + # sglang serves /metrics only with --enable-metrics; without it the + # endpoint 404s. Graceful shutdown reads the in-flight request count + # from there, so leaving it off silently downgrades every scale-down + # and rolling update to "kill in-flight generations". Cheap enough to + # always enable, and the caller can still have passed it explicitly. + if not any(a == "--enable-metrics" for a in argv): + argv.append("--enable-metrics") + if self.enable_kv_events: dp_size = int(getattr(self.server_args, "dp_size", 1) or 1) self._kv_events_port = free_tcp_port_block(dp_size) if dp_size > 1 else free_tcp_port() diff --git a/infera/engine/vllm/__main__.py b/infera/engine/vllm/__main__.py index b8f33c48..fe9b2bc5 100644 --- a/infera/engine/vllm/__main__.py +++ b/infera/engine/vllm/__main__.py @@ -24,6 +24,7 @@ from infera.common.registration import RegistrationClient from infera.common.worker_pool import DisaggMode, KvRegistrationMetadata from infera.engine.base import watch_engine_death +from infera.engine.drain import drain_engine_inflight from infera.engine.vllm.args import VllmWorkerArgs, parse_vllm_args from infera.engine.vllm.worker import VllmEngine @@ -354,15 +355,47 @@ async def main() -> None: except asyncio.CancelledError: pass + # Stop the heartbeat before touching the record: it re-asserts registration + # from config, so a refresh landing after deregistration would put the + # worker straight back into the pool. hb_task.cancel() try: await hb_task except asyncio.CancelledError: pass - await reg_client.deregister() - if nats_req_server is not None: - await nats_req_server.stop(drain=True, drain_timeout=args.drain_timeout) + async def _drain() -> None: + if nats_req_server is not None: + await nats_req_server.stop(drain=True, drain_timeout=args.drain_timeout) + else: + # HTTP transport: the router talks straight to the engine, so infera + # never saw these requests and has to ask the engine what is still + # in flight. + await drain_engine_inflight( + host=config.host, + port=config.port, + engine=config.engine, + timeout=args.drain_timeout, + ) + + # Deregister before draining, on every backend: removing the record is what + # stops new work arriving, and waiting on in-flight work while still being + # dispatched to just races arrivals. + # + # On Kubernetes the registry does drop a Pod on its deletionTimestamp, well + # before this process is signalled -- but only when the Pod is being + # deleted. A liveness-probe restart, a node graceful shutdown or a manual + # kill all deliver SIGTERM with the Pod object untouched, and on those paths + # the annotation is still there and still parsed, so this worker stays + # routable until it clears it. Draining first would hand it new work for the + # whole drain window. + # + # The cost is that the worker is gone from /v1/workers while it finishes, + # rather than visibly draining. + if not await reg_client.deregister(): + # deregister() already logged why, including whether it matters here. + logger.warning("draining anyway") + await _drain() if kv_relay is not None: await kv_relay.stop() await engine.stop() diff --git a/infera/kv/snapshot.py b/infera/kv/snapshot.py index 16301c74..5d0c71d7 100644 --- a/infera/kv/snapshot.py +++ b/infera/kv/snapshot.py @@ -243,9 +243,25 @@ def register_target( model: str, compat_key: str, ) -> None: - """Tell the reconciler to periodically pull this (publisher, tree).""" + """Tell the reconciler to periodically pull this (publisher, tree). + + Also pulls it now rather than on the next tick. The loop waits out the + full interval between sweeps, so without this a worker that joins while + the reconciler is running is invisible to kv-aware routing for up to + ``interval_s`` -- 30 s in production. + + For a genuinely new worker that would be harmless, since its cache is + empty and an empty view is accurate. It is not harmless on a router + restart or a rolling upgrade: every existing worker arrives through this + same path with a warm cache, and until its snapshot lands the policy + scores them all as holding nothing. + """ key = (publisher_id, endpoint, model, compat_key) + new = key not in self._targets self._targets[key] = None + if new: + self._urgent.add(key) + self._kick.set() def unregister_target( self, diff --git a/infera/router/auto.py b/infera/router/auto.py index 918884b9..0db4a494 100644 --- a/infera/router/auto.py +++ b/infera/router/auto.py @@ -5,20 +5,28 @@ ############################################################################### from __future__ import annotations +import logging + from fastapi import Response +from fastapi.responses import JSONResponse from infera.common.worker_pool import DisaggMode from infera.router.base import BaseRouter from infera.router.disagg import DisaggRouter from infera.router.mixed import MixedRouter +logger = logging.getLogger(__name__) + class AutoRouter(BaseRouter): """Per-request router selector. - Selection policy (v0.1: PD-preferred with mixed fallback): - - If the model has BOTH prefill and decode workers → DisaggRouter - - Otherwise (only mixed, partial PD, or empty) → MixedRouter + Selection policy (PD-preferred with mixed fallback): + - BOTH prefill and decode workers → DisaggRouter + - Exactly one PD pool, and no mixed workers → 503 naming the + empty pool (half a PD deployment cannot serve, and saying "no mixed + worker" would describe something the operator never deployed) + - Otherwise (mixed workers present, or nothing at all) → MixedRouter (MixedRouter itself returns 503 if no mixed worker is available) This supports mixed deployments (some models PD, others mixed) and rolling @@ -38,11 +46,17 @@ def __init__(self, *args, **kwargs) -> None: self.policy, nats_client=self.nats_client, request_max_retries=self.request_max_retries, + # One breaker shared by both sub-routers: otherwise each would build + # its own default and the configured thresholds would never reach + # them, since AutoRouter is what the server actually constructs. + breaker=self.breaker, ) # Pass the NATS request client to the PD router too, so disaggregated # (prefill/decode) dispatch uses the per-instance NATS transport when # configured (it falls back to HTTP only when nats_client is None). - self._disagg = DisaggRouter(self.pool, self.policy, nats_client=self.nats_client) + self._disagg = DisaggRouter( + self.pool, self.policy, nats_client=self.nats_client, breaker=self.breaker + ) async def aclose(self) -> None: await self._mixed.aclose() @@ -60,4 +74,30 @@ async def dispatch( has_d = self.pool.list_active(model=model, mode=DisaggMode.DECODE) if has_p and has_d: return await self._disagg.dispatch(body, stream=stream, path=path) + # Exactly one PD pool populated: the deployment is disaggregated but + # half of it is gone. Falling through to the mixed router would be + # correct-but-useless -- there are no mixed workers either, so it + # answers "no active mixed worker", which sends the reader looking for + # something they never deployed while a decode (or prefill) pool sits + # right there. Scaling either side to zero is the usual cause. + if bool(has_p) != bool(has_d) and not self.pool.list_active( + model=model, mode=DisaggMode.MIXED + ): + present, missing = ("prefill", "decode") if has_p else ("decode", "prefill") + logger.warning( + "model=%r has %d %s worker(s) but no %s worker: PD dispatch needs both", + model, + len(has_p or has_d), + present, + missing, + ) + return JSONResponse( + content={ + "error": ( + f"model={model!r} has {len(has_p or has_d)} {present} worker(s) " + f"but no {missing} worker; PD dispatch requires both pools" + ) + }, + status_code=503, + ) return await self._mixed.dispatch(body, stream=stream, path=path) diff --git a/infera/router/base.py b/infera/router/base.py index ad2f91d9..5553dbe3 100644 --- a/infera/router/base.py +++ b/infera/router/base.py @@ -10,6 +10,7 @@ from fastapi import Response from infera.common.worker_pool import WorkerPool +from infera.router.breaker import CircuitBreaker from infera.router.policy.base import Policy @@ -28,6 +29,7 @@ def __init__( policy: Policy, nats_client=None, request_max_retries: int = 1, + breaker: CircuitBreaker | None = None, ) -> None: self.pool = pool self.policy = policy @@ -40,6 +42,11 @@ def __init__( # disables retries (single attempt). Mid-stream failures are never # retried (output already partially sent). self.request_max_retries = max(0, request_max_retries) + # Per-worker failure memory across requests. Failover alone forgets a + # bad worker the moment the request ends, so the next one re-picks it. + # Subclasses that select their own target consult this when filtering + # candidates; DirectRouter does not select and leaves it unused. + self.breaker = breaker if breaker is not None else CircuitBreaker() @abstractmethod async def dispatch( diff --git a/infera/router/breaker.py b/infera/router/breaker.py new file mode 100644 index 00000000..1cd1af29 --- /dev/null +++ b/infera/router/breaker.py @@ -0,0 +1,304 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# SPDX-License-Identifier: MIT +############################################################################### +"""Per-worker circuit breaker for routers that select their own target. + +Failover alone is not enough. It retries a failed dispatch on another worker, +but the memory of that failure lives in a per-request ``tried`` set that is +discarded when the request returns -- so the next request scores the same broken +worker as if nothing happened, picks it again if cache locality says so, and +pays the failover cost again. A worker that is healthy to the platform and +broken for inference therefore taxes *every* request, indefinitely. + +That worker is not hypothetical: it accepts the connection, answers ``/health``, +stays ``ACTIVE`` in discovery, and fails before the first byte. Kubernetes +cannot see it and neither can discovery. The router is the only component that +knows, and without this it forgets immediately. + +Scope, deliberately narrow: + +* **Only pre-first-byte failures trip it.** A failure after bytes have been + streamed is already non-retryable by design -- the worker demonstrably served + part of the request, and treating that as a health signal would open the + breaker on ordinary client disconnects. +* **It never touches ``WorkerStatus``.** That field is owned by discovery; this + is the router's private opinion, applied when filtering candidates. +* **Only routers that select use it.** ``direct.py`` has no failover because the + gateway owns selection there, so a breaker would be wrong. + +States are the usual three. ``closed`` routes normally. After +``failure_threshold`` consecutive failures the breaker goes ``open`` and the +worker is excluded for ``cooldown`` seconds. It then becomes ``half_open`` and +admits one probe at a time: success closes it and clears the count, failure +reopens it with the cooldown doubled, up to ``max_cooldown``. Backing off +matters because the common cause -- a worker wedged on a bad KV handoff -- does +not resolve on the first retry, and a fixed cooldown turns into a probe every +``cooldown`` seconds forever. + +"One at a time" is bounded by ``probe_timeout`` rather than by waiting for an +outcome, because the outcome may never arrive: the slot is claimed while +filtering candidates, and only the one the policy dispatches to reports back. +""" + +from __future__ import annotations + +import logging +import time +from dataclasses import dataclass, field +from enum import Enum + +from infera.server import metrics + +logger = logging.getLogger(__name__) + +_STATE_VALUE = {"closed": 0, "half_open": 1, "open": 2} + + +def _observe(worker_id: str, state) -> None: + """Mirror a state change to Prometheus. Never allowed to fail a request -- + an unregistered collector or a duplicate registry must not take out the + data plane.""" + try: + metrics.worker_breaker_state.labels(worker_id=worker_id).set(_STATE_VALUE[state.value]) + except Exception: # pragma: no cover + pass + + +def is_worker_fault(status: int) -> bool: + """True if an HTTP status is evidence about the *worker*, not the request. + + Failover retries on any pre-first-byte error, including 4xx -- that is + correct, since a 400 costs nothing to re-ask. Feeding 4xx to the breaker is + not: a malformed request returns 400 from every worker it touches, so the + breaker would trip the entire healthy fleet on one bad client. 429 is + excluded for a different reason -- it means "full right now", which the + policy's load accounting already routes around, and a 5s cooldown with + doubling is far too heavy a response to transient backpressure. + """ + return status >= 500 or status == 0 + + +class BreakerState(str, Enum): + CLOSED = "closed" + OPEN = "open" + HALF_OPEN = "half_open" + + +@dataclass +class _Entry: + consecutive_failures: int = 0 + state: BreakerState = BreakerState.CLOSED + # Wall time after which an open breaker becomes half-open. + opens_until: float = 0.0 + # Cooldown applied on the *next* trip; doubles each time a probe fails. + next_cooldown: float = 0.0 + # When the outstanding half-open probe was admitted, so only one is in + # flight at a time. None means the slot is free. + probe_started_at: float | None = None + trips: int = 0 + + +@dataclass +class CircuitBreaker: + """Tracks per-worker health as seen by dispatch outcomes. + + Not thread-safe by design: routers drive it from a single asyncio loop, and + a lock here would sit on the hot path of every request for no benefit. + """ + + failure_threshold: int = 3 + cooldown: float = 5.0 + max_cooldown: float = 60.0 + #: How long a claimed probe slot is honoured before it is reclaimed. + #: + #: Claiming and releasing the slot are not paired: ``filter`` claims one for + #: every candidate it lets through, and the policy then dispatches exactly + #: one of them, so the rest are never told how they did. A 4xx, a client + #: disconnect or a request that never returns leaves the slot held too. + #: Bounding the claim keeps any of those from wedging a healthy worker out + #: of rotation permanently. The cost of reclaiming too early is one extra + #: probe; the cost of never reclaiming is a worker lost until restart. + probe_timeout: float = 60.0 + #: Injectable clock, so tests do not sleep. + now: object = field(default=time.monotonic) + _entries: dict[str, _Entry] = field(default_factory=dict, init=False) + + # --- queries ------------------------------------------------------------ + + def _entry(self, worker_id: str) -> _Entry: + e = self._entries.get(worker_id) + if e is None: + e = _Entry(next_cooldown=self.cooldown) + self._entries[worker_id] = e + return e + + @property + def enabled(self) -> bool: + """A threshold of 0 or less turns the breaker off entirely, so an + operator can fall back to plain failover without a code change.""" + return self.failure_threshold > 0 + + def allows(self, worker_id: str) -> bool: + """True if this worker may be dispatched to right now. + + Transitions open -> half_open as a side effect when the cooldown has + elapsed, because the alternative is a separate timer whose only job is + to flip a flag that this function already has to check. + """ + if not self.enabled: + return True + e = self._entries.get(worker_id) + if e is None or e.state is BreakerState.CLOSED: + return True + now = self.now() + if e.state is BreakerState.OPEN: + if now < e.opens_until: + return False + e.state = BreakerState.HALF_OPEN + e.probe_started_at = None + _observe(worker_id, e.state) + logger.info("breaker: worker %s half-open, admitting one probe", worker_id) + # half_open: one probe at a time, and only for as long as a probe could + # plausibly still be running -- see probe_timeout for why the claim has + # to expire rather than wait for an outcome that may never come. + if e.probe_started_at is not None: + if now - e.probe_started_at < self.probe_timeout: + return False + logger.info( + "breaker: worker %s probe slot unclaimed after %.0fs; admitting another", + worker_id, + self.probe_timeout, + ) + e.probe_started_at = now + return True + + def filter(self, workers): + """Drop workers whose breaker is open. Returns a list. + + If every candidate is open, returns them all rather than nothing: a + request served by a probably-bad worker beats a 503 when there is no + alternative, and refusing to route would turn a partial outage into a + total one. + """ + allowed = [w for w in workers if self.allows(self._id_of(w))] + if allowed: + return allowed + if workers: + logger.warning( + "breaker: all %d candidate(s) open; routing anyway rather than failing", + len(workers), + ) + return list(workers) + + @staticmethod + def _id_of(w) -> str: + # Accepts a WorkerInfo or anything exposing .worker_id. + return getattr(w, "worker_id", None) or str(w) + + def state_of(self, worker_id: str) -> BreakerState: + e = self._entries.get(worker_id) + return e.state if e else BreakerState.CLOSED + + # --- outcomes ----------------------------------------------------------- + + def record_success(self, worker_id: str) -> None: + e = self._entries.get(worker_id) + if e is None: + return + if e.state is not BreakerState.CLOSED: + logger.info("breaker: worker %s recovered, closing", worker_id) + e.consecutive_failures = 0 + e.state = BreakerState.CLOSED + e.probe_started_at = None + e.next_cooldown = self.cooldown + _observe(worker_id, e.state) + + def record_neutral(self, worker_id: str) -> None: + """Release the probe slot without scoring the worker either way. + + For an outcome that says nothing about worker health -- a 4xx, which + every worker would answer identically, or a 429, which is backpressure + the policy already routes around. Counting it as recovery is as wrong + as counting it as failure: it would reset the failure count and close + an open breaker, so a worker alternating 500s and 400s would never + accumulate the consecutive failures needed to trip. But the slot such a + request consumed must still come back, or one bad client can wedge a + recovering worker out of rotation. + """ + e = self._entries.get(worker_id) + if e is None: + return + e.probe_started_at = None + + def forget(self, worker_id: str) -> None: + """Drop everything remembered about a worker that has left the fleet. + + Worker ids are addresses and a rebuilt Pod never reuses one, so without + this every rollout strands another entry -- and another pair of + Prometheus series, since both are labelled by worker id. + """ + if self._entries.pop(worker_id, None) is None: + return + for collector in (metrics.worker_breaker_state, metrics.worker_breaker_trips_total): + try: + collector.remove(worker_id) + except Exception: # noqa: BLE001 - never registered, or already gone + pass + + def record_failure(self, worker_id: str) -> None: + """Record a pre-first-byte dispatch failure.""" + if not self.enabled: + return + e = self._entry(worker_id) + e.consecutive_failures += 1 + was_probe = e.state is BreakerState.HALF_OPEN + e.probe_started_at = None + + if was_probe: + # A failed probe reopens immediately and backs off further, without + # waiting for the threshold again -- we already know it is bad. + e.next_cooldown = min(e.next_cooldown * 2, self.max_cooldown) + self._open(worker_id, e) + return + if e.consecutive_failures >= self.failure_threshold: + self._open(worker_id, e) + + def _open(self, worker_id: str, e: _Entry) -> None: + # A trip is an edge into exclusion, not every failure that lands while + # the worker is already excluded. A failed probe counts: it is a fresh + # verdict on a worker that was given another chance, and the doubling + # cooldown bounds how often one can happen. A failure while already + # open does not -- those arrive at the request rate, via the all-open + # fallback, and counting them turns the metric into a request counter + # that drowns out real trips and prints the warning on every request. + newly_tripped = e.state is not BreakerState.OPEN + e.state = BreakerState.OPEN + e.opens_until = self.now() + e.next_cooldown + _observe(worker_id, e.state) + if not newly_tripped: + return + e.trips += 1 + try: + metrics.worker_breaker_trips_total.labels(worker_id=worker_id).inc() + except Exception: # pragma: no cover - metrics must never break routing + pass + logger.warning( + "breaker: worker %s open for %.1fs after %d consecutive failure(s)", + worker_id, + e.next_cooldown, + e.consecutive_failures, + ) + + # --- introspection for metrics / tests ----------------------------------- + + def snapshot(self) -> dict[str, dict]: + return { + wid: { + "state": e.state.value, + "consecutive_failures": e.consecutive_failures, + "trips": e.trips, + } + for wid, e in self._entries.items() + } diff --git a/infera/router/disagg.py b/infera/router/disagg.py index a340e744..80f846d9 100644 --- a/infera/router/disagg.py +++ b/infera/router/disagg.py @@ -17,6 +17,7 @@ from infera.common.nats_request import TYPE_DATA, TYPE_DONE, TYPE_ERROR from infera.common.worker_pool import DisaggMode from infera.router.base import BaseRouter +from infera.router.breaker import is_worker_fault from infera.router.cache_control import parse_cache_hints from infera.router.disagg_protocols import ( ProtocolMismatch, @@ -119,6 +120,11 @@ async def dispatch( model = body.get("model") prefills = self.pool.list_active(model=model, mode=DisaggMode.PREFILL) decodes = self.pool.list_active(model=model, mode=DisaggMode.DECODE) + # Independently per role: a wedged prefill and a wedged decode are + # different events against different pools, and one open breaker + # must not remove the other role's healthy workers. + prefills = self.breaker.filter(prefills) + decodes = self.breaker.filter(decodes) if not prefills or not decodes: obs["outcome"] = "503" metrics.pd_bootstrap_failures_total.labels(reason="no_pd_workers").inc() @@ -131,7 +137,7 @@ async def dispatch( # skips prefill) vs D (load-heavy) differently. p_target, p_blocks = self.policy.pick(prefills, body, role_hint="prefill") d_target, d_blocks = self.policy.pick(decodes, body, role_hint="decode") - return await self._run_pd( + return await self._dispatch_pd( obs, p_target, d_target, p_blocks, d_blocks, body, stream, path ) @@ -162,11 +168,50 @@ async def dispatch_direct( }, status_code=503, ) - return await self._run_pd( + return await self._dispatch_pd( obs, RouteTarget(p), RouteTarget(d), [], [], body, stream, path ) - async def _run_pd( + def _score_leg_headers(self, worker_id: str, status: int) -> None: + """Score a streaming leg from its response headers. + + Headers are the pre-first-byte moment, so a failure recorded here is + exactly what the breaker wants. Success is not: a 200 header says the + request was accepted, nothing more, and the worker this class exists to + catch is the one that accepts a request and then produces nothing. That + profile would otherwise be scored as recovery -- resetting the failure + count and closing an open breaker -- so a 2xx is neutral here and the + success is recorded once the stream has actually delivered something. + """ + if status < 400: + self.breaker.record_neutral(worker_id) + else: + self._score_leg(worker_id, status) + + def _score_leg(self, worker_id: str, status: int) -> None: + """Record one PD leg's HTTP outcome against the worker that produced it. + + The two legs are two different workers whose health is independent, so + each has to be scored from its own response. A decode that answers + cannot vouch for a prefill that did not: scoring both off one status + code let a prefill 500-ing every request be reset to healthy by the + decode leg beside it, and there is no status code at all for the + client-facing streaming response, whose 200 is a framework default set + before either leg has been dispatched. + + A 5xx is the worker's fault. A 4xx is the request's -- every worker + would answer the same, so it frees the probe slot without counting + either way. Anything below 400 is the evidence of health that resets + the consecutive-failure count. + """ + if is_worker_fault(status): + self.breaker.record_failure(worker_id) + elif status < 400: + self.breaker.record_success(worker_id) + else: + self.breaker.record_neutral(worker_id) + + async def _dispatch_pd( self, obs, p_target: RouteTarget, @@ -333,16 +378,34 @@ async def _post(url, leg, worker_id, leg_body, leg_headers): with metrics.track_pd_leg(leg=leg, worker_id=worker_id): return await self._client.post(url, json=leg_body, headers=leg_headers) - try: - p_resp, d_resp = await asyncio.gather( - _post(p_url, "prefill", p.worker_id, p_body, p_headers), - _post(d_url, "decode", d.worker_id, d_body, d_headers), - ) - except httpx.HTTPError as exc: + # Gathered with return_exceptions so a failure can be attributed to + # the leg that produced it. Letting gather raise surfaces whichever + # one failed first with no way to tell which that was, and blaming a + # fixed leg means a decode outage evicts the healthy prefill worker + # while the broken decode is never scored at all. + p_resp, d_resp = await asyncio.gather( + _post(p_url, "prefill", p.worker_id, p_body, p_headers), + _post(d_url, "decode", d.worker_id, d_body, d_headers), + return_exceptions=True, + ) + failed = None + for worker_id, leg, result in ( + (p.worker_id, "prefill", p_resp), + (d.worker_id, "decode", d_resp), + ): + if isinstance(result, BaseException): + self.breaker.record_failure(worker_id) + if failed is None: + failed = (leg, result) + else: + self._score_leg(worker_id, result.status_code) + if failed is not None: + leg, exc = failed + if not isinstance(exc, httpx.HTTPError): + raise exc obs["outcome"] = "502" metrics.pd_bootstrap_failures_total.labels(reason="worker_unreachable").inc() - return _sanitized_error("PD request failed", exc, status_code=502) - + return _sanitized_error(f"PD {leg} leg failed", exc, status_code=502) if p_resp.status_code >= 400: logger.warning( "prefill worker %s returned %d (decode may fail)", @@ -375,16 +438,22 @@ def _start_prefill_drain_nats(self, p, p_payload): like the HTTP path. Strong ref guards against GC mid-flight.""" async def _drain(): + # Scored here for the same reason the HTTP legs are scored at their + # own responses: this transport never touches the HTTP client, so + # nothing else observes how this worker did. try: async for kind, _st, data in self.nats_client.stream(p.worker_id, p_payload): if kind == TYPE_ERROR: logger.warning("prefill leg (nats) %s failed: %s", p.worker_id, data[:200]) metrics.pd_bootstrap_failures_total.labels(reason="prefill_exception").inc() + self.breaker.record_failure(p.worker_id) return if kind == TYPE_DONE: + self.breaker.record_success(p.worker_id) return except Exception as exc: logger.warning("prefill nats drain %s failed: %s", p.worker_id, exc) + self.breaker.record_failure(p.worker_id) task = asyncio.create_task(_drain(), name="nats-prefill-drain") self._pending_prefill_tasks.add(task) @@ -431,6 +500,7 @@ async def _concurrent_nats( elif kind == TYPE_ERROR: # st carries 504 on inactivity timeout; worker errors -> 502. code = st or 502 + self._score_leg(d.worker_id, code) obs["outcome"] = str(code) return JSONResponse( content={ @@ -454,6 +524,7 @@ async def _concurrent_nats( }, status_code=502, ) + self._score_leg(d.worker_id, status) obs["outcome"] = "ok" if status < 400 else f"{status // 100}xx" return JSONResponse(content=payload, status_code=status) finally: @@ -467,14 +538,22 @@ async def _concurrent_nats( async def _stream_dual_nats(self, p_target, p_blocks, d_target, d_blocks, d_payload, p_task): """Stream decode's reply over NATS while prefill drains in background.""" d = d_target.worker + served = False try: async for kind, _st, data in self.nats_client.stream(d.worker_id, d_payload): if kind == TYPE_DATA: if data: + if not served: + # Bytes are flowing, so this worker is doing the + # work; an accepted request alone would not show it. + self.breaker.record_success(d.worker_id) + served = True yield data elif kind == TYPE_ERROR: logger.warning("decode (nats) %s stream failed: %s", d.worker_id, data[:200]) metrics.pd_bootstrap_failures_total.labels(reason="decode_stream_broken").inc() + if not served: + self.breaker.record_failure(d.worker_id) yield ( f'data: {{"error":"decode {d.worker_id} nats stream failed"}}\n\n' ).encode() @@ -547,8 +626,10 @@ async def _dispatch_serial( p_failed = True obs["outcome"] = "502" metrics.pd_bootstrap_failures_total.labels(reason="prefill_unreachable").inc() + self.breaker.record_failure(p.worker_id) return _sanitized_error("prefill leg failed", exc, status_code=502) + self._score_leg(p.worker_id, p_resp.status_code) if p_resp.status_code >= 400: p_failed = True obs["outcome"] = f"{p_resp.status_code // 100}xx" @@ -632,8 +713,10 @@ async def _dispatch_serial( except httpx.HTTPError as exc: obs["outcome"] = "502" metrics.pd_bootstrap_failures_total.labels(reason="decode_unreachable").inc() + self.breaker.record_failure(d.worker_id) return _sanitized_error("decode leg failed", exc, status_code=502) + self._score_leg(d.worker_id, d_resp.status_code) try: d_payload = d_resp.json() except ValueError: @@ -677,10 +760,12 @@ async def _stream_decode_only( exc or "", ) metrics.pd_bootstrap_failures_total.labels(reason="decode_unreachable").inc() + self.breaker.record_failure(d_target.worker.worker_id) err = json.dumps({"error": "decode unreachable"}) yield f"data: {err}\n\n".encode() return + self._score_leg_headers(d_target.worker.worker_id, d_resp.status_code) if d_resp.status_code >= 400: try: body_bytes = await d_resp.aread() @@ -709,8 +794,14 @@ async def _stream_decode_only( _DONE_NEEDLE = b"data: [DONE]" _TAIL_KEEP = len(_DONE_NEEDLE) - 1 tail = b"" + served = False try: async for chunk in d_resp.aiter_raw(): + if not served and chunk: + # Bytes are flowing, so the worker is doing the work -- + # which the headers alone did not establish. + self.breaker.record_success(d_target.worker.worker_id) + served = True if not done_seen: window = tail + chunk if _DONE_NEEDLE in window: @@ -831,11 +922,13 @@ async def _stream_dual( exc or "", ) metrics.pd_bootstrap_failures_total.labels(reason="decode_unreachable").inc() + self.breaker.record_failure(d_target.worker.worker_id) # json.dumps: exc text may contain chars that break SSE. err = json.dumps({"error": "decode unreachable"}) yield f"data: {err}\n\n".encode() return + self._score_leg_headers(d_target.worker.worker_id, d_resp.status_code) if d_resp.status_code >= 400: # Engine accepted but rejected; surface its body verbatim. try: @@ -867,7 +960,13 @@ async def _stream_dual( _DONE_NEEDLE = b"data: [DONE]" _TAIL_KEEP = len(_DONE_NEEDLE) - 1 tail = b"" + served = False async for chunk in d_resp.aiter_raw(): + if not served and chunk: + # Bytes are flowing, so the worker is doing the work -- + # which the headers alone did not establish. + self.breaker.record_success(d_target.worker.worker_id) + served = True if not done_seen: window = tail + chunk if _DONE_NEEDLE in window: @@ -908,6 +1007,9 @@ async def _stream_dual( try: p_resp = await asyncio.shield(p_task) except asyncio.CancelledError: + # The request was torn down from above, so the prefill worker + # was never given the chance to answer. That is not evidence + # about it either way, and scoring it would be inventing one. logger.debug("prefill task cancelled (parent torn down)") except Exception as exc: logger.warning( @@ -918,8 +1020,15 @@ async def _stream_dual( exc or "", ) metrics.pd_bootstrap_failures_total.labels(reason="prefill_exception").inc() + self.breaker.record_failure(p.worker_id) else: - if getattr(p_resp, "status_code", 0) >= 400: + # The prefill leg is scored here rather than beside the decode + # leg above because this is where its own answer arrives: it + # runs concurrently, so nothing about it is known until now. + p_status = getattr(p_resp, "status_code", None) + if p_status is not None: + self._score_leg(p.worker_id, p_status) + if p_status is not None and p_status >= 400: logger.warning( "prefill leg %s returned %d (decode will hang on KVPoll)", p.worker_id, diff --git a/infera/router/dp_routing.py b/infera/router/dp_routing.py index ffd905f2..c91214b7 100644 --- a/infera/router/dp_routing.py +++ b/infera/router/dp_routing.py @@ -15,6 +15,12 @@ from infera.common.worker_pool import EngineType from infera.router.policy.target import RouteTarget +#: Both engines honour this (SGLang ``DataParallelController``, vLLM +#: ``_get_data_parallel_rank``, case-insensitively). Named rather than inlined +#: so anything asserting on it breaks loudly on a rename instead of silently +#: never matching. Mirrors Rust's DP_RANK_HEADER. +DP_RANK_HEADER = "X-Data-Parallel-Rank" + def dp_rank_header(target: RouteTarget) -> dict[str, str] | None: """``X-Data-Parallel-Rank`` header pinning the request to a DP rank. Both @@ -22,7 +28,7 @@ def dp_rank_header(target: RouteTarget) -> dict[str, str] | None: case-insensitive) honour it; no-op when the target carries no rank.""" if target.dp_rank is None: return None - return {"X-Data-Parallel-Rank": str(target.dp_rank)} + return {DP_RANK_HEADER: str(target.dp_rank)} def inject_disagg_prefill_dp_rank( diff --git a/infera/router/mixed.py b/infera/router/mixed.py index 6f1836a1..9f2fe4fa 100644 --- a/infera/router/mixed.py +++ b/infera/router/mixed.py @@ -16,6 +16,7 @@ from infera.common.nats_request import TYPE_DATA, TYPE_DONE, TYPE_ERROR from infera.common.worker_pool import DisaggMode from infera.router.base import BaseRouter +from infera.router.breaker import is_worker_fault from infera.router.cache_control import parse_cache_hints from infera.router.dp_routing import dp_rank_header from infera.router.engine_priority import inject_engine_priority @@ -82,13 +83,35 @@ async def dispatch( for w in self.pool.list_active(model=model, mode=DisaggMode.MIXED) if w.worker_id not in tried ] + # Drop workers the breaker has open. Falls back to the unfiltered + # list when every candidate is open -- a request served by a + # probably-bad worker beats turning a partial outage into a 503. + candidates = self.breaker.filter(candidates) if not candidates: break target, blocks = self.policy.pick(candidates, body) tried.add(target.worker.worker_id) try: - return await self._attempt(target, blocks, body, hints, path, stream, obs) + resp = await self._attempt(target, blocks, body, hints, path, stream, obs) + # Only a clean response is evidence the worker is healthy. + # A 4xx says the request was bad -- every worker would + # answer the same -- so scoring it as recovery would reset + # the failure count and reopen a breaker that should stay + # shut. It still has to free the probe slot it took. + if getattr(resp, "status_code", 200) < 400: + self.breaker.record_success(target.worker.worker_id) + else: + self.breaker.record_neutral(target.worker.worker_id) + return resp except _Retry as r: + # Pre-first-byte only: _Retry is never raised once bytes have + # been streamed, so a mid-stream failure cannot trip this. + # 4xx is retried but not held against the worker -- see + # is_worker_fault(). + if is_worker_fault(getattr(r.response, "status_code", 0)): + self.breaker.record_failure(target.worker.worker_id) + else: + self.breaker.record_neutral(target.worker.worker_id) last_error = r.response logger.info( "failover: worker %s failed before first byte; %d worker(s) tried", @@ -242,6 +265,10 @@ async def _attempt_unary( ) ) from None obs["outcome"] = "ok" if status < 400 else f"{status // 100}xx" + # Same rule as the HTTP path below: a 5xx before any data is a + # worker fault and retryable; a 4xx belongs to the request. + if is_worker_fault(status): + raise _Retry(JSONResponse(content=payload_json, status_code=status)) return JSONResponse(content=payload_json, status_code=status) # Direct HTTP forward. @@ -275,6 +302,16 @@ async def _attempt_unary( ) ) from None obs["outcome"] = "ok" if resp.status_code < 400 else f"{resp.status_code // 100}xx" + # A 5xx here is the worker failing before a single byte reached the + # client, which is exactly the case failover exists for -- and until now + # this path returned it verbatim instead, so a unary request over HTTP + # never failed over and never fed the circuit breaker. The streaming + # path and the Rust router both retry it; this brings the third one into + # line. 4xx still passes straight through: the request itself is bad and + # every worker would say the same, so retrying only triples the latency + # of an error the client needs to see. + if is_worker_fault(resp.status_code): + raise _Retry(JSONResponse(content=payload_json, status_code=resp.status_code)) return JSONResponse(content=payload_json, status_code=resp.status_code) async def _normalized_stream( diff --git a/infera/server/__main__.py b/infera/server/__main__.py index f9400561..499d346f 100644 --- a/infera/server/__main__.py +++ b/infera/server/__main__.py @@ -21,6 +21,7 @@ from infera.kv.subscriber import KvEventSubscriberPool from infera.kv.writer import KvIndexWriter from infera.router.auto import AutoRouter +from infera.router.breaker import CircuitBreaker from infera.router.direct import DirectRouter from infera.router.policy.factory import build_policy from infera.server.app import init_app @@ -92,6 +93,18 @@ async def main(args) -> None: # older than what's already in the index. writer.set_reconciler(reconciler) + # Per-worker failure memory, shared by every router this process builds. + # DirectRouter never selects, so it holds one but does not consult it. + # + # Built here rather than beside the routers because on_worker_removed + # closes over it and the registry starts first, so a worker that leaves + # during startup would otherwise hit an unbound name. + breaker = CircuitBreaker( + failure_threshold=args.breaker_failure_threshold, + cooldown=args.breaker_cooldown_s, + max_cooldown=args.breaker_max_cooldown_s, + ) + # Per-worker snapshot of fields we need at removal time. The Registry # has already evicted the WorkerInfo from its pool by the time # on_worker_removed fires, so we stash the kv block at registration. @@ -152,6 +165,11 @@ def on_worker_removed(worker_id: str) -> None: except Exception: logger.exception("policy.on_worker_removed failed for %s", worker_id) + # A worker id is an address, and a rebuilt Pod never reuses one, so an + # entry left here outlives the fleet member it describes -- one more + # per rollout, each pinning a Prometheus series labelled by that id. + breaker.forget(worker_id) + # Phase 1 reconciler/subscriber cleanup. snap = kv_snapshots.pop(worker_id, None) if snap is None: @@ -231,6 +249,7 @@ def on_worker_removed(worker_id: str) -> None: policy, nats_client=nats_request_client, request_max_retries=args.request_max_retries, + breaker=breaker, ) logger.info("router-mode=direct (honouring GAIE EPP x-worker-instance-id)") else: @@ -239,6 +258,7 @@ def on_worker_removed(worker_id: str) -> None: policy, nats_client=nats_request_client, request_max_retries=args.request_max_retries, + breaker=breaker, ) app = init_app( registry, diff --git a/infera/server/args.py b/infera/server/args.py index 4551d6cb..bddfa807 100644 --- a/infera/server/args.py +++ b/infera/server/args.py @@ -206,4 +206,28 @@ def parse_server_args(argv: list[str] | None = None) -> argparse.Namespace: "backlog). Mid-stream failures are never retried. Default 1; 0 disables. " "Overrides $INFERA_REQUEST_MAX_RETRIES.", ) + parser.add_argument( + "--breaker-failure-threshold", + type=int, + default=int(os.environ.get("INFERA_BREAKER_FAILURE_THRESHOLD", "3") or 3), + help="Consecutive pre-first-byte worker faults (5xx / unreachable; 4xx " + "and 429 excluded) before the router takes a worker out of rotation. " + "Failover alone forgets between requests, so a worker that is ACTIVE in " + "discovery but broken for inference is otherwise re-picked forever. " + "0 disables the breaker. Overrides $INFERA_BREAKER_FAILURE_THRESHOLD.", + ) + parser.add_argument( + "--breaker-cooldown-s", + type=float, + default=float(os.environ.get("INFERA_BREAKER_COOLDOWN_S", "5") or 5), + help="Seconds a tripped worker is excluded before one probe request is " + "admitted. Overrides $INFERA_BREAKER_COOLDOWN_S.", + ) + parser.add_argument( + "--breaker-max-cooldown-s", + type=float, + default=float(os.environ.get("INFERA_BREAKER_MAX_COOLDOWN_S", "60") or 60), + help="Ceiling for the cooldown, which doubles on each failed probe. " + "Overrides $INFERA_BREAKER_MAX_COOLDOWN_S.", + ) return parser.parse_args(argv) diff --git a/infera/server/launch_rust.py b/infera/server/launch_rust.py index fbc23bce..16ea9a1d 100644 --- a/infera/server/launch_rust.py +++ b/infera/server/launch_rust.py @@ -80,6 +80,12 @@ def exec_rust(args: argparse.Namespace) -> None: args.request_transport, "--request-max-retries", str(args.request_max_retries), + "--breaker-failure-threshold", + str(args.breaker_failure_threshold), + "--breaker-cooldown-s", + str(args.breaker_cooldown_s), + "--breaker-max-cooldown-s", + str(args.breaker_max_cooldown_s), ] # kv-aware needs the tokenizer + overlap weights, or it degrades to # load-only routing (no cache locality). Resolve HF ids to a local path diff --git a/infera/server/metrics.py b/infera/server/metrics.py index ebd8af9b..b130bcd5 100644 --- a/infera/server/metrics.py +++ b/infera/server/metrics.py @@ -133,6 +133,24 @@ ) +worker_breaker_state = Gauge( + "infera_router_worker_breaker_state", + "Router-side circuit breaker per worker: 0=closed, 1=half_open, 2=open. " + "Non-zero means the router is routing around a worker that discovery still " + "reports ACTIVE — the gap this metric exists to make visible.", + labelnames=("worker_id",), + registry=REGISTRY, +) + +worker_breaker_trips_total = Counter( + "infera_router_worker_breaker_trips_total", + "Times a worker's breaker has opened. A worker tripping repeatedly while " + "staying ACTIVE is broken for inference but healthy to the platform.", + labelnames=("worker_id",), + registry=REGISTRY, +) + + # ---------------------------------------------------------------------- # KV-aware policy internals # ---------------------------------------------------------------------- diff --git a/manual/features/graceful_shutdown.md b/manual/features/graceful_shutdown.md new file mode 100644 index 00000000..9559bd4c --- /dev/null +++ b/manual/features/graceful_shutdown.md @@ -0,0 +1,71 @@ +# Graceful shutdown + +```{admonition} One-pager +:class: tip +**What:** a worker being removed stops receiving new requests immediately, then +finishes the generations it already accepted before the process exits. +**Why:** a severed generation cannot be retried — the tokens already streamed +cannot be un-sent — so without this, every rolling upgrade or scale-down +produces a burst of client errors. **Requires:** nothing, for finishing in-flight +work; Kubernetes with the default `kubernetes` discovery backend for the advance +notice described below. +``` + +```{important} +Finishing in-flight work happens on every backend, bounded by `--drain-timeout`. +What needs **Kubernetes with the default `kubernetes` discovery backend** is the +*advance* notice — the router learning a worker is leaving before the process is +signalled. That is not an implementation gap: it relies on the orchestrator +knowing a Pod is being removed, which nothing outside Kubernetes can tell the +router. `discoveryBackend: etcd` is rejected by the operator for in-cluster +deployments. +``` + +## What happens + +Removing a worker — a rolling update, a scale-down, draining a node — separates +two things that would otherwise happen at once: + +1. **It stops receiving.** Kubernetes marks the Pod the moment its removal is + requested, which is *before* the worker process is signalled. The router sees + that mark and stops choosing the worker within milliseconds, so new requests + go elsewhere while it is still running and long before it is told to stop. +2. **It keeps serving.** The worker finishes the generations it already + accepted, bounded by `--drain-timeout`, and only then exits. + +The early mark is what makes this different from simply stopping a process. +The `preStop` delay that follows is not spent waiting for the router to notice +— that already happened — but letting work in progress finish before the +process is signalled at all. + +Deploying through the operator needs no configuration: it injects the `preStop` +delay and sizes the termination grace period to cover the whole sequence. For +hand-written manifests and the per-stage timings, see +[Scaling a fleet](scaling.md). + +## When the Pod is not being deleted + +A worker can also be stopped without its Pod going anywhere — a liveness probe +failing and restarting the container, a node being shut down gracefully, someone +killing the process. There is no deletion, so there is no early mark, and the +router has no way to know until the worker says so. + +On those paths the worker removes its own registration as its first act on +`SIGTERM`, which stops new requests arriving, and drains after. In-flight +generations still finish. What is missing is the head start: from the moment the +decision is made to the moment the process is signalled, the router is still +sending work, because nothing has told it otherwise. + +## Elsewhere + +Deployments outside Kubernetes use an external etcd for discovery, where a +worker record is simply present or absent and nothing observes that a process is +leaving. Shutdown behaves as in the section above — deregister, then drain — +with in-flight work finished either way. The early notice is the part that needs +Kubernetes. + +In both cases the worker is absent from `/v1/workers` while it drains rather +than shown as draining, since removing the record is what stops new work +arriving. A Pod being deleted does show as `draining`, but earlier — between its +deletion being requested and the process being signalled, before the drain +itself begins. diff --git a/manual/features/routing_and_transport.md b/manual/features/routing_and_transport.md index 7ec103d9..a1437c92 100644 --- a/manual/features/routing_and_transport.md +++ b/manual/features/routing_and_transport.md @@ -113,6 +113,47 @@ idle-timeout-before-first-token, or a 429 admission reject. It never retries mid-stream (once tokens flow, a failure surfaces to the client). Raise it for more resilience; set `0` to fail fast. +### Circuit breaker + +Failover on its own has no memory. Its `tried` set lives for one request, so a +worker that is broken for inference but healthy to discovery — it accepts the +connection, answers `/health`, stays `ACTIVE` — gets re-picked by the *next* +request, and every request after that pays the failover cost again. + +The breaker is that missing memory. After `--breaker-failure-threshold` +consecutive faults a worker is dropped from the candidate list for +`--breaker-cooldown-s`, then one probe request is admitted: if it succeeds the +worker is restored, if it fails the cooldown doubles, up to +`--breaker-max-cooldown-s`. + +| Flag | Env | Default | Meaning | +|---|---|---|---| +| `--breaker-failure-threshold` | `INFERA_BREAKER_FAILURE_THRESHOLD` | `3` | consecutive faults before removal; `0` disables | +| `--breaker-cooldown-s` | `INFERA_BREAKER_COOLDOWN_S` | `5` | exclusion window before a probe | +| `--breaker-max-cooldown-s` | `INFERA_BREAKER_MAX_COOLDOWN_S` | `60` | cap on the doubling backoff | + +Two exclusions are deliberate. **4xx never counts** — a malformed request returns +400 from every worker it reaches, so counting it would trip the entire healthy +fleet on one bad client. **429 never counts** either: it means "full right now", +which the policy's load accounting already routes around, and a doubling cooldown +is far too heavy a response to transient backpressure. + +If *every* candidate is open the router dispatches anyway rather than returning +503 — a request served by a probably-bad worker beats turning a partial outage +into a total one. + +The breaker never writes `WorkerStatus`; that field belongs to discovery. This is +the router's private opinion, and it is visible as +`infera_router_worker_breaker_state` (0 closed / 1 half-open / 2 open) and +`infera_router_worker_breaker_trips_total`. A worker tripping repeatedly while +discovery still reports it `ACTIVE` is the signal worth alerting on. + +Both the Python and Rust routers implement this identically, with the same flags. + +The breaker is the router's view of a worker that is failing. For the orderly +case — a worker being removed on purpose — see [Scaling a fleet](scaling.md), +which covers draining in-flight generations before shutdown. + ## KV-event transport Powers [KV-aware routing](kv_aware_routing.md). `--kv-event-transport`: diff --git a/manual/features/scaling.md b/manual/features/scaling.md new file mode 100644 index 00000000..4c4653cf --- /dev/null +++ b/manual/features/scaling.md @@ -0,0 +1,480 @@ +# Scaling a fleet + +Adding and removing workers while traffic is flowing. Every number on this page +was measured on the hardware described in [Measurements](#measurements) — none +of it is projected. + +## How it works + +There is no scaling controller. Workers **self-register** into discovery when +they are ready and **deregister** when they shut down, and the router routes to +whatever is registered at that instant. Scaling is therefore just starting and +stopping worker processes; nothing has to be told about it. + +``` +worker ready ──► register (etcd lease / Pod annotation) ──► router's watch fires + ──► receives traffic +leaving ──► stop being routed to ──► drain in-flight ──► deregister ──► exit +``` + +What triggers "stop being routed to" depends on the backend: under Kubernetes +the orchestrator marks the Pod before the process is even signalled, elsewhere +removing the record is what does it. See [Who says the worker is +leaving](#who-says-the-worker-is-leaving), and +[Graceful shutdown](graceful_shutdown.md) for the feature as a whole. + +That shape is why scale-up and scale-down have very different costs. Scale-up is +bounded by **model load**, which is minutes. Scale-down is bounded by the +**longest in-flight generation**, which is seconds — and the router stops +choosing the worker in milliseconds, long before it stops serving. + +## Scaling up + +Start another worker with the same `--model-name` and the same discovery +settings. It joins when it is ready, and not before: registration happens after +the engine has loaded weights, so a worker in the pool is always a worker that +can serve. + +```bash +infera-worker ... --port 20002 --etcd-endpoint http://etcd:2379 +``` + +On Kubernetes, raise `replicas` on the worker service in the `InferaDeployment`. + +**Budget minutes, not seconds.** Measured cold start for an 8B model on one +MI355X was **140 s** from `docker run` to appearing in `/v1/workers`, almost all +of it weight loading. Anything that reacts to load by starting a worker has to +tolerate that delay — a rule that scales up when a queue is deep will still be +scaling up long after the queue drained. + +The corollary matters more than it looks: for a burst shorter than the cold +start, **adding workers cannot help**. Either keep headroom, or shift traffic +between roles that are already running (see +[PD disaggregation](pd_disaggregation.md)). + +## Scaling down + +Send `SIGTERM`. Do not `SIGKILL`, and do not simply delete the Pod without a +grace period. + +The worker then, in this order: + +1. **Stops being routed to.** The worker removes its registration, which is + what stops new work arriving. Under Kubernetes a Pod being *deleted* has + already left routing well before this — the registry acts on its + deletionTimestamp, before the process is even signalled — so this is only + cleanup there (see [Who says the worker is + leaving](#who-says-the-worker-is-leaving)). +2. **Drains.** On the NATS transport infera tracks in-flight requests directly. + On HTTP the router talks straight to the engine, so infera asks the engine + instead, polling its `/metrics` until running, queued, and PD-handoff queues + all reach zero. Bounded by `--drain-timeout` (default 30 s). +3. **Stops the engine.** + +Requests already in flight run to completion. Requests that arrive during the +drain go to other workers. + +**Two different timings, easily conflated.** A worker stops *receiving* new +requests within a second of the shutdown starting — the router's watch picking +up either the `deletionTimestamp` or the record's removal — and it is the number +that decides whether traffic is still being sent somewhere that is about to die. How long the *process* then lives is +a separate and much larger number, set by the longest generation it was already +serving. Measured: under a second to stop receiving, while a 40-second +generation ran to completion afterwards. + +Watching `/v1/workers` measures neither. The record now goes when the drain +*starts*, not when it ends, so its disappearance marks the beginning of the +in-flight work rather than the end of it — a worker finishing a long generation +is absent from that list for all of it. A Pod being deleted shows as `draining` +only for the window between its deletion being requested and the process being +signalled, which is the preStop hook's 15 s. + +```{note} +`--drain-timeout` is a **ceiling, not a delay** — a worker with nothing in flight +exits in about six seconds regardless. Set it above your p99 generation time. +Anything still running when it expires is cut, with a warning naming the count. +``` + +### The transport decides how well this works + +Draining is only as good as the router's view of what is in flight, and that +differs by transport — not by implementation quality, but by where the +information lives. + +| | who knows what is in flight | drain | +|---|---|---| +| **NATS** (`--request-transport nats`) | infera — it owns the request path and holds the in-flight set | exact, no polling | +| **HTTP** (default in the recipes) | only the engine — the router dials it directly and never sees the request | poll the engine's `/metrics`, behind a settle window | + +Measured with a GPU-free stand-in worker, same generation: + +- **NATS, one in-flight generation**: the log reads `draining 1 in-flight NATS + request(s)` — it knows the count — the 300-chunk generation completed in full, + and the worker deregistered **21.3 s** later, which is just the remaining + generation time with no overhead. +- **NATS, nothing in flight**: leaving rotation to exit in **3 ms**. +- **HTTP with a real engine, nothing in flight**: at least the **6 s** settle + window, because a single zero reading cannot be told apart from a gauge that + has not refreshed yet. + +So NATS costs a broker and buys a drain that is exact rather than inferred. It +also buys request cancellation the HTTP path does not have — a timeout or client +disconnect publishes to `infera.cancel.` and the worker tears down the +engine connection, instead of leaving it generating. + +### Admission control + +Setting `INFERA_NATS_REQ_MAX_PENDING` (or `--nats-req-max-pending`) above zero +on **both** the server and the workers makes the request path JetStream-backed: +a WorkQueue stream with one durable consumer per worker. The router reads that +consumer's backlog before dispatching and refuses a worker over the limit. + +This matters for scaling because it covers the window scaling cannot: a burst +shorter than a 140 s cold start cannot be answered by adding workers, so the +choice is between queueing behind a saturated worker and steering away from it. + +**Look at the distribution, not the status codes.** A refusal raises the same +retryable failure as any other pre-first-byte error, so the request fails over +to a freer worker and the client sees `200`. Only when every worker is over the +limit and retries are exhausted does a `429` reach the client. Measured with one +deliberately saturated worker (concurrency 1) and one fast one, limit 3: + +| | saturated worker | fast worker | +|---|---|---| +| 20 requests under backlog | **+0** | **+20** | +| round-robin without the throttle | +10 | +10 | + +The worker's consumer showed `num_ack_pending = 10` against a limit of 3 at the +time — the ack happens after the request is fully proxied, precisely so the +backlog gauge reflects genuinely in-flight work rather than mere delivery. + +```{note} +The check is per dispatch, so it steers *new* requests. Requests already +dispatched are unaffected, and a simultaneous burst is all admitted — every +admission check runs before any of them has built backlog. +``` + +```{note} +The Rust router does not implement the NATS transport (`lib.rs`: "Configs +outside this set (NATS transport, ...) are served by the Python backend"), so +the Rust data plane and the NATS drain are currently an either/or. +``` + +### Why in-flight work is visible at all + +The engine's own gauges are the only source of truth on the HTTP path, and they +have three properties worth knowing: + +- **SGLang serves `/metrics` only with `--enable-metrics`.** Without it the + endpoint 404s and the drain has nothing to read. The worker entrypoint injects + the flag, so this is handled — but a hand-rolled deployment that bypasses it + will silently lose the drain. +- **The gauges lag.** Measured on SGLang: `num_running_reqs` stayed at 12 for + 5–15 s after the last response completed. The drain therefore requires the + count to read zero continuously for a settle window before believing it, + which also protects against a request accepted moments before `SIGTERM` that + has not been counted yet. +- **PD handoff queues count as in-flight.** A prefill worker can show no running + and no queued requests while KV transfers are still outstanding. Stopping it + there strands the decode workers waiting on that KV, so + `num_prefill_bootstrap_queue_reqs`, `num_prefill_inflight_queue_reqs`, + `num_decode_prealloc_queue_reqs` and `num_decode_transfer_queue_reqs` are + included in the count. + +If the in-flight count cannot be read at all — an unknown engine, a renamed +series, a dead HTTP server — the worker logs a warning naming the metric it +looked for and shuts down **without** draining rather than blocking. A rolling +update that stalls on a parse failure is worse than one that cuts a request, and +a silent full-timeout wait would be indistinguishable from a genuinely busy +worker. + +### On Kubernetes + +The recipes deploy with `discoveryBackend: kubernetes` and +`--request-transport http`, so the shutdown path differs from a bare +etcd deployment in two ways — and gains one stage. + +**Discovery is a Pod annotation, not an etcd lease.** Registering writes +`infera.amd.com/worker-info` on the worker's own Pod; deregistering clears it. +The registry additionally marks a Pod `DRAINING` the moment it carries a +`deletionTimestamp`, without waiting for the container to exit. That matters +because a terminating Pod keeps `phase: Running` — without the check it would +stay a routing candidate for the whole `preStop` delay, turning a hook meant to +make shutdown graceful into extra seconds of accepting work about to be killed. + +The mark, rather than an outright removal, is what keeps the two timings above +distinguishable on this backend too: the worker leaves routing immediately and +its record stays until it clears its own annotation at the end of the drain, so +`/v1/workers` shows a rollout in progress instead of a worker that vanished. + +**There is a `preStop` delay before `SIGTERM`.** The operator injects +`sleep 15`, so the full sequence is: + +``` +deletion requested ──► deletionTimestamp set ──► registry drops the worker + ──► preStop sleep 15 (still serving what it has) + ──► SIGTERM ──► drain ──► deregister ──► engine.stop() + ──► [kubelet SIGKILL at terminationGracePeriodSeconds] +``` + +### Who says the worker is leaving + +The two discovery backends learn this in different ways, and only one of them +needs the worker to say anything. The difference is not an inconsistency to be +smoothed over — it is what each backend can actually observe. + +**Kubernetes: the orchestrator says so.** A condemned Pod carries +`deletionTimestamp` from the moment deletion is requested, which is before the +`preStop` hook runs and therefore before the process is signalled at all. The +registry reads it and drops the worker from routing immediately — measured at +under 100 ms against the 15 s `preStop` delay. The worker announcing the same +thing later would add nothing: routing has already stopped, and the terminating +check returns before the annotation is even parsed. + +So on this backend the annotation carries **identity only** — worker id, URL, +model, engine, role, KV endpoints — all of it fixed for the life of the +process. That is deliberate. The heartbeat re-asserts the annotation to +self-heal, rebuilding it from config; if state lived there too, a refresh +landing mid-drain would overwrite it with a payload that omits the status, +which parses as `ACTIVE`, and the worker would be handed new work it is about +to refuse. + +**Everywhere else: removing the record says so.** On etcd there is no +orchestrator at all — a record is either present with an unexpired lease or it +is gone, with no third state to put it in. The same is true on Kubernetes +whenever the Pod is *not* being deleted: a liveness probe restarting the +container, a node shutting down gracefully, someone killing the process. No +deletionTimestamp is set, so the registry reads the annotation as usual and the +worker stays routable until it clears it. + +So every shutdown deregisters *first* and drains after. In-flight generations +are finished either way; the cost is that the worker is absent from +`/v1/workers` while it drains rather than shown as draining. The head start — +leaving routing before the process is signalled at all — is what deleting a Pod +buys, and only that. + +```{warning} +`discoveryBackend: etcd` **is not supported for in-cluster deployments** and the +operator refuses it. The combination keeps the `preStop` delay while losing the +early notice it exists to provide: the server no longer watches Pods, so nothing +reads the `deletionTimestamp`, and the only signal left arrives after `SIGTERM` +— that is, after the delay has already elapsed. For its whole duration the +router keeps handing new work to a Pod that is already condemned. Use the +default `kubernetes` backend in Kubernetes; external etcd is for deployments +outside it. +``` + +### Worst case, and the budget + +Every stage is individually bounded: + +| Stage | Bound | Set by | +|---|---|---| +| `preStop` | 15 s | operator | +| drain | `--drain-timeout` (default 30 s) | flag | +| deregister | 10 s | registration HTTP client timeout | +| `engine.stop()` | 30 s | `SIGTERM` to the engine's process group, then `SIGKILL` | +| **total** | **≈95 s at defaults** | | + +`terminationGracePeriodSeconds` has to cover that whole sum, because the kubelet +`SIGKILL`s the moment it expires — mid-drain if that is where things are. The +operator now **derives** it as `preStop + --drain-timeout + 50 s` of teardown +headroom, with a 120 s floor, reading the flag from `ServiceSpec.Args` or from +the container directly when an `extraPodSpec` template supplies it. + +```{warning} +This used to be a fixed 120 s with a comment saying it "must exceed preStop + +the worker `--drain-timeout`" — and nothing parsed that flag, so the invariant +was documented and unenforced. Raising `--drain-timeout` for long generations +(the only reason anyone raises it) pushed shutdown past the grace and turned the +drain back into a kill. Measured on a live cluster before the change: a worker +declaring `--drain-timeout 300` still received `terminationGracePeriodSeconds: +120`, i.e. 365 s of budget granted 120. +``` + +Measured on a live k3s cluster: `kubectl delete pod` took the worker out of +routing in **87–93 ms**, against the 15 000 ms `preStop` delay. That gap is the +whole point of reading `deletionTimestamp` — the alternatives (the `DELETE` +event, or `phase` leaving `Running`) only fire once the container has already +exited, so without it the router would keep assigning work for the entire +`preStop` window and then have it killed. + +Two runs, both with a Pod deleted while holding in-flight work: + +- **Real SGLang Qwen3-8B** deployed by the operator (`InferaDeployment`, two + workers, one MI355X each, Kubernetes discovery, HTTP transport): four + concurrent 2500-token generations in flight, **4/4 completed with HTTP 200** + and full-length output (6.5–13.3 kB), replacement Pod registered before the + drain finished. +- **GPU-free stand-in workers**, same path: a 300-chunk generation completed + in full across the drain. + +```{note} +`spec.services..resources` is **ignored when `extraPodSpec` is set** — the +template is passed through verbatim, so the GPU request has to live on your own +container. A worker that omits it schedules, starts, and then fails with "No +accelerator available". +``` + +If you set the grace period yourself it is respected as long as it is **larger** +than the derived value; it is only ever raised, never lowered. + +## PD and DP + +Prefill and decode register into separate pools and are selected per request, so +they scale **independently** — add prefill for longer inputs, decode for more +concurrent users. Two constraints: + +- **Neither pool can go to zero.** PD dispatch fails closed when either side is + empty — `minReplicas: 0` on either is an outage, not an idle saving. The 503 + names the empty pool (`has 1 decode worker(s) but no prefill worker`), so the + cause is visible without reading the fleet. +- **A DP worker's shape decides who picks the rank.** A worker registering + `dp_size > 1` with **no** `dp_rank` is rank-multiplexed: the router fans it + out into one target per rank and pins `X-Data-Parallel-Rank`. A worker that + registers its own `dp_rank` is a plain endpoint and opts out — its address + already selects the rank. Both are valid; only the first involves the router. + +## Across machines + +Nothing about scaling changes when workers live on different hosts — discovery +is already the coordination point, so a worker on another machine joins the same +way. Two things do change, and both are configuration rather than mechanism: + +**`--advertise-host` must be the node's routable address.** It is the URL peers +dial, and the single-node habit of leaving it at `127.0.0.1` registers an +address that resolves to the wrong machine everywhere else. The failure is +quiet in the worst way: the router *lists* the worker and cannot reach it, so it +looks like a broken worker rather than a misconfiguration. On Kubernetes, take +it from the downward API (`POD_IP`). + +**Discovery must be reachable from every node.** An etcd bound only to loopback, +or advertising a loopback client URL, works perfectly on the node running it and +is invisible from the others. + +Checking both before deploying costs nothing: + +```bash +# from each worker node +curl -s -o /dev/null -w '%{http_code}\n' -X POST http://:2379/v3/kv/range -d '{"key":"Lw=="}' +# from the router node, once a worker has registered +curl -s http://:8000/v1/workers | jq -r '.workers[].url' # must be dialable +``` + +Measured on two nodes (chi2800 / chi2866, one MI355X each, workers advertising +their own IPs, etcd and router on the first node): both workers registered with +distinct addresses, 12 requests distributed 7/7 across the machines, and a +`SIGTERM` to the **remote** worker drained cleanly — its three in-flight +3000-token generations all completed (13.7–14.4 k characters), its record +disappeared after 30 s, and 100 requests flowing through the router during the +whole transition saw **0 failures**. (As above, the record surviving 30 s is the +generations finishing, not 30 s of continuing to receive work.) + +```{warning} +This covers workers on separate machines. It does **not** cover a single worker +*spanning* machines (`numberOfNodes > 1`, LeaderWorkerSet) or PD over RDMA +between nodes — neither has been exercised here. Note also that on this cluster +`rdma/hca` is not advertised as an allocatable resource, so a PD deployment +would need host networking and direct device access rather than a device plugin. +``` + +## Measurements + +SGLang 0.5.15 and vLLM 0.1.dev19253, Qwen3-8B, one MI355X per instance, HTTP +transport, etcd discovery, real router. + +| | | +|---|---| +| Cold start (`docker run` → in `/v1/workers`) | **140 s** | +| Scale-down: `SIGTERM` → stops receiving new requests | **< 1 s** | +| Scale-down: `SIGTERM` → record gone from `/v1/workers` | **30–38 s** | +| Router reaction to a worker's record being deleted | **15 ms** | +| Drain settle window | 6 s | + +Two runs, both with traffic flowing throughout: + +**Drain under load.** Six concurrent 4000-token generations in flight at +`SIGTERM`. Both engines: **6/6 completed with HTTP 200** and full-length output +(15–19 k characters). SGLang 22 s, vLLM 19 s from signal to last response. + +**Scale up then down.** Two instances, continuous traffic, a third added and +then one removed. **260 requests, 0 failures**, including in the 5-second +windows around each transition. The removed instance left rotation, drained the +one generation it was holding (`engine idle for 6s, 1 request(s) completed`), +and only then exited. + +```{note} +This run predates the change that made the shutdown order backend-specific, so +its logs show the worker announcing `DRAINING` before draining. On etcd the two +steps are now the other way round — deregister, then drain — which is what stops +new work arriving on a backend where nothing else can. The request counts are +unaffected: both orderings stop new work before waiting on in-flight work. +``` + +**PD scaling, measured.** A 1P1D fake fleet grown to 2P2D and shrunk back under +continuous traffic: **200 requests, 0 failures**, both pools scaling +independently and the drained workers finishing their in-flight work. Taking the +last prefill away then returns 503 naming the empty pool. + +```{warning} +**Not measured:** multi-node workers, TP > 1, PD scaling with a *real* engine +(the run above used GPU-free stand-ins, so no KV moved), and scale-down during an +active KV transfer. The PD handoff queues are counted in the drain, but that +path has not been exercised on hardware. +``` + +## Scaling a deployment + +Edit the service's `replicas` in the `InferaDeployment`. That is the only +supported way in, and it is the only write that survives: + +```bash +kubectl patch inferadeployment qwen --type=merge \ + -p '{"spec":{"services":{"decode":{"replicas":5}}}}' +``` + +For a multi-node service the count is **groups**, not pods: `replicas: 5` with +`numberOfNodes: 3` is fifteen pods and five servable instances, since only +node-rank 0 of each group registers. + +Pods removed by a scale-down drain first — the operator injects the `preStop` +delay and a grace period sized from `--drain-timeout`, so the sequence is the +same one `kubectl delete pod` follows. + +```{warning} +Do **not** scale the generated `Deployment` or `LeaderWorkerSet` directly. Both +carry a real `/scale` subresource, so the write succeeds and nothing reports an +error — and then the next reconcile reverts it, because this reconciler assigns +the whole child `.Spec` on every pass. Measured: a `kubectl scale` to 3 went +back to 1 in under 3 seconds. The only symptom is a replica count that keeps +snapping back. +``` + +## Autoscaling + +Infera ships no autoscaler, and there is currently no `/scale` surface for an +external one to drive. + +An `InferaDeployment` cannot carry `/scale` itself, and that is a property of +its shape rather than an omission: `spec.services` is a map with user-chosen +keys, while the scale subresource requires `specReplicasPath` to be a *static* +dot-notation JSONPath, and a CRD may declare only one. A single path could name +one service — hardcoding `decode`, say — which leaves every other pool, and in +a PD deployment specifically the prefill pool, with no handle at all. + +Pointing an autoscaler at the generated workload does not work either, for the +reason in the warning above: those objects are derived state and are rewritten +every pass. + +Two harder problems sit behind the plumbing anyway: + +- **A 140-second cold start sits inside a control loop that ticks every 15 + seconds.** A burst shorter than the cold start cannot be answered by adding + workers at all. +- **Nothing in Kubernetes lets a scaler choose *which* replica to remove**, so + the one holding the warmest KV cache is as likely to go as any other. Upstream + has declined to fix this (k8s#123541, closed as not planned). + +The signals worth scaling on (`vllm:num_requests_waiting`, +`sglang:num_queue_reqs`, KV utilisation) are exposed by the engines and already +read by the drain path, but nothing polls them continuously yet. diff --git a/manual/getting_started/quickstart.md b/manual/getting_started/quickstart.md index 081d6783..7f5c8f73 100644 --- a/manual/getting_started/quickstart.md +++ b/manual/getting_started/quickstart.md @@ -80,6 +80,12 @@ kubernetes` (needs a k8s API + label selector) and `--request-transport nats` `etcd` discovery + `http` request transport + `zmq` KV events. Set the same three flags on **every** server and worker, or they won't find each other. See [Routing & transport](../features/routing_and_transport.md). + +One behaviour differs on this path: stopping a worker still lets its in-flight +generations finish, but nothing takes it out of routing until it is signalled. +Leaving rotation *before* the process is told to stop needs Kubernetes to +report the Pod as going away — see +[Graceful shutdown](../features/graceful_shutdown.md). ``` ```{tip} diff --git a/manual/reference/cli.md b/manual/reference/cli.md index ed4afd56..184171c5 100644 --- a/manual/reference/cli.md +++ b/manual/reference/cli.md @@ -22,6 +22,9 @@ the **same** on the server and every worker. See | `--kv-prefill-overlap-weight` | (unset) | KV-aware PD: prefill-side weight (typical `20.0`); overrides the global | | `--kv-decode-overlap-weight` | (unset) | KV-aware PD: decode-side weight (typical `2.0`); overrides the global | | `--request-max-retries` | `1` | retry on an alternate worker on pre-response failure (never mid-stream); `0` disables | +| `--breaker-failure-threshold` | `3` | consecutive worker faults (5xx / unreachable) before a worker leaves rotation; `0` disables the breaker | +| `--breaker-cooldown-s` | `5` | how long a tripped worker is excluded before one probe request is admitted | +| `--breaker-max-cooldown-s` | `60` | ceiling for that cooldown, which doubles on each failed probe | | `--discovery-backend` | `kubernetes` | `kubernetes` \| `etcd` | | `--etcd-endpoint` | — | required for `--discovery-backend etcd` | | `--etcd-prefix` | `/infera/workers/` | etcd key prefix the fleet registers under | diff --git a/manual/sphinx/_toc.yml.in b/manual/sphinx/_toc.yml.in index cb76e25c..a4b79ab2 100644 --- a/manual/sphinx/_toc.yml.in +++ b/manual/sphinx/_toc.yml.in @@ -45,6 +45,10 @@ subtrees: title: KV-aware routing - file: features/routing_and_transport.md title: Routing and transport + - file: features/scaling.md + title: Scaling a fleet + - file: features/graceful_shutdown.md + title: Graceful shutdown - caption: Serving entries: diff --git a/pyproject.toml b/pyproject.toml index e71ed3ed..6119a9c7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -41,7 +41,6 @@ infera-kvd-probe = "infera.kvd.bench.probe:main" infera-kvd-l3-bench = "infera.kvd.bench.l3_bench:main" # node + PD preflight suite (gpu / network / storage / firmware / host probes) infera-preflight = "infera.tools.preflight.cli:main" - [build-system] requires = ["setuptools>=69", "setuptools_scm[toml]>=8", "wheel"] build-backend = "setuptools.build_meta" diff --git a/rust/router/src/breaker.rs b/rust/router/src/breaker.rs new file mode 100644 index 00000000..26303ae8 --- /dev/null +++ b/rust/router/src/breaker.rs @@ -0,0 +1,715 @@ +/////////////////////////////////////////////////////////////////////////////// +// Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +// +// SPDX-License-Identifier: MIT +/////////////////////////////////////////////////////////////////////////////// +//! Per-worker circuit breaker. Mirrors `infera/router/breaker.py` — same three +//! states, same thresholds, same all-open fallback — so the two data planes +//! behave identically under a wedged worker. +//! +//! Failover on its own is not enough. It retries a failed dispatch elsewhere, +//! but the memory of that failure lives in a per-request `tried` set that is +//! dropped when the request returns, so the next request scores the same broken +//! worker as if nothing had happened. A worker that answers `/health`, stays +//! ACTIVE in etcd, and fails before the first byte therefore taxes *every* +//! request, indefinitely. +//! +//! Unlike the Python side this is shared across tokio worker threads, so the +//! map lives behind a `Mutex`. The critical sections are a hash lookup and a +//! few integer writes; contention is not a concern at any plausible request +//! rate, and a lock-free design here would buy nothing for the complexity. + +use std::collections::{HashMap, HashSet}; +use std::sync::Mutex; +use std::time::{Duration, Instant}; + +/// Statuses that are evidence about the *worker*, not the request. +/// +/// Failover retries on any pre-first-byte error including 4xx, which is +/// correct — re-asking costs nothing. Feeding 4xx to the breaker is not: a +/// malformed request returns 400 from every worker it touches, so one bad +/// client would trip the whole healthy fleet. 429 is excluded for a different +/// reason: it means "full right now", which load accounting already routes +/// around, and a doubling cooldown is far too heavy for transient backpressure. +pub fn is_worker_fault(status: u16) -> bool { + status >= 500 || status == 0 +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum BreakerState { + Closed, + Open, + HalfOpen, +} + +impl BreakerState { + pub fn as_str(self) -> &'static str { + match self { + BreakerState::Closed => "closed", + BreakerState::Open => "open", + BreakerState::HalfOpen => "half_open", + } + } +} + +#[derive(Debug)] +struct Entry { + consecutive_failures: u32, + state: BreakerState, + /// Instant after which an open breaker becomes half-open. + opens_until: Instant, + /// Cooldown applied on the *next* trip; doubles each time a probe fails. + next_cooldown: Duration, + /// When the outstanding half-open probe was admitted, so only one runs at + /// a time. `None` means the slot is free. + probe_started_at: Option, + trips: u64, +} + +/// How long a claimed probe slot is honoured before it is reclaimed. +/// +/// Claiming and releasing are not paired: `filter` claims a slot for every +/// candidate it lets through, and the policy dispatches to exactly one of them, +/// so the rest are never told how they did. A 4xx records neither outcome +/// either, and a cancelled request -- the client hung up, or the worker never +/// answered, which is the very condition this guards against -- unwinds without +/// reaching any record call. Any of those would otherwise hold the slot +/// forever, leaving a recovered worker permanently out of rotation. +const PROBE_TIMEOUT: Duration = Duration::from_secs(60); + +pub struct CircuitBreaker { + failure_threshold: u32, + cooldown: Duration, + max_cooldown: Duration, + probe_timeout: Duration, + entries: Mutex>, +} + +impl CircuitBreaker { + pub fn new(failure_threshold: u32, cooldown: Duration, max_cooldown: Duration) -> Self { + Self { + failure_threshold, + cooldown, + max_cooldown, + probe_timeout: PROBE_TIMEOUT, + entries: Mutex::new(HashMap::new()), + } + } + + /// Whether this worker may be dispatched to right now. + /// + /// Transitions Open -> HalfOpen as a side effect once the cooldown has + /// elapsed, because the alternative is a background timer whose only job is + /// to flip a flag this function already has to read. Call it once per + /// candidate per request: in HalfOpen it *consumes* the single probe slot. + pub fn allows(&self, worker_id: &str) -> bool { + self.allows_at(worker_id, Instant::now()) + } + + /// A threshold of 0 turns the breaker off entirely, so an operator can fall + /// back to plain failover without a code change. + /// Take the entry lock, recovering it if a previous holder panicked. + /// + /// Every request passes through the breaker, so `unwrap` here would turn + /// one panic into a panic on every subsequent request -- a local fault + /// escalated into total unavailability, which is the outcome this whole + /// class exists to avoid. The guarded state is a plain map of counters; a + /// panic mid-update can leave one worker's entry stale, never the process + /// unusable. + fn entries(&self) -> std::sync::MutexGuard<'_, HashMap> { + self.entries + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + } + + fn enabled(&self) -> bool { + self.failure_threshold > 0 + } + + fn allows_at(&self, worker_id: &str, now: Instant) -> bool { + if !self.enabled() { + return true; + } + let mut map = self.entries(); + let Some(e) = map.get_mut(worker_id) else { + return true; + }; + match e.state { + BreakerState::Closed => return true, + BreakerState::Open => { + if now < e.opens_until { + return false; + } + e.state = BreakerState::HalfOpen; + e.probe_started_at = None; + tracing::info!(worker = worker_id, "breaker half-open, admitting one probe"); + } + BreakerState::HalfOpen => {} + } + // One probe at a time, but only for as long as one could plausibly + // still be running: the claim expires rather than waiting for an + // outcome that may never arrive. See PROBE_TIMEOUT. + if let Some(started) = e.probe_started_at { + if now.duration_since(started) < self.probe_timeout { + return false; + } + tracing::info!( + worker = worker_id, + "breaker: probe slot unclaimed, admitting another" + ); + } + e.probe_started_at = Some(now); + true + } + + /// Drop workers whose breaker is open. + /// + /// If *every* candidate is open the full list is returned instead of an + /// empty one: a request served by a probably-bad worker beats a guaranteed + /// 503, and refusing to route would turn a partial outage into a total one. + pub fn filter(&self, workers: &[W], id_of: impl Fn(&W) -> &str) -> Vec { + let allowed: Vec = workers + .iter() + .filter(|w| self.allows(id_of(w))) + .cloned() + .collect(); + if !allowed.is_empty() { + return allowed; + } + if !workers.is_empty() { + tracing::warn!( + candidates = workers.len(), + "breaker: all candidates open; routing anyway rather than failing" + ); + } + workers.to_vec() + } + + pub fn record_success(&self, worker_id: &str) { + let mut map = self.entries(); + if let Some(e) = map.get_mut(worker_id) { + if e.state != BreakerState::Closed { + tracing::info!(worker = worker_id, "breaker: worker recovered, closing"); + } + e.consecutive_failures = 0; + e.state = BreakerState::Closed; + e.probe_started_at = None; + e.next_cooldown = self.cooldown; + } + } + + /// Release the probe slot without scoring the worker either way. + /// + /// For an outcome that says nothing about worker health: a 4xx, which every + /// worker would answer identically, or a 429, which is backpressure the + /// policy already routes around. Counting either as recovery is as wrong as + /// counting it as failure -- it would reset the failure count and close an + /// open breaker, so a worker alternating 500s and 400s could never reach + /// the consecutive failures needed to trip, and one 429 from a worker the + /// all-open fallback reached would undo its backoff. The slot such a + /// request consumed still has to come back. + pub fn record_neutral(&self, worker_id: &str) { + let mut map = self.entries(); + if let Some(e) = map.get_mut(worker_id) { + e.probe_started_at = None; + } + } + + /// Drop everything remembered about workers no longer in the fleet. + /// + /// Called with the full active set on each discovery snapshot, mirroring + /// `Policy::sync_workers`. Worker ids are addresses and a rebuilt Pod never + /// reuses one, so without this every rollout strands another entry -- and + /// another pair of Prometheus series, since /metrics exports one per entry + /// labelled by worker id. + pub fn retain_workers(&self, active: &HashSet) { + let mut map = self.entries(); + map.retain(|id, _| active.contains(id)); + } + + /// Record a pre-first-byte dispatch failure. Callers must gate this on + /// [`is_worker_fault`] when the failure carries an HTTP status. + pub fn record_failure(&self, worker_id: &str) { + self.record_failure_at(worker_id, Instant::now()); + } + + fn record_failure_at(&self, worker_id: &str, now: Instant) { + if !self.enabled() { + return; + } + let mut map = self.entries(); + let e = map.entry(worker_id.to_string()).or_insert_with(|| Entry { + consecutive_failures: 0, + state: BreakerState::Closed, + opens_until: now, + next_cooldown: self.cooldown, + probe_started_at: None, + trips: 0, + }); + e.consecutive_failures += 1; + let was_probe = e.state == BreakerState::HalfOpen; + e.probe_started_at = None; + + if was_probe { + // A failed probe reopens immediately and backs off further, without + // waiting out the threshold again — we already know it is bad. The + // common cause (a worker wedged on a bad KV handoff) does not clear + // on the first retry, and a fixed cooldown would probe it at a + // constant rate forever. + e.next_cooldown = (e.next_cooldown * 2).min(self.max_cooldown); + } else if e.consecutive_failures < self.failure_threshold { + return; + } + // A trip is an edge into exclusion, not every failure that lands while + // the worker is already excluded. A failed probe counts: it is a fresh + // verdict on a worker that was given another chance, and the doubling + // cooldown bounds how often one can happen. A failure while already + // open does not -- those arrive at the request rate, via the all-open + // fallback, and counting them turns the metric into a request counter + // that drowns out real trips and prints the warning on every request. + let newly_tripped = e.state != BreakerState::Open; + e.state = BreakerState::Open; + e.opens_until = now + e.next_cooldown; + if newly_tripped { + e.trips += 1; + tracing::warn!( + worker = worker_id, + cooldown_s = e.next_cooldown.as_secs_f64(), + failures = e.consecutive_failures, + "breaker: worker open" + ); + } + } + + pub fn state_of(&self, worker_id: &str) -> BreakerState { + self.entries() + .get(worker_id) + .map(|e| e.state) + .unwrap_or(BreakerState::Closed) + } + + /// `(worker_id, state, trips)` for metrics export. + pub fn snapshot(&self) -> Vec<(String, BreakerState, u64)> { + let map = self.entries(); + let mut out: Vec<_> = map + .iter() + .map(|(k, e)| (k.clone(), e.state, e.trips)) + .collect(); + out.sort_by(|a, b| a.0.cmp(&b.0)); + out + } +} + +impl Default for CircuitBreaker { + fn default() -> Self { + Self::new(3, Duration::from_secs(5), Duration::from_secs(60)) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + /// Time is driven forward explicitly rather than slept, so the cooldown + /// behaviour is tested at full speed and deterministically. + fn cb() -> CircuitBreaker { + CircuitBreaker::new(3, Duration::from_secs(5), Duration::from_secs(20)) + } + + #[derive(Clone)] + struct W(&'static str); + + #[test] + fn threshold_zero_disables_it() { + // Without the guard, `failures >= 0` would trip on the first failure — + // the exact opposite of what --breaker-failure-threshold=0 promises. + let b = CircuitBreaker::new(0, Duration::from_secs(5), Duration::from_secs(20)); + for _ in 0..20 { + b.record_failure("w1"); + } + assert!(b.allows("w1")); + assert_eq!(b.state_of("w1"), BreakerState::Closed); + let ws = vec![W("a"), W("b")]; + assert_eq!(b.filter(&ws, |w| w.0).len(), 2); + } + + #[test] + fn unknown_worker_is_allowed() { + let b = cb(); + assert!(b.allows("w1")); + assert_eq!(b.state_of("w1"), BreakerState::Closed); + } + + #[test] + fn failures_below_threshold_do_not_open() { + let b = cb(); + b.record_failure("w1"); + b.record_failure("w1"); + assert!(b.allows("w1")); + assert_eq!(b.state_of("w1"), BreakerState::Closed); + } + + #[test] + fn opens_at_threshold_and_excludes() { + let b = cb(); + for _ in 0..3 { + b.record_failure("w1"); + } + assert_eq!(b.state_of("w1"), BreakerState::Open); + assert!(!b.allows("w1")); + } + + #[test] + fn success_resets_the_count() { + let b = cb(); + b.record_failure("w1"); + b.record_failure("w1"); + b.record_success("w1"); + b.record_failure("w1"); + b.record_failure("w1"); + assert_eq!(b.state_of("w1"), BreakerState::Closed); + } + + #[test] + fn half_open_admits_exactly_one_probe() { + let b = cb(); + let t0 = Instant::now(); + for _ in 0..3 { + b.record_failure_at("w1", t0); + } + let t1 = t0 + Duration::from_millis(5_100); + assert!(b.allows_at("w1", t1), "cooldown elapsed -> one probe"); + assert_eq!(b.state_of("w1"), BreakerState::HalfOpen); + assert!( + !b.allows_at("w1", t1), + "a second concurrent request must not also probe" + ); + } + + #[test] + fn successful_probe_closes() { + let b = cb(); + let t0 = Instant::now(); + for _ in 0..3 { + b.record_failure_at("w1", t0); + } + let t1 = t0 + Duration::from_millis(5_100); + b.allows_at("w1", t1); + b.record_success("w1"); + assert_eq!(b.state_of("w1"), BreakerState::Closed); + assert!(b.allows_at("w1", t1)); + } + + #[test] + fn failed_probe_reopens_with_doubled_cooldown() { + let b = cb(); + let t0 = Instant::now(); + for _ in 0..3 { + b.record_failure_at("w1", t0); + } + let t1 = t0 + Duration::from_millis(5_100); + b.allows_at("w1", t1); + b.record_failure_at("w1", t1); // probe fails -> reopen, 5s -> 10s + assert_eq!(b.state_of("w1"), BreakerState::Open); + + let t2 = t1 + Duration::from_millis(5_100); // old cooldown would be up + assert!(!b.allows_at("w1", t2), "backoff must have doubled"); + let t3 = t1 + Duration::from_millis(10_100); + assert!(b.allows_at("w1", t3)); + } + + #[test] + fn cooldown_is_capped() { + let b = cb(); + let mut t = Instant::now(); + for _ in 0..3 { + b.record_failure_at("w1", t); + } + for _ in 0..6 { + t += Duration::from_secs(1000); + b.allows_at("w1", t); + b.record_failure_at("w1", t); + } + assert!( + b.allows_at("w1", t + Duration::from_millis(20_100)), + "cooldown must not grow without bound" + ); + } + + #[test] + fn filter_drops_open_workers() { + let b = cb(); + for _ in 0..3 { + b.record_failure("bad"); + } + let ws = vec![W("good"), W("bad")]; + let got = b.filter(&ws, |w| w.0); + assert_eq!(got.len(), 1); + assert_eq!(got[0].0, "good"); + } + + #[test] + fn filter_returns_all_when_every_worker_is_open() { + let b = cb(); + for id in ["a", "b"] { + for _ in 0..3 { + b.record_failure(id); + } + } + let ws = vec![W("a"), W("b")]; + assert_eq!(b.filter(&ws, |w| w.0).len(), 2); + } + + #[test] + fn filter_of_empty_is_empty() { + let b = cb(); + let ws: Vec = vec![]; + assert!(b.filter(&ws, |w| w.0).is_empty()); + } + + #[test] + fn workers_are_independent() { + let b = cb(); + for _ in 0..3 { + b.record_failure("bad"); + } + assert!(b.allows("good")); + } + + #[test] + fn success_on_unknown_worker_is_harmless() { + let b = cb(); + b.record_success("never-seen"); + assert!(b.allows("never-seen")); + } + + #[test] + fn snapshot_reports_trips() { + let b = cb(); + let t0 = Instant::now(); + for _ in 0..3 { + b.record_failure_at("w1", t0); + } + let t1 = t0 + Duration::from_millis(5_100); + b.allows_at("w1", t1); + b.record_failure_at("w1", t1); + let snap = b.snapshot(); + assert_eq!(snap.len(), 1); + assert_eq!(snap[0].1, BreakerState::Open); + assert_eq!(snap[0].2, 2, "initial trip plus the failed probe"); + } + + #[test] + fn client_errors_are_not_worker_faults() { + for s in [400u16, 404, 422, 429] { + assert!(!is_worker_fault(s), "{s} must not trip the breaker"); + } + for s in [0u16, 500, 502, 503, 504] { + assert!(is_worker_fault(s), "{s} must trip the breaker"); + } + } + + #[test] + fn a_bad_client_cannot_trip_the_fleet() { + let b = cb(); + let ws = vec![W("a"), W("b"), W("c")]; + for _ in 0..10 { + for w in b.filter(&ws, |w| w.0) { + if is_worker_fault(400) { + b.record_failure(w.0); + } + } + } + for w in &ws { + assert_eq!(b.state_of(w.0), BreakerState::Closed); + } + } + + #[test] + fn the_regression_this_exists_for() { + // A worker that fails every dispatch must stop being selected. Before + // this type existed `tried` was per-request, so `bad` was offered on + // all ten requests. + let b = cb(); + let ws = vec![W("good"), W("bad")]; + let mut offered_bad = 0; + for _ in 0..10 { + let cands = b.filter(&ws, |w| w.0); + if cands.iter().any(|w| w.0 == "bad") { + offered_bad += 1; + b.record_failure("bad"); + } + b.record_success("good"); + } + assert_eq!(offered_bad, 3, "bad worker must stop being offered"); + } + + #[test] + fn concurrent_probes_admit_only_one() { + // The Python breaker is single-loop; this one is shared across tokio + // threads, so the half-open slot has to be safe under real contention. + use std::sync::Arc; + let b = Arc::new(cb()); + for _ in 0..3 { + b.record_failure("w1"); + } + // Force half-open by driving the clock through the private hook. + let t = Instant::now() + Duration::from_secs(6); + assert!(b.allows_at("w1", t)); + b.record_failure_at("w1", t); // back to open, then reopen at 10s + let t2 = t + Duration::from_secs(11); + + let admitted = Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let hs: Vec<_> = (0..16) + .map(|_| { + let b = b.clone(); + let admitted = admitted.clone(); + std::thread::spawn(move || { + if b.allows_at("w1", t2) { + admitted.fetch_add(1, std::sync::atomic::Ordering::SeqCst); + } + }) + }) + .collect(); + for h in hs { + h.join().unwrap(); + } + assert_eq!( + admitted.load(std::sync::atomic::Ordering::SeqCst), + 1, + "exactly one of 16 racing threads may probe" + ); + } + + // filter() claims the probe slot for every candidate it lets through, but + // the policy dispatches to exactly one of them, so the others are never + // told how they did. Without a bound on the claim those workers sit in + // half-open holding a slot nothing will ever release -- healthy, and + // permanently unroutable until the process restarts. + #[test] + fn a_probe_slot_taken_but_never_dispatched_is_reclaimed() { + let b = CircuitBreaker::new(3, Duration::from_secs(5), Duration::from_secs(60)); + let t = Instant::now(); + for _ in 0..3 { + b.record_failure_at("w1", t); + } + let due = t + Duration::from_secs(6); + assert!(b.allows_at("w1", due), "cooldown elapsed -> a probe is due"); + + // The request went to another worker; nothing reports back for w1. + let later = due + PROBE_TIMEOUT + Duration::from_secs(1); + assert!( + b.allows_at("w1", later), + "an unused probe claim must not be permanent" + ); + } + + // 4xx says the request was bad, not the worker. Scoring it as recovery + // would reset the failure count and close an open breaker, so a worker + // alternating 500s and 400s could never reach three consecutive failures. + #[test] + fn a_neutral_outcome_frees_the_slot_without_scoring_it() { + let b = CircuitBreaker::new(3, Duration::from_secs(5), Duration::from_secs(60)); + let t = Instant::now(); + for _ in 0..3 { + b.record_failure_at("w1", t); + } + let due = t + Duration::from_secs(6); + assert!(b.allows_at("w1", due)); + + b.record_neutral("w1"); + assert_eq!( + b.state_of("w1"), + BreakerState::HalfOpen, + "a 4xx is not a recovery" + ); + assert!( + b.allows_at("w1", due), + "but the slot is free for a real probe" + ); + } + + // Worker ids are addresses and a rebuilt Pod never reuses one, so entries + // for departed workers accumulate for the process lifetime -- each also + // pinning a Prometheus series, which is the part that actually hurts. + #[test] + fn workers_gone_from_the_fleet_are_forgotten() { + let b = CircuitBreaker::new(3, Duration::from_secs(5), Duration::from_secs(60)); + let t = Instant::now(); + for _ in 0..3 { + b.record_failure_at("gone", t); + b.record_failure_at("stay", t); + } + assert_eq!(b.snapshot().len(), 2); + + let active: HashSet = ["stay".to_string()].into_iter().collect(); + b.retain_workers(&active); + + let snap = b.snapshot(); + assert_eq!(snap.len(), 1); + assert!(snap.iter().any(|(id, _, _)| id == "stay")); + assert_eq!(b.state_of("gone"), BreakerState::Closed); + } + + // `trips` answers "how often did this worker go bad", which is what an + // alert on rate(trips_total) is asking. Counting every failure that lands + // while the breaker is already open answers "how many requests hit a bad + // worker" instead -- a different and much larger number, dominated by + // whichever worker the all-open fallback keeps feeding. + #[test] + fn trips_count_outages_not_requests() { + let b = CircuitBreaker::new(3, Duration::from_secs(5), Duration::from_secs(60)); + let t = Instant::now(); + for _ in 0..3 { + b.record_failure_at("w1", t); + } + assert_eq!(b.snapshot()[0].2, 1, "three failures are one outage"); + + // Still open, still failing -- the all-open fallback keeps dispatching. + for _ in 0..20 { + b.record_failure_at("w1", t); + } + assert_eq!( + b.snapshot()[0].2, + 1, + "failures while already open are not new trips" + ); + + // A failed probe does count: a fresh verdict on a worker that was given + // another chance, bounded by the doubling cooldown rather than the + // request rate. + let due = t + Duration::from_secs(6); + assert!(b.allows_at("w1", due)); + b.record_failure_at("w1", due); + assert_eq!(b.snapshot()[0].2, 2, "a failed probe is a new verdict"); + } + + // The breaker sits on every request, so a poisoned lock must not be able to + // take the data plane with it: the guarded state is a map of counters, and + // a stale entry for one worker beats refusing to route at all. + #[test] + fn a_poisoned_lock_does_not_wedge_the_router() { + use std::sync::Arc; + + let b = Arc::new(CircuitBreaker::new( + 3, + Duration::from_secs(5), + Duration::from_secs(60), + )); + let t = Instant::now(); + b.record_failure_at("w1", t); + + let poisoner = Arc::clone(&b); + let _ = std::thread::spawn(move || { + let _guard = poisoner.entries(); + panic!("poison the lock"); + }) + .join(); + + // Every accessor must still work rather than propagating the panic. + assert!(b.allows_at("w2", t)); + b.record_failure_at("w2", t); + b.record_success("w2"); + assert!(!b.snapshot().is_empty()); + } +} diff --git a/rust/router/src/config.rs b/rust/router/src/config.rs index 123ab838..2e25f01b 100644 --- a/rust/router/src/config.rs +++ b/rust/router/src/config.rs @@ -27,6 +27,19 @@ pub struct Config { #[arg(long, default_value_t = 1)] pub request_max_retries: usize, + /// Consecutive pre-first-byte worker faults before a worker is taken out + /// of rotation. Failover alone forgets between requests; this remembers. + #[arg(long, default_value_t = 3)] + pub breaker_failure_threshold: u32, + + /// Seconds a tripped worker is excluded before one probe is admitted. + #[arg(long, default_value_t = 5.0)] + pub breaker_cooldown_s: f64, + + /// Ceiling for the cooldown, which doubles on each failed probe. + #[arg(long, default_value_t = 60.0)] + pub breaker_max_cooldown_s: f64, + /// `round-robin` or `kv-aware` (DP-attention cache-locality routing). #[arg(long, default_value = "round-robin")] pub router_policy: String, diff --git a/rust/router/src/disagg.rs b/rust/router/src/disagg.rs index fcf7b9d2..d2371532 100644 --- a/rust/router/src/disagg.rs +++ b/rust/router/src/disagg.rs @@ -12,6 +12,7 @@ //! KVPoll until a ~300s timeout. A detached `tokio::spawn` gives us exactly //! that: it outlives the client connection. +use std::sync::Arc; use std::time::Duration; use axum::body::{Body, Bytes}; @@ -19,6 +20,7 @@ use axum::http::{header, StatusCode}; use axum::response::Response; use serde_json::{Map, Value}; +use crate::breaker::{is_worker_fault, CircuitBreaker}; use crate::dp; use crate::handlers::AppState; use crate::policy::{ActiveGuard, Role}; @@ -41,16 +43,21 @@ pub async fn dispatch( ) -> Response { // role_hint lets a cost-aware policy weight P (cache-heavy: a hit skips a // whole prefill pass) differently from D (route by load). - let p_pick = state.policy.pick( - snap.list_active(model, DisaggMode::Prefill), - request, - Role::Prefill, - ); - let d_pick = state.policy.pick( - snap.list_active(model, DisaggMode::Decode), - request, - Role::Decode, - ); + // Each pool is filtered against the breaker independently: a wedged prefill + // and a wedged decode are different events against different pools, and one + // open breaker must not remove the other role's healthy workers. + let p_avail = state + .breaker + .filter(snap.list_active(model, DisaggMode::Prefill), |w| { + w.worker_id.as_str() + }); + let d_avail = state + .breaker + .filter(snap.list_active(model, DisaggMode::Decode), |w| { + w.worker_id.as_str() + }); + let p_pick = state.policy.pick(&p_avail, request, Role::Prefill); + let d_pick = state.policy.pick(&d_avail, request, Role::Decode); let p = p_pick.target; let d = d_pick.target; // One guard for both legs; dropped when the decode body finishes streaming @@ -121,7 +128,14 @@ async fn stream_dual( d_body: Map, guard: ActiveGuard, ) -> Response { - spawn_prefill_drain(state.http.clone(), p_url, p_body, p.dp_rank); + spawn_prefill_drain( + state.http.clone(), + state.breaker.clone(), + p.worker.worker_id.clone(), + p_url, + p_body, + p.dp_rank, + ); match open_decode(state, d, &d_url, &d_body).await { Ok(resp) => Response::builder() @@ -167,13 +181,30 @@ async fn unary_dual( st.as_u16() ); } + if is_worker_fault(st.as_u16()) { + state.breaker.record_failure(&p.worker.worker_id); + } else if st.is_success() { + state.breaker.record_success(&p.worker.worker_id); + } else { + state.breaker.record_neutral(&p.worker.worker_id); + } + } + Err(e) => { + tracing::warn!("prefill {} failed: {e}", p_url); + state.breaker.record_failure(&p.worker.worker_id); } - Err(e) => tracing::warn!("prefill {} failed: {e}", p_url), } match d_res { Ok(resp) => { let st = resp.status(); + if is_worker_fault(st.as_u16()) { + state.breaker.record_failure(&d.worker.worker_id); + } else if st.is_success() { + state.breaker.record_success(&d.worker.worker_id); + } else { + state.breaker.record_neutral(&d.worker.worker_id); + } let ct = content_type(&resp); match resp.bytes().await { Ok(bytes) => Response::builder() @@ -187,10 +218,13 @@ async fn unary_dual( ), } } - Err(e) => json_error( - StatusCode::BAD_GATEWAY, - &format!("decode {} unreachable: {e}", d.worker.worker_id), - ), + Err(e) => { + state.breaker.record_failure(&d.worker.worker_id); + json_error( + StatusCode::BAD_GATEWAY, + &format!("decode {} unreachable: {e}", d.worker.worker_id), + ) + } } } @@ -198,6 +232,8 @@ async fn unary_dual( /// Never awaited by the request path, so a client disconnect can't cancel it. fn spawn_prefill_drain( http: reqwest::Client, + breaker: Arc, + worker_id: String, url: String, body: Map, dp_rank: Option, @@ -217,8 +253,21 @@ fn spawn_prefill_drain( st.as_u16() ); } + // This leg is detached, so its outcome never reaches the client + // — but a prefill that 5xx's still leaves decode hanging on + // KVPoll, which is exactly the failure worth remembering. + if is_worker_fault(st.as_u16()) { + breaker.record_failure(&worker_id); + } else if st.is_success() { + breaker.record_success(&worker_id); + } else { + breaker.record_neutral(&worker_id); + } + } + Err(e) => { + tracing::warn!("prefill {url} failed: {e} (decode may hang on KVPoll)"); + breaker.record_failure(&worker_id); } - Err(e) => tracing::warn!("prefill {url} failed: {e} (decode may hang on KVPoll)"), } }); } @@ -237,6 +286,11 @@ async fn open_decode( Ok(resp) => { let st = resp.status(); if st.is_client_error() || st.is_server_error() { + if is_worker_fault(st.as_u16()) { + state.breaker.record_failure(&d.worker.worker_id); + } else { + state.breaker.record_neutral(&d.worker.worker_id); + } let txt = resp.text().await.unwrap_or_default(); return Err(format!( "decode {} error {}: {}", @@ -245,6 +299,7 @@ async fn open_decode( &txt[..txt.len().min(300)] )); } + state.breaker.record_success(&d.worker.worker_id); return Ok(resp); } Err(e) if attempt < DECODE_OPEN_RETRIES => { @@ -255,7 +310,12 @@ async fn open_decode( tokio::time::sleep(backoff).await; backoff = (backoff * 2).min(Duration::from_millis(500)); } - Err(e) => return Err(format!("decode {} unreachable: {e}", d.worker.worker_id)), + Err(e) => { + // Exhausted the in-request retries: this worker is not merely + // slow to accept a connection. + state.breaker.record_failure(&d.worker.worker_id); + return Err(format!("decode {} unreachable: {e}", d.worker.worker_id)); + } } } unreachable!("loop returns on the final attempt") diff --git a/rust/router/src/discovery.rs b/rust/router/src/discovery.rs index fa357568..6852e55c 100644 --- a/rust/router/src/discovery.rs +++ b/rust/router/src/discovery.rs @@ -17,10 +17,17 @@ use futures::StreamExt; use serde::Serialize; use serde_json::Value; +use crate::breaker::CircuitBreaker; use crate::policy::Policy; use crate::pool::{SharedPool, Snapshot, Worker}; -pub async fn run(base: String, prefix: String, pool: SharedPool, policy: Arc) { +pub async fn run( + base: String, + prefix: String, + pool: SharedPool, + policy: Arc, + breaker: Arc, +) { let prefix = if prefix.ends_with('/') { prefix } else { @@ -28,7 +35,7 @@ pub async fn run(base: String, prefix: String, pool: SharedPool, policy: Arc backoff = 1, Err(e) => { tracing::warn!("etcd discovery error: {e}; retry in {backoff}s"); @@ -64,6 +71,7 @@ async fn discover_once( prefix: &str, pool: &SharedPool, policy: &Arc, + breaker: &Arc, ) -> anyhow::Result<()> { let client = reqwest::Client::builder().build()?; let re = range_end(prefix); @@ -95,7 +103,7 @@ async fn discover_once( workers.len(), prefix ); - publish(pool, policy, &workers); + publish(pool, policy, breaker, &workers); // 2. watch for changes (long-lived NDJSON stream) let create = serde_json::json!({ @@ -121,7 +129,7 @@ async fn discover_once( } if let Ok(msg) = serde_json::from_slice::(line) { if apply_watch(prefix, &msg, &mut workers) { - publish(pool, policy, &workers); + publish(pool, policy, breaker, &workers); } } } @@ -129,11 +137,20 @@ async fn discover_once( Ok(()) } -fn publish(pool: &SharedPool, policy: &Arc, workers: &HashMap>) { +fn publish( + pool: &SharedPool, + policy: &Arc, + breaker: &Arc, + workers: &HashMap>, +) { let all: Vec> = workers.values().cloned().collect(); // Let cost-aware policies reconcile per-worker state (kv-event subscriptions, // load bookkeeping) against the new fleet before we swap the snapshot in. policy.sync_workers(&all); + // Same reconcile for the breaker: a worker id is an address no rebuilt Pod + // reuses, so entries for departed workers would otherwise accumulate for + // the process lifetime, each pinning a Prometheus series. + breaker.retain_workers(&workers.keys().cloned().collect()); pool.store(Arc::new(Snapshot::build(all))); } diff --git a/rust/router/src/handlers.rs b/rust/router/src/handlers.rs index d731267f..1031f205 100644 --- a/rust/router/src/handlers.rs +++ b/rust/router/src/handlers.rs @@ -15,6 +15,7 @@ use axum::routing::{get, post}; use axum::{Json, Router}; use serde_json::json; +use crate::breaker::CircuitBreaker; use crate::policy::Policy; use crate::pool::SharedPool; use crate::proxy; @@ -26,6 +27,9 @@ pub struct AppState { pub http: reqwest::Client, pub started: Instant, pub retries: usize, + /// Per-worker failure memory. Shared across threads and across requests — + /// that persistence across requests is the whole point (see breaker.rs). + pub breaker: Arc, } pub fn app(state: AppState) -> Router { @@ -75,11 +79,44 @@ async fn models(State(st): State) -> impl IntoResponse { async fn metrics(State(st): State) -> impl IntoResponse { let snap = st.pool.load(); - format!( + let mut out = format!( "# infera-router (rust)\n\ infera_router_active_workers {}\n\ infera_router_uptime_seconds {}\n", snap.active_count(), st.started.elapsed().as_secs() - ) + ); + // Non-zero state means the router is routing around a worker that + // discovery still reports ACTIVE — the gap this metric exists to show. + for (worker_id, state, trips) in st.breaker.snapshot() { + let v = match state { + crate::breaker::BreakerState::Closed => 0, + crate::breaker::BreakerState::HalfOpen => 1, + crate::breaker::BreakerState::Open => 2, + }; + // Escaped: worker ids come from discovery records, and a stray quote, + // backslash or newline in one would not corrupt a single line but end + // the whole exposition, failing every scrape of this endpoint. + let worker_id = escape_label_value(&worker_id); + out.push_str(&format!( + "infera_router_worker_breaker_state{{worker_id=\"{worker_id}\"}} {v}\n\ + infera_router_worker_breaker_trips_total{{worker_id=\"{worker_id}\"}} {trips}\n" + )); + } + out +} + +/// Escape a Prometheus label value: backslash, double quote and newline, per +/// the text exposition format. +fn escape_label_value(v: &str) -> String { + let mut out = String::with_capacity(v.len()); + for c in v.chars() { + match c { + '\\' => out.push_str("\\\\"), + '"' => out.push_str("\\\""), + '\n' => out.push_str("\\n"), + _ => out.push(c), + } + } + out } diff --git a/rust/router/src/lib.rs b/rust/router/src/lib.rs index 2555db29..0983afca 100644 --- a/rust/router/src/lib.rs +++ b/rust/router/src/lib.rs @@ -15,6 +15,7 @@ //! Modules are `pub` so the binary and the `tests/` suite share one API. pub mod block_hasher; +pub mod breaker; pub mod cache_control; pub mod config; pub mod disagg; diff --git a/rust/router/src/main.rs b/rust/router/src/main.rs index 3b728175..35c123dd 100644 --- a/rust/router/src/main.rs +++ b/rust/router/src/main.rs @@ -7,12 +7,13 @@ //! crate (see `lib.rs`); this just wires config → discovery → server. use std::sync::Arc; -use std::time::Instant; +use std::time::{Duration, Instant}; use arc_swap::ArcSwap; use tracing_subscriber::EnvFilter; use infera_router::block_hasher::BlockHasher; +use infera_router::breaker; use infera_router::config::Config; use infera_router::handlers::{app, AppState}; use infera_router::kv_event::KvEventClient; @@ -55,12 +56,21 @@ async fn main() -> anyhow::Result<()> { // without locking, so reads scale across cores. let pool = Arc::new(ArcSwap::from_pointee(Snapshot::empty())); + // Built before discovery starts: the reconcile loop prunes its entries + // against the live fleet, the same way it reconciles the policy's. + let breaker = Arc::new(breaker::CircuitBreaker::new( + cfg.breaker_failure_threshold, + Duration::from_secs_f64(cfg.breaker_cooldown_s), + Duration::from_secs_f64(cfg.breaker_max_cooldown_s), + )); + { let pool = pool.clone(); let policy = policy.clone(); + let breaker = breaker.clone(); let base = cfg.etcd_base(); let prefix = cfg.etcd_prefix.clone(); - tokio::spawn(async move { discovery::run(base, prefix, pool, policy).await }); + tokio::spawn(async move { discovery::run(base, prefix, pool, policy, breaker).await }); } let state = AppState { @@ -69,6 +79,7 @@ async fn main() -> anyhow::Result<()> { http: proxy::build_upstream_client()?, started: Instant::now(), retries: cfg.request_max_retries, + breaker, }; let addr = format!("{}:{}", cfg.host, cfg.port); diff --git a/rust/router/src/proxy.rs b/rust/router/src/proxy.rs index 7bab99df..29924349 100644 --- a/rust/router/src/proxy.rs +++ b/rust/router/src/proxy.rs @@ -18,6 +18,7 @@ use axum::response::Response; use futures::Stream; use serde_json::Value; +use crate::breaker::is_worker_fault; use crate::dp; use crate::handlers::AppState; use crate::policy::{ActiveGuard, Role}; @@ -108,6 +109,10 @@ async fn mixed_dispatch( if avail.is_empty() { break; } + // Drop workers the breaker has open. Falls back to the unfiltered list + // when every candidate is open — a request served by a probably-bad + // worker beats turning a partial outage into a 503. + let avail = state.breaker.filter(&avail, |w| w.worker_id.as_str()); let pick = state.policy.pick(&avail, request, Role::Mixed); tried.insert(pick.target.worker.worker_id.clone()); // Load guard: started here, dropped when this attempt fails (fail-over) @@ -116,9 +121,27 @@ async fn mixed_dispatch( state.policy.clone(), vec![(pick.target.route_key(), pick.blocks.clone())], ); + let wid = pick.target.worker.worker_id.clone(); match attempt(state, &pick.target, &raw, stream, path, guard).await { - Ok(resp) => return resp, - Err(err_resp) => last_err = Some(err_resp), + Ok(resp) => { + state.breaker.record_success(&wid); + return resp; + } + Err(err_resp) => { + // `attempt` only returns Err before any byte reached the client, + // so a mid-stream failure can never trip the breaker. 4xx is + // failed over but not held against the worker — see + // is_worker_fault(). + if is_worker_fault(err_resp.status().as_u16()) { + state.breaker.record_failure(&wid); + } else { + // A 4xx is not held against the worker, but the probe slot + // it consumed has to come back or one bad client wedges a + // recovering worker out of rotation. + state.breaker.record_neutral(&wid); + } + last_err = Some(err_resp); + } } } last_err.unwrap_or_else(|| json_error(StatusCode::SERVICE_UNAVAILABLE, "all workers failed")) diff --git a/rust/router/tests/functional.rs b/rust/router/tests/functional.rs index 86b0fac9..8cdebda9 100644 --- a/rust/router/tests/functional.rs +++ b/rust/router/tests/functional.rs @@ -21,6 +21,7 @@ use axum::{Json, Router}; use serde_json::{json, Value}; use infera_router::block_hasher::BlockHasher; +use infera_router::breaker::CircuitBreaker; use infera_router::handlers::{app, AppState}; use infera_router::kv_event::KvEventClient; use infera_router::policy::{KvEventAwarePolicy, RoundRobin}; @@ -106,6 +107,7 @@ fn make_state(workers: Vec>, retries: usize) -> AppState { http: proxy::build_upstream_client().unwrap(), started: Instant::now(), retries, + breaker: Arc::new(CircuitBreaker::default()), } } @@ -199,6 +201,7 @@ fn make_kv_state(workers: Vec>, retries: usize) -> AppState { http: proxy::build_upstream_client().unwrap(), started: Instant::now(), retries, + breaker: Arc::new(CircuitBreaker::default()), } } @@ -287,6 +290,86 @@ async fn mixed_failover_to_healthy_worker() { assert_eq!(ok.hit_count(), 1); } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn breaker_stops_reselecting_a_dead_worker() { + // The regression behind issue #82, end to end through the real router + // rather than against the breaker in isolation. + // + // Failover already made every one of these ten requests succeed, so a test + // that only checked status codes passed before the fix and after it. What + // was broken is the *cost*: `tried` is per-request, so RoundRobin kept + // offering the dead worker its turn -- 5 of 10 requests paid a wasted + // upstream round trip. The assertion that matters is bad.hit_count(), which + // is 5 without the breaker and 3 with it. + let (url_bad, bad) = spawn_mock(500, false, json!(null)).await; + let (url_ok, ok) = spawn_mock(200, false, json!({"ok": true})).await; + let state = make_state( + vec![ + worker( + json!({"worker_id": "bad", "url": url_bad, "model_name": "m", "disagg_mode": "mixed"}), + ), + worker( + json!({"worker_id": "ok", "url": url_ok, "model_name": "m", "disagg_mode": "mixed"}), + ), + ], + 1, + ); + let router = spawn_router(state).await; + + for _ in 0..10 { + let resp = client() + .post(format!("{router}/v1/chat/completions")) + .json(&json!({"model": "m"})) + .send() + .await + .unwrap(); + assert_eq!( + resp.status(), + 200, + "failover must still serve every request" + ); + } + + // Default threshold is 3. RoundRobin offers `bad` on every other request, + // so it takes 3 of its turns to trip; after that it is out of rotation and + // the 5s cooldown does not elapse within the test. + assert_eq!( + bad.hit_count(), + 3, + "dead worker must stop being re-picked after the threshold (was 5 before the fix)" + ); + assert_eq!(ok.hit_count(), 10, "healthy worker still serves everything"); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn breaker_ignores_client_errors() { + // A 400 comes from the request, not the worker: it would be returned by + // every worker in the fleet, so counting it would circuit-break all of + // them. Ten bad requests must leave the worker in rotation. + let (url, w) = spawn_mock(400, false, json!(null)).await; + let state = make_state( + vec![worker( + json!({"worker_id": "w1", "url": url, "model_name": "m", "disagg_mode": "mixed"}), + )], + 0, + ); + let router = spawn_router(state).await; + + for _ in 0..10 { + let _ = client() + .post(format!("{router}/v1/chat/completions")) + .json(&json!({"model": "m"})) + .send() + .await + .unwrap(); + } + assert_eq!( + w.hit_count(), + 10, + "4xx must not take a healthy worker out of rotation" + ); +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn mixed_streaming_relays_sse() { let (url, _mock) = spawn_mock(200, true, json!(null)).await; diff --git a/scripts/check-gofmt.sh b/scripts/check-gofmt.sh new file mode 100755 index 00000000..e4387b20 --- /dev/null +++ b/scripts/check-gofmt.sh @@ -0,0 +1,45 @@ +#!/usr/bin/env bash +# pre-commit hook: refuse Go sources that gofmt would rewrite. +# +# Worth a script rather than an inline entry, because `gofmt -l` prints the +# offending files and *still exits 0*. A bare `entry: gofmt -l` therefore passes +# on every commit and enforces nothing — which is how this tree accumulated Go +# files main had never formatted, discovered only when something unrelated ran +# gofmt over the directory. +# +# Formatting is checked, never applied: a hook that rewrites files mid-commit +# leaves the staged and working copies disagreeing, and the diff you reviewed is +# not the diff you commit. +set -euo pipefail + +if ! command -v gofmt >/dev/null 2>&1; then + cat >&2 <<'MSG' +gofmt not found, but this commit touches Go sources. + +Install the Go toolchain (https://go.dev/dl/), or skip this one check with: + + SKIP=gofmt git commit ... +MSG + exit 1 +fi + +# pre-commit passes the staged files matching `files:`; nothing to do otherwise. +[ "$#" -eq 0 ] && exit 0 + +unformatted="$(gofmt -l "$@")" +[ -z "$unformatted" ] && exit 0 + +{ + echo + echo "Refusing the commit: gofmt would rewrite these files." + echo + printf ' %s\n' $unformatted + echo + echo "Format them in place, then re-stage:" + echo + # Unquoted on purpose: word-splitting rejoins the newline-separated list into + # a single space-separated command line. + echo " gofmt -w" $unformatted + echo +} >&2 +exit 1 diff --git a/tests/unit/common/test_discovery_k8s_terminating.py b/tests/unit/common/test_discovery_k8s_terminating.py new file mode 100644 index 00000000..05d49727 --- /dev/null +++ b/tests/unit/common/test_discovery_k8s_terminating.py @@ -0,0 +1,225 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# SPDX-License-Identifier: MIT +############################################################################### +"""A condemned Pod must leave the pool before its process is killed. + +Kubernetes keeps ``phase: Running`` on a terminating Pod until its containers +exit, so liveness alone cannot tell a healthy worker from one that is seconds +from SIGTERM. The operator makes that window long on purpose -- it injects a +``preStop sleep`` so in-flight work has time to finish -- which means that +without a ``deletionTimestamp`` check the router spends the entire drain window +assigning new requests to a worker that is guaranteed to be killed. + +These tests pin the removal rules rather than the implementation: what matters +is which observable Pod states take a worker out of rotation. +""" + +from __future__ import annotations + +import json + +from infera.common.discovery_k8s import WORKER_INFO_ANNOTATION, KubernetesRegistry +from infera.common.worker_pool import WorkerStatus + + +def _payload(worker_id: str = "10.0.0.1:8080") -> str: + host, port = worker_id.split(":") + return json.dumps( + { + "worker_id": worker_id, + "url": f"http://{worker_id}", + "model_name": "m", + "engine": "sglang", + "disagg_mode": "mixed", + "disagg_meta": {}, + "kv_events_endpoint": None, + "kv_block_size": None, + "dp_rank": None, + "dp_size": None, + "request_transport": "http", + } + ) + + +def _pod(name="w-0", *, annotated=True, phase="Running", terminating=False): + meta: dict = {"name": name} + if annotated: + meta["annotations"] = {WORKER_INFO_ANNOTATION: _payload()} + if terminating: + meta["deletionTimestamp"] = "2026-08-04T03:00:00Z" + return {"metadata": meta, "status": {"phase": phase}} + + +def _registry(): + removed: list[str] = [] + reg = KubernetesRegistry( + "app=infera", + namespace="infera", + on_worker_removed=removed.append, + ) + return reg, removed + + +def _ids(reg): + return [w.worker_id for w in reg.pool.list_active()] + + +def test_running_pod_registers(): + reg, _ = _registry() + reg._handle_pod(_pod(), deleted=False) + assert _ids(reg) == ["10.0.0.1:8080"] + + +def test_terminating_pod_is_removed_while_still_running(): + """The case this exists for. `phase` is still Running -- only the deletion + timestamp distinguishes a healthy worker from one inside its preStop delay, + and every request routed to it in that window is work that gets cut.""" + reg, removed = _registry() + reg._handle_pod(_pod(), deleted=False) + assert _ids(reg) == ["10.0.0.1:8080"] + + reg._handle_pod(_pod(terminating=True, phase="Running"), deleted=False) + assert _ids(reg) == [], "a condemned Pod must not stay a routing candidate" + assert removed == ["10.0.0.1:8080"] + + +def test_terminating_pod_never_enters_the_pool(): + """A relist during a rolling update can surface an already-terminating Pod + the registry has never seen. It must not be admitted.""" + reg, _ = _registry() + reg._handle_pod(_pod(terminating=True), deleted=False) + assert _ids(reg) == [] + + +def test_removal_is_idempotent(): + """Watch events are re-delivered after a 410/relist, so the same + terminating Pod arrives more than once.""" + reg, removed = _registry() + reg._handle_pod(_pod(), deleted=False) + for _ in range(3): + reg._handle_pod(_pod(terminating=True), deleted=False) + assert removed == ["10.0.0.1:8080"], "must not fire the removal callback repeatedly" + + +def test_other_removal_rules_still_hold(): + for label, kwargs, deleted in ( + ("explicit DELETE", {}, True), + ("annotation cleared", {"annotated": False}, False), + ("no longer Running", {"phase": "Failed"}, False), + ): + reg, _ = _registry() + reg._handle_pod(_pod(), deleted=False) + reg._handle_pod(_pod(**kwargs), deleted=deleted) + assert _ids(reg) == [], f"{label} must still deregister" + + +def _all(reg): + return {w.worker_id: w.status for w in reg.pool.list_all()} + + +def test_a_draining_worker_stays_visible(): + """Out of routing, still on the books. + + Dropping the record entirely makes a worker finishing its in-flight + generations look exactly like one that crashed, so `/v1/workers` cannot + tell an orderly rollout from a fleet losing workers -- at exactly the + moment someone is watching one happen. + """ + reg, _ = _registry() + reg._handle_pod(_pod(), deleted=False) + reg._handle_pod(_pod(terminating=True), deleted=False) + + assert _ids(reg) == [], "still must not be a routing candidate" + assert _all(reg) == {"10.0.0.1:8080": WorkerStatus.DRAINING} + + +def test_the_record_goes_when_the_worker_deregisters(): + """The worker clears its own annotation on SIGTERM, before draining, which + lands here as 'annotation gone'. Without that the draining record would be + immortal.""" + reg, removed = _registry() + reg._handle_pod(_pod(), deleted=False) + reg._handle_pod(_pod(terminating=True), deleted=False) + assert _all(reg), "precondition: the record survives the terminating event" + + reg._handle_pod(_pod(terminating=True, annotated=False), deleted=False) + assert _all(reg) == {}, "a drained worker must leave the pool" + assert removed == ["10.0.0.1:8080"], "announced once, not once per stage" + + +def test_a_worker_killed_without_its_pod_being_deleted_is_announced_once(): + """The path deregistering-before-draining exists for. + + A liveness probe restarting the container, a node shutting down gracefully, + someone killing the process -- all send SIGTERM with the Pod object + untouched. There is no deletionTimestamp, so the annotation the worker + clears on its way out is the only signal, and the callback behind it is what + stops the KV subscriber. A later DELETE for the same Pod must not repeat it. + """ + reg, removed = _registry() + reg._handle_pod(_pod(), deleted=False) + + reg._handle_pod(_pod(annotated=False), deleted=False) + assert _ids(reg) == [], "clearing the annotation must stop new work arriving" + assert removed == ["10.0.0.1:8080"], "the only signal on this path must announce" + + reg._handle_pod(_pod(annotated=False), deleted=True) + assert removed == ["10.0.0.1:8080"], "the eventual DELETE must not announce again" + + +def test_announced_once_across_draining_then_delete(): + """The callbacks stop a KV subscriber and clear block accounting. Firing + them twice for one worker is not free, and firing them late (only at the + DELETE) would leave the router accounting for a worker it no longer routes + to for the whole drain.""" + reg, removed = _registry() + reg._handle_pod(_pod(), deleted=False) + reg._handle_pod(_pod(terminating=True), deleted=False) + assert removed == ["10.0.0.1:8080"], "must announce as soon as it leaves routing" + + reg._handle_pod(_pod(terminating=True), deleted=True) + assert _all(reg) == {} + assert removed == ["10.0.0.1:8080"], "the DELETE must not re-announce" + + +def test_a_worker_that_never_drained_still_announces_on_delete(): + """The drain path is not the only way out: a crash or an evicted Pod goes + straight to removal, and that still has to reach the callbacks.""" + reg, removed = _registry() + reg._handle_pod(_pod(), deleted=False) + reg._handle_pod(_pod(), deleted=True) + assert _all(reg) == {} + assert removed == ["10.0.0.1:8080"] + + +def test_a_list_reconciles_away_a_worker_whose_delete_was_missed(): + """A list is a complete snapshot, so a tracked Pod it does not mention is + gone. + + This is not hypothetical. Keeping a draining record alive means its removal + now depends on a later event, and the watch drops out routinely -- the + re-list exists precisely because etcd compaction expires the + resourceVersion every few minutes. A Pod deleted inside that window + produces no event anyone sees, so without reconciling against the list the + record is immortal: `/v1/workers` reports a worker that does not exist, and + the model's canary is never forgotten because a phantom still holds it. + """ + reg, removed = _registry() + reg._handle_pod(_pod(), deleted=False) + reg._handle_pod(_pod(terminating=True), deleted=False) + assert _all(reg), "precondition: the draining record is being kept" + + # The Pod is gone; a fresh list simply does not contain it. + reg._reconcile_absent(seen=set()) + + assert _all(reg) == {}, "a tracked Pod missing from a full list must be dropped" + assert removed == ["10.0.0.1:8080"], "already announced when it started draining" + + +def test_a_list_keeps_workers_it_still_sees(): + reg, _ = _registry() + reg._handle_pod(_pod(name="w-0"), deleted=False) + reg._reconcile_absent(seen={"w-0"}) + assert _ids(reg) == ["10.0.0.1:8080"], "a Pod present in the list must survive" diff --git a/tests/unit/common/test_draining_status.py b/tests/unit/common/test_draining_status.py new file mode 100644 index 00000000..fca252a2 --- /dev/null +++ b/tests/unit/common/test_draining_status.py @@ -0,0 +1,108 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# SPDX-License-Identifier: MIT +############################################################################### +"""Where DRAINING comes from, and why the registration record has no state. + +``WorkerStatus.DRAINING`` takes a worker out of ``list_active`` while leaving it +visible, which is what distinguishes an orderly rollout from a crash. What sets +it is the Kubernetes registry, from the Pod's ``deletionTimestamp`` -- not the +worker. + +That is a deliberate split. Kubernetes knows a Pod is condemned before the +process is signalled, so the orchestrator answers "is this worker leaving" +earlier and more authoritatively than the worker could. etcd has no such signal: +a record is present or absent, so there removing it is what stops new work, and +a shutdown deregisters before it drains. Either way the record itself carries +identity only, which is what makes the heartbeat -- which rebuilds it from +config -- safe to run at any point. +""" + +from __future__ import annotations + +import pytest + +from infera.common.discovery import worker_info_from_json +from infera.common.registration import build_worker_payload +from infera.common.worker_pool import EngineType, WorkerPool, WorkerStatus +from infera.engine.base import EngineConfig + + +def _cfg(): + return EngineConfig(model_name="m", host="10.0.0.1", port=8080, engine=EngineType.SGLANG) + + +def test_the_record_carries_identity_and_no_state(): + """Every field is fixed for the life of the process, so two builds are + byte-identical. State written here would be erased by the next heartbeat, + which rebuilds the payload from the same config.""" + payload = build_worker_payload(_cfg()) + assert "status" not in payload + assert payload == build_worker_payload(_cfg()) + + +def test_a_record_without_a_status_reads_as_active(): + """Registration says nothing about state, so the parser's default is what + every healthy worker resolves to.""" + info = worker_info_from_json(build_worker_payload(_cfg())) + assert info.status is WorkerStatus.ACTIVE + + +def test_a_draining_worker_is_excluded_but_still_visible(): + """The value DRAINING adds over deleting the record is not routing -- both + stop new work -- it is that the worker stays observable while it finishes.""" + pool = WorkerPool() + pool.add(worker_info_from_json(build_worker_payload(_cfg()))) + assert [w.worker_id for w in pool.list_active()] == ["10.0.0.1:8080"] + + worker = pool.get("10.0.0.1:8080") + worker.status = WorkerStatus.DRAINING + assert pool.list_active() == [], "a draining worker must not be routed to" + assert pool.get("10.0.0.1:8080") is not None, "but it must still be observable" + + +# --- which step stops new work arriving --------------------------------------- + + +def test_no_client_announces_a_status(): + """Neither backend writes state into the record any more. + + Under Kubernetes it was never read -- the registry acts on deletionTimestamp + and returns before parsing the annotation -- while the heartbeat rebuilt the + payload and erased it. On etcd it was read, but deregistering already stops + new work, so announcing first only added a second mechanism for the same + thing. + """ + from infera.common.registration import RegistrationClient + from infera.common.registration_k8s import K8sRegistrationClient + + for client in (RegistrationClient, K8sRegistrationClient): + assert not hasattr(client, "announce_draining"), client.__name__ + + +@pytest.mark.asyncio +async def test_the_k8s_registry_marks_a_condemned_pod_draining(): + """The one remaining producer of DRAINING, and the reason the worker does + not need to announce anything under Kubernetes.""" + import json as _json + + from infera.common.discovery_k8s import WORKER_INFO_ANNOTATION, KubernetesRegistry + + reg = KubernetesRegistry(label_selector="x=y", namespace="ns") + + def _pod(*, deleting: bool): + meta = { + "name": "worker-1", + "annotations": {WORKER_INFO_ANNOTATION: _json.dumps(build_worker_payload(_cfg()))}, + } + if deleting: + meta["deletionTimestamp"] = "2026-08-10T00:00:00Z" + return {"metadata": meta, "status": {"phase": "Running"}} + + reg._handle_pod(_pod(deleting=False), deleted=False) + assert [w.worker_id for w in reg._pool.list_active()] == ["10.0.0.1:8080"] + + reg._handle_pod(_pod(deleting=True), deleted=False) + assert reg._pool.list_active() == [], "a condemned Pod must leave routing" + assert reg._pool.get("10.0.0.1:8080").status is WorkerStatus.DRAINING diff --git a/tests/unit/common/test_nats_drain.py b/tests/unit/common/test_nats_drain.py index 39203a70..b075ce7f 100644 --- a/tests/unit/common/test_nats_drain.py +++ b/tests/unit/common/test_nats_drain.py @@ -10,10 +10,11 @@ from __future__ import annotations import asyncio +from types import SimpleNamespace import pytest -from infera.common.nats_request import NatsRequestServer +from infera.common.nats_request import REQUEST_STREAM, NatsRequestServer, request_durable @pytest.mark.asyncio @@ -40,3 +41,123 @@ async def test_no_drain_cancels_in_flight_immediately(): await asyncio.gather(t, return_exceptions=True) assert t.cancelled() + + +# --- JetStream (admission throttle) path ------------------------------------- +# +# Under the throttle the router publishes into a WorkQueue stream and this +# worker pulls from it, so a request can be accepted, queued, and invisible to +# `_inflight` -- which only tracks what has already been delivered here. + + +class _FakeJs: + """JetStream stand-in: a scripted num_pending sequence + a call log.""" + + def __init__(self, pending: list[int] | None = None, *, fail: bool = False) -> None: + self._pending = list(pending or []) + self._fail = fail + self.deleted: list[tuple[str, str]] = [] + self.info_calls = 0 + + async def consumer_info(self, stream, consumer, timeout=None): + self.info_calls += 1 + if self._fail: + raise RuntimeError("broker hiccup") + value = self._pending.pop(0) if self._pending else 0 + return SimpleNamespace(num_pending=value, num_ack_pending=0) + + async def delete_consumer(self, stream, consumer): + self.deleted.append((stream, consumer)) + return True + + +class _FakeSub: + """Records when the subject was unsubscribed, in units of backlog polls.""" + + def __init__(self, js: _FakeJs) -> None: + self._js = js + self.unsubscribed_after_polls: int | None = None + + async def unsubscribe(self): + self.unsubscribed_after_polls = self._js.info_calls + + +@pytest.mark.asyncio +async def test_the_door_closes_only_after_the_backlog_clears(): + """`unsubscribe()` discards whatever is left in the stream, so it must not + run until the backlog has been handed over -- otherwise a request the + router already accepted is silently dropped and the client waits out the + full idle timeout for a reply nobody will send. + + Pinned by *when* the unsubscribe happened, not just that the backlog was + polled: polling and then closing the door anyway would pass a weaker check. + """ + js = _FakeJs([2, 1, 0]) + sub = _FakeSub(js) + srv = NatsRequestServer("w:1", 30000) + srv._js = js + srv._sub = sub + + await srv.stop(drain=True, drain_timeout=5) + + assert sub.unsubscribed_after_polls is not None, "unsubscribe never ran" + assert sub.unsubscribed_after_polls >= 3, ( + "unsubscribed while the stream still held work: only " + f"{sub.unsubscribed_after_polls} poll(s) had happened, backlog clears on the 3rd" + ) + + +@pytest.mark.asyncio +async def test_a_backlog_that_never_clears_is_bounded_by_the_deadline(): + """A stuck backlog must not hold the process past its drain budget -- the + kubelet's SIGKILL does not wait, and everything after this still has to + run.""" + js = _FakeJs([5] * 1000) + srv = NatsRequestServer("w:1", 30000) + srv._js = js + + start = asyncio.get_running_loop().time() + await srv.stop(drain=True, drain_timeout=0.5) + elapsed = asyncio.get_running_loop().time() - start + + assert elapsed < 3.0, f"stop() overran its 0.5s budget by far: {elapsed:.1f}s" + + +@pytest.mark.asyncio +async def test_an_unreadable_backlog_does_not_block_shutdown(): + """A stalled rollout is worse than a dropped queued request, and a broker + that cannot answer looks identical to a genuinely busy one.""" + js = _FakeJs(fail=True) + srv = NatsRequestServer("w:1", 30000) + srv._js = js + + await asyncio.wait_for(srv.stop(drain=True, drain_timeout=30), timeout=2.0) + + assert js.info_calls == 1, "should give up after the first failed read" + + +@pytest.mark.asyncio +async def test_stop_deletes_the_durable_consumer(): + """The durable outlives the subscription by definition, and its name comes + from worker_id -- which a rebuilt Pod never reuses. Left behind, every + rollout adds an orphan holding WorkQueue quota nothing will consume.""" + js = _FakeJs([0]) + srv = NatsRequestServer("10.0.0.1:30000", 30000) + srv._js = js + + await srv.stop(drain=True, drain_timeout=1) + + assert len(js.deleted) == 1, f"expected one delete_consumer call, got {js.deleted}" + stream, consumer = js.deleted[0] + assert stream == REQUEST_STREAM + assert consumer == request_durable("10.0.0.1:30000") + + +@pytest.mark.asyncio +async def test_core_nats_path_touches_no_jetstream(): + """Without the throttle there is no stream and no durable; the shutdown + path must not assume otherwise.""" + srv = NatsRequestServer("w:1", 30000) + assert srv._js is None + + await srv.stop(drain=True, drain_timeout=1) # must not raise diff --git a/tests/unit/common/test_shutdown_order.py b/tests/unit/common/test_shutdown_order.py new file mode 100644 index 00000000..f244959b --- /dev/null +++ b/tests/unit/common/test_shutdown_order.py @@ -0,0 +1,134 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# SPDX-License-Identifier: MIT +############################################################################### +"""The shutdown sequence: deregister, then drain. + +Removing the record is what stops new work arriving, so it has to happen before +waiting on work already in flight. These pin that order in the entrypoints +themselves and pin that a deregistration which did not happen says so. +""" + +from __future__ import annotations + +import ast +import inspect +from pathlib import Path + +import pytest + +from infera.common.registration import RegistrationClient +from infera.common.registration_k8s import K8sRegistrationClient + +ENTRYPOINTS = ( + "infera/engine/vllm/__main__.py", + "infera/engine/sglang/__main__.py", +) + + +def _shutdown_call_order(path: Path) -> list[str]: + """The two calls that matter, in source order. + + Reduced to deregistering and draining, so reordering anything else in the + shutdown sequence does not make this fail. Calling `_drain` is a Call node + while defining it is not, so the definition is not counted. + """ + seen = [] + for call in (n for n in ast.walk(ast.parse(path.read_text())) if isinstance(n, ast.Call)): + func = call.func + if isinstance(func, ast.Attribute) and func.attr == "deregister": + seen.append(("deregister", call.lineno)) + elif isinstance(func, ast.Name) and func.id == "_drain": + seen.append(("drain", call.lineno)) + # ast.walk is breadth-first, so sort back into source order. + return [name for name, _ in sorted(seen, key=lambda p: p[1])] + + +@pytest.mark.parametrize("entrypoint", ENTRYPOINTS) +def test_every_entrypoint_deregisters_before_it_drains(entrypoint): + """Draining while still registered means the router keeps assigning work + for the whole drain window, so the drain never converges. + + On etcd removing the record is the only way to express departure. On + Kubernetes the registry does drop a Pod on its deletionTimestamp, but only + when the Pod is being deleted -- a liveness-probe restart, a node graceful + shutdown or a manual kill all send SIGTERM with the Pod object untouched, + leaving the annotation present and the worker routable until it clears it. + """ + root = Path(inspect.getfile(RegistrationClient)).parents[2] + order = _shutdown_call_order(root / entrypoint) + + assert "deregister" in order, f"{entrypoint} never deregisters" + assert "drain" in order, f"{entrypoint} never drains" + assert order.index("deregister") < order.index("drain"), ( + f"{entrypoint} drains before deregistering, so new work keeps arriving " + "for the whole drain window" + ) + + +class _Resp: + def __init__(self, code: int): + self.status_code = code + + def raise_for_status(self): + if self.status_code >= 400: + raise RuntimeError(f"HTTP {self.status_code}") + + +@pytest.mark.asyncio +async def test_etcd_reports_a_revoke_the_server_refused(): + """httpx does not raise on 4xx/5xx, so a refused revoke reaches the same + code path as a successful one. etcd answering 'lease not found', or a proxy + answering 503, is the likeliest way this fails.""" + + class _Http: + async def post(self, path, json=None): # noqa: A002 - mirrors httpx + return _Resp(500) + + async def aclose(self): + pass + + # __new__, so the real httpx client is never built and never leaked. + c = RegistrationClient.__new__(RegistrationClient) + c._http = _Http() + c._lease_id, c._key, c._worker_id, c._lease_ttl = 1, "/k", "w", 30 + + assert await c.deregister() is False, ( + "etcd refused the revoke, so the lease is still alive and the worker is " + "still routable -- that cannot report success" + ) + + +@pytest.mark.asyncio +async def test_etcd_reports_an_unreachable_server(): + class _Http: + async def post(self, path, json=None): # noqa: A002 - mirrors httpx + raise RuntimeError("etcd unreachable") + + async def aclose(self): + pass + + # __new__, so the real httpx client is never built and never leaked. + c = RegistrationClient.__new__(RegistrationClient) + c._http = _Http() + c._lease_id, c._key, c._worker_id, c._lease_ttl = 1, "/k", "w", 30 + + assert await c.deregister() is False + + +@pytest.mark.asyncio +async def test_kubernetes_reports_a_patch_that_failed(): + """Clearing the annotation is what takes the worker out of the pool on the + paths where the Pod object is untouched.""" + c = K8sRegistrationClient.__new__(K8sRegistrationClient) + c._worker_id = "w" + c._pod_name = "p" + c._namespace = "ns" + + async def _patch(*_args, **_kwargs): + raise RuntimeError("apiserver unreachable") + + c._patch_annotation = _patch + + assert await c.deregister() is False diff --git a/tests/unit/engine/test_drain.py b/tests/unit/engine/test_drain.py new file mode 100644 index 00000000..0acb5823 --- /dev/null +++ b/tests/unit/engine/test_drain.py @@ -0,0 +1,217 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# SPDX-License-Identifier: MIT +############################################################################### +"""Graceful drain on the HTTP transport. + +On NATS infera owns the request path and knows what is in flight. On HTTP the +router talks straight to the engine, so infera has to ask the engine — which +means every failure mode is a *measurement* failure, and the interesting cases +are all about what happens when the number cannot be trusted. + +The rule these tests pin down: never treat "unknown" as "idle". An unreadable +metric that defaults to zero would make the drain pass instantly and cut live +generations, and it would do so silently. +""" + +from __future__ import annotations + +import asyncio +import logging + +import httpx +import pytest + +from infera.common.engine_metrics import inflight_from_metrics, metric_name, parse_metric +from infera.common.worker_pool import EngineType +from infera.engine import drain as drain_mod +from infera.engine.drain import drain_engine_inflight + +# --- reading the engine's numbers --------------------------------------------- + + +def test_parses_plain_and_labelled_gauges(): + assert parse_metric("vllm:num_requests_running 3.0\n", "vllm:num_requests_running") == 3.0 + assert parse_metric('x:g{a="b"} 7\n', "x:g") == 7.0 + + +def test_missing_series_is_none_not_zero(): + """The distinction the whole drain rests on.""" + assert parse_metric("something_else 1\n", "vllm:num_requests_running") is None + + +def test_inflight_counts_running_plus_waiting(): + """A queued request is work a client is waiting on; killing the process + loses it just as surely as one mid-generation.""" + text = "vllm:num_requests_running 2\nvllm:num_requests_waiting 5\n" + assert inflight_from_metrics(text, EngineType.VLLM) == 7.0 + + +def test_unknown_engine_yields_none(): + """ATOM has no mapping. Guessing one would report an idle engine, and an + idle engine is exactly the answer that makes a drain cut live requests.""" + assert metric_name("requests_running", EngineType.ATOM) is None + assert inflight_from_metrics("vllm:num_requests_running 4\n", EngineType.ATOM) is None + + +def test_partial_metrics_still_count(): + text = "sglang:num_running_reqs 1\n" + assert inflight_from_metrics(text, EngineType.SGLANG) == 1.0 + + +def test_a_missing_series_says_so(caplog): + """Counting only what is readable is the right call -- refusing a partial + page would let one renamed series stop the drain waiting at all. But the + absent series contributes zero, so queued work can be cut while the drain + reports itself finished. That has to be audible: these names do drift, and + the vLLM KV gauge was renamed under this very module. + """ + text = "sglang:num_running_reqs 1\n" # num_queue_reqs absent + with caplog.at_level(logging.WARNING, logger="infera.common.engine_metrics"): + assert inflight_from_metrics(text, EngineType.SGLANG) == 1.0 + assert "sglang:num_queue_reqs" in caplog.text + caplog.clear() + + # A complete page must stay quiet, or the warning is noise on every poll. + both = "sglang:num_running_reqs 1\nsglang:num_queue_reqs 2\n" + with caplog.at_level(logging.WARNING, logger="infera.common.engine_metrics"): + assert inflight_from_metrics(both, EngineType.SGLANG) == 3.0 + assert caplog.text == "" + + +# --- the drain loop ----------------------------------------------------------- + + +def _patch_client(monkeypatch, handler): + """Route drain's httpx client at a mock transport.""" + real = httpx.AsyncClient + + def factory(*a, **kw): + kw.pop("timeout", None) + return real(transport=httpx.MockTransport(handler), timeout=5.0) + + monkeypatch.setattr(drain_mod.httpx, "AsyncClient", factory) + + +@pytest.mark.asyncio +async def test_returns_when_engine_goes_idle(monkeypatch): + counts = iter([3, 2] + [0] * 100) + + def handler(request): + return httpx.Response(200, text=f"vllm:num_requests_running {next(counts)}\n") + + _patch_client(monkeypatch, handler) + drained = await drain_engine_inflight( + host="1.2.3.4", + port=8000, + engine=EngineType.VLLM, + timeout=5, + poll_interval=0.01, + settle=0.05, + ) + assert drained is True + + +@pytest.mark.asyncio +async def test_a_single_zero_reading_is_not_enough(monkeypatch): + """Measured on SGLang: the gauge lags the work by 5-15s, so one zero can + mean "idle" or "not counted yet". A late non-zero must restart the window + rather than being ignored.""" + counts = iter([0, 0, 4] + [0] * 100) + seen: list[float] = [] + + def handler(request): + v = next(counts) + seen.append(v) + return httpx.Response(200, text=f"vllm:num_requests_running {v}\n") + + _patch_client(monkeypatch, handler) + drained = await drain_engine_inflight( + host="1.2.3.4", + port=8000, + engine=EngineType.VLLM, + timeout=5, + poll_interval=0.01, + settle=0.05, + ) + assert drained is True + assert 4 in seen, "the late non-zero reading must have been observed" + # It must have kept polling well past the point where the first two zeros + # would have satisfied a naive implementation. + assert len(seen) > 3 + + +@pytest.mark.asyncio +async def test_times_out_while_still_busy(monkeypatch): + def handler(request): + return httpx.Response(200, text="vllm:num_requests_running 4\n") + + _patch_client(monkeypatch, handler) + drained = await drain_engine_inflight( + host="1.2.3.4", port=8000, engine=EngineType.VLLM, timeout=0.2, poll_interval=0.01 + ) + assert drained is False, "a busy engine must not report a clean drain" + + +@pytest.mark.asyncio +async def test_unreadable_metric_does_not_hang(monkeypatch, caplog): + """A rolling update that stalls on a parse failure is worse than one that + cuts a request -- and a silent full-timeout wait is indistinguishable from + a genuinely busy worker.""" + + def handler(request): + return httpx.Response(200, text="totally_different_metric 1\n") + + _patch_client(monkeypatch, handler) + started = asyncio.get_running_loop().time() + drained = await drain_engine_inflight( + host="1.2.3.4", port=8000, engine=EngineType.VLLM, timeout=30, poll_interval=0.01 + ) + elapsed = asyncio.get_running_loop().time() - started + assert drained is False + assert elapsed < 1.0, f"returned after {elapsed:.1f}s; must not wait out the timeout" + assert "WITHOUT draining" in caplog.text + + +@pytest.mark.asyncio +async def test_unreachable_engine_does_not_hang(monkeypatch): + def handler(request): + raise httpx.ConnectError("refused", request=request) + + _patch_client(monkeypatch, handler) + drained = await drain_engine_inflight( + host="1.2.3.4", port=8000, engine=EngineType.VLLM, timeout=30, poll_interval=0.01 + ) + assert drained is False + + +@pytest.mark.asyncio +async def test_zero_timeout_is_a_no_op(monkeypatch): + calls = [] + + def handler(request): + calls.append(1) + return httpx.Response(200, text="vllm:num_requests_running 0\n") + + _patch_client(monkeypatch, handler) + assert ( + await drain_engine_inflight(host="1.2.3.4", port=8000, engine=EngineType.VLLM, timeout=0) + is False + ) + assert calls == [], "--drain-timeout 0 must not even probe" + + +@pytest.mark.asyncio +async def test_never_raises_on_the_shutdown_path(monkeypatch): + """This runs immediately before engine.stop(); an exception here would skip + the teardown that follows.""" + + def factory(*a, **kw): + raise RuntimeError("boom") + + monkeypatch.setattr(drain_mod.httpx, "AsyncClient", factory) + assert ( + await drain_engine_inflight(host="1.2.3.4", port=8000, engine=EngineType.VLLM, timeout=1) + is False + ) diff --git a/tests/unit/kv/test_snapshot.py b/tests/unit/kv/test_snapshot.py index 3eef1970..9413f509 100644 --- a/tests/unit/kv/test_snapshot.py +++ b/tests/unit/kv/test_snapshot.py @@ -289,12 +289,14 @@ def test_producer_empty_snapshot_for_unknown_stream() -> None: # ---------------------------------------------------------------------- -async def _make_reconciler(pull_fn) -> tuple[KVIndex, KvIndexWriter, SnapshotReconciler]: +async def _make_reconciler( + pull_fn, interval_s: float = 10_000 +) -> tuple[KVIndex, KvIndexWriter, SnapshotReconciler]: index = KVIndex() queue: asyncio.Queue = asyncio.Queue(maxsize=100) writer = KvIndexWriter(index=index, queue=queue) await writer.start() - rec = SnapshotReconciler(index=index, writer=writer, pull_fn=pull_fn, interval_s=10_000) + rec = SnapshotReconciler(index=index, writer=writer, pull_fn=pull_fn, interval_s=interval_s) return index, writer, rec @@ -685,3 +687,112 @@ def test_index_drop_tree() -> None: matches_b = index.find_matches(model="m", compat_key="ckB", chain=[chain[0]], candidates=["w1"]) assert matches_a["w1"] == OverlapBlocks() assert matches_b["w1"].device == 1 + + +async def test_target_registered_while_running_is_not_left_for_a_whole_interval() -> None: + """A worker that joins after the reconciler is already running must have its + snapshot pulled promptly, not on the next periodic tick. + + This is the router-restart / rolling-upgrade case, and it is where the delay + actually costs something. A *newly started* worker has an empty cache, so an + empty routing view of it is correct. But when the router restarts, every + existing worker arrives through the same path with a cache that is genuinely + warm -- and until its snapshot lands, kv-aware routing scores all of them as + holding nothing. + + The loop sits in `wait_for(self._kick.wait(), timeout=interval_s)`, so a + registration that does not set the kick waits out the full interval (30 s in + production). + """ + chain = hash_token_blocks(list(range(4)), block_size=4) + snap = Snapshot( + publisher_id="w-late", + publisher_type="worker", + model_name="m", + compat_key="ck", + index_block_size=4, + batch_id=0, + blocks=( + SnapshotBlock( + sequence_hash=chain[0].sequence_hash, + parent_sequence_hash=None, + block_hash=chain[0].block_hash, + tiers=("device",), + ), + ), + ) + pulled: list[str] = [] + + async def pull_fn(publisher_id, *args, **kwargs): + pulled.append(publisher_id) + return snap + + # A long interval stands in for production's 30 s: if registration relies on + # the periodic tick, this test waits for it and fails on the assertion. + index, writer, rec = await _make_reconciler(pull_fn, interval_s=5.0) + try: + await rec.start() + await asyncio.sleep(0.05) + pulled.clear() + + rec.register_target(publisher_id="w-late", endpoint="ignored", model="m", compat_key="ck") + await asyncio.sleep(0.3) + assert "w-late" in pulled, ( + "a worker registered while the reconciler was running was not pulled " + "within 0.3s; it is waiting out the periodic interval" + ) + finally: + await rec.stop() + await writer.stop() + + +async def test_reregistering_a_known_target_does_not_re_pull() -> None: + """Registration is re-asserted routinely -- the Kubernetes backend rewrites + its Pod annotation every 30 s, and etcd watch redelivers on relist. Kicking + on every one of those would turn a self-heal into a snapshot stampede + proportional to fleet size. + """ + chain = hash_token_blocks(list(range(4)), block_size=4) + snap = Snapshot( + publisher_id="w1", + publisher_type="worker", + model_name="m", + compat_key="ck", + index_block_size=4, + batch_id=0, + blocks=( + SnapshotBlock( + sequence_hash=chain[0].sequence_hash, + parent_sequence_hash=None, + block_hash=chain[0].block_hash, + tiers=("device",), + ), + ), + ) + pulls = 0 + + async def pull_fn(*args, **kwargs): + nonlocal pulls + pulls += 1 + return snap + + index, writer, rec = await _make_reconciler(pull_fn, interval_s=5.0) + try: + await rec.start() + await asyncio.sleep(0.05) + kw = dict(publisher_id="w1", endpoint="ignored", model="m", compat_key="ck") + rec.register_target(**kw) + await asyncio.sleep(0.2) + after_first = pulls + assert after_first > 0, "the first registration must pull" + + for _ in range(5): + rec.register_target(**kw) + await asyncio.sleep(0.2) + assert pulls == after_first, ( + f"re-registering pulled {pulls - after_first} more time(s); " + "only a new target should kick" + ) + finally: + await rec.stop() + await writer.stop() diff --git a/tests/unit/router/test_breaker.py b/tests/unit/router/test_breaker.py new file mode 100644 index 00000000..d6ace31d --- /dev/null +++ b/tests/unit/router/test_breaker.py @@ -0,0 +1,301 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# SPDX-License-Identifier: MIT +############################################################################### +"""Circuit breaker: a failing worker must stop being re-picked. + +The bug this guards against is not "failover is broken" -- failover works. It is +that failover's memory is per-request, so the *next* request scores a wedged +worker as if nothing happened. Every test here is written against that: the +assertions are about what happens on the second and third request, not the +first. + +Time is injected rather than slept, so the cooldown behaviour is tested at full +speed and deterministically. +""" + +from __future__ import annotations + +import pytest + +from infera.router.breaker import BreakerState, CircuitBreaker, is_worker_fault + + +class FakeClock: + def __init__(self) -> None: + self.t = 1000.0 + + def __call__(self) -> float: + return self.t + + def advance(self, dt: float) -> None: + self.t += dt + + +class W: + """Minimal stand-in for WorkerInfo -- the breaker only reads worker_id.""" + + def __init__(self, wid: str) -> None: + self.worker_id = wid + + +@pytest.fixture +def clock(): + return FakeClock() + + +@pytest.fixture +def cb(clock): + return CircuitBreaker(failure_threshold=3, cooldown=5.0, max_cooldown=20.0, now=clock) + + +def test_unknown_worker_is_allowed(cb): + """A worker never seen must route; the breaker is not an allowlist.""" + assert cb.allows("w1") is True + assert cb.state_of("w1") is BreakerState.CLOSED + + +def test_failures_below_threshold_do_not_open(cb): + for _ in range(2): + cb.record_failure("w1") + assert cb.allows("w1") is True + assert cb.state_of("w1") is BreakerState.CLOSED + + +def test_opens_at_threshold_and_excludes(cb): + """The actual bug: after N failures the worker must stop being offered.""" + for _ in range(3): + cb.record_failure("w1") + assert cb.state_of("w1") is BreakerState.OPEN + assert cb.allows("w1") is False + + +def test_success_resets_the_count(cb): + """Intermittent failures must not accumulate into a trip.""" + cb.record_failure("w1") + cb.record_failure("w1") + cb.record_success("w1") + cb.record_failure("w1") + cb.record_failure("w1") + assert cb.state_of("w1") is BreakerState.CLOSED + assert cb.allows("w1") is True + + +def test_half_open_admits_exactly_one_probe(cb, clock): + for _ in range(3): + cb.record_failure("w1") + clock.advance(5.1) + assert cb.allows("w1") is True, "cooldown elapsed -> one probe admitted" + assert cb.state_of("w1") is BreakerState.HALF_OPEN + assert cb.allows("w1") is False, "a second concurrent request must not also probe" + + +def test_successful_probe_closes(cb, clock): + for _ in range(3): + cb.record_failure("w1") + clock.advance(5.1) + cb.allows("w1") + cb.record_success("w1") + assert cb.state_of("w1") is BreakerState.CLOSED + assert cb.allows("w1") is True + + +def test_failed_probe_reopens_with_doubled_cooldown(cb, clock): + """A wedged worker does not recover on the first retry. A fixed cooldown + would probe it forever at a constant rate; this asserts the backoff.""" + for _ in range(3): + cb.record_failure("w1") + clock.advance(5.1) + cb.allows("w1") + cb.record_failure("w1") # probe fails -> reopen, cooldown 5 -> 10 + assert cb.state_of("w1") is BreakerState.OPEN + + clock.advance(5.1) # old cooldown would have elapsed + assert cb.allows("w1") is False, "backoff must have doubled" + clock.advance(5.0) # now past 10s + assert cb.allows("w1") is True + + +def test_cooldown_is_capped(cb, clock): + for _ in range(3): + cb.record_failure("w1") + for _ in range(6): # drive the doubling past max_cooldown + clock.advance(1000.0) + cb.allows("w1") + cb.record_failure("w1") + clock.advance(20.1) # max_cooldown = 20 + assert cb.allows("w1") is True, "cooldown must not grow without bound" + + +def test_filter_drops_open_workers(cb): + ws = [W("good"), W("bad")] + for _ in range(3): + cb.record_failure("bad") + assert [w.worker_id for w in cb.filter(ws)] == ["good"] + + +def test_filter_returns_all_when_every_worker_is_open(cb): + """Refusing to route would turn a partial outage into a total one. A + request served by a probably-bad worker beats a guaranteed 503.""" + ws = [W("a"), W("b")] + for wid in ("a", "b"): + for _ in range(3): + cb.record_failure(wid) + assert {w.worker_id for w in cb.filter(ws)} == {"a", "b"} + + +def test_filter_of_empty_is_empty(cb): + assert cb.filter([]) == [] + + +def test_workers_are_independent(cb): + for _ in range(3): + cb.record_failure("bad") + assert cb.allows("good") is True + assert cb.state_of("good") is BreakerState.CLOSED + + +def test_success_on_unknown_worker_is_harmless(cb): + cb.record_success("never-seen") + assert cb.allows("never-seen") is True + + +def test_snapshot_reports_trips(cb, clock): + for _ in range(3): + cb.record_failure("w1") + clock.advance(5.1) + cb.allows("w1") + cb.record_failure("w1") + snap = cb.snapshot()["w1"] + assert snap["state"] == "open" + assert snap["trips"] == 2, "initial trip plus the failed probe" + + +@pytest.mark.parametrize("status", [500, 502, 503, 504, 0]) +def test_server_errors_are_worker_faults(status): + assert is_worker_fault(status) is True + + +@pytest.mark.parametrize("status", [400, 404, 422, 429]) +def test_client_errors_are_not_worker_faults(status): + """A malformed request 400s on every worker it reaches. Counting that as a + health signal would trip the breaker across an entirely healthy fleet. + 429 is excluded separately: it means "full now", which load accounting + already routes around, and a doubling cooldown is far too heavy for it.""" + assert is_worker_fault(status) is False + + +def test_a_bad_client_cannot_trip_the_fleet(cb): + """Ten malformed requests against three healthy workers must leave all + three closed.""" + ws = [W("a"), W("b"), W("c")] + for _ in range(10): + for w in cb.filter(ws): + if is_worker_fault(400): + cb.record_failure(w.worker_id) + assert all(cb.state_of(w.worker_id) is BreakerState.CLOSED for w in ws) + + +def test_the_regression_this_exists_for(cb): + """End to end in breaker terms: a worker that fails every dispatch stops + being selected, instead of being re-picked on every subsequent request. + + Before this class existed, `tried` was per-request, so the loop below would + have offered `bad` on all ten requests. + """ + ws = [W("good"), W("bad")] + offered_bad = 0 + for _ in range(10): + candidates = cb.filter(ws) + if any(w.worker_id == "bad" for w in candidates): + offered_bad += 1 + cb.record_failure("bad") + cb.record_success("good") + assert offered_bad == 3, f"bad worker offered {offered_bad} times, expected 3 (the threshold)" + + +def test_threshold_zero_disables_it(): + """The documented off switch. Without this, threshold=0 would satisfy + `failures >= threshold` on the very first failure and trip immediately -- + the exact opposite of what --breaker-failure-threshold=0 promises.""" + off = CircuitBreaker(failure_threshold=0) + for _ in range(20): + off.record_failure("w1") + assert off.allows("w1") is True + assert off.state_of("w1") is BreakerState.CLOSED + ws = [W("a"), W("b")] + assert len(off.filter(ws)) == 2 + + +def test_a_probe_slot_taken_but_never_dispatched_is_reclaimed(cb, clock): + """The wedge: filter() claims the probe slot for *every* candidate it lets + through, while the policy dispatches exactly one of them. + + So a recovering worker routinely has its slot taken by a request that then + went elsewhere, and nothing records an outcome for it. Without a time bound + on the claim, that worker sits in half_open with the slot held forever -- + permanently out of rotation while perfectly healthy, and only a router + restart brings it back. + """ + good, bad = W("good"), W("bad") + for _ in range(3): + cb.record_failure("bad") + clock.advance(5.1) + + # A request arrives, both are offered, the policy picks `good`. + assert bad in cb.filter([good, bad]), "cooldown elapsed -> bad is due a probe" + cb.record_success("good") + + # `bad` never learns how its probe went, because it never got one. + clock.advance(cb.probe_timeout + 1.0) + assert bad in cb.filter([good, bad]), "an unused probe claim must not be permanent" + + +def test_a_neutral_outcome_frees_the_slot_without_scoring_it(cb, clock): + """4xx is evidence about the request, not the worker, so it must neither + trip the breaker nor count as a recovery -- but the probe slot it consumed + still has to come back, or the worker is wedged by one bad client.""" + for _ in range(3): + cb.record_failure("w1") + clock.advance(5.1) + assert cb.allows("w1") is True, "cooldown elapsed -> one probe admitted" + + cb.record_neutral("w1") + assert cb.state_of("w1") is BreakerState.HALF_OPEN, "a 4xx is not a recovery" + assert cb.allows("w1") is True, "but the slot is free for a real probe" + + +def test_forgetting_a_worker_drops_its_entry(cb): + """Worker ids are addresses, and a rebuilt pod never reuses one. Entries + kept for workers discovery has dropped grow without bound, and each also + pins a Prometheus series.""" + for _ in range(3): + cb.record_failure("gone") + assert "gone" in cb.snapshot() + cb.forget("gone") + assert "gone" not in cb.snapshot() + assert cb.state_of("gone") is BreakerState.CLOSED + + +def test_trips_count_outages_not_requests(cb, clock): + """`trips` answers "how often did this worker go bad", which is what an + alert on its rate is asking. Counting every failure that lands while the + breaker is already open answers "how many requests hit a bad worker" + instead -- a much larger number, dominated by whichever worker the all-open + fallback keeps feeding.""" + for _ in range(3): + cb.record_failure("w1") + assert cb.snapshot()["w1"]["trips"] == 1, "three failures are one outage" + + for _ in range(20): + cb.record_failure("w1") + assert cb.snapshot()["w1"]["trips"] == 1, "failures while already open are not new trips" + + # A failed probe does count: it is a fresh verdict on a worker that was + # given another chance, and the doubling cooldown bounds how often one can + # happen -- unlike the failures above, which arrive at the request rate. + clock.advance(cb.cooldown + 1) + assert cb.allows("w1") is True + cb.record_failure("w1") + assert cb.snapshot()["w1"]["trips"] == 2, "a failed probe is a new verdict" diff --git a/tests/unit/router/test_disagg_breaker.py b/tests/unit/router/test_disagg_breaker.py new file mode 100644 index 00000000..fc35dd35 --- /dev/null +++ b/tests/unit/router/test_disagg_breaker.py @@ -0,0 +1,404 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# SPDX-License-Identifier: MIT +############################################################################### +"""The breaker's PD call sites, exercised on the paths that actually reach them. + +These are the streaming generators in ``disagg.py``. They only run when a +decode leg is unreachable -- which is exactly the condition the breaker exists +for, and exactly what no existing test drove. Two of the five record sites were +written against a local name (``d``) that does not exist in these scopes; the +result would have been a ``NameError`` raised *while handling* a decode outage, +turning a clean SSE error into a 500. Lint caught it, but nothing executed it. +""" + +from __future__ import annotations + +import json + +import httpx +import pytest + +from infera.common.nats_request import TYPE_DATA, TYPE_DONE, TYPE_ERROR +from infera.common.worker_pool import DisaggMode, EngineType, WorkerInfo +from infera.router.disagg import DisaggRouter +from infera.router.policy.target import RouteTarget + + +class _FakePolicy: + def pick(self, candidates, body, role_hint=None): + return RouteTarget(candidates[0]), [] + + def on_request_started(self, route_key, blocks): + pass + + def on_request_finished(self, route_key, blocks): + pass + + +class _FakePool: + def __init__(self, workers): + self._workers = workers + + def list_active(self, model=None, mode=None): + return list(self._workers) + + +def _w(wid): + return WorkerInfo( + worker_id=wid, + url=f"http://{wid}", + model_name="m", + engine=EngineType.SGLANG, + request_transport="http", + ) + + +def _router(): + r = DisaggRouter(_FakePool([_w("p1"), _w("d1")]), _FakePolicy()) + # Every send fails at the transport layer: the decode leg is unreachable. + r._client = httpx.AsyncClient( + transport=httpx.MockTransport( + lambda request: (_ for _ in ()).throw(httpx.ConnectError("refused", request=request)) + ) + ) + # The pre-flight retry loop sleeps between attempts; not worth the wall time. + r._DECODE_OPEN_MAX_RETRIES = 0 + return r + + +async def _drain(agen) -> bytes: + out = b"" + async for chunk in agen: + out += chunk if isinstance(chunk, bytes) else chunk.encode() + return out + + +@pytest.mark.asyncio +async def test_decode_only_stream_records_failure_on_unreachable(): + r = _router() + d_target = RouteTarget(_w("d1")) + body = await _drain( + r._stream_decode_only(d_target, [], "http://d1/v1/chat/completions", {"model": "m"}) + ) + assert b"decode unreachable" in body, "client must get a clean SSE error, not a traceback" + assert r.breaker.state_of("d1").value == "closed", "one failure is below the threshold" + for _ in range(2): + await _drain( + r._stream_decode_only(d_target, [], "http://d1/v1/chat/completions", {"model": "m"}) + ) + assert r.breaker.state_of("d1").value == "open", "three unreachable decodes must trip it" + await r.aclose() + + +@pytest.mark.asyncio +async def test_dual_stream_records_failure_on_unreachable_decode(): + """Only the decode leg is broken here. Failing both -- which is what a + transport that refuses everything does -- cannot show that the pools are + scored independently, because then prefill deserves to trip too. + """ + r = _router() + + def _only_decode_is_down(request): + if "d1" in str(request.url): + raise httpx.ConnectError("refused", request=request) + return httpx.Response(200, json={"id": "x"}) + + r._client = httpx.AsyncClient(transport=httpx.MockTransport(_only_decode_is_down)) + + p_target = RouteTarget(_w("p1")) + d_target = RouteTarget(_w("d1")) + for _ in range(3): + body = await _drain( + r._stream_dual( + p_target, + [], + d_target, + [], + "http://p1/v1/chat/completions", + "http://d1/v1/chat/completions", + {"model": "m"}, + {"model": "m"}, + ) + ) + assert b"decode unreachable" in body + + assert r.breaker.state_of("d1").value == "open" + # The prefill leg is a separate pool: a wedged decode must not evict it. + assert r.breaker.state_of("p1").value == "closed" + await r.aclose() + + +class _RolePool: + """Unlike _FakePool, hands back the pool the caller actually asked for, so + a dispatch gets a real prefill/decode pair rather than the same worker twice.""" + + def __init__(self, prefill, decode): + self._by_mode = {DisaggMode.PREFILL: [prefill], DisaggMode.DECODE: [decode]} + + def list_active(self, model=None, mode=None): + return list(self._by_mode.get(mode, [])) + + +def _pd_worker(wid, mode, transport="http"): + meta = {"protocol": "sglang-bootstrap"} + if mode is DisaggMode.PREFILL: + meta["params"] = {"bootstrap_addr": f"{wid}:9000"} + return WorkerInfo( + worker_id=wid, + url=f"http://{wid}", + model_name="m", + engine=EngineType.SGLANG, + request_transport=transport, + disagg_mode=mode, + disagg_meta=meta, + ) + + +def _nats_router(*, fail_decode): + """A PD pair that both registered for the NATS transport, which is what + selects the NATS dispatch path.""" + return DisaggRouter( + _RolePool( + _pd_worker("p1", DisaggMode.PREFILL, transport="nats"), + _pd_worker("d1", DisaggMode.DECODE, transport="nats"), + ), + _FakePolicy(), + nats_client=_FakeNatsPD(fail_decode=fail_decode), + ) + + +def _ok_router(): + """A PD router whose every leg answers 200.""" + r = DisaggRouter( + _RolePool(_pd_worker("p1", DisaggMode.PREFILL), _pd_worker("d1", DisaggMode.DECODE)), + _FakePolicy(), + ) + r._client = httpx.AsyncClient( + transport=httpx.MockTransport( + lambda request: httpx.Response( + 200, json={"id": "x", "choices": [{"message": {"content": "hi"}}]} + ) + ) + ) + return r + + +@pytest.mark.asyncio +async def test_a_served_request_clears_the_failure_count(): + """Without a success recorded anywhere, "three consecutive failures" decays + into "three failures ever": the counter only climbs, so a healthy worker + that fails once a day trips on day three.""" + r = _ok_router() + for _ in range(2): + r.breaker.record_failure("d1") + assert r.breaker.snapshot()["d1"]["consecutive_failures"] == 2 + + await r.dispatch({"model": "m"}, stream=False) + + assert r.breaker.snapshot()["d1"]["consecutive_failures"] == 0 + assert r.breaker.state_of("d1").value == "closed" + await r.aclose() + + +@pytest.mark.asyncio +async def test_a_healthy_leg_does_not_launder_a_broken_one(): + """The two legs are two workers whose health is independent, so each is + scored from its own response. + + Scoring both off the single client-facing status code let the decode leg's + 200 count as evidence for a prefill that had just 500'd -- resetting its + failure count, and reopening a breaker that was already open. + """ + + def _prefill_is_broken(request): + if "p1" in str(request.url): + return httpx.Response(500, json={"error": "prefill exploded"}) + return httpx.Response(200, json={"id": "x", "choices": [{"message": {"content": "hi"}}]}) + + r = DisaggRouter( + _RolePool(_pd_worker("p1", DisaggMode.PREFILL), _pd_worker("d1", DisaggMode.DECODE)), + _FakePolicy(), + ) + r._client = httpx.AsyncClient(transport=httpx.MockTransport(_prefill_is_broken)) + + for _ in range(3): + resp = await r.dispatch({"model": "m"}, stream=False) + assert resp.status_code == 200, "decode answers, so the client still gets 200" + + assert r.breaker.state_of("p1").value == "open", ( + "a prefill that 500s every request must trip, even though the decode leg beside it succeeds" + ) + assert r.breaker.state_of("d1").value == "closed", "the healthy leg is untouched" + await r.aclose() + + +@pytest.mark.asyncio +async def test_a_streaming_dispatch_is_not_scored_before_it_runs(): + """A StreamingResponse is returned before its generator is touched: the + decode leg has not been POSTed and its 200 is Starlette's default, not an + outcome. Scoring it there records success for both roles before either was + dispatched and resets the count the legs are about to raise -- a decode + worker failing every request would never trip. + + Drives dispatch(stream=True), which is the production path; the older tests + call the generators directly and so cannot see this. + """ + r = DisaggRouter( + _RolePool(_pd_worker("p1", DisaggMode.PREFILL), _pd_worker("d1", DisaggMode.DECODE)), + _FakePolicy(), + ) + r._client = httpx.AsyncClient( + transport=httpx.MockTransport( + lambda request: (_ for _ in ()).throw(httpx.ConnectError("refused", request=request)) + ) + ) + r._DECODE_OPEN_MAX_RETRIES = 0 + + for i in range(3): + resp = await r.dispatch({"model": "m"}, stream=True) + assert await _drain(resp.body_iterator), f"request {i} produced no body" + + assert r.breaker.state_of("d1").value == "open", ( + "three unreachable decodes must trip the breaker; if this is closed the " + "wrapper scored the stream before the decode leg ran" + ) + await r.aclose() + + +@pytest.mark.asyncio +async def test_a_tripped_pd_worker_recovers_after_a_good_probe(): + """The probe is dispatched and succeeds; if nothing records that, the + worker stays half-open with its probe slot held and never routes again.""" + r = _ok_router() + for _ in range(3): + r.breaker.record_failure("d1") + assert r.breaker.state_of("d1").value == "open" + + # Let the cooldown lapse so the next dispatch is the half-open probe. + r.breaker._entries["d1"].opens_until = 0.0 + + await r.dispatch({"model": "m"}, stream=False) + + assert r.breaker.state_of("d1").value == "closed", "a good probe must close it" + await r.aclose() + + +@pytest.mark.asyncio +async def test_a_decode_outage_does_not_trip_the_prefill_worker(): + """asyncio.gather raises whichever leg failed, without saying which. Blaming + a fixed one means a decode that refuses connections evicts the healthy + prefill worker from rotation while the broken decode is never scored at + all -- the exact inversion of what the breaker is for.""" + + def _only_decode_is_down(request): + if "d1" in str(request.url): + raise httpx.ConnectError("refused", request=request) + return httpx.Response(200, json={"id": "x"}) + + r = DisaggRouter( + _RolePool(_pd_worker("p1", DisaggMode.PREFILL), _pd_worker("d1", DisaggMode.DECODE)), + _FakePolicy(), + ) + r._client = httpx.AsyncClient(transport=httpx.MockTransport(_only_decode_is_down)) + + for _ in range(3): + await r.dispatch({"model": "m"}, stream=False) + + assert r.breaker.state_of("d1").value == "open", "the leg that refused must be the one scored" + assert r.breaker.state_of("p1").value == "closed", "the healthy leg must not be evicted" + await r.aclose() + + +@pytest.mark.asyncio +async def test_a_stream_that_dies_after_its_headers_is_not_a_success(): + """A 200 header only means the request was accepted. Treating it as + recovery resets the failure count, so a decode worker that answers 200 and + then sends nothing -- which is precisely the "healthy to the platform, + broken for inference" profile the breaker exists for -- can never trip, and + erases real failures on its way.""" + r = DisaggRouter( + _RolePool(_pd_worker("p1", DisaggMode.PREFILL), _pd_worker("d1", DisaggMode.DECODE)), + _FakePolicy(), + ) + + def _headers_then_nothing(request): + if "d1" in str(request.url): + # 200, then the body raises as soon as it is read. + return httpx.Response(200, stream=_DyingStream()) + return httpx.Response(200, json={"id": "x"}) + + r._client = httpx.AsyncClient(transport=httpx.MockTransport(_headers_then_nothing)) + + for _ in range(2): + r.breaker.record_failure("d1") + before = r.breaker.snapshot()["d1"]["consecutive_failures"] + assert before == 2 + + resp = await r.dispatch({"model": "m"}, stream=True) + await _drain(resp.body_iterator) + + after = r.breaker.snapshot()["d1"]["consecutive_failures"] + assert after >= before, ( + f"consecutive_failures went {before} -> {after}: a stream that produced " + "no output was scored as evidence of health" + ) + await r.aclose() + + +class _DyingStream(httpx.AsyncByteStream): + """Headers arrive, then the body fails -- no bytes ever reach the client.""" + + async def __aiter__(self): + raise httpx.ReadError("connection died after headers") + yield b"" # unreachable; makes this an async generator + + +@pytest.mark.asyncio +async def test_the_nats_transport_scores_its_legs_too(): + """The NATS paths do not use the HTTP client, so scoring placed at HTTP + response sites misses them entirely. A decode worker failing every request + over NATS would be invisible to the breaker, and -- worse -- one already + open could never be closed again, since nothing would ever record the + success that ends its half-open state.""" + r = _nats_router(fail_decode=True) + + for _ in range(3): + await r.dispatch({"model": "m"}, stream=False) + + assert r.breaker.state_of("d1").value == "open", "a failing NATS decode leg must trip" + await r.aclose() + + +@pytest.mark.asyncio +async def test_a_good_nats_probe_closes_the_breaker(): + r = _nats_router(fail_decode=False) + for _ in range(3): + r.breaker.record_failure("d1") + r.breaker._entries["d1"].opens_until = 0.0 + + await r.dispatch({"model": "m"}, stream=False) + + assert r.breaker.state_of("d1").value == "closed", ( + "a NATS worker that answers cleanly must be able to recover; otherwise " + "it stays half-open forever, throttled to one request per probe window" + ) + await r.aclose() + + +class _FakeNatsPD: + """Scripted NATS transport: decode either answers or errors, prefill is fine.""" + + def __init__(self, *, fail_decode: bool): + self.fail_decode = fail_decode + + async def admit(self, worker_id): + return True + + async def stream(self, worker_id, payload): + if worker_id == "d1" and self.fail_decode: + yield (TYPE_ERROR, 502, b"decode exploded") + return + yield (TYPE_DATA, None, json.dumps({"id": "x"}).encode()) + yield (TYPE_DONE, 200, b"") diff --git a/tests/unit/router/test_failover.py b/tests/unit/router/test_failover.py index 27ebe9d9..65a2b99f 100644 --- a/tests/unit/router/test_failover.py +++ b/tests/unit/router/test_failover.py @@ -14,7 +14,7 @@ import pytest from infera.common.nats_request import TYPE_DATA, TYPE_DONE, TYPE_ERROR -from infera.common.worker_pool import EngineType, WorkerInfo +from infera.common.worker_pool import DisaggMode, EngineType, WorkerInfo from infera.router.mixed import MixedRouter from infera.router.policy.target import RouteTarget @@ -208,3 +208,198 @@ def handler(request: httpx.Request) -> httpx.Response: assert resp.status_code == 200 assert json.loads(bytes(resp.body))["id"] == "ok-http" await r.aclose() + + +# --- circuit breaker: failure memory ACROSS requests (issue #82) --------------- + + +@pytest.mark.asyncio +async def test_breaker_stops_reselecting_a_dead_worker(): + """The regression behind issue #82, through the real MixedRouter. + + Failover already made all ten of these requests succeed, so a test that + only checked status codes passed both before and after the fix. What was + broken is the cost: ``tried`` is per-request, so a worker that is broken + for inference but healthy to discovery was re-picked on *every* request and + every one of them paid a wasted round trip. ``nats.streamed`` is the + assertion that matters. + """ + scripts = { + "w1": [(TYPE_ERROR, 502, b"wedged")], + "w2": [(TYPE_DATA, None, b"ok"), (TYPE_DONE, 200, b"")], + } + nats = _FakeNats(scripts) + r = _router([_w("w1"), _w("w2")], nats, retries=1) + for _ in range(10): + resp = await r.dispatch({"model": "m"}, stream=True) + assert await _drain_stream(resp) == b"ok", "failover must still serve every request" + + # _FakePolicy always picks candidates[0], so w1 is offered every request + # until the breaker (threshold 3) takes it out; its 5s cooldown does not + # elapse during the test. + assert nats.streamed.count("w1") == 3, ( + f"w1 dispatched {nats.streamed.count('w1')} times; expected 3 " + "(it was 10 before the breaker existed)" + ) + assert nats.streamed.count("w2") == 10 + await r.aclose() + + +@pytest.mark.asyncio +async def test_breaker_ignores_client_errors(): + """A 400 is the request's fault, and every worker would return it. Counting + it would circuit-break an entirely healthy fleet on one bad client.""" + nats = _FakeNats({"w1": [(TYPE_ERROR, 400, b"bad request")]}) + r = _router([_w("w1")], nats, retries=0) + for _ in range(10): + await r.dispatch({"model": "m"}, stream=True) + assert nats.streamed.count("w1") == 10, "4xx must not take a healthy worker out of rotation" + await r.aclose() + + +@pytest.mark.asyncio +async def test_breaker_recovers_after_cooldown(): + """A worker that comes back must be picked up again, not stay excluded.""" + scripts = { + "w1": [(TYPE_ERROR, 502, b"wedged")], + "w2": [(TYPE_DATA, None, b"ok"), (TYPE_DONE, 200, b"")], + } + nats = _FakeNats(scripts) + r = _router([_w("w1"), _w("w2")], nats, retries=1) + + clock = [1000.0] + r.breaker.now = lambda: clock[0] + + for _ in range(3): + await _drain_stream(await r.dispatch({"model": "m"}, stream=True)) + assert nats.streamed.count("w1") == 3 # tripped + + await _drain_stream(await r.dispatch({"model": "m"}, stream=True)) + assert nats.streamed.count("w1") == 3, "still open during the cooldown" + + scripts["w1"] = [(TYPE_DATA, None, b"back"), (TYPE_DONE, 200, b"")] + clock[0] += 5.1 + body = await _drain_stream(await r.dispatch({"model": "m"}, stream=True)) + assert body == b"back", "the half-open probe must reach the recovered worker" + assert r.breaker.state_of("w1").value == "closed" + await r.aclose() + + +# --- unary 5xx must fail over too (the gap the breaker fell through) --------- + + +@pytest.mark.asyncio +async def test_unary_http_fails_over_on_5xx(): + """A worker 500 before any byte reached the client is exactly what failover + is for. This path used to return it verbatim, so a non-streaming request + over HTTP -- the default in every k8s example -- never failed over and never + reached the circuit breaker.""" + + def handler(request: httpx.Request) -> httpx.Response: + if request.url.host == "w1": + return httpx.Response(500, json={"error": "boom"}) + return httpx.Response(200, json={"id": "ok"}) + + r = _router([_w("w1", transport="http"), _w("w2", transport="http")], nats=None, retries=1) + r._client = httpx.AsyncClient(transport=httpx.MockTransport(handler)) + resp = await r.dispatch({"model": "m"}, stream=False) + assert resp.status_code == 200 + assert json.loads(bytes(resp.body))["id"] == "ok" + await r.aclose() + + +@pytest.mark.asyncio +async def test_unary_http_does_not_fail_over_on_4xx(): + """The request is bad, not the worker. Every worker would answer the same, + so retrying only triples the latency of an error the client must see.""" + hits: list[str] = [] + + def handler(request: httpx.Request) -> httpx.Response: + hits.append(request.url.host) + return httpx.Response(400, json={"error": "bad request"}) + + r = _router([_w("w1", transport="http"), _w("w2", transport="http")], nats=None, retries=1) + r._client = httpx.AsyncClient(transport=httpx.MockTransport(handler)) + resp = await r.dispatch({"model": "m"}, stream=False) + assert resp.status_code == 400 + assert hits == ["w1"], f"4xx must not be retried, but hit {hits}" + await r.aclose() + + +@pytest.mark.asyncio +async def test_unary_5xx_trips_the_breaker(): + """The consequence that made this worth fixing: without failover the + breaker never saw a unary failure, so a wedged worker was re-picked + forever on the most common configuration.""" + + def handler(request: httpx.Request) -> httpx.Response: + if request.url.host == "bad": + return httpx.Response(503, json={"error": "wedged"}) + return httpx.Response(200, json={"id": "ok"}) + + r = _router([_w("bad", transport="http"), _w("good", transport="http")], nats=None, retries=1) + r._client = httpx.AsyncClient(transport=httpx.MockTransport(handler)) + for _ in range(3): + assert (await r.dispatch({"model": "m"}, stream=False)).status_code == 200 + assert r.breaker.state_of("bad").value == "open" + await r.aclose() + + +# --- half a PD deployment must say so -------------------------------------- + + +class _ModePool: + """Pool that can answer per-disagg-mode, unlike _FakePool.""" + + def __init__(self, workers): + self._w = workers + + def list_active(self, model=None, mode=None): + return [w for w in self._w if mode is None or w.disagg_mode == mode] + + +def _pd_worker(wid, mode): + return WorkerInfo( + worker_id=wid, + url=f"http://{wid}", + model_name="m", + engine=EngineType.SGLANG, + disagg_mode=mode, + request_transport="http", + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "present,missing", + [(DisaggMode.PREFILL, "decode"), (DisaggMode.DECODE, "prefill")], +) +async def test_half_a_pd_deployment_names_the_empty_pool(present, missing): + """Scaling either PD pool to zero fails closed -- correctly -- but used to + report "no active mixed worker", which points at something the operator + never deployed while the surviving pool sits right there.""" + from infera.router.auto import AutoRouter + + r = AutoRouter(_ModePool([_pd_worker("w1", present)]), _FakePolicy()) + resp = await r.dispatch({"model": "m"}, stream=False) + assert resp.status_code == 503 + body = json.loads(bytes(resp.body))["error"] + assert missing in body and "PD dispatch requires both pools" in body, body + assert "mixed" not in body, f"must not blame mixed workers: {body}" + await r.aclose() + + +@pytest.mark.asyncio +async def test_a_mixed_worker_still_absorbs_a_half_pd_fleet(): + """A mixed worker alongside half a PD pool can serve, so the 503 must not + fire -- this is the rolling-upgrade case.""" + from infera.router.auto import AutoRouter + + pool = _ModePool([_pd_worker("p", DisaggMode.PREFILL), _pd_worker("m1", DisaggMode.MIXED)]) + r = AutoRouter(pool, _FakePolicy(), nats_client=None, request_max_retries=0) + r._mixed._client = httpx.AsyncClient( + transport=httpx.MockTransport(lambda req: httpx.Response(200, json={"id": "ok"})) + ) + resp = await r.dispatch({"model": "m"}, stream=False) + assert resp.status_code == 200 + await r.aclose()