[RFC]: Stage aware saturation in Flow control

[asharkhan3101]

Summary

Tracked in #2585. Flow control saturation gating is stage-aware since #2221, but it still uses a single detector instance for every stage. In this RFC, I present design that lets each pipeline stage use its own detector instance, keeps the existing single-detector configuration working unchanged, and leaves room for future stages and combination modes.

Background and problem

In a prefill/decode deployment the stages have different resource profiles: prefill is token-bound and decode is often request-bound. The concurrency-detector supports hybrid mode to cover both dimensions with one instance, but one instance has one set of limits (maxConcurrency, maxTokenConcurrency). A deployment cannot today say "prefill gates on tokens with limit X, decode gates on requests with limit Y" because flowControl.saturationDetector takes exactly one pluginRef.

The flow control processor already partitions endpoints per stage and calls the detector once per stage. What is missing is the ability to select a different detector for each call.

Non-goals

• Combining per-stage values other than max or a single selected stage (no weighted sums, no min).
• Per-stage detectors in the legacy admission controller (flow control gate disabled). The legacy path keeps the default detector.
• Per-stage stale-metrics accounting; that is tracked separately in #2475.

Configuration

Extend SaturationDetectorConfig in
apix/config/v1alpha1/endpointpickerconfig_types.go:

// SaturationDetectorConfig contains the configuration for a saturation detector.
type SaturationDetectorConfig struct {
	// PluginRef specifies the default saturation detector instance. It is
	// used for every stage that has no specific override and for the legacy
	// admission path. If unspecified, "utilization-detector" is used.
	PluginRef string `json:"pluginRef,omitempty"`

	// Stages overrides the detector per pipeline stage. Keys are stage names
	// ("prefill", "decode"); values reference plugin instances declared in
	// the top-level Plugins section. A stage without an override uses
	// PluginRef.
	Stages map[string]string `json:"stages,omitempty"`

	// GateOn selects which stage's signal gates dispatch. "max" (default)
	// gates on the highest stage signal. Any other value must be a stage
	// name; dispatch then gates on that stage alone.
	GateOn string `json:"gateOn,omitempty"`
}

Example:

plugins:
  - type: concurrency-detector
    name: prefill-concurrency
    parameters:
      concurrencyMode: "tokens"
      maxTokenConcurrency: 300000
  - type: concurrency-detector
    name: decode-concurrency
    parameters:
      concurrencyMode: "requests"
      maxConcurrency: 64
  - type: utilization-detector

flowControl:
  saturationDetector:
    pluginRef: utilization-detector
    stages:
      prefill: prefill-concurrency
      decode: decode-concurrency
    gateOn: max

gateOn: prefill is how a deployment restricts gating to the prefill stage.

Resolution

The config loader (pkg/epp/config/loader/configloader.go) resolves each stage reference to a fwkfc.SaturationDetector at load time, type-asserting every referenced plugin, and stores the map on flowcontrol.Config. The existing single-detector field stays and remains the default fallback.

Validations:

• stages keys must belong to the known stage set. Today the set is {"prefill", "decode"}; the set lives in one place so a future stage is added by editing that set plus partitionEndpoints.
• Every stage reference must name an instantiated plugin that implements fwkfc.SaturationDetector.
• gateOn, when not empty, must be "max" or a known stage name.

Evaluation semantics

The processor keeps partitioning endpoints per stage. For each non-empty stage partition it selects stageDetectors[stage] when configured, otherwise the default detector, and records the stage metric as today. The effective
saturation is:

• max(prefill, decode) when gateOn is max (the default);
• the selected stage's saturation when gateOn names a stage.

Empty partitions keep their current behavior: the stage metric series is deleted and the stage is skipped. If gateOn names a stage and that stage partition is empty, the effective saturation is 1.0 (fail closed), because gating on a stage that has no capacity must not open the gate. If both partitions are empty, the effective saturation falls back to the default detector over the whole pool, as today.

Where this does not apply

• Legacy admission controller (pkg/epp/requestcontrol/admission.go):
  uses the default detector over all candidates, unchanged. When the flow
  control gate is disabled and stages is set, the loader warns that the
  per-stage settings are ignored (mirrors the existing warning for other
  flowControl settings).
• Scheduling filters: the concurrency-detector/utilization-detector
  filter path inside profiles is per-endpoint and stage-agnostic; it is not
  changed.

Metrics

No new metrics. flow_control_pool_saturation{stage="prefill"} and {stage="decode"} keep their series, each now sourced from its stage detector. flow_control_stale_endpoints continues to be written by whichever utilization detector is invoked last; per-stage stale accounting remains tracked in #2475.