Skip to content

feat(kv-router): replicate session bindings with one schema and one applier - #86

Draft
ishandhanani wants to merge 1 commit into
idhanani/session-affinity-unification-p4from
idhanani/session-affinity-unification-p5
Draft

ishandhanani wants to merge 1 commit into
idhanani/session-affinity-unification-p4from
idhanani/session-affinity-unification-p5

Conversation

@ishandhanani

Copy link
Copy Markdown
Owner

Summary

Fifth PR of the session-affinity unification stack (base: idhanani/session-affinity-unification-p4). Replicated session bindings get one wire schema and one applier, carried by both existing transports, with a one-release dual-publish window.

Before. Two formats and two appliers: the selection service's AffinityBindingEvent on its ZMQ mesh (with the partition; writer = random process id; applied if the worker is in the partition's catalog) and the frontend's SessionAffinityUpdate on the runtime event plane (no partition; writer = discovery instance id; applied if the worker is discovered). They could not interoperate, and the frontend payload could not scope a binding to a partition (a latent bug for multi-partition frontends).

After

lib/kv-router/src/services/selection/affinity/replication.rs:

  • AffinityBindingEvent: the one schema (partition flattened as model_name and routing_group, session_id, worker_id, dp_rank, sequence, writer_id).
  • AffinityEventSink: the one table sink, building the event for its partition and handing it to the host's transport closure.
  • AffinityResolver::apply_replica_event(partition, event): the one applier. Skips the replica's own writer id, skips other partitions, advances the replica clock, and applies the binding only if the host can schedule its worker (TargetLiveness). Returns a ReplicaEventDisposition.
  • One writer-id rule: the id installed with the table's replication sink. Frontends use their discovery instance id (stable per instance); the standalone service keeps a random non-zero process id, since it has no discovery.

Transports (unchanged in shape, one schema each):

Transport v2 v1 (one release)
Selection service ZMQ mesh topic dynamo.session-affinity.v2 topic dynamo.session-affinity.v1, same payload
Frontend runtime event plane subject session_affinity_events_v2, payload AffinityBindingEvent subject session_affinity_events, payload SessionAffinityUpdate (no partition); a receiver attributes it to its own partition, as every v1.6 frontend did

Both transports publish and apply both versions for one release so mixed-version replicas converge; applying the same binding twice is idempotent (the second apply refreshes the same version). Both v1 forms carry a TODO(v1.8) removal note. The frontend's HostAffinity records its partition for this: the embedded partition's key for KV routing, the default key for the builtin modes.

Behavior changes

Decisions with defaults needing sign-off:

  1. Dual publish on both transports for one release. This doubles affinity replication traffic (one event per binding change and per lease release) on the mesh and the event plane for that release. Alternative: publish v2 only and subscribe to both, which leaves older replicas blind to newer replicas' bindings during a rolling upgrade. I chose convergence over bandwidth; flagging it.
  2. Apply-time liveness is TargetLiveness. The selection service previously applied a binding if the worker was in the partition's catalog in any lifecycle state; it now requires the worker to be schedulable. The frontend KV host previously required the worker to be discovered; it now also requires the bound rank to be within the worker's known DP range (when configs are known). Both are stricter and only drop bindings the resolver would have dropped at the next request anyway.
  3. Partition scoping on the frontend. A v2 event for another partition is ignored; v1 events keep today's behavior (applied to the receiving partition).

Validation

Run with RUSTC_WRAPPER= and a per-worktree CARGO_TARGET_DIR.

Command Result
cargo test -p dynamo-kv-router --features standalone-selection --lib 1290 passed, 2 ignored (+3 replication tests)
cargo test -p dynamo-llm --lib kv_router 412 passed
cargo test -p dynamo-llm --lib session_affinity 50 passed (+1 dual-apply test)
cargo test -p dynamo-custom-policy-builtin -p dynamo-mocker passed
cargo test -p dynamo-ext-proc passed
cargo clippy -p dynamo-kv-router --features standalone-selection,bench --all-targets -- -D warnings clean
cargo clippy -p dynamo-llm -p dynamo-ext-proc -p dynamo-custom-policy-builtin -p dynamo-mocker --all-targets -- -D warnings clean
cargo fmt --all -- --check clean

New tests: schema_round_trips_with_the_partition_flattened (JSON and msgpack), applier_scopes_by_partition_writer_and_liveness, replicas_converge_through_the_shared_schema (two resolvers with their own liveness sources, standing in for a frontend and a selection-service replica, agree on a binding through the shared schema in both directions, including a re-bind after the worker departs), and on the frontend v1_and_v2_updates_both_apply_to_the_partition (v1 converted and v2 both applied, another partition ignored, duplicate apply idempotent, applier stops when the resolver is gone). rejected_worker_update_still_advances_replica_clock and replica_update_backpressure_is_nonfatal are kept over the shared types.

Not run: cross-process replication over a real ZMQ mesh or event plane (no multi-replica fixture in the unit suites); the convergence test is at the schema and applier level.

Related Issues

@ishandhanani
ishandhanani had a problem deploying to external_collaborator October 6, 2026 09:19 — with GitHub Actions Failure
@github-actions github-actions Bot added feat documentation Improvements or additions to documentation router labels Oct 6, 2026
@ishandhanani
ishandhanani added this pull request to stack #89 October 6, 2026 20:45
…pplier

Replicated session bindings had two wire formats and two appliers: the
selection service's `AffinityBindingEvent` on its ZMQ mesh (with the
partition, writer = random process id) and the frontend's
`SessionAffinityUpdate` on the runtime event plane (no partition, writer =
discovery instance id, applied against discovered instances). The two could
not interoperate and the frontend payload could not scope a binding to a
partition.

`services/selection/affinity/replication.rs` now holds the one schema,
`AffinityBindingEvent`, the one sink that produces it for a partition
(`AffinityEventSink`), and the one applier,
`AffinityResolver::apply_replica_event`: skip the replica's own writer id,
skip other partitions, advance the replica clock, apply only a binding whose
worker the host can schedule (`TargetLiveness`). The writer id is the one
installed with the table's replication sink: the discovery instance id where
a host has one, otherwise a random process id.

Both transports carry the shared schema on a new v2 topic or subject and
keep publishing and applying v1 for one release, so mixed-version replicas
converge: the mesh on `dynamo.session-affinity.v2` beside `.v1` (same
payload), the frontend on `session_affinity_events_v2` beside the
partition-less `session_affinity_events`, which a receiver attributes to
its own partition as before. The frontend host records its partition for
that: the embedded partition's key for KV routing, the default key for the
builtin modes. Applying a binding twice is idempotent.

Tests: schema round trip through JSON and msgpack with the partition
flattened; the applier's partition, writer, and liveness scoping; v1 and v2
updates both applied to one partition and an event for another partition
ignored; and a convergence test in which two replicas with their own
liveness sources agree on a binding through the shared schema in both
directions.

Signed-off-by: Ishan Dhanani <ishandhanani@gmail.com>
@ishandhanani
ishandhanani force-pushed the idhanani/session-affinity-unification-p5 branch from 676ff33 to 95b8018 Compare October 6, 2026 21:42
@ishandhanani
ishandhanani had a problem deploying to external_collaborator October 6, 2026 21:42 — with GitHub Actions Failure

This branch had an error being deployed

1 failed deployment
external_collaborator — 95b80181 Deployed Oct 6, 2026 by ishandhanani via ok-to-test #271
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

documentation Improvements or additions to documentation feat router

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant