Feat/provider load balancing

[nerdalert]

Summary

This PR adds local provider selection policies to the AI gateway's intelligent_route filter.

It is rebased on #386, which provides the authenticated provider-routing boundary used by this implementation.

  Client
    |
    v
  Consumer Gateway
    |
    v
  intelligent_route
    |
    +--> Provider Gateway A
    +--> Provider Gateway B
    `--> Provider Gateway C

Grid remains the asynchronous control plane. It determines provider eligibility, admission state, ordering, and selection-group membership, then publishes that information in the routing overlay.

The AI gateway consumes the accepted overlay from memory. When the overlay is loaded, it validates the group metadata and builds an immutable index for request-time selection.

The supported picker policies are:

• deterministic: select the first viable provider, preserving existing behavior.
• roundRobin: rotate evenly among viable providers in the active group.
• random: select uniformly among viable providers in the active group.

Session affinity is evaluated before the picker. A valid existing binding continues to use its provider when admission permits it. New, unbound requests use the configured picker.

The request path does not call Grid, Kubernetes, the operator, overlay-sync, ConfigMaps, EPP, Prometheus, or a remote scoring service.

Responsibility Boundary

This change keeps control-plane and request-time responsibilities separate:

  Grid
    -> discovers providers
    -> evaluates health, freshness, and admission
    -> applies routing and scoring policy
    -> assigns selection groups
    -> publishes an immutable overlay

  Praxis AI intelligent_route
    -> reads the accepted in-memory snapshot
    -> resolves session affinity
    -> identifies the best viable group
    -> applies the configured local picker
    -> selects a provider cluster

  Praxis core load_balancer
    -> selects an endpoint within that provider cluster

intelligent_route does not recompute scores, geography, health, or group membership.

Overlay Contract

The routing overlay may include an optional selection policy and producer-assigned groups:

  selection_policy:
    picker: roundRobin

  candidates:
    - cluster: provider-a
      selection_group: 0
    - cluster: provider-b
      selection_group: 0
    - cluster: provider-c
      selection_group: 0

Group 0 is preferred. If a group contains no candidate eligible for a new request, selection proceeds to the next group.

Provider-level retries after a forwarding failure are not introduced by this PR. Praxis core may retry endpoints within the selected cluster, but retrying a different Grid provider cluster remains follow-up work.

Compatibility

Existing overlays without selection_policy or selection_group retain the legacy deterministic, ordered-first behavior.

An upgraded deployment therefore does not begin balancing traffic unless its overlay explicitly enables a picker.

Unknown picker values and malformed group metadata fail validation. Invalid reloads retain the last-known-good snapshot, and in-flight requests continue using the immutable snapshot with which they started.

The picker is not inferred from:

• scoring policy;
• candidate score or rank;
• provider count;
• the presence of selection groups;
• queue-depth metrics;
• KV-cache metrics.

Scores and ranks do not implicitly become traffic weights. Weighted selection is a separate future extension.

Snapshot And Hot-Path Behavior

Selection-group indexes are constructed when a routing snapshot is loaded, not while processing requests.

Round-robin counters are:

• local to each gateway process;
• scoped to a capability and selection group;
• stored with the immutable routing snapshot;
• concurrency-safe;
• preserved when an overlay rewrite has the same semantic revision;
• replaced when a genuinely changed snapshot is accepted.

The request path performs bounded in-memory lookups and one local picker operation. It does not parse configuration, rebuild groups, scan unrelated capabilities, scrape metrics, or make control-plane calls.

This provides balanced behavior per gateway. It does not attempt to coordinate one global round-robin counter across a gateway fleet.

Security And Tenancy

The picker only selects among candidates already present in the accepted routing overlay and eligible under their published admission state. It does not discover providers, expand authorization, resolve credentials, or override Grid's trust and policy decisions.

Selection groups must not cross security boundaries unintentionally. Grid remains responsible for ensuring that candidates grouped together are valid for the same capability and applicable policy scope. Praxis AI validates the group structure but does not reinterpret tenancy, locality, credential, or provider trust rules.

Provider credentials and authorization data are not used as picker inputs or exposed through selection metadata. The authenticated provider-routing boundary remains owned by #386 (#386).

Operational Observability

Successful route decisions expose bounded in-process metadata for:

• picker policy;
• selection group;
• selected provider site, cluster, and stable ID;
• admission state;
• rank and selection tier when present;
• serving overlay revision;
• session-affinity creation, reuse, or failover state.

This evidence supports logs and tracing without adding a request-time
control-plane dependency.

Provider IDs, session identifiers, request IDs, and other high-cardinality values should not become unbounded metric labels. Aggregate metrics should use controlled dimensions. Rich per-request routing details are better represented as bounded trace attributes when OpenTelemetry support is enabled separately.

Multiple Consumer Gateways

This design also works when several consumer or edge gateways accept traffic. Each gateway loads the same Grid overlay and independently applies the configured picker:

                      +--> Consumer Gateway East --+
  Client / GTM / DNS -+                            +--> Providers A, B, C
                      +--> Consumer Gateway West --+

No consumer gateway coordinates with another gateway during request handling. Each process maintains its own round-robin counters and session-affinity state. This preserves the local, low-latency request path and avoids introducing a shared database or remote selection service.

With enough new requests, each gateway should independently distribute traffic across the active provider group. The combined distribution will generally approach the expected split, but a strict global sequence or exact fleet-wide ratio is not guaranteed.

Limitations include:

• Round-robin counters are process-local, not synchronized across gateways.
• A gateway restart or genuinely changed overlay snapshot resets that gateway's picker state.
• Uneven DNS, GTM, or load-balancer traffic across consumer gateways can produce uneven aggregate provider traffic.
• Small request samples may not produce an exact fleet-wide split.
• Session affinity intentionally overrides picker distribution for established sessions.
• Long-lived or unevenly sized requests can create unequal backend load even when request counts are equal.
• The picker balances provider selections, not tokens, latency, concurrency, or inference cost.
• Provider health and admission changes reach gateways asynchronously through new Grid overlays; AI does not probe providers on the request path.
• This PR does not add cross-gateway counter coordination, provider-level retries, metric-derived weights, or globally exact traffic allocation.
• External DNS, Anycast, or GTM remains responsible for distributing incoming clients across consumer gateways. Grid does not provide that entry-point balancing.

Functional Proof

The focused
provider-traffic demo (https://github.com/praxis-proxy/demos/tree/main/demos/grid-provider-traffic) exercises the intended topology:

  One consumer gateway
    |
    +--> Provider Gateway A --> backend A
    +--> Provider Gateway B --> backend B
    `--> Provider Gateway C --> backend C

The clean runtime validation demonstrated:

• one consumer gateway receiving the complete measured workload;
• three provider gateways advertising the same capability;
• all three candidates accepted in one active selection group;
• 60 serial requests without session-affinity headers;
• HTTP 200 for all 60 requests;
• exact 20/20/20 provider distribution;
• a repeating three-provider round-robin sequence;
• an unchanged serving overlay during the measurement window;
• no gateway restart during selection;
• complete environment teardown.

The demo proves equal provider selection from a stable snapshot. It does not claim to prove weighted routing, provider-level retry, lower-group runtime failover, or fleet-global round-robin coordination.

Other Grid demos cover separate concerns:

• llm-d pool metrics (https://github.com/praxis-proxy/demos/tree/main/demos/grid-llmd-pool-metrics) demonstrates asynchronous EPP metric scoring and rank changes.

• Grid GLB (https://github.com/praxis-proxy/demos/tree/main/demos/grid-glb-demo) demonstrates entry-point, edge, and provider routing across a larger topology.

OpenTelemetry instrumentation and the tracing UI are separate work and are not modified by this PR.

Dependency

This branch is based on the authenticated provider-routing work in #386 (#386).

The load-balancing implementation is isolated in one commit and should be rebased onto main after #386 merges.

Follow-Up Work

This PR establishes equal local selection within producer-defined groups.

Separate follow-up work includes:

• provider-level retry within the selected group;
• explicit lower-group fallback and recovery runtime proof;
• metric-derived traffic_weight publication by Grid;
• weighted request or new-session selection in Praxis AI;
• normalization across providers with different capacities;
• affinity behavior when weights or membership change;
• protection against noisy metrics and excessive overlay churn;
• throughput and latency benchmarks;
• multi-consumer runtime validation;
• bounded aggregate picker metrics;
• globally coordinated selection, if a concrete requirement emerges.

Weighted routing must remain explicit. Scores, ranks, queue depth, and KV-cache pressure should not silently become traffic percentages.

Related Issue

Related to
praxis-proxy/grid#31.

The research spike includes the focused provider-traffic demo narrative, runtime evidence, and a video showing one consumer gateway distributing requests across three provider gateways.

This PR implements the AI request-time selection portion of the spike. It does not close the Grid-side overlay producer work or the remaining retry, fallback, and weighted-selection research.

Validation

• Unit tests
• Integration or functional tests
• make lint

Validated with:

cargo +nightly fmt --all -- --check
cargo test -p praxis-ai-filters
cargo test -p praxis-ai-filters --no-default-features
cargo test -p praxis-ai-filters --all-features
cargo clippy --workspace --all-targets -- -D warnings
make lint
make test
RUSTDOCFLAGS="-D warnings" cargo doc --workspace --no-deps
git diff --check

Focused coverage includes:

• deterministic selection;
• random selection with deterministic ordinal coverage;
• round-robin distribution;
• concurrent round-robin selection;
• counter wraparound;
• counter isolation across capabilities;
• multiple and interleaved capabilities;
• contiguous group validation;
• mixed grouped and ungrouped rejection;
• admission-aware group fallback;
• legacy overlay compatibility;
• semantically unchanged reloads;
• changed-snapshot picker state;
• duplicate candidates as explicit independent selection slots.

Checklist

• I reviewed every changed line and can explain the change.
• New capabilities include an example configuration and functional demo.
• User-facing behavior and generated documentation are updated.
• Performance-sensitive changes include benchmark or load-test evidence.
• Commits are signed and include a Signed-off-by trailer.

The runtime demo validates functional distribution. It is not presented as a throughput or latency benchmark.

Breaking Changes

None intended.

Existing overlays retain deterministic behavior. New deployments can opt into roundRobin or random through an explicit selection policy. This PR does not change scoring semantics, add weighted routing, or modify OpenTelemetry instrumentation.

/cc @jordigilh