feat(disaggregation): gating at head; skip when header pins or ≤1 revision

Moves the gating filter to the head of the scheduler filter chain
(prepended after strict, so it ends up in front of it), making the
revision-axis decision the very first step per request. Chain order
becomes:

  gating → strict → operator → prefer

Only one of gating/strict actually makes the revision decision on any
given request; the other detects the situation and short-circuits.

Gating skips (passthrough, no cache List, no RNG) when:

  1. The revision-axis strict header is set — strict downstream will
     pin the revision. Gating firing here would either waste work or
     stochastic-pick a revision that strict then rejects (503).
  2. The candidate pool contains ≤1 unique revision — nothing to
     shape, either because the fleet has one revision or an upstream
     filter already narrowed.

Otherwise gating does its weighted-random-pick as before (weight per
revision = Σ Ready pod count over gating.requireRoles.values), drops
uncovered revisions, and keeps only the chosen revision's pods.

Everything downstream of the head now sees a single-revision candidate
pool. Operator-declared filters and prefer no longer walk multi-
revision pools when the decision has already been made.

Also renames scanCoverage → readyPodsByRevisionRole. The old name
implied "checks coverage" but the function also returns raw per-role
counts used by the gating weights; the new name describes the data
structure directly.

Tests:
  - wire_test.go: order assertion flipped to gating < strict < operator
    < prefer.
  - New: TestGatingFilter_SkipsWhenRevisionHeaderPresent — panicking
    rand01 asserts the fast path fires when the header is set.
  - New: TestGatingFilter_SkipsWhenSingleRevisionInPool — same, for
    the ≤1-revision fast path.

Live-verified on kind against the three shapes covered by
scripts/verify-slow-transition.sh (10p10d, 2p20d, 20p2d) — all steps
land within ±10pp of the expected cross-role weighted share; zero
cross-revision mismatches, zero request failures during transitions.

Signed-off-by: Mathis Felardos <mathis@mistral.ai>

Author: hasB4K

pkg/epp/disaggregation/controller.go

@@ -151,11 +151,11 @@ func (c *Controller) ResponseHeader(_ context.Context, _ *fwksched.InferenceRequ
 	}
 }
 
-// scanCoverage lists Ready pods matching the controller's scope and groups
+// readyPodsByRevisionRole lists Ready pods matching the controller's scope and groups
 // them by (revision → role → count). One cache-backed List per call. Used
 // by the gating filter (per request, for both coverage AND revision-weight
 // computation) and by boot validation.
-func (c *Controller) scanCoverage(ctx context.Context) (revisions map[string]struct{}, roleCounts map[string]map[string]int, err error) {
+func (c *Controller) readyPodsByRevisionRole(ctx context.Context) (revisions map[string]struct{}, roleCounts map[string]map[string]int, err error) {
 	var podList corev1.PodList
 	if err := c.reader.List(ctx, &podList,
 		client.InNamespace(c.namespace),

pkg/epp/disaggregation/controller_test.go

@@ -433,6 +433,49 @@ func TestGatingFilter_NoGatingConfigIsNoop(t *testing.T) {
 	}
 }
 
+func TestGatingFilter_SkipsWhenRevisionHeaderPresent(t *testing.T) {
+	// A request with the revision-axis strict header means strict downstream
+	// will pin the revision; gating firing here would either waste work or,
+	// worse, weight-pick a revision that strict then rejects (503). The
+	// filter must passthrough without reading the cache — assert by using a
+	// panicking rand01 as a canary.
+	c := gatingFixture(t, validConfig(), map[string]map[string]int{
+		"v1": {"prefill": 2, "decode": 18},
+		"v2": {"prefill": 1, "decode": 2},
+	})
+	f := newGatingFilter(c)
+	f.rand01 = func() float64 { panic("rand01 must not be called when the revision header is set") }
+	pods := []fwksched.Endpoint{
+		endpoint("p1a", revLabels("v1")),
+		endpoint("p2", revLabels("v2")),
+	}
+	req := &fwksched.InferenceRequest{Headers: map[string]string{"x-disagg-revision": "v1"}}
+	got := f.Filter(context.Background(), req, pods)
+	if len(got) != len(pods) {
+		t.Fatalf("expected passthrough of all %d pods, got %d", len(pods), len(got))
+	}
+}
+
+func TestGatingFilter_SkipsWhenSingleRevisionInPool(t *testing.T) {
+	// If the candidate pool already has only one unique revision — e.g.
+	// because the fleet has one revision, or an upstream filter narrowed —
+	// gating has nothing to shape and must skip. Panicking rand01 asserts
+	// the fast path fires.
+	c := gatingFixture(t, validConfig(), map[string]map[string]int{
+		"v1": {"prefill": 3, "decode": 3},
+	})
+	f := newGatingFilter(c)
+	f.rand01 = func() float64 { panic("rand01 must not be called with a single-revision pool") }
+	pods := []fwksched.Endpoint{
+		endpoint("p1", revLabels("v1")),
+		endpoint("p2", revLabels("v1")),
+	}
+	got := f.Filter(context.Background(), nil, pods)
+	if len(got) != len(pods) {
+		t.Fatalf("expected passthrough of all %d pods, got %d", len(pods), len(got))
+	}
+}
+
 func TestGatingFilter_DisabledModeIsNoop(t *testing.T) {
 	cfg := validConfig()
 	cfg.Gating.Mode = GatingModeDisabled

pkg/epp/disaggregation/filter.go

@@ -48,8 +48,9 @@ func (f *modeSelectorsFilter) Filter(ctx context.Context, request *fwksched.Infe
 	return f.controller.filterSelectors(ctx, request, pods, f.keepFn)
 }
 
-// gatingFilter is where the disaggregation controller actually shapes
-// traffic. Two things happen per Filter call:
+// gatingFilter is where the disaggregation controller decides the revision
+// axis when no client header has done so. Runs at the HEAD of the filter
+// chain. Two things happen per Filter call:
 //
 //  1. Drop any candidate whose revision fails the coverage check —
 //     revisions missing pods on any listed role (rollout drift) are
@@ -65,6 +66,14 @@ func (f *modeSelectorsFilter) Filter(ctx context.Context, request *fwksched.Infe
 // the picker choice in the operator YAML: prefill traffic to revision R
 // converges on (crossRolePods(R) / Σ crossRolePods) even when prefill and
 // decode replica counts differ per revision (2p+18d vs. 1p+2d etc.).
+//
+// The filter short-circuits (passthrough, no cache read) when:
+//   - The request carries the revision-axis strict header — the strict
+//     filter downstream will pin the revision, we shouldn't stochastic-
+//     pick something that would then get dropped as a mismatch (503).
+//   - The candidate pool contains at most one unique revision — nothing
+//     to shape, either because the fleet has one revision or because an
+//     upstream filter already narrowed things.
 type gatingFilter struct {
 	controller *Controller
 	typedName  fwkplugin.TypedName
@@ -89,35 +98,37 @@ func newGatingFilter(controller *Controller) *gatingFilter {
 
 func (f *gatingFilter) TypedName() fwkplugin.TypedName { return f.typedName }
 
-func (f *gatingFilter) Filter(ctx context.Context, _ *fwksched.InferenceRequest, pods []fwksched.Endpoint) []fwksched.Endpoint {
+func (f *gatingFilter) Filter(ctx context.Context, request *fwksched.InferenceRequest, pods []fwksched.Endpoint) []fwksched.Endpoint {
 	gating := f.controller.config.Gating
 	revisionLabelKey := f.controller.revisionLabelKey
 	if !gating.Active() || revisionLabelKey == "" {
 		return append(make([]fwksched.Endpoint, 0, len(pods)), pods...)
 	}
+
+	// Fast path: header-pinned. Strict downstream will do the narrowing.
+	if f.hasRevisionHeader(request) {
+		return append(make([]fwksched.Endpoint, 0, len(pods)), pods...)
+	}
+
+	// Fast path: pool already has ≤1 unique revision. Nothing to shape.
+	seenRevisions := uniqueRevisions(pods, revisionLabelKey)
+	if len(seenRevisions) <= 1 {
+		return append(make([]fwksched.Endpoint, 0, len(pods)), pods...)
+	}
+
 	requiredRoles := gating.RequireRoles.Values
 
-	_, roleCounts, err := f.controller.scanCoverage(ctx)
+	_, roleCounts, err := f.controller.readyPodsByRevisionRole(ctx)
 	if err != nil {
 		// Cache-backed List failing is unexpected; be pessimistic and drop
 		// everything rather than route to potentially-broken revisions.
 		return nil
 	}
 
-	// Collect the revisions present in the candidate pool and their
-	// cross-role weights (0 if any required role has no Ready pod → gated).
-	weights := make(map[string]int)
-	for _, endpoint := range pods {
-		if endpoint == nil || endpoint.GetMetadata() == nil {
-			continue
-		}
-		revision := endpoint.GetMetadata().Labels[revisionLabelKey]
-		if revision == "" {
-			continue
-		}
-		if _, already := weights[revision]; already {
-			continue
-		}
+	// Weight per revision in the pool (0 if any required role has no Ready
+	// pod → gated).
+	weights := make(map[string]int, len(seenRevisions))
+	for revision := range seenRevisions {
 		weights[revision] = crossRoleWeight(roleCounts[revision], requiredRoles)
 	}
 
@@ -162,6 +173,41 @@ func crossRoleWeight(perRole map[string]int, required []string) int {
 	return total
 }
 
+// hasRevisionHeader reports whether the request carries the header that
+// pins the revision axis. Convention: the revision axis is the first
+// entry in Selectors — same source as controller.revisionLabelKey — and
+// only counts when it's a strict-mode selector (prefer mode is a hint,
+// not a pin).
+func (f *gatingFilter) hasRevisionHeader(request *fwksched.InferenceRequest) bool {
+	if request == nil {
+		return false
+	}
+	selectors := f.controller.config.Selectors
+	if len(selectors) == 0 {
+		return false
+	}
+	revSelector := selectors[0]
+	if revSelector.Mode != ModeStrict {
+		return false
+	}
+	return request.Headers[revSelector.HeaderName] != ""
+}
+
+// uniqueRevisions returns the set of distinct revision-label values in
+// the candidate pool. Used only by the gating fast path.
+func uniqueRevisions(pods []fwksched.Endpoint, revisionLabelKey string) map[string]struct{} {
+	seen := make(map[string]struct{})
+	for _, endpoint := range pods {
+		if endpoint == nil || endpoint.GetMetadata() == nil {
+			continue
+		}
+		if revision := endpoint.GetMetadata().Labels[revisionLabelKey]; revision != "" {
+			seen[revision] = struct{}{}
+		}
+	}
+	return seen
+}
+
 // pickWeightedRevision returns one revision from weights, chosen with
 // probability weights[r] / Σ weights. Iteration order is stabilised via
 // sort so a fixed rand01 gives repeatable output in tests.

pkg/epp/disaggregation/register.go

@@ -92,7 +92,7 @@ func gatingForLog(g *Gating) string {
 //     single revision — the filter would otherwise silently drop every
 //     endpoint at request time.
 func validateRolesObserved(ctx context.Context, c *Controller, roles []string) error {
-	revisions, roleCounts, err := c.scanCoverage(ctx)
+	revisions, roleCounts, err := c.readyPodsByRevisionRole(ctx)
 	if err != nil {
 		return err
 	}

pkg/epp/disaggregation/wire.go

@@ -6,20 +6,26 @@ import (
 )
 
 // WireInto attaches a Controller to a SchedulerConfig and a request-control
-// Config in the correct positions:
+// Config. The scheduler filter chain, in execution order per request:
 //
-//   - strict-selector filter prepended to every scheduler profile's filter
-//     chain, so no downstream filter ever sees a wrong-revision candidate.
-//   - prefer-selector filter appended to every profile's filter chain, so
-//     its fallback captures the fully-narrowed pool.
-//   - gating filter appended after the two selector wrappers (only when the
-//     config declares an active Gating block), so revisions that fail the
-//     gating check are dropped before the picker sees them.
-//   - response-header processor attached to the request-control pipeline.
+//	gating → strict → operator filters → prefer
 //
-// No scorer is registered: pick behaviour is left to the operator's YAML
-// picker choice. The gating filter handles cross-role liveness; any
-// traffic-shaping across surviving revisions is a follow-on concern.
+// The revision axis is decided at the head of the chain — either by the
+// gating filter (stochastic, cache-driven weighted pick, fires when no
+// client header is present) or by the strict filter (deterministic, from
+// the request header, fires when the header is set). Only one of the two
+// makes the revision decision on any given request; the other detects the
+// situation and fast-paths.
+//
+// Everything downstream of the head therefore sees a candidate pool
+// already narrowed to a single revision. Operator-declared filters run in
+// the middle. The prefer filter is appended at the tail so that its
+// "keep-if-empty" fallback captures the fully-narrowed pool — an operator
+// filter that trims after prefer would violate prefer's fallback set.
+//
+// The response-header processor is attached to the request-control
+// pipeline. No scorer is registered: pick behaviour is left to the
+// operator's YAML picker choice.
 func WireInto(schedulerConfig *scheduling.SchedulerConfig, requestControlConfig *requestcontrol.Config, controller *Controller) error {
 	if controller == nil {
 		return nil
@@ -38,15 +44,18 @@ func WireInto(schedulerConfig *scheduling.SchedulerConfig, requestControlConfig
 	}
 
 	for _, profile := range schedulerConfig.Profiles() {
+		// Prepend order matters: the LAST PrependFilter call ends up at
+		// position 0 of the chain. Strict is prepended first so gating
+		// ends up in front of it — gating is the first thing to run.
 		if hasStrict {
 			profile.PrependFilter(strictFilter)
 		}
+		if gatingActive {
+			profile.PrependFilter(gating)
+		}
 		if hasPrefer {
 			profile.AppendFilter(preferFilter)
 		}
-		if gatingActive {
-			profile.AppendFilter(gating)
-		}
 	}
 	requestControlConfig.AddPlugins(controller)
 	return nil

pkg/epp/disaggregation/wire_test.go

@@ -23,10 +23,12 @@ func buildSchedulerConfigForTest(t *testing.T) *scheduling.SchedulerConfig {
 	return scheduling.NewSchedulerConfig(singleprofilehandler.NewSingleProfileHandler(), profiles)
 }
 
-// TestWireInto_StrictAtHead_PreferAtTail proves the placement contract: an
-// existing operator-declared filter ends up sandwiched between our strict and
-// prefer filters.
-func TestWireInto_StrictAtHead_PreferAtTail(t *testing.T) {
+// TestWireInto_GatingAtHead_PreferAtTail proves the placement contract:
+// the revision decision is made at the head (gating first, then strict —
+// only one of them fires per request), any operator-declared filters run
+// on the resulting single-revision pool, and prefer runs at the tail so
+// its keep-if-empty fallback sees the fully-narrowed pool.
+func TestWireInto_GatingAtHead_PreferAtTail(t *testing.T) {
 	schedulerConfig := buildSchedulerConfigForTest(t)
 	profile := schedulerConfig.Profiles()["default"]
 	profile.AppendFilter(&namedTestFilter{name: "operator-filter"})
@@ -47,18 +49,18 @@ func TestWireInto_StrictAtHead_PreferAtTail(t *testing.T) {
 
 	// Read back the filter chain via profile.String() — the framework does
 	// not expose the private slice directly. Expected order in the string:
-	// strict, operator, prefer, gating.
+	// gating, strict, operator, prefer.
 	got := profile.String()
+	gatingAt := indexOf(got, gatingFilterType)
 	strictAt := indexOf(got, strictFilterType)
 	operatorAt := indexOf(got, "operator-filter")
 	preferAt := indexOf(got, preferFilterType)
-	gatingAt := indexOf(got, gatingFilterType)
-	if strictAt < 0 || operatorAt < 0 || preferAt < 0 || gatingAt < 0 {
+	if gatingAt < 0 || strictAt < 0 || operatorAt < 0 || preferAt < 0 {
 		t.Fatalf("expected all four filters in profile string, got %q", got)
 	}
-	if strictAt >= operatorAt || operatorAt >= preferAt || preferAt >= gatingAt {
-		t.Fatalf("expected order strict < operator < prefer < gating in %q (positions %d, %d, %d, %d)",
-			got, strictAt, operatorAt, preferAt, gatingAt)
+	if gatingAt >= strictAt || strictAt >= operatorAt || operatorAt >= preferAt {
+		t.Fatalf("expected order gating < strict < operator < prefer in %q (positions %d, %d, %d, %d)",
+			got, gatingAt, strictAt, operatorAt, preferAt)
 	}
 }