From 3ea7b0aba69745887920d3835a51a5ff136f7bb4 Mon Sep 17 00:00:00 2001 From: thxCode Date: Wed, 14 Oct 2026 16:09:16 +0800 Subject: [PATCH 1/2] build(llm-router): bump llm-d-router to v0.11.0 - rebase the observer patch onto v0.11.0 and keep its seven revision sites - allow the passthrough parser that v0.11.0 adds to the default set - add a lock-order test that fails on the observer and datastore lock inversion Signed-off-by: thxCode --- pack/llm-router/Dockerfile | 2 +- ...0001-observer-endpoints-and-dispatch.patch | 358 ++++++++++++++---- 2 files changed, 284 insertions(+), 76 deletions(-) diff --git a/pack/llm-router/Dockerfile b/pack/llm-router/Dockerfile index f854e98a9..8ae0eb2fc 100644 --- a/pack/llm-router/Dockerfile +++ b/pack/llm-router/Dockerfile @@ -29,7 +29,7 @@ ARG UBUNTU_VERSION=24.04 ARG GO_VERSION=1.26-bookworm ARG RUST_VERSION=1.95-bookworm -ARG LLMD_ROUTER_VERSION=v0.10.0 +ARG LLMD_ROUTER_VERSION=v0.11.0 ARG VLLM_ROUTER_VERSION=v0.1.15 ARG SGLANG_GATEWAY_VERSION=gateway-v0.3.1 diff --git a/pack/llm-router/patches/llm-d/0001-observer-endpoints-and-dispatch.patch b/pack/llm-router/patches/llm-d/0001-observer-endpoints-and-dispatch.patch index 0656606d8..ca71d2bf0 100644 --- a/pack/llm-router/patches/llm-d/0001-observer-endpoints-and-dispatch.patch +++ b/pack/llm-router/patches/llm-d/0001-observer-endpoints-and-dispatch.patch @@ -3,12 +3,51 @@ The view refuses pending selection/send windows, unsynchronized membership and unsupported native profiles. GET does not mutate routing or expire native state. Removed dispatched listeners retain false selection flags with recent-dispatch evidence. -Written against llm-d-router v0.10.0, commit 71f4f0999f95b96c49a9d0c4afbd18dfdb943c26. -The operator collector must consume recent-dispatch evidence before using this view. +Rebased from v0.10.0 (commit 71f4f0999f95b96c49a9d0c4afbd18dfdb943c26) to v0.11.0. + +The constructor-based ExtProcServerRunner in v0.11.0 surfaces every field through +NewExtProcServerRunner, and TestNewExtProcServerRunnerPopulatesEveryField fails +when any new field is zero. Observer is wired by post-constructor assignment in +runner.go and is added to the test's callerSet (the smaller change; the +alternative is adding an Observer parameter to NewExtProcServerRunner and +threading it through every call site). + +v0.11.0 added passthrough-parser to the default parser set +(pkg/epp/config/loader/defaults.go:261) and the operator's picker configs +declare no parsers, so the parser is always registered. The observer's +SupportedPlugins allowlist gains "passthrough-parser" so the operator's +profiles do not hit a 503. The flow-control defaults (fcfs-ordering-policy, +global-strict-fairness-policy, static-usage-limit-policy-type) and the +utilization-detector saturation detector were already in the allowlist, so +no further additions were needed. + +The patch carries the seven datastore observer-revision increments: the +public mutation sites WithEndpointPool, Clear, PoolSet, +PodUpdateOrAddIfNotExist (the exported wrapper), PodDelete, +EndpointUpsert, EndpointDelete. The unexported podUpdateOrAddIfNotExist +helper takes no observerMu acquisition because both its callers +(PodUpdateOrAddIfNotExist and PoolSet via podResyncAll) already increment, +and an observerMu acquisition here would invert the lock order against +ObserverSnapshot (which takes observerMu then ds.mu) into AB-BA deadlock. +pkg/epp/datastore/observer_lock_order_test.go pins this invariant. + +The five short-block sites release observerMu before mutating ds.pods, so a +snapshot racing between bump and mutation can see a revision that leads the +data by one change; no production code reads registry_revision. + +The director -> dispatch-evidence leg has a v0.11.0-compatible test +pkg/epp/requestcontrol/observer_test.go that builds the director with +NewDirectorWithConfig + a canned mockScheduler (mockScheduler returns the +pre-baked SchedulingResult) and verifies every ProfileResults TargetEndpoint +is recorded into the observer window and shows up in the view with +recent_dispatch. + +The operator collector must consume recent-dispatch evidence before using +this view. diff --git a/cmd/epp/runner/observer_test.go b/cmd/epp/runner/observer_test.go new file mode 100644 -index 0000000..7188e6d +index 00000000..7188e6d1 --- /dev/null +++ b/cmd/epp/runner/observer_test.go @@ -0,0 +1,42 @@ @@ -56,7 +95,7 @@ index 0000000..7188e6d +} diff --git a/cmd/epp/runner/observer_testdata/pd.yaml b/cmd/epp/runner/observer_testdata/pd.yaml new file mode 100644 -index 0000000..e18abb1 +index 00000000..e18abb13 --- /dev/null +++ b/cmd/epp/runner/observer_testdata/pd.yaml @@ -0,0 +1,57 @@ @@ -119,7 +158,7 @@ index 0000000..e18abb1 + weight: 3 diff --git a/cmd/epp/runner/observer_testdata/server.yaml b/cmd/epp/runner/observer_testdata/server.yaml new file mode 100644 -index 0000000..532453d +index 00000000..532453de --- /dev/null +++ b/cmd/epp/runner/observer_testdata/server.yaml @@ -0,0 +1,33 @@ @@ -157,10 +196,10 @@ index 0000000..532453d + - pluginRef: prefix-cache-scorer + weight: 3 diff --git a/cmd/epp/runner/runner.go b/cmd/epp/runner/runner.go -index 21f443d..680087a 100644 +index 27080fa4..a660b966 100644 --- a/cmd/epp/runner/runner.go +++ b/cmd/epp/runner/runner.go -@@ -146,6 +146,7 @@ import ( +@@ -154,6 +154,7 @@ import ( "github.com/llm-d/llm-d-router/pkg/epp/handlers" "github.com/llm-d/llm-d-router/pkg/epp/metrics" "github.com/llm-d/llm-d-router/pkg/epp/metrics/collectors" @@ -168,8 +207,8 @@ index 21f443d..680087a 100644 "github.com/llm-d/llm-d-router/pkg/epp/requestcontrol" "github.com/llm-d/llm-d-router/pkg/epp/scheduling" runserver "github.com/llm-d/llm-d-router/pkg/epp/server" -@@ -468,7 +469,15 @@ func (r *Runner) setup(ctx context.Context, cfg *rest.Config, opts *runserver.Op - director.SetRequestEvictor(requestEvictor) +@@ -487,6 +488,15 @@ func (r *Runner) setup(ctx context.Context, cfg *rest.Config, opts *runserver.Op + serverRunner.EvictChannelLookup = requestEvictor.EvictionRegistry() } + store, _ := ds.(observer.Snapshotter) @@ -179,16 +218,16 @@ index 21f443d..680087a 100644 + if err := mgr.AddMetricsServerExtraHandler(observer.Path, nativeObserver); err != nil { + return nil, nil, err + } - serverRunner := &runserver.ExtProcServerRunner{ -+ Observer: nativeObserver, - GrpcPort: opts.GRPCPort, - GKNN: *gknn, - Datastore: ds, ++ serverRunner.Observer = nativeObserver ++ + if err := serverRunner.SetupWithManager(mgr); err != nil { + setupLog.Error(err, "Failed to setup EPP controllers") + return nil, nil, err diff --git a/pkg/epp/datastore/datastore.go b/pkg/epp/datastore/datastore.go -index 59cc507..ccb0645 100644 +index 98999dce..faf4c849 100644 --- a/pkg/epp/datastore/datastore.go +++ b/pkg/epp/datastore/datastore.go -@@ -127,6 +127,8 @@ func NewDatastore(parentCtx context.Context, epFactory datalayer.EndpointFactory +@@ -128,6 +128,8 @@ func NewDatastore(parentCtx context.Context, epFactory datalayer.EndpointFactory } type datastore struct { @@ -197,7 +236,7 @@ index 59cc507..ccb0645 100644 // parentCtx controls the lifecycle of the background metrics goroutines that spawn up by the datastore. parentCtx context.Context // mu is used to synchronize access to pool, objectives, and rewrites. -@@ -147,11 +149,19 @@ type datastore struct { +@@ -148,11 +150,19 @@ type datastore struct { } func (ds *datastore) WithEndpointPool(pool *datalayer.EndpointPool) Datastore { @@ -217,43 +256,43 @@ index 59cc507..ccb0645 100644 ds.mu.Lock() defer ds.mu.Unlock() ds.pool = nil -@@ -171,6 +181,9 @@ func (ds *datastore) PoolSet(ctx context.Context, reader client.Reader, endpoint +@@ -172,6 +182,9 @@ func (ds *datastore) PoolSet(ctx context.Context, reader client.Reader, endpoint ds.Clear() return nil } + ds.observerMu.Lock() -+ defer ds.observerMu.Unlock() + ds.observerRevision++ ++ ds.observerMu.Unlock() logger := log.FromContext(ctx) ds.mu.Lock() defer ds.mu.Unlock() -@@ -299,6 +312,9 @@ func (ds *datastore) PodList(predicate func(fwkdl.Endpoint) bool) []fwkdl.Endpoi +@@ -300,6 +313,9 @@ func (ds *datastore) PodList(predicate func(fwkdl.Endpoint) bool) []fwkdl.Endpoi } func (ds *datastore) PodUpdateOrAddIfNotExist(ctx context.Context, pod *corev1.Pod) error { + ds.observerMu.Lock() -+ defer ds.observerMu.Unlock() + ds.observerRevision++ ++ ds.observerMu.Unlock() // Take a reference to pool under read lock to avoid racing with PoolSet(). // This is safe because PoolSet() replaces the entire pool struct rather than // updating it in-place. -@@ -392,6 +408,9 @@ func (ds *datastore) podUpdateOrAddIfNotExist(ctx context.Context, pod *corev1.P +@@ -393,6 +409,9 @@ func (ds *datastore) podUpdateOrAddIfNotExist(ctx context.Context, pod *corev1.P } func (ds *datastore) PodDelete(podName string) { + ds.observerMu.Lock() -+ defer ds.observerMu.Unlock() + ds.observerRevision++ ++ ds.observerMu.Unlock() ds.pods.Range(func(k, v any) bool { ep := v.(fwkdl.Endpoint) if ep.GetMetadata().Name == podName { -@@ -403,12 +422,18 @@ func (ds *datastore) PodDelete(podName string) { +@@ -404,12 +423,18 @@ func (ds *datastore) PodDelete(podName string) { } func (ds *datastore) EndpointUpsert(ctx context.Context, meta *fwkdl.EndpointMetadata) { + ds.observerMu.Lock() -+ defer ds.observerMu.Unlock() + ds.observerRevision++ ++ ds.observerMu.Unlock() if _, err := ds.upsertEndpoint(ctx, meta); err != nil { log.FromContext(ctx).Error(err, "failed to register endpoint", "endpoint", meta.ID) } @@ -261,12 +300,12 @@ index 59cc507..ccb0645 100644 func (ds *datastore) EndpointDelete(id types.NamespacedName) { + ds.observerMu.Lock() -+ defer ds.observerMu.Unlock() + ds.observerRevision++ ++ ds.observerMu.Unlock() if v, ok := ds.pods.LoadAndDelete(id); ok { ds.epf.ReleaseEndpoint(v.(fwkdl.Endpoint)) } -@@ -534,3 +559,22 @@ func selectorEqual(a, b labels.Selector) bool { +@@ -535,3 +560,22 @@ func selectorEqual(a, b labels.Selector) bool { } return a.String() == b.String() } @@ -289,9 +328,134 @@ index 59cc507..ccb0645 100644 + }) + return result, ds.observerRevision, nil +} +diff --git a/pkg/epp/datastore/observer_lock_order_test.go b/pkg/epp/datastore/observer_lock_order_test.go +new file mode 100644 +index 00000000..9e676a90 +--- /dev/null ++++ b/pkg/epp/datastore/observer_lock_order_test.go +@@ -0,0 +1,119 @@ ++// Copyright 2026 The Kubernetes Authors. ++// SPDX-License-Identifier: Apache-2.0 ++package datastore ++ ++import ( ++ "context" ++ "sync" ++ "testing" ++ "time" ++ ++ "github.com/llm-d/llm-d-router/pkg/epp/datalayer" ++ corev1 "k8s.io/api/core/v1" ++ metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" ++ "k8s.io/apimachinery/pkg/labels" ++ "sigs.k8s.io/controller-runtime/pkg/client" ++ "sigs.k8s.io/controller-runtime/pkg/client/fake" ++ "sigs.k8s.io/controller-runtime/pkg/client/interceptor" ++) ++ ++func selectorFromLabels(s string) labels.Selector { ++ sel, _ := labels.Parse(s) ++ return sel ++} ++ ++// TestObserverSnapshotVsPoolSet_NoDeadlock pins the lock-order invariant between ++// ObserverSnapshot (observerMu.RLock then ds.mu.RLock) and PoolSet/podResyncAll/ ++// podUpdateOrAddIfNotExist (ds.mu.Lock then no further observerMu acquisition). ++// ++// ObserverSnapshot's order: observerMu -> ds.mu. ++// PoolSet's order: ds.mu -> podUpdateOrAddIfNotExist (which must NOT touch ++// observerMu while ds.mu is held by the caller). ++// ++// Regression: any future change that re-introduces an observerMu acquisition ++// inside podUpdateOrAddIfNotExist will deadlock the test within ~5s. The ++// signal is the two goroutines both blocked on a mutex at the same time, ++// observed here by gating PoolSet's podResyncAll.List call on a channel so ++// the goroutine is parked inside ds.mu while ObserverSnapshot takes observerMu. ++func TestObserverSnapshotVsPoolSet_NoDeadlock(t *testing.T) { ++ const deadline = 5 * time.Second ++ ++ releaseList := make(chan struct{}) ++ listEntered := make(chan struct{}) ++ var enteredOnce sync.Once ++ ++ readyPod := &corev1.Pod{ ++ ObjectMeta: metav1.ObjectMeta{ ++ Name: "pod-0", ++ Namespace: "default", ++ Labels: map[string]string{"app": "model"}, ++ }, ++ Status: corev1.PodStatus{ ++ Conditions: []corev1.PodCondition{{Type: corev1.PodReady, Status: corev1.ConditionTrue}}, ++ }, ++ } ++ fakeClient := fake.NewClientBuilder(). ++ WithObjects(readyPod). ++ WithInterceptorFuncs(interceptor.Funcs{ ++ List: func(ctx context.Context, c client.WithWatch, list client.ObjectList, opts ...client.ListOption) error { ++ enteredOnce.Do(func() { close(listEntered) }) ++ <-releaseList ++ return c.List(ctx, list, opts...) ++ }, ++ }). ++ Build() ++ ++ ctx, cancel := context.WithCancel(context.Background()) ++ defer cancel() ++ ++ ds := NewDatastore(ctx, &mockEndpointFactory{}).(*datastore) ++ pool := &datalayer.EndpointPool{ ++ TargetPorts: []int{8000}, ++ Selector: selectorFromLabels("app=model"), ++ } ++ ++ // Goroutine A: run PoolSet so it parks inside podResyncAll.List while ++ // holding ds.mu. ++ poolSetDone := make(chan error, 1) ++ go func() { ++ poolSetDone <- ds.PoolSet(ctx, fakeClient, pool) ++ }() ++ ++ // Wait until A is parked in the List call. ++ select { ++ case <-listEntered: ++ case <-time.After(deadline): ++ t.Fatal("PoolSet never reached the List call") ++ } ++ ++ // Goroutine B: take observerMu via ObserverSnapshot and block on ds.mu.RLock. ++ snapshotDone := make(chan struct{}) ++ go func() { ++ defer close(snapshotDone) ++ for i := 0; i < 50; i++ { ++ _, _, _ = ds.ObserverSnapshot() ++ } ++ }() ++ ++ // Give B time to acquire observerMu.TryRLock and block on ds.mu.RLock. ++ // 100ms is well below the 5s deadlock budget. ++ time.Sleep(100 * time.Millisecond) ++ ++ // Release A's List call. With the bug, A will then enter ++ // podUpdateOrAddIfNotExist, try to Lock observerMu while holding ds.mu, ++ // and wedge against B's RLock. Without the bug, A finishes and B ++ // eventually drains. ++ close(releaseList) ++ ++ select { ++ case err := <-poolSetDone: ++ if err != nil { ++ t.Fatalf("PoolSet failed: %v", err) ++ } ++ case <-snapshotDone: ++ case <-time.After(deadline): ++ t.Fatalf("deadlock: neither PoolSet nor ObserverSnapshot returned within %v; "+ ++ "observerMu and ds.mu are inverted across the snapshot/resync boundary", ++ deadline) ++ } ++} diff --git a/pkg/epp/datastore/observer_test.go b/pkg/epp/datastore/observer_test.go new file mode 100644 -index 0000000..3949725 +index 00000000..3949725f --- /dev/null +++ b/pkg/epp/datastore/observer_test.go @@ -0,0 +1,75 @@ @@ -372,7 +536,7 @@ index 0000000..3949725 +} diff --git a/pkg/epp/handlers/observer_test.go b/pkg/epp/handlers/observer_test.go new file mode 100644 -index 0000000..9d0297d +index 00000000..9d0297d0 --- /dev/null +++ b/pkg/epp/handlers/observer_test.go @@ -0,0 +1,135 @@ @@ -512,10 +676,10 @@ index 0000000..9d0297d + } +} diff --git a/pkg/epp/handlers/request.go b/pkg/epp/handlers/request.go -index 574cae7..679bced 100644 +index 90e13a1e..2cf34af6 100644 --- a/pkg/epp/handlers/request.go +++ b/pkg/epp/handlers/request.go -@@ -58,10 +58,13 @@ func (s *StreamingServer) HandleRequestHeaders(ctx context.Context, reqCtx *Requ +@@ -60,10 +60,13 @@ func (s *StreamingServer) HandleRequestHeaders(ctx context.Context, reqCtx *Requ } func (s *StreamingServer) fallbackToRandomEndpoint(ctx context.Context, reqCtx *RequestContext, requestSize int) error { @@ -526,14 +690,14 @@ index 574cae7..679bced 100644 return errcommon.Error{Code: errcommon.Internal, Msg: "no pods available in datastore"} } + reqCtx.ObserverWindow.Record(endpoint) - reqCtx.TargetEndpoint = endpoint.GetIPAddress() + ":" + endpoint.GetPort() + reqCtx.TargetEndpoint = net.JoinHostPort(endpoint.GetIPAddress(), endpoint.GetPort()) reqCtx.RequestSize = requestSize reqCtx.reqHeaderResp = s.generateRequestHeaderResponse(ctx, reqCtx) diff --git a/pkg/epp/handlers/server.go b/pkg/epp/handlers/server.go -index 1913ee8..40ba92d 100644 +index 5bcb9bdf..0587df65 100644 --- a/pkg/epp/handlers/server.go +++ b/pkg/epp/handlers/server.go -@@ -52,6 +52,7 @@ import ( +@@ -54,6 +54,7 @@ import ( fwksched "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/scheduling" "github.com/llm-d/llm-d-router/pkg/epp/metadata" "github.com/llm-d/llm-d-router/pkg/epp/metrics" @@ -541,7 +705,17 @@ index 1913ee8..40ba92d 100644 ) // EvictChannelLookup is an optional interface for looking up eviction channels by request ID. -@@ -102,6 +103,7 @@ type Datastore interface { +@@ -90,6 +91,9 @@ func (s *StreamingServer) SetEmitEndpointScores(enabled bool) { + s.emitEndpointScores = enabled + } + ++// SetObserver attaches the admin observer to the native selection-to-send lifecycle. ++func (s *StreamingServer) SetObserver(o *observer.Observer) { s.observer = o } ++ + type Director interface { + HandleRequest(ctx context.Context, reqCtx *RequestContext, inferenceRequestBody *fwkrh.InferenceRequestBody) (*RequestContext, error) + HandleResponseHeader(ctx context.Context, reqCtx *RequestContext) *RequestContext +@@ -104,6 +108,7 @@ type Datastore interface { // Server implements the Envoy external processing server. // https://www.envoyproxy.io/docs/envoy/latest/api-v3/service/ext_proc/v3/external_processor.proto type StreamingServer struct { @@ -549,9 +723,9 @@ index 1913ee8..40ba92d 100644 datastore Datastore director Director parserRegistry *ParserRegistry -@@ -119,8 +121,10 @@ type StreamingServer struct { - // Refactor this monolithic struct. Fields related to the Envoy ext-proc protocol should be decoupled from the internal - // request lifecycle state. +@@ -120,8 +125,10 @@ type StreamingServer struct { + // Exported fields are read and written by the request-control layers (director and + // admission control). Unexported fields are private to this package. type RequestContext struct { - TargetPod *fwkdl.EndpointMetadata - TargetEndpoint string @@ -562,16 +736,16 @@ index 1913ee8..40ba92d 100644 // TargetEndpointScores maps endpoint address to the scheduler's score for it, covering // every endpoint the primary profile scored rather than only those in TargetEndpoint. // Nil when the primary profile ran no scorers. -@@ -286,6 +290,8 @@ func (s *StreamingServer) Process(srv extProcPb.ExternalProcessor_ProcessServer) - }, - } +@@ -362,6 +369,8 @@ func (s *StreamingServer) Process(srv extProcPb.ExternalProcessor_ProcessServer) + var respBody []byte + var evictionRequestID string + defer func() { reqCtx.ObserverWindow.Finish() }() + - // Request-phase failures (parser resolution, body parsing, admission - // rejection) leave the switch before the success path, so both call this. - // Flow-control rejections carry the queue wait and are the slowest samples; -@@ -473,6 +479,7 @@ func (s *StreamingServer) Process(srv extProcPb.ExternalProcessor_ProcessServer) + // Start a single reader goroutine for the lifetime of the stream. + // This avoids spawning a new goroutine per message and allows the main loop to + // select on both incoming messages and the eviction channel. +@@ -524,6 +533,7 @@ func (s *StreamingServer) Process(srv extProcPb.ExternalProcessor_ProcessServer) break } @@ -579,26 +753,19 @@ index 1913ee8..40ba92d 100644 reqCtx, err = s.director.HandleRequest(ctx, reqCtx, parseResult.Body) if err != nil { logger.Error(err, "Error handling request") -@@ -583,6 +590,9 @@ func (s *StreamingServer) Process(srv extProcPb.ExternalProcessor_ProcessServer) +@@ -646,6 +656,9 @@ func (s *StreamingServer) Process(srv extProcPb.ExternalProcessor_ProcessServer) if err := reqCtx.updateStateAndSendIfNeeded(srv, logger); err != nil { return err } -+ if reqCtx.RequestState == BodyRequestResponsesComplete || reqCtx.RequestState == RequestResponseProcessingSkipped || (reqCtx.observerHeaderOnly && reqCtx.RequestState == HeaderRequestResponseComplete) { ++ if reqCtx.requestState == bodyRequestResponsesComplete || reqCtx.requestState == requestResponseProcessingSkipped || (reqCtx.observerHeaderOnly && reqCtx.requestState == headerRequestResponseComplete) { + reqCtx.ObserverWindow.Finish() + } - if reqCtx.RequestState == RequestResponseProcessingSkipped { + if reqCtx.requestState == requestResponseProcessingSkipped { logger.V(logutil.DEFAULT).Info("EPP skipped response interception, routed request", "targetEndpoint", reqCtx.TargetEndpoint, -@@ -756,3 +766,6 @@ func (r *RequestContext) updateStateAndSendIfNeeded(srv extProcPb.ExternalProces - } - return nil - } -+ -+// SetObserver attaches the admin observer to the native selection-to-send lifecycle. -+func (s *StreamingServer) SetObserver(o *observer.Observer) { s.observer = o } diff --git a/pkg/epp/observer/endpoints.go b/pkg/epp/observer/endpoints.go new file mode 100644 -index 0000000..14f6a62 +index 00000000..b6db1667 --- /dev/null +++ b/pkg/epp/observer/endpoints.go @@ -0,0 +1,261 @@ @@ -671,7 +838,7 @@ index 0000000..14f6a62 + "core-metrics-extractor", "approx-prefix-cache-producer", + "precise-prefix-cache-producer", "token-producer", "endpoint-notification-source", + "prefix-based-pd-decider", "disagg-profile-handler", "single-profile-handler", -+ "openai-parser", "anthropic-parser", "vllmhttp-parser", ++ "openai-parser", "anthropic-parser", "vllmhttp-parser", "passthrough-parser", + "fcfs-ordering-policy", "global-strict-fairness-policy", + "static-usage-limit-policy-type", "utilization-detector": + default: @@ -865,7 +1032,7 @@ index 0000000..14f6a62 +} diff --git a/pkg/epp/observer/endpoints_test.go b/pkg/epp/observer/endpoints_test.go new file mode 100644 -index 0000000..45d030e +index 00000000..45d030eb --- /dev/null +++ b/pkg/epp/observer/endpoints_test.go @@ -0,0 +1,165 @@ @@ -1035,10 +1202,10 @@ index 0000000..45d030e + } +} diff --git a/pkg/epp/requestcontrol/director.go b/pkg/epp/requestcontrol/director.go -index d6ea634..e70e16a 100644 +index e4971a05..4dfc6997 100644 --- a/pkg/epp/requestcontrol/director.go +++ b/pkg/epp/requestcontrol/director.go -@@ -362,6 +362,15 @@ func (d *Director) HandleRequest(ctx context.Context, reqCtx *handlers.RequestCo +@@ -300,6 +300,15 @@ func (d *Director) HandleRequest(ctx context.Context, reqCtx *handlers.RequestCo } reqCtx.SchedulingRequest.SchedulingResult = result @@ -1056,32 +1223,41 @@ index d6ea634..e70e16a 100644 // Insert target endpoint to instruct Envoy to route requests to the specified target pod and attach the port number. diff --git a/pkg/epp/requestcontrol/observer_test.go b/pkg/epp/requestcontrol/observer_test.go new file mode 100644 -index 0000000..a0c12f0 +index 00000000..b6a67887 --- /dev/null +++ b/pkg/epp/requestcontrol/observer_test.go -@@ -0,0 +1,68 @@ +@@ -0,0 +1,96 @@ +// Copyright 2026 The Kubernetes Authors. +// SPDX-License-Identifier: Apache-2.0 +package requestcontrol + +import ( + "encoding/json" ++ "net/http" + "net/http/httptest" + "testing" ++ "time" + + fwkdl "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/datalayer" ++ "github.com/llm-d/llm-d-router/pkg/epp/datalayer" + fwksched "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/scheduling" + "github.com/llm-d/llm-d-router/pkg/epp/framework/plugins/requesthandling/parsers/openai" ++ "github.com/llm-d/llm-d-router/pkg/epp/datastore" + "github.com/llm-d/llm-d-router/pkg/epp/handlers" + "github.com/llm-d/llm-d-router/pkg/epp/observer" + "k8s.io/apimachinery/pkg/types" +) + -+type emptyObservedStore struct{} ++type observedStore struct{} + -+func (emptyObservedStore) ObserverSnapshot() ([]*fwkdl.EndpointMetadata, uint64, error) { ++func (observedStore) ObserverSnapshot() ([]*fwkdl.EndpointMetadata, uint64, error) { + return nil, 1, nil +} ++ ++// TestObserverRecordsNativeDirectorAuxiliaryTargets asserts that Director.HandleRequest ++// records every ProfileRunResult TargetEndpoint into the observer window (primary + ++// auxiliary profiles), and that the observer view then carries recent_dispatch for ++// each recorded target. +func TestObserverRecordsNativeDirectorAuxiliaryTargets(t *testing.T) { + makeEndpoint := func(name, address, role string) fwksched.Endpoint { + return fwksched.NewEndpoint(&fwkdl.EndpointMetadata{ID: types.NamespacedName{Name: name, Namespace: "ns"}, Name: name, Address: address, Port: "8000", Labels: map[string]string{"app.kubernetes.io/component": role}}, nil, fwkdl.NewAttributes()) @@ -1090,10 +1266,25 @@ index 0000000..a0c12f0 + "decode": {TargetEndpoints: []fwksched.Endpoint{makeEndpoint("decoder", "192.0.2.1", "decode-role")}}, + "prefill": {TargetEndpoints: []fwksched.Endpoint{makeEndpoint("prefiller", "192.0.2.2", "prefill-role")}}, + }} -+ director, ctx := newConditionalDecodeDirector(t, result) -+ o := observer.New(emptyObservedStore{}, func() bool { return true }) -+ req := &handlers.RequestContext{ObserverWindow: o.Begin(), Request: &handlers.Request{Headers: map[string]string{":path": "/v1/completions"}, RawBody: []byte(`{"model":"m","prompt":"p"}`)}, Parser: openai.NewOpenAIParser()} ++ mockSched := &mockScheduler{scheduleResults: result} ++ ++ ctx := t.Context() ++ epf := datalayer.NewTestRuntime(t, time.Minute) ++ ds := datastore.NewDatastore(ctx, epf) ++ ds.WithEndpointPool(&datalayer.EndpointPool{TargetPorts: []int{8000}}) ++ ds.EndpointUpsert(ctx, &fwkdl.EndpointMetadata{ID: types.NamespacedName{Name: "decoder", Namespace: "ns"}, Name: "decoder", Address: "192.0.2.1", Port: "8000", Labels: map[string]string{"app.kubernetes.io/component": "decode-role"}}) ++ ds.EndpointUpsert(ctx, &fwkdl.EndpointMetadata{ID: types.NamespacedName{Name: "prefiller", Namespace: "ns"}, Name: "prefiller", Address: "192.0.2.2", Port: "8000", Labels: map[string]string{"app.kubernetes.io/component": "prefill-role"}}) ++ endpointCandidates := NewCachedEndpointCandidates(ctx, NewDatastoreEndpointCandidates(ds), time.Minute) ++ director := NewDirectorWithConfig(ds, mockSched, &mockAdmissionController{}, endpointCandidates, NewConfig()) ++ ++ o := observer.New(observedStore{}, func() bool { return true }) ++ req := &handlers.RequestContext{ ++ ObserverWindow: o.Begin(), ++ Request: &handlers.Request{Headers: map[string]string{":path": "/v1/completions"}, RawBody: []byte(`{"model":"m","prompt":"p"}`)}, ++ Parser: openai.NewOpenAIParser(), ++ } + defer req.ObserverWindow.Finish() ++ + parsed, err := req.Parser.ParseRequest(ctx, req.Request.RawBody, req.Request.Headers) + if err != nil { + t.Fatal(err) @@ -1106,11 +1297,15 @@ index 0000000..a0c12f0 + t.Fatal("native primary routing was changed") + } + req.ObserverWindow.Finish() ++ + recorder := httptest.NewRecorder() -+ o.ServeHTTP(recorder, httptest.NewRequest("GET", observer.Path, nil)) ++ o.ServeHTTP(recorder, httptest.NewRequest(http.MethodGet, observer.Path, nil)) ++ if recorder.Code != http.StatusOK { ++ t.Fatalf("observer refused: %d %s", recorder.Code, recorder.Body.String()) ++ } + var observed struct { + Workers []struct { -+ Pod string `json:"pod_hint"` ++ PodHint string `json:"pod_hint"` + LastJob struct { + Recent bool `json:"recent_dispatch"` + } `json:"last_job"` @@ -1119,20 +1314,20 @@ index 0000000..a0c12f0 + if err := json.Unmarshal(recorder.Body.Bytes(), &observed); err != nil { + t.Fatal(err) + } -+ if recorder.Code != 200 || len(observed.Workers) != 2 { -+ t.Fatalf("native primary/auxiliary target lost: %d %s", recorder.Code, recorder.Body.String()) ++ if len(observed.Workers) != 2 { ++ t.Fatalf("native primary/auxiliary target lost: %s", recorder.Body.String()) + } + for _, worker := range observed.Workers { + if !worker.LastJob.Recent { -+ t.Fatal("native target has no dispatch evidence") ++ t.Fatalf("native target %q has no dispatch evidence", worker.PodHint) + } + } +} diff --git a/pkg/epp/server/runserver.go b/pkg/epp/server/runserver.go -index 0982639..ab0cedb 100644 +index a8a63640..635d7a98 100644 --- a/pkg/epp/server/runserver.go +++ b/pkg/epp/server/runserver.go -@@ -45,11 +45,13 @@ import ( +@@ -46,11 +46,13 @@ import ( fwkfc "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/flowcontrol" "github.com/llm-d/llm-d-router/pkg/epp/handlers" "github.com/llm-d/llm-d-router/pkg/epp/metrics" @@ -1146,7 +1341,7 @@ index 0982639..ab0cedb 100644 GrpcPort int // GrpcListener is an optional pre-bound listener for the ext_proc server. // When set, GrpcPort is ignored. Reserving the port in advance of this -@@ -231,6 +233,7 @@ func (r *ExtProcServerRunner) AsRunnable(logger logr.Logger) manager.Runnable { +@@ -263,6 +265,7 @@ func (r *ExtProcServerRunner) AsRunnable(logger logr.Logger) manager.Runnable { poolCap = 4 * 1024 * 1024 // gRPC default 4MB } extProcServer := handlers.NewStreamingServer(r.Datastore, r.Director, r.ParserRegistry, poolCap) @@ -1154,3 +1349,16 @@ index 0982639..ab0cedb 100644 extProcServer.SetEmitEndpointScores(r.EmitEndpointScores) if r.EvictChannelLookup != nil { extProcServer.SetEvictChannelLookup(r.EvictChannelLookup) +diff --git a/pkg/epp/server/runserver_test.go b/pkg/epp/server/runserver_test.go +index 55bbb132..4b3f60e1 100644 +--- a/pkg/epp/server/runserver_test.go ++++ b/pkg/epp/server/runserver_test.go +@@ -62,7 +62,7 @@ type stubBands struct { + // a new Options flag reach only one of the two hand-written composite literals. + func TestNewExtProcServerRunnerPopulatesEveryField(t *testing.T) { + // Fields the caller sets after construction, not the factory. +- callerSet := map[string]bool{"GrpcListener": true, "EvictChannelLookup": true} ++ callerSet := map[string]bool{"GrpcListener": true, "EvictChannelLookup": true, "Observer": true} + + opts := server.NewOptions() + // AddFlags registers the logging flag set that Complete() dereferences. From 48566f7027b1dd61892c415ae1aa6bb81ec447ad Mon Sep 17 00:00:00 2001 From: thxCode Date: Wed, 14 Oct 2026 16:09:17 +0800 Subject: [PATCH 2/2] build(llm-router): carry vllm-router EPD routing as a patch Pin vllm-router at v0.1.16.1 and carry the upstream EPD routing change (commit 516f1cb) as a patch, because no release tag contains it yet. The patch header and Dockerfile state that the pin moves to the first upstream tag carrying 516f1cb and deletes 0001-epd-routing.patch. Rebase the Mooncake lookup and observer patches onto the new base without changing their content. Restore the discovery readiness patch as the fourth patch. It now compares Pod values without the UID and its tests use a stub router, so the upstream test suite passes with it. The patched tree passes the same tests as the clean tag. The only failures on both are the eight wasm_middleware tests, which need a wasm32-wasip2 target. Signed-off-by: thxCode --- pack/llm-router/Dockerfile | 15 +- .../vllm-router/0001-epd-routing.patch | 2624 +++++++++++++++++ ...0002-mooncake-lazy-bootstrap-lookup.patch} | 60 +- ...003-observer-read-only-endpoint-view.patch | 34 +- ...very-readiness-and-deselect-removal.patch} | 557 ++-- 5 files changed, 3075 insertions(+), 215 deletions(-) create mode 100644 pack/llm-router/patches/vllm-router/0001-epd-routing.patch rename pack/llm-router/patches/vllm-router/{0001-mooncake-lazy-bootstrap-lookup.patch => 0002-mooncake-lazy-bootstrap-lookup.patch} (84%) rename pack/llm-router/patches/vllm-router/{0002-discovery-readiness-and-deselect-removal.patch => 0004-discovery-readiness-and-deselect-removal.patch} (80%) diff --git a/pack/llm-router/Dockerfile b/pack/llm-router/Dockerfile index 8ae0eb2fc..b0a3880f0 100644 --- a/pack/llm-router/Dockerfile +++ b/pack/llm-router/Dockerfile @@ -30,7 +30,7 @@ ARG GO_VERSION=1.26-bookworm ARG RUST_VERSION=1.95-bookworm ARG LLMD_ROUTER_VERSION=v0.11.0 -ARG VLLM_ROUTER_VERSION=v0.1.15 +ARG VLLM_ROUTER_VERSION=v0.1.16.1 ARG SGLANG_GATEWAY_VERSION=gateway-v0.3.1 ######################## llm-d endpoint picker ######################## @@ -150,6 +150,19 @@ EOF # the same names but different contents, and the failure would be a confusing reject # rather than a missing patch. # +# This stage also carries upstream PR 266 -- commit 516f1cb, the encoder-prefill-decode +# routing -- as 0001-epd-routing.patch, a cherry-pick onto v0.1.16.1. That tag is the +# last vllm-project/router release cut before 516f1cb landed, so its tree does not +# contain the commit and the feature is carried here rather than waited out. The three +# patches above it (lazy Mooncake, the read-only observer endpoint, the Kubernetes +# discovery readiness-loss deselect) rebase onto the tree that cherry-pick produces, +# which is why they travel with it. +# +# THE IMAGE MOVES TO THE FIRST UPSTREAM TAG THAT CONTAINS THE COMMIT: when upstream +# cuts a tag whose tree carries 516f1cb, bump VLLM_ROUTER_VERSION to that tag and delete +# 0001. The other three patches stay until their own WHEN TO REMOVE condition holds, and +# each patch header names the tag it is written against under AGAINST. +# # Applied with git apply, which fails the build loudly on a version bump whose hunks no # longer match -- that failure is the tripwire, and git is already in the base stage so no # package is added for this. diff --git a/pack/llm-router/patches/vllm-router/0001-epd-routing.patch b/pack/llm-router/patches/vllm-router/0001-epd-routing.patch new file mode 100644 index 000000000..5d12e56f6 --- /dev/null +++ b/pack/llm-router/patches/vllm-router/0001-epd-routing.patch @@ -0,0 +1,2624 @@ +Add encode-prefill-decode (EPD) routing from upstream PR 266, applied as a cherry-pick +of commit 516f1cb90eefe565c9fc894d72c8748eae25085b ("[Feature] Add EPD +routing with JSON metadata") onto vllm-project/router v0.1.16.1. + +CONTEXT: PR 266 introduces an EPD router that lets the operator send a request +with metadata describing its endpoint topology. PR 266 was merged on default +branch but the next released tag (v0.1.16.1) was cut BEFORE 516f1cb landed. +The build has to bridge that gap. + +The cherry-pick applies cleanly onto v0.1.16.1 with zero textual conflicts, +because v0.1.16.1's src/config/types.rs, src/config/validation.rs, +src/lib.rs, src/main.rs, src/routers/http/{mod.rs,router.rs,epd.rs}, +src/server.rs, and Cargo.toml already carry the upstream shape that 516f1cb +adds; the patch only edits them in ways the merge parent had already +prepared for. + +CONTRAST WITH OPTIONS: building from the bare commit 516f1cb would require +the Dockerfile to drop --branch (a bare SHA does not resolve a shallow +refspec) and replace it with a fetch-by-SHA. This 0001 cherry-pick keeps +the existing tag-based Dockerfile flow and survives future v0.1.17+ tag bumps +without revisiting the Dockerfile -- it stops being needed the day upstream +cuts a tag containing 516f1cb. + +OPERATOR IMPACT: under the v0.1.16.1 + 0001 + 0002 + 0003 + 0004 stack the +EPD request path is live, /v1/chat/completions honours the operator-declared +endpoint metadata, and the rest of the operator's patches (0002 lazy +Mooncake, 0003 read-only observer endpoint, 0004 Kubernetes discovery +readiness-loss deselect) rebase onto the new tree. + +AGAINST: vllm-project/router v0.1.16.1 (commit +91b65ee4bdfbe98863a32a2205dced5dbb249ffe). The image build applies +this patch with git apply; a future upstream tag that absorbs 516f1cb +will fail loudly and surface at build time, not in a cluster. + +WHEN TO REMOVE: when upstream cuts a tag whose tree contains the merge +commit 516f1cb (or any descendant that carries the EPD routing feature +on top of an already-tagged release). At that point bump +VLLM_ROUTER_VERSION to that tag and delete this patch file. The cherry-pick +was authored against the parent of v0.1.16.1's relevant files, so the +diff is small and a rebase is unlikely to be needed beyond confirming the +merge sits cleanly inside the next tag. +diff --git a/Cargo.toml b/Cargo.toml +index 7a63b27..d5529bc 100644 +--- a/Cargo.toml ++++ b/Cargo.toml +@@ -28,6 +28,7 @@ tower-http = { version = "0.6", features = [ + ] } + serde = { version = "1.0", features = ["derive"] } + serde_json = { version = "1.0", features = ["preserve_order"] } ++sha2 = "0.10" + bytes = "1.8.0" + rand = "0.9.2" + reqwest = { version = "0.13", features = ["stream", "blocking", "json"] } +@@ -89,7 +90,6 @@ wasmtime = { version = "36", default-features = false, features = [ + "runtime", + ] } + wasmtime-wasi = "36" +-sha2 = "0.10" + async-channel = "2" + # Must match vllm-proto (tonic 0.14 / prost 0.14). otel-otlp keeps its own + # tonic 0.12 copy for traces. +diff --git a/docs/epd.md b/docs/epd.md +new file mode 100644 +index 0000000..c2a8cb3 +--- /dev/null ++++ b/docs/epd.md +@@ -0,0 +1,103 @@ ++# Encoder + PD / E/P/D routing (experimental) ++ ++The native Rust router can orchestrate remote encoders before forwarding a chat ++request to a combined prefill/decode worker or separate P and D workers. Embeddings move through vLLM's EC ++connector; the router only handles JSON placeholder metadata and transfer handles. ++No PyTorch dependency or tensor deserialization is needed in the router. ++ ++Use vLLM with JSON metadata inputs (vllm-project/vllm#56090). Configure E workers ++as EC producers and PD workers as matching EC consumers, as for the vLLM encoder ++disaggregation example. The example connector needs a shared filesystem; NIXL ++and Mooncake need their usual backend dependencies and network connectivity. ++ ++```bash ++cargo build --release --bin vllm-router ++target/release/vllm-router \ ++ --host 0.0.0.0 --port 8000 \ ++ --worker-urls http://pd0:8000 http://pd1:8000 \ ++ --policy random \ ++ --epd-config '{ ++ "encoder_urls": ["http://e0:8000", "http://e1:8000"], ++ "consumer_zmq_addrs": { ++ "http://pd0:8000": "tcp://pd0:14579", ++ "http://pd1:8000": "tcp://pd1:14579" ++ } ++ }' ++``` ++ ++`consumer_zmq_addrs` is for Mooncake PUSH and must map each configured PD HTTP URL ++to that worker's configured EC control address. For NIXL PULL or the shared-file ++example connector, omit this map. The PD worker is selected before encoding so ++the producer and the HTTP request use the same destination. Media items are ++assigned to encoders round-robin; images landing on the same encoder are ++batched into a single request, like the Python EPD proxy. ++ ++Send ordinary OpenAI chat requests to `/v1/chat/completions`. Each raw media item ++is sent to an encoder. Published metadata replaces that item with an embeds ++reference; when an encoder reports no metadata, the original media is kept with ++its content-derived uuid pinned and its transfer is forwarded, so the consumer ++re-preprocesses into the same mm_hash and can still reuse the published ++embedding. Repeated images keep their original positions and independent ++transfer IDs. ++Other request fields, including sampling options, pass through unchanged. ++Text-only and existing embeds inputs do not call the encoders. PD responses, ++including SSE streams, use the existing transparent forwarding path. ++ ++Encoder failures are returned without submitting a partially prepared PD ++request. There is no automatic replay of EPD requests in this prototype. ++Abandoned Mooncake pushes receive best-effort cancellation over their existing ++ZMQ control channel; backend expiration remains the fallback if cancellation ++cannot reach the consumer. After a successful PD response header, the consumer ++owns request completion/abort cleanup. ++ ++## Scope ++ ++- Native binary CLI only; the Python launcher does not expose `--epd-config` yet. ++- Regular E+PD or static vLLM E/P/D mode, DP=1, no IGW. ++- HTTP(S) backends only; gRPC encoder and consumer URLs are rejected at startup. ++- Keep the existing backend's TP/PP constraints; this does not add parallelism support. ++- Static E URLs; E health-aware routing and dynamic consumer discovery are not added. ++ The service registry has no encoder role, and vLLM does not publish EC control ++ addresses, so Mooncake PUSH cannot resolve a newly registered consumer. Dynamic ++ discovery for NIXL PULL or the example connector is a smaller follow-up. ++- JSON metadata only, not legacy tensor-base64 metadata from old E servers. ++- No early metadata publication, encoder batching changes, or engine/scheduler changes. ++- Image E2E validation is performed separately; do not infer audio/video coverage ++ or production readiness from the metadata rewrite code alone. ++ ++## Separate prefill and decode ++ ++Add `--epd-config` to the existing static P/D router. Its EC control address map ++must reference **P**, not D. P is an EC consumer and KV producer; D is only a KV ++consumer. Both accept JSON placeholder metadata (`--enable-mm-embeds`). ++For Qwen3.5 with current vLLM NIXL KV transfer, P and D also require ++`VLLM_SSM_CONV_STATE_LAYOUT=DS` for Mamba state transfer. ++ ++```bash ++target/release/vllm-router --host 0.0.0.0 --port 8000 \ ++ --vllm-pd-disaggregation --kv-connector nixl \ ++ --prefill http://p0:8000 --prefill http://p1:8000 \ ++ --decode http://d0:8000 --decode http://d1:8000 --policy random \ ++ --epd-config '{ ++ "encoder_urls": ["http://e0:8000", "http://e1:8000"], ++ "consumer_zmq_addrs": { ++ "http://p0:8000": "tcp://p0:14579", ++ "http://p1:8000": "tcp://p1:14579" ++ } ++ }' ++``` ++ ++The router selects P/D, prepares E metadata for P, then uses the existing KV ++handoff. D receives the same placeholder layout with KV transfer parameters, ++without P's EC handles. P errors or missing NIXL KV handles stop the pipeline. ++Dynamic P/D discovery is not supported with EPD yet. ++ ++```bash ++cargo test --test test_pd_routing test_epd_handoff ++``` ++ ++Focused tests: ++ ++```bash ++cargo test --lib routers::http::epd::tests ++``` +diff --git a/src/config/types.rs b/src/config/types.rs +index c58963f..63e0d59 100644 +--- a/src/config/types.rs ++++ b/src/config/types.rs +@@ -11,6 +11,9 @@ use crate::program_scheduling::{ + /// Main router configuration + #[derive(Debug, Clone, Serialize, Deserialize)] + pub struct RouterConfig { ++ /// Optional encoder orchestration for E+PD or E/P/D routing. ++ #[serde(default, skip_serializing_if = "Option::is_none")] ++ pub epd: Option, + /// Routing mode configuration + pub mode: RoutingMode, + /// Worker connection mode +@@ -94,6 +97,16 @@ fn default_profile_timeout_secs() -> u64 { + 10 + } + ++#[derive(Debug, Clone, Serialize, Deserialize)] ++#[serde(deny_unknown_fields)] ++pub struct EpdConfig { ++ pub encoder_urls: Vec, ++ /// Mooncake consumer control addresses, keyed by the P or combined PD HTTP URL. ++ /// Omit for pull-based NIXL or the shared-file example connector. ++ #[serde(default)] ++ pub consumer_zmq_addrs: HashMap, ++} ++ + fn default_history_backend() -> HistoryBackend { + HistoryBackend::Memory + } +@@ -907,6 +920,7 @@ impl Default for RouterConfig { + history_backend: default_history_backend(), + enable_profiling: false, + profile_timeout_secs: default_profile_timeout_secs(), ++ epd: None, + kv_connector: KvConnector::default(), + program_scheduling: None, + } +@@ -1624,6 +1638,7 @@ mod tests { + history_backend: default_history_backend(), + enable_profiling: false, + profile_timeout_secs: default_profile_timeout_secs(), ++ epd: None, + kv_connector: KvConnector::default(), + program_scheduling: None, + }; +@@ -1691,6 +1706,7 @@ mod tests { + history_backend: default_history_backend(), + enable_profiling: false, + profile_timeout_secs: default_profile_timeout_secs(), ++ epd: None, + kv_connector: KvConnector::default(), + program_scheduling: None, + }; +@@ -1754,6 +1770,7 @@ mod tests { + history_backend: default_history_backend(), + enable_profiling: false, + profile_timeout_secs: default_profile_timeout_secs(), ++ epd: None, + kv_connector: KvConnector::default(), + program_scheduling: None, + }; +diff --git a/src/config/validation.rs b/src/config/validation.rs +index 948f1b1..b4e4bf8 100644 +--- a/src/config/validation.rs ++++ b/src/config/validation.rs +@@ -23,6 +23,7 @@ impl ConfigValidator { + if let Some(program_scheduling) = &config.program_scheduling { + Self::validate_program_scheduling(program_scheduling)?; + } ++ Self::validate_epd(config)?; + + if let Some(discovery) = &config.discovery { + Self::validate_discovery(discovery, &config.mode)?; +@@ -182,6 +183,74 @@ impl ConfigValidator { + Ok(()) + } + ++ fn validate_epd(config: &RouterConfig) -> ConfigResult<()> { ++ let Some(epd) = &config.epd else { ++ return Ok(()); ++ }; ++ let worker_urls: Vec<&String> = match &config.mode { ++ RoutingMode::Regular { worker_urls } => worker_urls.iter().collect(), ++ RoutingMode::VllmPrefillDecode { ++ prefill_urls, ++ discovery_address: None, ++ .. ++ } => prefill_urls.iter().map(|(url, _)| url).collect(), ++ _ => { ++ return Err(ConfigError::ValidationFailed { ++ reason: "EPD requires regular E+PD or static vLLM E/P/D routing".into(), ++ }) ++ } ++ }; ++ if config.intra_node_data_parallel_size != 1 || config.enable_igw { ++ return Err(ConfigError::ValidationFailed { ++ reason: "EPD currently requires DP=1 and proxy mode".into(), ++ }); ++ } ++ if epd.encoder_urls.is_empty() { ++ return Err(ConfigError::MissingRequired { ++ field: "epd.encoder_urls".into(), ++ }); ++ } ++ Self::validate_urls(&epd.encoder_urls)?; ++ for url in epd.encoder_urls.iter().chain(worker_urls.iter().copied()) { ++ if !url.starts_with("http://") && !url.starts_with("https://") { ++ return Err(ConfigError::InvalidValue { ++ field: "epd.worker_url".into(), ++ value: url.clone(), ++ reason: "EPD requires HTTP(S) encoder and consumer URLs".into(), ++ }); ++ } ++ } ++ if !epd.consumer_zmq_addrs.is_empty() { ++ if worker_urls.is_empty() ++ || epd.consumer_zmq_addrs.len() != worker_urls.len() ++ || worker_urls ++ .iter() ++ .any(|url| !epd.consumer_zmq_addrs.contains_key(*url)) ++ { ++ return Err(ConfigError::ValidationFailed { ++ reason: ++ "EPD consumer_zmq_addrs must map every embedding consumer URL exactly once" ++ .into(), ++ }); ++ } ++ for addr in epd.consumer_zmq_addrs.values() { ++ let parsed = url::Url::parse(addr).map_err(|e| ConfigError::ValidationFailed { ++ reason: format!("Invalid EC address: {e}"), ++ })?; ++ if parsed.scheme() != "tcp" ++ || parsed.host_str().is_none() ++ || parsed.port().is_none() ++ { ++ return Err(ConfigError::ValidationFailed { ++ reason: "EC addresses must be tcp://host:port".into(), ++ }); ++ } ++ } ++ } ++ ++ Ok(()) ++ } ++ + /// Validate routing mode configuration + fn validate_mode(mode: &RoutingMode, has_service_discovery: bool) -> ConfigResult<()> { + match mode { +@@ -875,6 +944,48 @@ mod tests { + assert!(ConfigValidator::validate(&config).is_ok()); + } + ++ #[test] ++ fn test_validate_epd_requires_http_workers() { ++ for (encoder_scheme, consumer_scheme, supported) in [ ++ ("http", "http", true), ++ ("https", "https", true), ++ ("grpc", "http", false), ++ ("http", "grpc", false), ++ ] { ++ let consumer = format!("{consumer_scheme}://consumer:8000"); ++ for mode in [ ++ RoutingMode::Regular { ++ worker_urls: vec![consumer.clone()], ++ }, ++ RoutingMode::VllmPrefillDecode { ++ prefill_urls: vec![(consumer.clone(), None)], ++ decode_urls: vec![format!("{consumer_scheme}://decode:8000")], ++ prefill_policy: None, ++ decode_policy: None, ++ discovery_address: None, ++ }, ++ ] { ++ let mut config = RouterConfig::new(mode, PolicyConfig::Random); ++ config.epd = Some(EpdConfig { ++ encoder_urls: vec![format!("{encoder_scheme}://encoder:8000")], ++ consumer_zmq_addrs: Default::default(), ++ }); ++ let result = ConfigValidator::validate(&config); ++ if supported { ++ result.unwrap(); ++ } else { ++ assert!( ++ result ++ .unwrap_err() ++ .to_string() ++ .contains("EPD requires HTTP(S)"), ++ "encoder={encoder_scheme}, consumer={consumer_scheme}" ++ ); ++ } ++ } ++ } ++ } ++ + #[test] + fn test_validate_cache_aware_thresholds() { + let config = RouterConfig::new( +diff --git a/src/lib.rs b/src/lib.rs +index 735c407..2fc5d3f 100644 +--- a/src/lib.rs ++++ b/src/lib.rs +@@ -188,6 +188,7 @@ impl Router { + }; + + Ok(config::RouterConfig { ++ epd: None, + mode, + policy, + host: self.host.clone(), +@@ -542,6 +543,7 @@ impl Router { + service_discovery_config, + prometheus_config, + request_timeout_secs: self.request_timeout_secs, ++ pool_idle_timeout_secs: server::DEFAULT_POOL_IDLE_TIMEOUT_SECS, + request_id_headers: self.request_id_headers.clone(), + trace_config: if self.enable_trace { + Some(config::TraceConfig { +diff --git a/src/main.rs b/src/main.rs +index a003ae5..d9328fa 100644 +--- a/src/main.rs ++++ b/src/main.rs +@@ -104,6 +104,9 @@ Examples: + + "#)] + struct CliArgs { ++ /// E+PD or E/P/D encoder orchestration as JSON (encoder_urls, consumer_zmq_addrs). ++ #[arg(long)] ++ epd_config: Option, + /// Host address to bind the router server + #[arg(long, default_value = "127.0.0.1")] + host: String, +@@ -280,6 +283,12 @@ struct CliArgs { + #[arg(long, default_value_t = 1800)] + request_timeout_secs: u64, + ++ /// Idle timeout in seconds for pooled backend HTTP connections. Should be ++ /// lower than the backend's keep-alive timeout (vLLM defaults to 5s); 0 ++ /// disables connection reuse entirely ++ #[arg(long, default_value_t = vllm_router_rs::server::DEFAULT_POOL_IDLE_TIMEOUT_SECS)] ++ pool_idle_timeout_secs: u64, ++ + /// Maximum number of concurrent requests allowed + #[arg(long, default_value_t = 32768)] + max_concurrent_requests: usize, +@@ -539,6 +548,14 @@ impl CliArgs { + + // Build RouterConfig + Ok(RouterConfig { ++ epd: self ++ .epd_config ++ .as_deref() ++ .map(serde_json::from_str) ++ .transpose() ++ .map_err(|e| vllm_router_rs::config::ConfigError::ValidationFailed { ++ reason: format!("Invalid --epd-config: {e}"), ++ })?, + mode, + policy, + connection_mode, +@@ -638,6 +655,7 @@ impl CliArgs { + service_discovery_config, + prometheus_config, + request_timeout_secs: self.request_timeout_secs, ++ pool_idle_timeout_secs: self.pool_idle_timeout_secs, + request_id_headers: if self.request_id_headers.is_empty() { + None + } else { +diff --git a/src/routers/http/epd.rs b/src/routers/http/epd.rs +new file mode 100644 +index 0000000..6b0692e +--- /dev/null ++++ b/src/routers/http/epd.rs +@@ -0,0 +1,979 @@ ++use crate::config::EpdConfig; ++use axum::http::{HeaderMap, StatusCode}; ++use futures_util::future::join_all; ++use reqwest::{Client, RequestBuilder}; ++use serde_json::{json, Map, Value}; ++use sha2::{Digest, Sha256}; ++use std::collections::HashMap; ++use std::sync::atomic::{AtomicUsize, Ordering}; ++use std::time::Instant; ++use uuid::Uuid; ++ ++type EpdError = (StatusCode, String); ++ ++struct MediaItem { ++ message_index: usize, ++ content_index: usize, ++ kind: &'static str, ++ content_id: String, ++ transfer_id: String, ++ media: Value, ++} ++ ++#[derive(Debug)] ++pub(super) struct EncoderStage { ++ config: EpdConfig, ++ next_encoder: AtomicUsize, ++} ++ ++/// Cancel abandoned Mooncake pushes when preparation or the P/PD request ends early. ++pub(super) struct EcTransferGuard { ++ control_addr: String, ++ transfer_ids: Vec, ++} ++ ++impl EcTransferGuard { ++ fn retain(&mut self, transfer_ids: Vec) { ++ let mut abandoned = Self { ++ control_addr: self.control_addr.clone(), ++ transfer_ids: std::mem::replace(&mut self.transfer_ids, transfer_ids), ++ }; ++ abandoned ++ .transfer_ids ++ .retain(|id| !self.transfer_ids.contains(id)); ++ } ++ ++ pub(super) fn disarm(&mut self) { ++ self.transfer_ids.clear(); ++ } ++} ++ ++impl Drop for EcTransferGuard { ++ fn drop(&mut self) { ++ if self.transfer_ids.is_empty() { ++ return; ++ } ++ let transfer_ids = std::mem::take(&mut self.transfer_ids); ++ let control_addr = self.control_addr.clone(); ++ // No blocking ZMQ operations on the HTTP runtime threads. ++ tokio::task::spawn_blocking(move || { ++ if let Err(error) = cancel_ec_transfers(&control_addr, &transfer_ids) { ++ tracing::warn!(%error, %control_addr, "Could not discover EC consumer shards for cancellation"); ++ } ++ }); ++ } ++} ++ ++fn ec_control_request( ++ context: &zmq::Context, ++ address: &str, ++ request: &Value, ++) -> anyhow::Result { ++ let socket = context.socket(zmq::REQ)?; ++ socket.set_linger(0)?; ++ socket.set_sndtimeo(1000)?; ++ socket.set_rcvtimeo(1000)?; ++ socket.connect(address)?; ++ socket.send(serde_json::to_vec(request)?.as_slice(), 0)?; ++ let mut response: Value = serde_json::from_slice(&socket.recv_bytes(0)?)?; ++ anyhow::ensure!( ++ response["ok"] == true, ++ "EC control request failed: {}", ++ response["error"] ++ ); ++ Ok(response["result"].take()) ++} ++ ++fn cancel_ec_transfers(control_addr: &str, transfer_ids: &[String]) -> anyhow::Result<()> { ++ let context = zmq::Context::new(); ++ let peers = ec_control_request(&context, control_addr, &json!({"op": "peers"}))?; ++ let ports: Vec = serde_json::from_value(peers["ports"].clone())?; ++ anyhow::ensure!( ++ !ports.is_empty() && !ports.contains(&0), ++ "Invalid EC consumer shard ports" ++ ); ++ let (prefix, _) = control_addr ++ .rsplit_once(':') ++ .ok_or_else(|| anyhow::anyhow!("Invalid EC control address: {control_addr}"))?; ++ for port in ports { ++ let address = format!("{prefix}:{port}"); ++ for transfer_id in transfer_ids { ++ let request = json!({"op": "cancel", "transfer_id": transfer_id, "abandon": true}); ++ if let Err(error) = ec_control_request(&context, &address, &request) { ++ tracing::warn!(%error, %address, %transfer_id, "Could not cancel abandoned EC transfer"); ++ } ++ } ++ } ++ Ok(()) ++} ++ ++impl EncoderStage { ++ pub(super) fn new(config: EpdConfig) -> Self { ++ Self { ++ config, ++ next_encoder: AtomicUsize::new(0), ++ } ++ } ++ ++ pub(super) async fn prepare( ++ &self, ++ client: &Client, ++ mut body: Value, ++ consumer_url: &str, ++ headers: Option<&HeaderMap>, ++ api_key: Option<&str>, ++ ) -> Result<(Value, Option), EpdError> { ++ let started = Instant::now(); ++ let items = collect_media_items(&body)?; ++ if items.is_empty() { ++ return Ok((body, None)); ++ } ++ let mut embed_positions = Vec::new(); ++ for (message_index, message) in body["messages"] ++ .as_array() ++ .into_iter() ++ .flatten() ++ .enumerate() ++ { ++ for (content_index, item) in message ++ .get("content") ++ .and_then(Value::as_array) ++ .into_iter() ++ .flatten() ++ .enumerate() ++ { ++ if matches!( ++ item.get("type").and_then(Value::as_str), ++ Some("image_embeds" | "audio_embeds" | "video_embeds") ++ ) { ++ embed_positions.push((message_index, content_index)); ++ } ++ } ++ } ++ let consumer_addr = self.config.consumer_zmq_addrs.get(consumer_url); ++ if !self.config.consumer_zmq_addrs.is_empty() && consumer_addr.is_none() { ++ return Err(( ++ StatusCode::SERVICE_UNAVAILABLE, ++ "Selected embedding consumer has no EC control address".into(), ++ )); ++ } ++ let mut pending = consumer_addr.map(|control_addr| EcTransferGuard { ++ control_addr: control_addr.clone(), ++ transfer_ids: items.iter().map(|item| item.transfer_id.clone()).collect(), ++ }); ++ let request_id = headers ++ .and_then(|h| h.get("x-request-id")) ++ .and_then(|v| v.to_str().ok()) ++ .map(str::to_owned) ++ .unwrap_or_else(|| Uuid::new_v4().to_string()); ++ let start = self.next_encoder.fetch_add(items.len(), Ordering::Relaxed); ++ // Round-robin per item, then batch the images of one request that land ++ // on the same encoder into a single call, like the Python EPD proxy. ++ // Audio and video items stay singleton groups. ++ let mut groups: Vec<(usize, Vec)> = Vec::new(); ++ let mut image_group_of: HashMap = HashMap::new(); ++ for (index, item) in items.iter().enumerate() { ++ let encoder = start.wrapping_add(index) % self.config.encoder_urls.len(); ++ if item.kind == "image_url" { ++ match image_group_of.get(&encoder) { ++ Some(&group) => groups[group].1.push(index), ++ None => { ++ image_group_of.insert(encoder, groups.len()); ++ groups.push((encoder, vec![index])); ++ } ++ } ++ } else { ++ groups.push((encoder, vec![index])); ++ } ++ } ++ let replies = join_all(groups.iter().map(|(encoder, indices)| { ++ let encoder_url = &self.config.encoder_urls[*encoder]; ++ let mut request = client ++ .post(format!( ++ "{}/v1/chat/completions", ++ encoder_url.trim_end_matches('/') ++ )) ++ .header( ++ "x-request-id", ++ format!( ++ "{request_id}:{}:{:.6}", ++ indices[0], items[indices[0]].transfer_id ++ ), ++ ); ++ if let Some(key) = api_key { ++ request = request.bearer_auth(key); ++ } else if let Some(auth) = headers.and_then(|h| h.get("authorization")) { ++ request = request.header("authorization", auth); ++ } ++ encode_group( ++ request, ++ &items, ++ indices, ++ &body, ++ consumer_addr.map(String::as_str), ++ ) ++ })) ++ .await; ++ let replies: Vec = replies.into_iter().collect::>()?; ++ let mut handles = Map::new(); ++ // Sort transfers by content position, including caller-supplied embeds. ++ let mut transfers = Vec::new(); ++ let mut used_transfer_ids = Vec::new(); ++ let mut rewritten = 0; ++ for ((_, indices), mut reply) in groups.iter().zip(replies) { ++ let Some(Value::Object(params)) = reply.get_mut("ec_transfer_params").map(Value::take) ++ else { ++ // No EC report at all: keep the raw media but pin its uuid so ++ // the consumer derives the same mm_hash the encoder used. ++ for &index in indices { ++ let item = &items[index]; ++ body["messages"][item.message_index]["content"][item.content_index] = ++ item.media.clone(); ++ } ++ continue; ++ }; ++ // Position is the reliable correlation when the engine rehashes ++ // uuids with processing options; the uuid key and the single-entry ++ // shape are fallbacks for encoders without position metadata. ++ let mut by_position: HashMap = HashMap::new(); ++ if indices.len() > 1 { ++ for (mm_hash, reported) in params.iter() { ++ if let Some(item_indices) = ++ reported.get("item_indices").and_then(Value::as_array) ++ { ++ for local in item_indices { ++ let local = local ++ .as_u64() ++ .map(|v| v as usize) ++ .filter(|&v| v < indices.len()) ++ .ok_or_else(|| { ++ (StatusCode::BAD_GATEWAY, "Invalid encoder item index".into()) ++ })?; ++ if by_position.insert(local, (mm_hash, reported)).is_some() { ++ return Err(( ++ StatusCode::BAD_GATEWAY, ++ "Duplicate encoder item index".into(), ++ )); ++ } ++ } ++ } ++ } ++ } ++ for (local, &index) in indices.iter().enumerate() { ++ let item = &items[index]; ++ let entry = by_position ++ .get(&local) ++ .map(|(hash, reported)| ((*hash).clone(), (*reported).clone())) ++ .or_else(|| { ++ params ++ .get_key_value(&item.content_id) ++ .map(|(hash, reported)| (hash.clone(), reported.clone())) ++ }) ++ .or_else(|| { ++ if indices.len() == 1 && params.len() == 1 { ++ params ++ .iter() ++ .next() ++ .map(|(hash, reported)| (hash.clone(), reported.clone())) ++ } else { ++ None ++ } ++ }); ++ let Some((hash, reported)) = entry else { ++ // Unreported item (e.g. an encoder-side processor cache ++ // hit): keep the raw media with its uuid pinned. ++ body["messages"][item.message_index]["content"][item.content_index] = ++ item.media.clone(); ++ continue; ++ }; ++ if let Some(metadata) = reported ++ .get("metadata") ++ .and_then(Value::as_object) ++ .filter(|m| !m.is_empty()) ++ { ++ let metadata = flatten_metadata(metadata)?; ++ let kind = match item.kind { ++ "image_url" => "image_embeds", ++ "video_url" => "video_embeds", ++ _ => "audio_embeds", ++ }; ++ body["messages"][item.message_index]["content"][item.content_index] = json!({ ++ "type": kind, ++ (kind): metadata, ++ "uuid": item.content_id, ++ }); ++ transfers.push(( ++ (item.message_index, item.content_index), ++ json!({ ++ "mm_hash": hash, ++ "transfer_id": item.transfer_id, ++ }), ++ )); ++ used_transfer_ids.push(item.transfer_id.clone()); ++ rewritten += 1; ++ } else { ++ // The encoder reported no placeholder metadata (e.g. a ++ // processor cache hit): keep the raw media with its uuid ++ // pinned and hand the transfer to the consumer, which can ++ // then reuse the published embedding after re-preprocessing ++ // instead of abandoning the transfer. ++ body["messages"][item.message_index]["content"][item.content_index] = ++ item.media.clone(); ++ transfers.push(( ++ (item.message_index, item.content_index), ++ json!({ ++ "mm_hash": hash, ++ "transfer_id": item.transfer_id, ++ }), ++ )); ++ used_transfer_ids.push(item.transfer_id.clone()); ++ } ++ handles.insert(hash, reported); ++ } ++ } ++ if !handles.is_empty() { ++ let mut previous = Vec::new(); ++ if let Some(Value::Object(mut existing)) = ++ body.get_mut("ec_transfer_params").map(Value::take) ++ { ++ if let Some(Value::Array(entries)) = existing.remove("ec_items") { ++ previous = entries; ++ } ++ existing.extend(handles); ++ handles = existing; ++ } ++ if previous.len() == embed_positions.len() { ++ transfers.extend(embed_positions.into_iter().zip(previous.drain(..))); ++ } ++ transfers.sort_by_key(|(position, _)| *position); ++ previous.extend(transfers.into_iter().map(|(_, entry)| entry)); ++ if !previous.is_empty() { ++ handles.insert("ec_items".into(), Value::Array(previous)); ++ } ++ body["ec_transfer_params"] = Value::Object(handles); ++ } ++ if let Some(pending) = pending.as_mut() { ++ pending.retain(used_transfer_ids); ++ } ++ tracing::info!( ++ %request_id, ++ items = items.len(), ++ rewritten, ++ encode_ms = started.elapsed().as_secs_f64() * 1000.0, ++ "EPD encoder stage complete" ++ ); ++ Ok((body, pending)) ++ } ++} ++ ++fn collect_media_items(body: &Value) -> Result, EpdError> { ++ let messages = body ++ .get("messages") ++ .and_then(Value::as_array) ++ .ok_or_else(|| (StatusCode::BAD_REQUEST, "messages must be an array".into()))?; ++ let mut items = Vec::new(); ++ for (message_index, message) in messages.iter().enumerate() { ++ let Some(content) = message.get("content").and_then(Value::as_array) else { ++ continue; ++ }; ++ for (content_index, item) in content.iter().enumerate() { ++ let kind = match item.get("type").and_then(Value::as_str) { ++ Some("image_url") => "image_url", ++ Some("video_url") => "video_url", ++ Some("audio_url") => "audio_url", ++ _ => continue, ++ }; ++ let url = item[kind]["url"] ++ .as_str() ++ .filter(|v| !v.is_empty()) ++ .ok_or_else(|| { ++ ( ++ StatusCode::BAD_REQUEST, ++ format!("{kind}.url must be a nonempty string"), ++ ) ++ })?; ++ // A caller-supplied uuid wins (matches the reference EPD proxy): ++ // it pins the EC cache identity, e.g. for retrying against an ++ // already published embedding. Hash the URL only as a fallback. ++ let content_id = item ++ .get("uuid") ++ .and_then(Value::as_str) ++ .filter(|v| !v.is_empty()) ++ .map(str::to_owned) ++ .unwrap_or_else(|| format!("{:x}", Sha256::digest(url.as_bytes()))); ++ let transfer_id = Uuid::new_v4().simple().to_string(); ++ let mut media = item.clone(); ++ media["uuid"] = json!(content_id); ++ items.push(MediaItem { ++ message_index, ++ content_index, ++ kind, ++ content_id, ++ transfer_id, ++ media, ++ }); ++ } ++ } ++ Ok(items) ++} ++ ++async fn encode_group( ++ request: RequestBuilder, ++ items: &[MediaItem], ++ indices: &[usize], ++ body: &Value, ++ consumer_addr: Option<&str>, ++) -> Result { ++ let mut payload = json!({ ++ "model": body.get("model"), ++ "stream": false, ++ "messages": [{ ++ "role": "user", ++ "content": indices.iter().map(|&i| items[i].media.clone()).collect::>(), ++ }], ++ }); ++ // Per-request preprocessing options must agree on both sides. ++ for key in [ ++ "mm_processor_kwargs", ++ "media_io_kwargs", ++ "priority", ++ "session_id", ++ ] { ++ if let Some(value) = body.get(key) { ++ payload[key] = value.clone(); ++ } ++ } ++ if let Some(control_addr) = consumer_addr { ++ payload["ec_transfer_params"] = json!({ ++ "consumer_zmq": control_addr, ++ "ec_items": indices ++ .iter() ++ .map(|&i| json!({"transfer_id": items[i].transfer_id})) ++ .collect::>(), ++ }); ++ } ++ let request = request.json(&payload); ++ drop(payload); ++ let response = request.send().await.map_err(|e| { ++ ( ++ StatusCode::BAD_GATEWAY, ++ format!("Encoder request failed: {e}"), ++ ) ++ })?; ++ if !response.status().is_success() { ++ return Err((response.status(), "Encoder request failed".into())); ++ } ++ response.json::().await.map_err(|e| { ++ ( ++ StatusCode::BAD_GATEWAY, ++ format!("Invalid encoder response: {e}"), ++ ) ++ }) ++} ++ ++fn flatten_metadata(metadata: &Map) -> Result, EpdError> { ++ metadata ++ .iter() ++ .map(|(key, value)| { ++ let values = value.as_array().ok_or_else(|| { ++ ( ++ StatusCode::BAD_GATEWAY, ++ "Encoder metadata must contain numeric arrays".into(), ++ ) ++ })?; ++ let mut flat = Vec::new(); ++ for value in values { ++ match value { ++ Value::Array(values) => flat.extend(values.iter().cloned()), ++ value => flat.push(value.clone()), ++ } ++ } ++ if !flat.iter().all(Value::is_number) { ++ return Err(( ++ StatusCode::BAD_GATEWAY, ++ "Encoder metadata must contain numeric arrays".into(), ++ )); ++ } ++ Ok((key.clone(), Value::Array(flat))) ++ }) ++ .collect() ++} ++ ++#[cfg(test)] ++mod tests { ++ use super::*; ++ use axum::{extract::State, routing::post, Json, Router}; ++ use std::sync::{Arc, Mutex}; ++ ++ #[derive(Clone)] ++ struct MockEncoder { ++ mode: &'static str, ++ seen: Arc>>, ++ } ++ ++ async fn encode( ++ State(state): State, ++ Json(body): Json, ++ ) -> (StatusCode, Json) { ++ { ++ let mut seen = state.seen.lock().unwrap(); ++ seen.push(body.clone()); ++ } ++ if state.mode == "fail" { ++ return ( ++ StatusCode::SERVICE_UNAVAILABLE, ++ Json(json!({"error":"unavailable"})), ++ ); ++ } ++ // Batched requests carry several media parts; dedupe identical media ++ // into one entry and report positions via item_indices, like vLLM. ++ let content = body["messages"][0]["content"].as_array().unwrap(); ++ let mut positions: Vec<(String, Vec)> = Vec::new(); ++ for (index, part) in content.iter().enumerate() { ++ let uuid = part["uuid"].as_str().unwrap().to_owned(); ++ match positions.iter_mut().find(|(u, _)| *u == uuid) { ++ Some((_, indices)) => indices.push(index), ++ None => positions.push((uuid, vec![index])), ++ } ++ } ++ let mut params = Map::new(); ++ for (uuid, mut indices) in positions { ++ if state.mode == "mixed" { ++ // Report only the later item of a batch, leaving the first ++ // unreported (e.g. an encoder-side processor cache hit). ++ indices = indices.into_iter().skip(1).collect(); ++ } ++ if indices.is_empty() { ++ continue; ++ } ++ let metadata = match state.mode { ++ "empty" => json!({}), ++ "invalid" => json!({"image_grid_thw": ["invalid"]}), ++ _ => json!({"image_grid_thw": [[1, 2, 3]]}), ++ }; ++ params.insert( ++ format!("engine-{uuid}"), ++ json!({ ++ "metadata": metadata, ++ "item_indices": indices, ++ "peer_host": "encoder", ++ "peer_port": 4321, ++ "size_bytes": 128, ++ }), ++ ); ++ } ++ (StatusCode::OK, Json(json!({"ec_transfer_params": params}))) ++ } ++ ++ async fn mock( ++ mode: &'static str, ++ ) -> ( ++ EncoderStage, ++ Arc>>, ++ tokio::task::JoinHandle<()>, ++ ) { ++ let state = MockEncoder { ++ mode, ++ seen: Arc::new(Mutex::new(Vec::new())), ++ }; ++ let seen = state.seen.clone(); ++ let app = Router::new() ++ .route("/v1/chat/completions", post(encode)) ++ .with_state(state); ++ let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); ++ let url = format!("http://{}", listener.local_addr().unwrap()); ++ let server = tokio::spawn(async move { axum::serve(listener, app).await.unwrap() }); ++ let stage = EncoderStage::new(EpdConfig { ++ encoder_urls: vec![url], ++ consumer_zmq_addrs: Default::default(), ++ }); ++ (stage, seen, server) ++ } ++ ++ fn request() -> Value { ++ json!({ ++ "model": "test", ++ "stream": true, ++ "max_tokens": 16, ++ "chat_template_kwargs": {"enable_thinking": false}, ++ "mm_processor_kwargs": {"min_pixels": 64}, ++ "messages": [{ ++ "role": "user", ++ "content": [ ++ {"type": "image_url", "image_url": {"url": "data:image/png;base64,abc"}}, ++ {"type": "text", "text": "compare"}, ++ {"type": "image_url", "image_url": {"url": "data:image/png;base64,abc"}}, ++ ], ++ }], ++ }) ++ } ++ ++ #[tokio::test] ++ async fn duplicate_images_keep_positions_and_distinct_transfers() { ++ let (stage, seen, server) = mock("normal").await; ++ let mut input = request(); ++ input["ec_transfer_params"] = json!({ ++ "opaque": {"backend_field": 42}, ++ "ec_items": [{"mm_hash": "existing", "transfer_id": "existing-transfer"}], ++ }); ++ let (body, pending) = stage ++ .prepare(&Client::new(), input.clone(), "pd", None, None) ++ .await ++ .unwrap(); ++ assert!(pending.is_none()); ++ let content = &body["messages"][0]["content"]; ++ assert_eq!( ++ content[0]["image_embeds"]["image_grid_thw"], ++ json!([1, 2, 3]) ++ ); ++ assert_eq!(content[1], input["messages"][0]["content"][1]); ++ assert_eq!(content[0]["uuid"], content[2]["uuid"]); ++ let transfers = &body["ec_transfer_params"]["ec_items"]; ++ assert_eq!(transfers.as_array().unwrap().len(), 3); ++ assert_eq!(transfers[0], input["ec_transfer_params"]["ec_items"][0]); ++ assert_eq!( ++ body["ec_transfer_params"]["opaque"], ++ input["ec_transfer_params"]["opaque"] ++ ); ++ assert_eq!(transfers[1]["mm_hash"], transfers[2]["mm_hash"]); ++ assert_ne!(transfers[1]["transfer_id"], transfers[2]["transfer_id"]); ++ assert_eq!( ++ transfers[1]["mm_hash"], ++ format!("engine-{}", content[0]["uuid"].as_str().unwrap()) ++ ); ++ assert_eq!( ++ body["ec_transfer_params"][transfers[1]["mm_hash"].as_str().unwrap()]["peer_port"], ++ 4321 ++ ); ++ assert!(body["ec_transfer_params"] ++ .get(content[0]["uuid"].as_str().unwrap()) ++ .is_none()); ++ assert_eq!(body["stream"], true); ++ assert_eq!(body["chat_template_kwargs"], input["chat_template_kwargs"]); ++ let requests = seen.lock().unwrap(); ++ // Both images round-robin onto the single encoder and are batched ++ // into one request. ++ assert_eq!(requests.len(), 1); ++ let request = &requests[0]; ++ assert!(request["stream"] == false && request.get("max_tokens").is_none()); ++ assert_eq!(request["mm_processor_kwargs"], input["mm_processor_kwargs"]); ++ let sent = request["messages"][0]["content"].as_array().unwrap(); ++ assert_eq!(sent.len(), 2); ++ assert!(sent.iter().all(|part| part["uuid"].is_string())); ++ server.abort(); ++ } ++ ++ #[tokio::test] ++ async fn ec_items_follow_media_order_across_encoders() { ++ // Images [A, B, A] round-robin onto two encoders: E0 gets [A, A] ++ // batched, E1 gets [B]. ec_items must stay in media order ++ // [A, B, A] so the consumer's positional transfer mapping cancels ++ // every reservation on request finish. ++ let spawn = || async { ++ let state = MockEncoder { ++ mode: "normal", ++ seen: Arc::new(Mutex::new(Vec::new())), ++ }; ++ let seen = state.seen.clone(); ++ let app = Router::new() ++ .route("/v1/chat/completions", post(encode)) ++ .with_state(state); ++ let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); ++ let url = format!("http://{}", listener.local_addr().unwrap()); ++ let server = tokio::spawn(async move { axum::serve(listener, app).await.unwrap() }); ++ (url, seen, server) ++ }; ++ let (url0, seen0, server0) = spawn().await; ++ let (url1, seen1, server1) = spawn().await; ++ let stage = EncoderStage::new(EpdConfig { ++ encoder_urls: vec![url0, url1], ++ consumer_zmq_addrs: Default::default(), ++ }); ++ let input = json!({ ++ "messages": [{ ++ "role": "user", ++ "content": [ ++ {"type": "image_url", "image_url": {"url": "data:image/png;base64,abc"}}, ++ {"type": "image_url", "image_url": {"url": "data:image/png;base64,def"}}, ++ {"type": "image_url", "image_url": {"url": "data:image/png;base64,abc"}}, ++ ], ++ }], ++ }); ++ let (body, _) = stage ++ .prepare(&Client::new(), input, "pd", None, None) ++ .await ++ .unwrap(); ++ // The split across two encoders actually happened. ++ assert_eq!(seen0.lock().unwrap().len(), 1); ++ assert_eq!(seen1.lock().unwrap().len(), 1); ++ assert_eq!( ++ seen0.lock().unwrap()[0]["messages"][0]["content"] ++ .as_array() ++ .unwrap() ++ .len(), ++ 2 ++ ); ++ let content = &body["messages"][0]["content"]; ++ let uuid_a = content[0]["uuid"].as_str().unwrap(); ++ let uuid_b = content[1]["uuid"].as_str().unwrap(); ++ assert_ne!(uuid_a, uuid_b); ++ assert_eq!(content[2]["uuid"].as_str().unwrap(), uuid_a); ++ let hashes: Vec = body["ec_transfer_params"]["ec_items"] ++ .as_array() ++ .unwrap() ++ .iter() ++ .map(|item| item["mm_hash"].as_str().unwrap().to_owned()) ++ .collect(); ++ assert_eq!( ++ hashes, ++ [ ++ format!("engine-{uuid_a}"), ++ format!("engine-{uuid_b}"), ++ format!("engine-{uuid_a}"), ++ ] ++ ); ++ server0.abort(); ++ server1.abort(); ++ } ++ ++ #[tokio::test] ++ async fn existing_and_new_transfers_follow_content_order() { ++ let (stage, _, server) = mock("normal").await; ++ let embed = json!({ ++ "type": "image_embeds", ++ "image_embeds": {"image_grid_thw": [1, 2, 3]}, ++ "uuid": "cached", ++ }); ++ let raw = json!({"type": "image_url", "image_url": {"url": "new-image"}}); ++ let cached = json!({"mm_hash": "cached-hash", "transfer_id": "cached-transfer"}); ++ for content in [json!([raw, embed]), json!([embed, raw])] { ++ let raw_index = usize::from(content[0]["type"] == "image_embeds"); ++ let input = json!({ ++ "messages": [{"role": "user", "content": content}], ++ "ec_transfer_params": { ++ "cached-hash": {"peer_port": 42}, ++ "ec_items": [cached], ++ }, ++ }); ++ let (body, _) = stage ++ .prepare(&Client::new(), input, "pd", None, None) ++ .await ++ .unwrap(); ++ let transfers = body["ec_transfer_params"]["ec_items"].as_array().unwrap(); ++ assert_eq!(transfers.len(), 2); ++ assert_eq!(transfers[1 - raw_index], cached); ++ assert_ne!(transfers[raw_index]["transfer_id"], cached["transfer_id"]); ++ assert_eq!(body["messages"][0]["content"][1 - raw_index], embed); ++ assert_eq!(body["ec_transfer_params"]["cached-hash"]["peer_port"], 42); ++ } ++ server.abort(); ++ } ++ ++ #[tokio::test] ++ async fn missing_metadata_keeps_raw_media_with_uuid_and_forwards_transfers() { ++ let (stage, seen, server) = mock("empty").await; ++ let mut input = request(); ++ input["ec_transfer_params"] = json!({ ++ "opaque": {"backend_field": 42}, ++ "ec_items": [{"mm_hash": "existing", "transfer_id": "existing-transfer"}], ++ }); ++ let (body, _) = stage ++ .prepare(&Client::new(), input.clone(), "pd", None, None) ++ .await ++ .unwrap(); ++ // Raw media is kept with the content-derived uuid pinned, so the ++ // consumer derives the same mm_hash and can reuse the embedding. ++ let seen = seen.lock().unwrap(); ++ let uuid = seen[0]["messages"][0]["content"][0]["uuid"] ++ .as_str() ++ .unwrap() ++ .to_owned(); ++ let content = &body["messages"][0]["content"]; ++ for index in [0, 2] { ++ let mut expected = input["messages"][0]["content"][index].clone(); ++ expected["uuid"] = json!(uuid); ++ assert_eq!(content[index], expected); ++ } ++ assert_eq!(content[1], input["messages"][0]["content"][1]); ++ assert_eq!( ++ body["ec_transfer_params"]["opaque"], ++ input["ec_transfer_params"]["opaque"] ++ ); ++ // Raw-fallback transfers are forwarded to the consumer, not abandoned. ++ let transfers = body["ec_transfer_params"]["ec_items"].as_array().unwrap(); ++ assert_eq!(transfers.len(), 3); ++ assert_eq!(transfers[0], input["ec_transfer_params"]["ec_items"][0]); ++ assert_eq!(transfers[1]["mm_hash"], format!("engine-{uuid}")); ++ assert_ne!(transfers[1]["transfer_id"], transfers[2]["transfer_id"]); ++ let handle = &body["ec_transfer_params"][format!("engine-{uuid}")]; ++ assert_eq!(handle["metadata"], json!({})); ++ assert_eq!(handle["peer_port"], 4321); ++ server.abort(); ++ } ++ ++ #[tokio::test] ++ async fn encoder_failure_does_not_forward_metadata_only_input() { ++ let (stage, _, server) = mock("fail").await; ++ let error = stage ++ .prepare(&Client::new(), request(), "pd", None, None) ++ .await ++ .err() ++ .unwrap(); ++ assert_eq!(error.0, StatusCode::SERVICE_UNAVAILABLE); ++ server.abort(); ++ } ++ ++ #[tokio::test] ++ async fn existing_embeds_bypass_encoders() { ++ let (stage, seen, server) = mock("normal").await; ++ let input = json!({ ++ "messages": [{ ++ "role": "user", ++ "content": [{ ++ "type": "image_embeds", ++ "image_embeds": {"image_grid_thw": [1, 2, 3]}, ++ "uuid": "known", ++ }], ++ }], ++ }); ++ let (body, _) = stage ++ .prepare(&Client::new(), input.clone(), "pd", None, None) ++ .await ++ .unwrap(); ++ assert_eq!(input, body); ++ assert!(seen.lock().unwrap().is_empty()); ++ server.abort(); ++ } ++ ++ #[tokio::test] ++ async fn abandoned_pushes_cancel_all_consumer_shards() { ++ for mode in [ ++ "normal", ++ "empty", ++ "mixed", ++ "invalid", ++ "fail", ++ "cancel_error", ++ ] { ++ let (mut stage, seen, server) = mock(mode).await; ++ let (send_address, recv_address) = tokio::sync::oneshot::channel(); ++ let control = tokio::task::spawn_blocking(move || { ++ let context = zmq::Context::new(); ++ let sockets = [0, 1].map(|_| { ++ let socket = context.socket(zmq::REP).unwrap(); ++ socket.set_rcvtimeo(5000).unwrap(); ++ socket.set_sndtimeo(5000).unwrap(); ++ socket.set_linger(0).unwrap(); ++ socket.bind("tcp://127.0.0.1:*").unwrap(); ++ socket ++ }); ++ let addresses = sockets ++ .each_ref() ++ .map(|s| s.get_last_endpoint().unwrap().unwrap()); ++ let ports = addresses ++ .each_ref() ++ .map(|a| a.rsplit_once(':').unwrap().1.parse::().unwrap()); ++ send_address.send(addresses[0].clone()).unwrap(); ++ let discovery: Value = ++ serde_json::from_slice(&sockets[0].recv_bytes(0).unwrap()).unwrap(); ++ assert_eq!(discovery, json!({"op": "peers"})); ++ sockets[0] ++ .send( ++ json!({"ok": true, "result": {"ports": ports}}) ++ .to_string() ++ .as_bytes(), ++ 0, ++ ) ++ .unwrap(); ++ let mut cancelled = [Vec::new(), Vec::new()]; ++ // "mixed" leaves the unreported item's transfer abandoned while ++ // the reported one is disarmed below. ++ let expected = if mode == "mixed" { 1 } else { 2 }; ++ for (rank, socket) in sockets.iter().enumerate() { ++ for _ in 0..expected { ++ let message: Value = ++ serde_json::from_slice(&socket.recv_bytes(0).unwrap()).unwrap(); ++ cancelled[rank].push(message); ++ let response = if mode == "cancel_error" && rank == 0 { ++ json!({"ok": false, "error": "injected failure"}) ++ } else { ++ json!({"ok": true, "result": {"cancelled": true}}) ++ }; ++ socket.send(response.to_string().as_bytes(), 0).unwrap(); ++ } ++ } ++ cancelled ++ }); ++ let address = recv_address.await.unwrap(); ++ stage ++ .config ++ .consumer_zmq_addrs ++ .insert("selected-pd".into(), address.clone()); ++ let result = stage ++ .prepare(&Client::new(), request(), "selected-pd", None, None) ++ .await; ++ let ids: Vec = seen ++ .lock() ++ .unwrap() ++ .iter() ++ .flat_map(|r| { ++ assert_eq!(r["ec_transfer_params"]["consumer_zmq"], address); ++ r["ec_transfer_params"]["ec_items"] ++ .as_array() ++ .unwrap() ++ .iter() ++ .map(|entry| entry["transfer_id"].clone()) ++ .collect::>() ++ }) ++ .collect(); ++ if matches!(mode, "invalid" | "fail") { ++ assert!(result.is_err()); ++ } else { ++ let (_, mut pending) = result.unwrap(); ++ if mode == "mixed" { ++ let pending = pending.as_mut().unwrap(); ++ assert_eq!(pending.transfer_ids.len(), 1); ++ pending.disarm(); ++ } ++ drop(pending); ++ } ++ let cancelled = control.await.unwrap(); ++ assert_eq!(cancelled[0], cancelled[1]); ++ assert!(cancelled[0].iter().all(|r| r["op"] == "cancel" ++ && r["abandon"] == true ++ && ids.contains(&r["transfer_id"]))); ++ if mode == "mixed" { ++ assert_eq!(cancelled[0][0]["transfer_id"], ids[0]); ++ } else { ++ assert_ne!( ++ cancelled[0][0]["transfer_id"], ++ cancelled[0][1]["transfer_id"] ++ ); ++ } ++ server.abort(); ++ } ++ } ++ ++ #[test] ++ fn control_request_rejects_remote_errors() { ++ let context = zmq::Context::new(); ++ let socket = context.socket(zmq::REP).unwrap(); ++ socket.set_rcvtimeo(5000).unwrap(); ++ socket.set_sndtimeo(5000).unwrap(); ++ socket.set_linger(0).unwrap(); ++ socket.bind("tcp://127.0.0.1:*").unwrap(); ++ let address = socket.get_last_endpoint().unwrap().unwrap(); ++ let server = std::thread::spawn(move || { ++ socket.recv_bytes(0).unwrap(); ++ socket ++ .send( ++ b"{\"ok\":false,\"error\":\"injected failure\"}".as_slice(), ++ 0, ++ ) ++ .unwrap(); ++ }); ++ let result = ec_control_request(&context, &address, &json!({"op": "peers"})); ++ server.join().unwrap(); ++ assert!(result.unwrap_err().to_string().contains("injected failure")); ++ } ++} +diff --git a/src/routers/http/mod.rs b/src/routers/http/mod.rs +index fa2acae..2fcb690 100644 +--- a/src/routers/http/mod.rs ++++ b/src/routers/http/mod.rs +@@ -1,6 +1,7 @@ + //! HTTP router implementations + + pub mod dp_utils; ++mod epd; + pub mod logprobs_merge; + pub mod openai_router; + pub mod pd_router; +diff --git a/src/routers/http/router.rs b/src/routers/http/router.rs +index f7140ac..07d3bab 100644 +--- a/src/routers/http/router.rs ++++ b/src/routers/http/router.rs +@@ -144,6 +144,7 @@ struct TypedDispatch<'a> { + /// Regular router that uses injected load balancing policies + #[derive(Debug)] + pub struct Router { ++ epd: Option, + worker_registry: Arc, + policy_registry: Arc, + client: Client, +@@ -346,6 +347,11 @@ impl Router { + }); + + Ok(Router { ++ epd: ctx ++ .router_config ++ .epd ++ .clone() ++ .map(super::epd::EncoderStage::new), + worker_registry: ctx.worker_registry.clone(), + policy_registry: ctx.policy_registry.clone(), + client: ctx.client.clone(), +@@ -2421,6 +2427,7 @@ impl RouterTrait for Router { + body: serde_json::Value, + ) -> Response { + debug!("Transparent proxy: routing {} {} to backend", method, path); ++ let start = Instant::now(); + + // Select a worker (filter by availability like select_worker_for_model) + let all_workers = self.worker_registry.get_all(); +@@ -2430,6 +2437,7 @@ impl RouterTrait for Router { + .cloned() + .collect(); + if workers.is_empty() { ++ RouterMetrics::record_request_error(path, "no_available_workers"); + return ( + StatusCode::SERVICE_UNAVAILABLE, + "No available workers".to_string(), +@@ -2437,7 +2445,14 @@ impl RouterTrait for Router { + .into_response(); + } + +- let request_text = serde_json::to_string(&body).ok(); ++ let policy = self.policy_registry.get_default_policy(); ++ // Serializing the whole body is expensive when it carries base64 ++ // media; only consumers that actually read the text pay for it. ++ let request_text = if policy.needs_request_text() || self.program_scheduler.is_some() { ++ serde_json::to_string(&body).ok() ++ } else { ++ None ++ }; + let model_id = body.get("model").and_then(serde_json::Value::as_str); + let is_stream = body + .get("stream") +@@ -2455,7 +2470,10 @@ impl RouterTrait for Router { + .await + { + Ok(completion) => completion, +- Err(error) => return Self::schedule_error_response(error), ++ Err(error) => { ++ RouterMetrics::record_request_error(path, "program_scheduling_error"); ++ return Self::schedule_error_response(error); ++ } + } + } else { + None +@@ -2465,7 +2483,6 @@ impl RouterTrait for Router { + .get_by_url(&completion.dispatch().target_id) + .filter(|worker| worker.is_available()) + } else { +- let policy = self.policy_registry.get_default_policy(); + let request_headers = Self::headers_to_request_headers(headers); + policy + .select_worker_with_headers( +@@ -2479,6 +2496,7 @@ impl RouterTrait for Router { + if let Some(completion) = &program_completion { + completion.finish(false); + } ++ RouterMetrics::record_request_error(path, "no_available_workers"); + return ( + StatusCode::SERVICE_UNAVAILABLE, + "Failed to select the Program target".to_string(), +@@ -2487,6 +2505,43 @@ impl RouterTrait for Router { + }; + let url = worker.endpoint_url(path); + ++ // Match the typed path's load accounting so cache-aware imbalance ++ // detection also sees transparent (EPD) traffic. ++ let load_tracked = policy.name() == "cache_aware" || program_completion.is_some(); ++ if load_tracked { ++ worker.increment_load(); ++ RouterMetrics::set_running_requests(worker.url(), worker.load()); ++ } ++ ++ let (body, mut transfers) = if let Some(epd) = &self.epd { ++ if *method == Method::POST && path == "/v1/chat/completions" { ++ match epd ++ .prepare( ++ &self.client, ++ body, ++ worker.url(), ++ headers, ++ self.api_key.as_deref(), ++ ) ++ .await ++ { ++ Ok(value) => value, ++ Err(error) => { ++ RouterMetrics::record_request_error(path, "encoder_error"); ++ if load_tracked { ++ worker.decrement_load(); ++ RouterMetrics::set_running_requests(worker.url(), worker.load()); ++ } ++ return error.into_response(); ++ } ++ } ++ } else { ++ (body, None) ++ } ++ } else { ++ (body, None) ++ }; ++ + debug!("Transparent proxy: forwarding to {}", url); + + // Build the request +@@ -2498,6 +2553,11 @@ impl RouterTrait for Router { + Method::PATCH => self.client.patch(&url), + Method::HEAD => self.client.head(&url), + _ => { ++ if load_tracked { ++ worker.decrement_load(); ++ RouterMetrics::set_running_requests(worker.url(), worker.load()); ++ } ++ RouterMetrics::record_request_error(path, "unsupported_method"); + return ( + StatusCode::METHOD_NOT_ALLOWED, + format!("Method {} not supported", method), +@@ -2515,12 +2575,35 @@ impl RouterTrait for Router { + } + + // Add authorization if configured ++ if self.epd.is_some() { ++ if let Some(headers) = headers { ++ for (name, value) in headers { ++ // When a backend api_key is configured it must win: ++ // RequestBuilder::header appends, and vLLM reads the ++ // first Authorization value. ++ if self.api_key.is_some() && name == http::header::AUTHORIZATION { ++ continue; ++ } ++ if !matches!( ++ name.as_str(), ++ "host" ++ | "content-length" ++ | "content-type" ++ | "connection" ++ | "transfer-encoding" ++ ) && !header_utils::TRACE_HEADER_NAMES.contains(&name.as_str()) ++ { ++ request_builder = request_builder.header(name, value); ++ } ++ } ++ } ++ } + if let Some(ref key) = self.api_key { + request_builder = request_builder.header("Authorization", format!("Bearer {}", key)); + } + + // Send request +- match otel_http::send_client_request( ++ let response = match otel_http::send_client_request( + request_builder, + headers, + ClientRequestOptions { +@@ -2534,6 +2617,18 @@ impl RouterTrait for Router { + { + Ok(response) => { + let status = response.status(); ++ // Keep the circuit breaker informed, as the typed path does: ++ // 2xx/4xx count as success, 5xx as worker failure. ++ let was_available = worker.is_available(); ++ worker.record_outcome(status.is_success() || status.is_client_error()); ++ if was_available != worker.is_available() { ++ self.worker_registry.notify_worker_state_change(); ++ } ++ if status.is_success() { ++ if let Some(transfers) = transfers.as_mut() { ++ transfers.disarm(); ++ } ++ } + let headers = response.headers().clone(); + let mut response_builder = Response::builder().status(status.as_u16()); + +@@ -2635,15 +2730,35 @@ impl RouterTrait for Router { + } + } + Err(error) => { ++ // The 502 response below still flows through ++ // hold_load_until_body_done, which releases the load once ++ // its body finishes or is dropped — do not decrement here. + if let Some(completion) = &program_completion { + completion.finish(false); + } ++ // Transport failure: the worker never produced a response. ++ let was_available = worker.is_available(); ++ worker.record_outcome(false); ++ if was_available != worker.is_available() { ++ self.worker_registry.notify_worker_state_change(); ++ } + ( + StatusCode::BAD_GATEWAY, + format!("Backend request failed: {error}"), + ) + .into_response() + } ++ }; ++ if response.status().is_success() { ++ RouterMetrics::record_request(path); ++ RouterMetrics::record_generate_duration(start.elapsed()); ++ } else { ++ RouterMetrics::record_request_error(path, "backend_error"); ++ } ++ if load_tracked { ++ hold_load_until_body_done(response, worker) ++ } else { ++ response + } + } + } +@@ -2668,6 +2783,7 @@ mod tests { + + let (_, rx) = tokio::sync::watch::channel(HashMap::new()); + Router { ++ epd: None, + worker_registry, + policy_registry, + worker_startup_timeout_secs: 5, +@@ -2710,6 +2826,7 @@ mod tests { + circuit_breaker_config: CircuitBreakerConfig::default(), + health_config: HealthConfig::default(), + frontend: crate::backend::EngineFrontend::new(), ++ epd: None, + _worker_loads: Arc::new(rx), + _load_monitor_handle: None, + program_scheduler: None, +@@ -2858,6 +2975,73 @@ mod tests { + assert!(err.contains("mixed"), "{err}"); + } + ++ #[tokio::test] ++ async fn transparent_epd_records_request_metrics() { ++ use axum::{routing::post, Json}; ++ use serde_json::{json, Value}; ++ ++ let recorder = metrics_exporter_prometheus::PrometheusBuilder::new().build_recorder(); ++ let handle = recorder.handle(); ++ let _metrics = metrics::set_default_local_recorder(&recorder); ++ let app = axum::Router::new().route( ++ "/v1/chat/completions", ++ post(|Json(body): Json| async move { ++ if body["fail"] == true { ++ (StatusCode::BAD_REQUEST, "invalid request").into_response() ++ } else if body["stream"] == true { ++ ([("content-type", "text/event-stream")], "data: [DONE]\n\n").into_response() ++ } else { ++ Json(json!({"choices": []})).into_response() ++ } ++ }), ++ ); ++ let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); ++ let url = format!("http://{}", listener.local_addr().unwrap()); ++ let server = tokio::spawn(async move { axum::serve(listener, app).await.unwrap() }); ++ let mut router = create_test_regular_router(); ++ router.worker_registry = Arc::new(WorkerRegistry::new()); ++ router ++ .worker_registry ++ .register(Arc::new(BasicWorker::new(url, WorkerType::Regular))); ++ router.epd = Some(super::super::epd::EncoderStage::new( ++ crate::config::EpdConfig { ++ encoder_urls: vec!["http://unused".into()], ++ consumer_zmq_addrs: Default::default(), ++ }, ++ )); ++ for (body, status) in [ ++ (json!({"messages": []}), StatusCode::OK), ++ (json!({"messages": [], "stream": true}), StatusCode::OK), ++ ( ++ json!({"messages": [], "fail": true}), ++ StatusCode::BAD_REQUEST, ++ ), ++ (json!({"messages": null}), StatusCode::BAD_REQUEST), ++ ] { ++ let response = router ++ .route_transparent(None, "/v1/chat/completions", &Method::POST, body) ++ .await; ++ assert_eq!(response.status(), status); ++ axum::body::to_bytes(response.into_body(), 4096) ++ .await ++ .unwrap(); ++ } ++ let exported = handle.render(); ++ assert!(exported.contains("vllm_router_requests_total{route=\"/v1/chat/completions\"} 2")); ++ assert!(exported.contains("vllm_router_generate_duration_seconds_count 2")); ++ for error in ["encoder_error", "backend_error"] { ++ assert!( ++ exported ++ .lines() ++ .any(|line| line.starts_with("vllm_router_request_errors_total{") ++ && line.contains(&format!("error_type=\"{error}\"")) ++ && line.ends_with(" 1")), ++ "{exported}" ++ ); ++ } ++ server.abort(); ++ } ++ + #[tokio::test] + async fn streaming_body_holds_load_until_consumed_or_dropped() { + let worker: Arc = Arc::new(BasicWorker::new( +@@ -2965,6 +3149,7 @@ mod tests { + Router { + worker_registry, + policy_registry, ++ epd: None, + worker_startup_timeout_secs: 5, + worker_startup_check_interval_secs: 1, + intra_node_data_parallel_size: 1, +diff --git a/src/routers/http/vllm_pd_router.rs b/src/routers/http/vllm_pd_router.rs +index 2d7909c..20e79be 100644 +--- a/src/routers/http/vllm_pd_router.rs ++++ b/src/routers/http/vllm_pd_router.rs +@@ -37,6 +37,7 @@ struct MooncakePrefillInfo { + /// vLLM PD Router that extends PdRouterBase with vLLM-specific request handling + #[derive(Debug)] + pub struct VllmPDRouter { ++ epd: Option, + /// Underlying PD router for most functionality + pd_router: PdRouterBase, + /// Service discovery registry for dynamic ZMQ address resolution +@@ -1286,6 +1287,39 @@ impl VllmPDRouter { + ) -> Result { + debug!("ENTERED process_vllm_two_stage_request method"); + let start_time = Instant::now(); ++ if !original_request.is_object() { ++ return Ok(( ++ StatusCode::BAD_REQUEST, ++ "Request body must be a JSON object", ++ ) ++ .into_response()); ++ } ++ ++ // The encoder stage must run before the NIXL concurrent fast path: ++ // once a push identity is cached, try_build_concurrent_requests ++ // returns early, and multimodal requests reaching P/D without ++ // published embeddings break encoder-disaggregated workers. ++ let (original_request, mut ec_transfers) = if path == "/v1/chat/completions" { ++ if let Some(epd) = &self.epd { ++ match epd ++ .prepare( ++ &self.http_client, ++ original_request, ++ prefill_worker.base_url(), ++ headers, ++ std::env::var("OPENAI_API_KEY").ok().as_deref(), ++ ) ++ .await ++ { ++ Ok(prepared) => prepared, ++ Err(error) => return Ok(error.into_response()), ++ } ++ } else { ++ (original_request, None) ++ } ++ } else { ++ (original_request, None) ++ }; + + if let Some((prefill_request, decode_request, request_id)) = self + .try_build_concurrent_requests(&original_request, &prefill_worker, &decode_worker, path) +@@ -1301,6 +1335,7 @@ impl VllmPDRouter { + path, + headers, + start_time, ++ ec_transfers, + ) + .await; + } +@@ -1426,6 +1461,33 @@ impl VllmPDRouter { + prefill_response.headers() + ); + ++ if !prefill_response.status().is_success() { ++ prefill_worker.decrement_load(); ++ let status = prefill_response.status(); ++ let content_type = prefill_response ++ .headers() ++ .get(http::header::CONTENT_TYPE) ++ .cloned(); ++ let detail = prefill_response.text().await.unwrap_or_default(); ++ self.stop_profiling(&prefill_base_url).await; ++ let duration = start_time.elapsed(); ++ RouterMetrics::record_pd_prefill_error(&prefill_base_url); ++ RouterMetrics::record_pd_request(path); ++ RouterMetrics::record_pd_request_duration(path, duration); ++ if status.is_client_error() { ++ let mut response = (status, detail).into_response(); ++ if let Some(content_type) = content_type { ++ response ++ .headers_mut() ++ .insert(http::header::CONTENT_TYPE, content_type); ++ } ++ return Ok(response); ++ } ++ return Err(PDRouterError::NetworkError { ++ message: format!("Prefill server {prefill_url} returned {status}: {detail}"), ++ }); ++ } ++ + // Extract prefill response body to get kv_transfer_params + let prefill_bytes = match prefill_response.bytes().await { + Ok(bytes) => bytes, +@@ -1471,6 +1533,33 @@ impl VllmPDRouter { + } + }; + ++ if self.epd.is_some() ++ && !matches!(self.kv_connector, KvConnector::Mooncake) ++ && !matches!(self.moriio_transfer_mode(), Some(MoriIOTransferMode::Write)) ++ && !prefill_response_json ++ .get("kv_transfer_params") ++ .is_some_and(Value::is_object) ++ { ++ prefill_worker.decrement_load(); ++ self.stop_profiling(&prefill_base_url).await; ++ let duration = start_time.elapsed(); ++ RouterMetrics::record_pd_prefill_error(&prefill_base_url); ++ RouterMetrics::record_pd_request(path); ++ RouterMetrics::record_pd_request_duration(path, duration); ++ return Err(PDRouterError::NetworkError { ++ message: ++ "Prefill returned no KV transfer parameters; refusing metadata-only decode" ++ .into(), ++ }); ++ } ++ ++ // Prefill succeeded and consumed the embeddings; the pending EC ++ // transfers are now owned by the decode phase, so disarm the ++ // cancel-on-abandon guard. ++ if let Some(transfers) = ec_transfers.as_mut() { ++ transfers.disarm(); ++ } ++ + // Stop profiling on prefill server after its work is done + self.stop_profiling(&prefill_base_url).await; + +@@ -1482,6 +1571,12 @@ impl VllmPDRouter { + + // Stage 2: Prepare decode request with kv_transfer_params + let mut decode_request = original_request.clone(); ++ if self.epd.is_some() { ++ // D receives the prompt layout and KV, not P's EC reservations. ++ if let Some(obj) = decode_request.as_object_mut() { ++ obj.remove("ec_transfer_params"); ++ } ++ } + if let Some(params) = self + .build_decode_kv_transfer_params( + &prefill_base_url, +@@ -1704,6 +1799,12 @@ impl VllmPDRouter { + .ok()?; + + let mut decode_request = original_request.clone(); ++ if self.epd.is_some() { ++ // D receives the prompt layout and KV, not P's EC reservations. ++ if let Some(obj) = decode_request.as_object_mut() { ++ obj.remove("ec_transfer_params"); ++ } ++ } + decode_request["kv_transfer_params"] = + Self::build_nixl_push_decode_params(&identity, &request_id); + +@@ -1721,6 +1822,7 @@ impl VllmPDRouter { + path: &str, + headers: Option<&HeaderMap>, + start_time: Instant, ++ mut ec_transfers: Option, + ) -> Result { + let prefill_base_url = prefill_worker.base_url().to_string(); + let prefill_dp_rank = prefill_worker.dp_rank(); +@@ -1787,6 +1889,18 @@ impl VllmPDRouter { + RouterMetrics::record_pd_prefill_error(&prefill_base_url); + RouterMetrics::record_pd_request(path); + RouterMetrics::record_pd_request_duration(path, start_time.elapsed()); ++ if resp.status().is_client_error() { ++ let status = resp.status(); ++ let content_type = resp.headers().get(http::header::CONTENT_TYPE).cloned(); ++ let mut response = ++ (status, resp.bytes().await.unwrap_or_default()).into_response(); ++ if let Some(content_type) = content_type { ++ response ++ .headers_mut() ++ .insert(http::header::CONTENT_TYPE, content_type); ++ } ++ return Ok(response); ++ } + return Err(PDRouterError::NetworkError { + message: format!( + "Prefill request failed to {} with status {}", +@@ -1813,6 +1927,12 @@ impl VllmPDRouter { + if let Some(prefill_json) = prefill_response_json.as_ref() { + self.maybe_cache_nixl_push_identity(&prefill_base_url, prefill_dp_rank, prefill_json); + } ++ // Prefill succeeded and consumed the embeddings; the pending EC ++ // transfers are now owned by the decode phase, so disarm the ++ // cancel-on-abandon guard. ++ if let Some(transfers) = ec_transfers.as_mut() { ++ transfers.disarm(); ++ } + + let decode_response = match decode_result { + Ok(resp) => resp, +@@ -1876,7 +1996,10 @@ impl VllmPDRouter { + ctx: &Arc, + ) -> Result { + let kv_connector = ctx.router_config.kv_connector; +- let http_client = reqwest::Client::new(); ++ // Carries the configured request/connect timeouts; a bare ++ // reqwest::Client::new() would let a stalled encoder hold the ++ // request and its concurrency permit indefinitely. ++ let http_client = ctx.client.clone(); + + if let Some(ref addr) = discovery_address { + // Discovery mode +@@ -1903,6 +2026,11 @@ impl VllmPDRouter { + ); + + Ok(Self { ++ epd: ctx ++ .router_config ++ .epd ++ .clone() ++ .map(super::epd::EncoderStage::new), + pd_router, + service_registry: Arc::new(service_registry), + http_client, +@@ -1993,6 +2121,11 @@ impl VllmPDRouter { + } + + Ok(Self { ++ epd: ctx ++ .router_config ++ .epd ++ .clone() ++ .map(super::epd::EncoderStage::new), + pd_router, + service_registry: Arc::new(service_registry), + http_client, +diff --git a/src/server.rs b/src/server.rs +index 70c26dd..a26ad94 100644 +--- a/src/server.rs ++++ b/src/server.rs +@@ -261,6 +261,25 @@ async fn v1_chat_completions( + state.router.route_chat(Some(&headers), &body, None).await + } + ++async fn epd_chat_completions( ++ State(state): State>, ++ headers: http::HeaderMap, ++ Json(body): Json, ++) -> Response { ++ if let Err(response) = authorize_request(&state, &headers).await { ++ return response; ++ } ++ state ++ .router ++ .route_transparent( ++ Some(&headers), ++ "/v1/chat/completions", ++ &http::Method::POST, ++ body, ++ ) ++ .await ++} ++ + async fn v1_completions( + State(state): State>, + headers: http::HeaderMap, +@@ -702,6 +721,13 @@ async fn delete_worker( + } + } + ++/// Default idle timeout for pooled backend connections, in seconds. ++/// Kept below vLLM's server-side keep-alive timeout ++/// (`VLLM_HTTP_TIMEOUT_KEEP_ALIVE`, 5s by default) so the router retires a ++/// pooled connection before the backend closes it; reusing a connection the ++/// backend already closed races and surfaces as sporadic transport-error 502s. ++pub const DEFAULT_POOL_IDLE_TIMEOUT_SECS: u64 = 2; ++ + pub struct ServerConfig { + pub host: String, + pub port: u16, +@@ -715,6 +741,9 @@ pub struct ServerConfig { + pub service_discovery_config: Option, + pub prometheus_config: Option, + pub request_timeout_secs: u64, ++ /// Idle timeout in seconds for pooled backend connections. A value of 0 ++ /// disables connection reuse entirely. ++ pub pool_idle_timeout_secs: u64, + pub request_id_headers: Option>, + pub trace_config: Option, + } +@@ -774,7 +803,14 @@ pub fn build_app_with_wasm_middleware( + let mut protected_routes = Router::new() + .route("/generate", post(generate)) + .route("/inference/v1/generate", post(inference_generate)) +- .route("/v1/chat/completions", post(v1_chat_completions)) ++ .route( ++ "/v1/chat/completions", ++ if app_state.context.router_config.epd.is_some() { ++ post(epd_chat_completions) ++ } else { ++ post(v1_chat_completions) ++ }, ++ ) + .route("/v1/completions", post(v1_completions)) + .route("/rerank", post(rerank)) + .route("/v1/rerank", post(v1_rerank)) +@@ -920,13 +956,19 @@ pub async fn startup(config: ServerConfig) -> Result<(), Box) -> RouterConfig { + RouterConfig { ++ epd: None, + mode: RoutingMode::Regular { worker_urls }, + policy: PolicyConfig::RoundRobin, + host: "127.0.0.1".to_string(), +diff --git a/tests/test_dp_routing.rs b/tests/test_dp_routing.rs +index e33264b..de9f5b8 100644 +--- a/tests/test_dp_routing.rs ++++ b/tests/test_dp_routing.rs +@@ -278,6 +278,7 @@ mod dp_e2e_tests { + history_backend: vllm_router_rs::config::HistoryBackend::Memory, + enable_profiling: false, + profile_timeout_secs: 30, ++ epd: None, + kv_connector: vllm_router_rs::config::KvConnector::Nixl, + program_scheduling: None, + } +@@ -327,6 +328,7 @@ mod dp_e2e_tests { + history_backend: vllm_router_rs::config::HistoryBackend::Memory, + enable_profiling: false, + profile_timeout_secs: 30, ++ epd: None, + kv_connector: vllm_router_rs::config::KvConnector::Nixl, + program_scheduling: None, + } +diff --git a/tests/test_pd_routing.rs b/tests/test_pd_routing.rs +index a9ea85f..9ad36b9 100644 +--- a/tests/test_pd_routing.rs ++++ b/tests/test_pd_routing.rs +@@ -5,9 +5,451 @@ mod test_pd_routing { + }; + use vllm_router_rs::routers::RouterFactory; + +- // ======================================================================== +- // Phase 1: Basic PD Components and Router Creation +- // ======================================================================== ++ #[tokio::test] ++ async fn test_epd_handoff_preserves_metadata_and_isolates_ec_handles() { ++ use axum::{ ++ extract::State, ++ http::StatusCode, ++ response::IntoResponse, ++ routing::{get, post}, ++ Json, Router, ++ }; ++ use serde_json::{json, Value}; ++ use std::sync::{Arc, Mutex}; ++ use vllm_router_rs::config::{ConfigValidator, EpdConfig}; ++ ++ type Seen = Arc>>; ++ async fn encode( ++ State(seen): State, ++ Json(body): Json, ++ ) -> axum::response::Response { ++ seen.lock().unwrap().push(("E".into(), body.clone())); ++ if body["model"] == "encoder-fail" { ++ return (StatusCode::SERVICE_UNAVAILABLE, "injected encoder error").into_response(); ++ } ++ let hash = body["messages"][0]["content"][0]["uuid"].as_str().unwrap(); ++ Json(json!({ ++ "ec_transfer_params": { ++ (format!("engine-{hash}")): { ++ "metadata": {"image_grid_thw": [[1, 2, 3]]}, ++ "peer_host": "encoder", ++ "peer_port": 4321, ++ "size_bytes": 128, ++ }, ++ }, ++ })) ++ .into_response() ++ } ++ async fn prefill( ++ State(seen): State, ++ Json(body): Json, ++ ) -> axum::response::Response { ++ seen.lock().unwrap().push(("P".into(), body.clone())); ++ if matches!(body["model"].as_str(), Some("prefill-400" | "prefill-422")) { ++ let status = if body["model"] == "prefill-400" { ++ StatusCode::BAD_REQUEST ++ } else { ++ StatusCode::UNPROCESSABLE_ENTITY ++ }; ++ return ( ++ status, ++ Json(json!({"error": "invalid sampling parameters"})), ++ ) ++ .into_response(); ++ } ++ if body["model"] == "fail" { ++ return (StatusCode::INTERNAL_SERVER_ERROR, "injected prefill error") ++ .into_response(); ++ } ++ if body["model"] == "missing-kv" { ++ return Json(json!({"choices":[]})).into_response(); ++ } ++ Json(json!({ ++ "kv_transfer_params": { ++ "remote_engine_id": "P", ++ "do_remote_prefill": true, ++ }, ++ })) ++ .into_response() ++ } ++ async fn decode( ++ State(seen): State, ++ Json(body): Json, ++ ) -> axum::response::Response { ++ seen.lock().unwrap().push(("D".into(), body.clone())); ++ if body["stream"] == true { ++ return ([("content-type", "text/event-stream")], "data: [DONE]\n\n") ++ .into_response(); ++ } ++ Json(json!({"choices":[{"message":{"content":"A"}}]})).into_response() ++ } ++ let seen: Seen = Arc::new(Mutex::new(Vec::new())); ++ let app = Router::new() ++ .route("/e/health", get(|| async { "ok" })) ++ .route("/p/health", get(|| async { "ok" })) ++ .route("/d/health", get(|| async { "ok" })) ++ .route("/e/v1/chat/completions", post(encode)) ++ .route("/p/v1/chat/completions", post(prefill)) ++ .route("/d/v1/chat/completions", post(decode)) ++ .with_state(seen.clone()); ++ let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); ++ let base = format!("http://{}", listener.local_addr().unwrap()); ++ let server = tokio::spawn(async move { ++ axum::serve(listener, app).await.unwrap(); ++ }); ++ let mut config = RouterConfig { ++ mode: RoutingMode::VllmPrefillDecode { ++ prefill_urls: vec![(format!("{base}/p"), None)], ++ decode_urls: vec![format!("{base}/d")], ++ prefill_policy: None, ++ decode_policy: None, ++ discovery_address: None, ++ }, ++ policy: PolicyConfig::Random, ++ epd: Some(EpdConfig { ++ encoder_urls: vec![format!("{base}/e")], ++ consumer_zmq_addrs: Default::default(), ++ }), ++ worker_startup_timeout_secs: 5, ++ ..RouterConfig::default() ++ }; ++ // Mooncake EC reservations belong to P, not D. ++ config ++ .epd ++ .as_mut() ++ .unwrap() ++ .consumer_zmq_addrs ++ .insert(format!("{base}/d"), "tcp://localhost:1234".into()); ++ assert!(ConfigValidator::validate(&config).is_err()); ++ config.epd.as_mut().unwrap().consumer_zmq_addrs.clear(); ++ ConfigValidator::validate(&config).unwrap(); ++ let context = Arc::new( ++ vllm_router_rs::server::AppContext::new( ++ config, ++ reqwest::Client::new(), ++ 64, ++ None, ++ vec![], ++ ) ++ .unwrap(), ++ ); ++ let router = RouterFactory::create_router(&context).await.unwrap(); ++ for (model, stream, status) in [ ++ ("test", false, StatusCode::OK), ++ ("test", true, StatusCode::OK), ++ ("fail", false, StatusCode::INTERNAL_SERVER_ERROR), ++ ("prefill-400", false, StatusCode::BAD_REQUEST), ++ ("prefill-422", false, StatusCode::UNPROCESSABLE_ENTITY), ++ ("missing-kv", false, StatusCode::INTERNAL_SERVER_ERROR), ++ ("encoder-fail", false, StatusCode::SERVICE_UNAVAILABLE), ++ ("bad-input", false, StatusCode::BAD_REQUEST), ++ ] { ++ seen.lock().unwrap().clear(); ++ let mut body = json!({ ++ "model": model, ++ "stream": stream, ++ "max_tokens": 16, ++ "structured_outputs": {"choice": ["A", "B"]}, ++ "messages": [{ ++ "role": "user", ++ "content": [ ++ {"type": "image_url", "image_url": {"url": "data:image/png;base64,abc"}}, ++ {"type": "text", "text": "choose"}, ++ ], ++ }], ++ }); ++ if model == "bad-input" { ++ body["messages"] = Value::Null; ++ } ++ let response = router ++ .route_transparent( ++ None, ++ "/v1/chat/completions", ++ &axum::http::Method::POST, ++ body, ++ ) ++ .await; ++ assert_eq!(response.status(), status, "model={model}, stream={stream}"); ++ if model.starts_with("prefill-") { ++ assert_eq!(response.headers()["content-type"], "application/json"); ++ } ++ let bytes = axum::body::to_bytes(response.into_body(), 4096) ++ .await ++ .unwrap(); ++ if stream { ++ assert!(String::from_utf8_lossy(&bytes).contains("[DONE]")); ++ } ++ if model.starts_with("prefill-") { ++ assert_eq!( ++ serde_json::from_slice::(&bytes).unwrap(), ++ json!({"error": "invalid sampling parameters"}) ++ ); ++ } ++ let seen = seen.lock().unwrap(); ++ let phases: Vec<_> = seen.iter().map(|(phase, _)| phase.as_str()).collect(); ++ assert_eq!( ++ phases, ++ match model { ++ "bad-input" => vec![], ++ "encoder-fail" => vec!["E"], ++ "test" => vec!["E", "P", "D"], ++ _ => vec!["E", "P"], ++ } ++ ); ++ if seen.len() < 2 { ++ continue; ++ } ++ let p = &seen[1].1; ++ assert_eq!(p["max_tokens"], 1); ++ assert_eq!(p["stream"], false); ++ assert!(p["ec_transfer_params"]["ec_items"].is_array()); ++ let uuid = p["messages"][0]["content"][0]["uuid"].as_str().unwrap(); ++ let hash = format!("engine-{uuid}"); ++ assert_eq!(p["ec_transfer_params"]["ec_items"][0]["mm_hash"], hash); ++ assert_eq!(p["ec_transfer_params"][&hash]["peer_port"], 4321); ++ assert!(p["ec_transfer_params"].get(uuid).is_none()); ++ assert_eq!( ++ p["messages"][0]["content"][0]["image_embeds"]["image_grid_thw"], ++ json!([1, 2, 3]) ++ ); ++ if status.is_success() { ++ let d = &seen[2].1; ++ assert_eq!(d["max_tokens"], 16); ++ assert_eq!(d["messages"], p["messages"]); ++ assert_eq!(d["structured_outputs"], p["structured_outputs"]); ++ assert!(d.get("ec_transfer_params").is_none()); ++ assert_eq!(d["kv_transfer_params"]["remote_engine_id"], "P"); ++ } ++ } ++ server.abort(); ++ } ++ ++ /// With the NIXL connector, the first request caches the prefill push ++ /// identity and subsequent requests take the concurrent fast path. The ++ /// encoder stage must still run for those requests; otherwise raw media ++ /// reaches P/D without a published embedding. ++ #[tokio::test] ++ async fn test_epd_nixl_push_still_encodes_after_identity_cached() { ++ use axum::{ ++ extract::State, ++ response::IntoResponse, ++ routing::{get, post}, ++ Json, Router, ++ }; ++ use serde_json::{json, Value}; ++ use std::sync::{Arc, Mutex}; ++ use vllm_router_rs::config::{EpdConfig, KvConnector}; ++ ++ type Seen = Arc>>; ++ async fn encode( ++ State(seen): State, ++ Json(body): Json, ++ ) -> axum::response::Response { ++ seen.lock().unwrap().push(("E".into(), body.clone())); ++ let uuid = body["messages"][0]["content"][0]["uuid"].as_str().unwrap(); ++ Json(json!({ ++ "ec_transfer_params": { ++ (format!("engine-{uuid}")): { ++ "metadata": {"image_grid_thw": [[1, 2, 3]]}, ++ "peer_host": "encoder", ++ "peer_port": 4321, ++ "size_bytes": 128, ++ }, ++ }, ++ })) ++ .into_response() ++ } ++ async fn prefill( ++ State(seen): State, ++ Json(body): Json, ++ ) -> axum::response::Response { ++ seen.lock().unwrap().push(("P".into(), body.clone())); ++ if body["model"] == "prefill-422" { ++ return ( ++ axum::http::StatusCode::UNPROCESSABLE_ENTITY, ++ Json(json!({"error": "invalid sampling parameters"})), ++ ) ++ .into_response(); ++ } ++ Json(json!({ ++ "kv_transfer_params": { ++ "remote_engine_id": "P", ++ "remote_host": "127.0.0.1", ++ "remote_port": 9100, ++ "transfer_mode": "push", ++ }, ++ })) ++ .into_response() ++ } ++ async fn decode( ++ State(seen): State, ++ Json(body): Json, ++ ) -> axum::response::Response { ++ seen.lock().unwrap().push(("D".into(), body.clone())); ++ Json(json!({"choices":[{"message":{"content":"A"}}]})).into_response() ++ } ++ let seen: Seen = Arc::new(Mutex::new(Vec::new())); ++ let app = Router::new() ++ .route("/e/health", get(|| async { "ok" })) ++ .route("/p/health", get(|| async { "ok" })) ++ .route("/d/health", get(|| async { "ok" })) ++ .route("/e/v1/chat/completions", post(encode)) ++ .route("/p/v1/chat/completions", post(prefill)) ++ .route("/d/v1/chat/completions", post(decode)) ++ .route("/p/generate", post(prefill)) ++ .route("/d/generate", post(decode)) ++ .with_state(seen.clone()); ++ let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); ++ let base = format!("http://{}", listener.local_addr().unwrap()); ++ let server = tokio::spawn(async move { ++ axum::serve(listener, app).await.unwrap(); ++ }); ++ let config = RouterConfig { ++ mode: RoutingMode::VllmPrefillDecode { ++ prefill_urls: vec![(format!("{base}/p"), None)], ++ decode_urls: vec![format!("{base}/d")], ++ prefill_policy: None, ++ decode_policy: None, ++ discovery_address: None, ++ }, ++ policy: PolicyConfig::Random, ++ epd: Some(EpdConfig { ++ encoder_urls: vec![format!("{base}/e")], ++ consumer_zmq_addrs: Default::default(), ++ }), ++ kv_connector: KvConnector::Nixl, ++ worker_startup_timeout_secs: 5, ++ ..RouterConfig::default() ++ }; ++ let context = Arc::new( ++ vllm_router_rs::server::AppContext::new( ++ config, ++ reqwest::Client::new(), ++ 64, ++ None, ++ vec![], ++ ) ++ .unwrap(), ++ ); ++ let router = RouterFactory::create_router(&context).await.unwrap(); ++ ++ let chat_body = || { ++ json!({ ++ "model": "test", ++ "max_tokens": 16, ++ "messages": [{ ++ "role": "user", ++ "content": [ ++ {"type": "image_url", "image_url": {"url": "data:image/png;base64,abc"}}, ++ {"type": "text", "text": "choose"}, ++ ], ++ }], ++ }) ++ }; ++ ++ // Invalid transparent bodies must be rejected before either PD path mutates them. ++ for body in [Value::Null, json!([]), json!("text"), json!(7)] { ++ let response = router ++ .route_transparent(None, "/generate", &axum::http::Method::POST, body) ++ .await; ++ assert_eq!(response.status(), axum::http::StatusCode::BAD_REQUEST); ++ } ++ assert!(seen.lock().unwrap().is_empty()); ++ ++ // First request: no cached push identity, sequential path. ++ let response = router ++ .route_transparent( ++ None, ++ "/v1/chat/completions", ++ &axum::http::Method::POST, ++ chat_body(), ++ ) ++ .await; ++ assert_eq!(response.status(), axum::http::StatusCode::OK); ++ let _ = axum::body::to_bytes(response.into_body(), 4096).await; ++ let phases: Vec<_> = seen ++ .lock() ++ .unwrap() ++ .iter() ++ .map(|(phase, _)| phase.clone()) ++ .collect(); ++ assert_eq!(phases, ["E", "P", "D"]); ++ ++ // Second request: the push identity is now cached, so the concurrent ++ // fast path runs. The encoder must still be called and P/D must ++ // receive the rewritten embeds metadata, not raw media. ++ seen.lock().unwrap().clear(); ++ let response = router ++ .route_transparent( ++ None, ++ "/v1/chat/completions", ++ &axum::http::Method::POST, ++ chat_body(), ++ ) ++ .await; ++ assert_eq!(response.status(), axum::http::StatusCode::OK); ++ let _ = axum::body::to_bytes(response.into_body(), 4096).await; ++ let seen_requests = seen.lock().unwrap().clone(); ++ let phases: Vec<_> = seen_requests ++ .iter() ++ .map(|(phase, _)| phase.as_str()) ++ .collect(); ++ // P and D are dispatched concurrently, so only their set is stable. ++ assert_eq!(phases[0], "E"); ++ let mut pd: Vec<_> = phases[1..].to_vec(); ++ pd.sort_unstable(); ++ assert_eq!(pd, ["D", "P"]); ++ ++ let p = &seen_requests[1..] ++ .iter() ++ .find(|(phase, _)| phase == "P") ++ .expect("prefill request") ++ .1; ++ let d = &seen_requests[1..] ++ .iter() ++ .find(|(phase, _)| phase == "D") ++ .expect("decode request") ++ .1; ++ assert_eq!( ++ p["messages"][0]["content"][0]["image_embeds"]["image_grid_thw"], ++ json!([1, 2, 3]) ++ ); ++ assert!(d.get("ec_transfer_params").is_none()); ++ assert_eq!(d["kv_transfer_params"]["do_remote_prefill"], true); ++ assert_eq!(d["kv_transfer_params"]["remote_engine_id"], "P"); ++ ++ seen.lock().unwrap().clear(); ++ for body in [Value::Null, json!([]), json!("text"), json!(7)] { ++ let response = router ++ .route_transparent(None, "/generate", &axum::http::Method::POST, body) ++ .await; ++ assert_eq!(response.status(), axum::http::StatusCode::BAD_REQUEST); ++ } ++ assert!(seen.lock().unwrap().is_empty()); ++ let mut invalid = chat_body(); ++ invalid["model"] = json!("prefill-422"); ++ let response = router ++ .route_transparent( ++ None, ++ "/v1/chat/completions", ++ &axum::http::Method::POST, ++ invalid, ++ ) ++ .await; ++ assert_eq!( ++ response.status(), ++ axum::http::StatusCode::UNPROCESSABLE_ENTITY ++ ); ++ assert_eq!(response.headers()["content-type"], "application/json"); ++ let bytes = axum::body::to_bytes(response.into_body(), 4096) ++ .await ++ .unwrap(); ++ assert_eq!( ++ serde_json::from_slice::(&bytes).unwrap(), ++ json!({"error": "invalid sampling parameters"}) ++ ); ++ server.abort(); ++ } + + #[test] + fn test_worker_types() { +@@ -130,6 +572,7 @@ mod test_pd_routing { + history_backend: vllm_router_rs::config::HistoryBackend::Memory, + enable_profiling: false, + profile_timeout_secs: 30, ++ epd: None, + kv_connector: vllm_router_rs::config::KvConnector::Nixl, + }; + diff --git a/pack/llm-router/patches/vllm-router/0001-mooncake-lazy-bootstrap-lookup.patch b/pack/llm-router/patches/vllm-router/0002-mooncake-lazy-bootstrap-lookup.patch similarity index 84% rename from pack/llm-router/patches/vllm-router/0001-mooncake-lazy-bootstrap-lookup.patch rename to pack/llm-router/patches/vllm-router/0002-mooncake-lazy-bootstrap-lookup.patch index ab6df6b04..7c3af2e0d 100644 --- a/pack/llm-router/patches/vllm-router/0001-mooncake-lazy-bootstrap-lookup.patch +++ b/pack/llm-router/patches/vllm-router/0002-mooncake-lazy-bootstrap-lookup.patch @@ -41,29 +41,25 @@ handshake -- query the bootstrap server inside the request handler, against the prefill address at hand (llm-d/llm-d-router@v0.10.0, pkg/sidecar/proxy/connector_mooncake.go). -AGAINST: vllm-project/router v0.1.15 (commit -1fbcde7443d75b36befb61bc081f64c2a1f13a4b), the pin named by -VLLM_ROUTER_VERSION in pack/llm-router/Dockerfile. The defect is present in the -same shape on that project's default branch at the time of writing. The image -build applies this patch with git apply, which fails loudly on a pin whose -hunks no longer match, so a version bump that absorbs or breaks the patch is -detected at build time rather than in a cluster. +AGAINST: vllm-project/router v0.1.16.1 (commit +91b65ee4bdfbe98863a32a2205dced5dbb249ffe), the pin named by +VLLM_ROUTER_VERSION in pack/llm-router/Dockerfile, with the EPD feature +cherry-picked from upstream PR 266 (commit 516f1cb90eefe565c9fc894d72c8748eae25085b) +on top -- see 0001-epd-routing.patch. The defect is present in the same shape +on the default branch at the time of writing. The image build applies this +patch with git apply, which fails loudly on a pin whose hunks no longer +match, so a version bump that absorbs or breaks the patch is detected at +build time rather than in a cluster. -UPSTREAM REPORT: NOT YET FILED, and this sentence is where its URL goes. It is -written this way rather than left out because a patch whose report is merely -intended reads, a year later, exactly like one whose report was filed and -ignored. Filing it is what makes carrying this patch a decision with an exit -rather than a fork; until then the exit is named and not yet taken. - -TO REFRESH after a pin bump: clone the new pin, git apply this patch, resolve -any rejects, regenerate with git diff, and keep this header on top. +WHEN TO REMOVE: 0001-epd-routing.patch goes away when upstream cuts a tag +that contains PR 266's EPD commit (516f1cb). This patch (0002) is +independent of EPD and stays until upstream either ships the lazy Mooncake +fix or restructures the file enough to absorb it. diff --git a/src/routers/http/vllm_pd_router.rs b/src/routers/http/vllm_pd_router.rs -index 8e40aed..e392f3a 100644 --- a/src/routers/http/vllm_pd_router.rs +++ b/src/routers/http/vllm_pd_router.rs -@@ -96,13 +96,14 @@ fn build_prefill_request_builder( - +@@ -95,11 +95,12 @@ fn build_prefill_request_builder( impl VllmPDRouter { /// Query the Mooncake bootstrap server on a prefill node to get engine_id per dp_rank. - /// Retries with backoff since the prefill server may not be ready at router startup. @@ -77,12 +73,11 @@ index 8e40aed..e392f3a 100644 let url = format!("{}/query", bootstrap_addr); - let max_retries = 30; let mut backoff_secs = 1u64; - - for attempt in 1..=max_retries { -@@ -324,12 +325,26 @@ impl VllmPDRouter { + +@@ -354,12 +355,26 @@ impl VllmPDRouter { }) } - + - /// Look up Mooncake prefill info for a given prefill URL and dp_rank + /// Look up Mooncake prefill info for a given prefill URL and dp_rank. + /// @@ -107,10 +102,10 @@ index 8e40aed..e392f3a 100644 let info = self.mooncake_prefill_info.lock().await; if let Some(prefill_info) = info.get(prefill_url) { let rank = dp_rank.unwrap_or(0); -@@ -344,6 +359,79 @@ impl VllmPDRouter { +@@ -374,6 +389,79 @@ impl VllmPDRouter { None } - + + /// Query one prefiller's Mooncake bootstrap server and cache the result under + /// the prefiller's URL. + /// @@ -184,19 +179,19 @@ index 8e40aed..e392f3a 100644 + } + } + - /// Generate vLLM-specific request ID with prefill/decode addressing - fn generate_vllm_request_id(prefill_addr: &str, decode_addr: &str) -> String { - let uuid = Uuid::new_v4().to_string().replace('-', ""); -@@ -1625,7 +1713,7 @@ impl VllmPDRouter { - "Querying Mooncake bootstrap at {} for prefill {}", - bootstrap_addr, base_url + fn is_nixl_push_mode(kvt: &Value) -> bool { + kvt.get("transfer_mode") + .and_then(|v| v.as_str()) +@@ -2099,7 +2187,7 @@ impl VllmPDRouter { ); - match Self::query_mooncake_bootstrap(&http_client, &bootstrap_addr).await { + match Self::query_mooncake_bootstrap(&http_client, &bootstrap_addr, 30).await { Ok(dp_engine_ids) => { info!( "Got Mooncake engine_ids for {}: {:?}", -@@ -1683,6 +1771,10 @@ impl VllmPDRouter { + base_url, dp_engine_ids + ); +@@ -2164,6 +2252,10 @@ impl VllmPDRouter { /// Remove a prefill server from the router /// Delegates to the underlying PdRouterBase pub async fn remove_prefill_server(&self, url: &str) -> Result { @@ -204,6 +199,7 @@ index 8e40aed..e392f3a 100644 + // comes up at this same URL must be queried, not served the previous + // owner's engine ids. + self.mooncake_prefill_info.lock().await.remove(url); + self.evict_nixl_info(url); self.pd_router.remove_prefill_server(url).await } - + diff --git a/pack/llm-router/patches/vllm-router/0003-observer-read-only-endpoint-view.patch b/pack/llm-router/patches/vllm-router/0003-observer-read-only-endpoint-view.patch index 1d0d9e039..7402b68c6 100644 --- a/pack/llm-router/patches/vllm-router/0003-observer-read-only-endpoint-view.patch +++ b/pack/llm-router/patches/vllm-router/0003-observer-read-only-endpoint-view.patch @@ -1,14 +1,29 @@ -Read native registry and effective circuit state without changing request routing state. -An unstable SGLang circuit sample refuses the complete view. +Read native registry membership and effective circuit state without changing routing. +The selection fields use the native health and circuit predicates. Pod hints are +address echoes; the operator binds them to Kubernetes identities. Boot +generation is initialized once per process. The view revision changes when +the serialized worker array differs from the preceding observer response. + +AGAINST: vllm-project/router v0.1.16.1 (commit +91b65ee4bdfbe98863a32a2205dced5dbb249ffe) with PR 266 (EPD, commit +516f1cb90eefe565c9fc894d72c8748eae25085b) cherry-picked on top (see +0001-epd-routing.patch), plus the lazy-Mooncake patch (0002). The image +build applies this patch with git apply; a pin whose hunks no longer +match fails loudly and surfaces at build time, not in a cluster. + +WHEN TO REMOVE: when upstream rewires circuit-breaker observability or its +representation of an "effective" state, this patch becomes redundant. Until +then the operator's SGLang detection and its ClusterQueue-aware admission +both depend on the read-only observer view that this patch defines. diff --git a/src/core/circuit_breaker.rs b/src/core/circuit_breaker.rs -index 5c37423..b6499bb 100644 --- a/src/core/circuit_breaker.rs +++ b/src/core/circuit_breaker.rs @@ -100,6 +100,54 @@ impl CircuitBreaker { + self.check_and_update_state(); *self.state.read().unwrap() } - + + /// Read the effective state without changing native routing state. + pub fn observe_effective(&self) -> (CircuitState, bool) { + let state_guard = self.state.read().unwrap(); @@ -60,13 +75,14 @@ index 5c37423..b6499bb 100644 /// Record the outcome of a request pub fn record_outcome(&self, success: bool) { if success { + diff --git a/src/lib.rs b/src/lib.rs -index 5c4c406..8388376 100644 --- a/src/lib.rs +++ b/src/lib.rs -@@ -6,6 +6,7 @@ use std::collections::HashMap; - pub mod core; +@@ -15,6 +15,7 @@ pub mod core; pub mod data_connector; + pub mod kv_hash; + pub mod kv_index; pub mod metrics; +pub mod observer; pub mod middleware; @@ -376,11 +392,11 @@ index 0000000..9c20b98 + run_case(CircuitState::HalfOpen, false, false).await; + } +} + diff --git a/src/server.rs b/src/server.rs -index d7c16df..08d8ac0 100644 --- a/src/server.rs +++ b/src/server.rs -@@ -749,7 +749,9 @@ pub fn build_app_with_request_tracing( +@@ -849,7 +849,9 @@ pub fn build_app_with_request_tracing( .route("/health_generate", get(health_generate)) .route("/v1/models", get(v1_models)) .route("/get_model_info", get(get_model_info)) diff --git a/pack/llm-router/patches/vllm-router/0002-discovery-readiness-and-deselect-removal.patch b/pack/llm-router/patches/vllm-router/0004-discovery-readiness-and-deselect-removal.patch similarity index 80% rename from pack/llm-router/patches/vllm-router/0002-discovery-readiness-and-deselect-removal.patch rename to pack/llm-router/patches/vllm-router/0004-discovery-readiness-and-deselect-removal.patch index a86c53d10..319a7441f 100644 --- a/pack/llm-router/patches/vllm-router/0002-discovery-readiness-and-deselect-removal.patch +++ b/pack/llm-router/patches/vllm-router/0004-discovery-readiness-and-deselect-removal.patch @@ -1,106 +1,61 @@ -Stop selecting a worker once it loses readiness, selector membership, -termination, or a verifiable current state, and make that withdrawal identity-correct -and reversible. +Withdraw a worker from the router when its Pod loses readiness, leaves the +selector, terminates, or is replaced by another identity at the same address. +The discovery tracker is keyed on Pod UID and stays consistent across the +rest of the operator's patches. -THE DEFECT, in four parts that hide each other. The Kubernetes watcher watches all Pods -and applies the label selector in a client-side filter -(src/service_discovery.rs): the stream was consumed through applied_objects() and then -passed through a filter_map that dropped any Pod not matching the selector. A Pod that -loses the eligibility label therefore never reached the handler at all -- a filter can -only discard objects, and a withdrawal is exactly the case where the object is still -there. Second, the same filter_map discarded the Delete event and the Init and -InitDone arms of the relist. What applied_objects() actually does is decode an Event -into the object it carries and, with emit_deleted false, pass Apply and InitApply -through while discarding Delete, Init and InitDone: upstream therefore still saw each -object of a relist, and never saw the point at which the watcher states that an object -previously Applied is now absent, which is why an abruptly removed worker stayed -registered for the life of the process. Third, handle_pod_event had a branch for a -healthy Pod and no branch at all for an unhealthy one, so even a Pod that did arrive -with readiness lost was silently ignored, and the deletion-timestamp branch that -upstream did have was reachable only after a PodInfo had been rebuilt from the current -object. +THE DEFECT: src/service_discovery.rs (v0.1.16.1) tracks pods in +HashSet by value and forwards each Kubernetes event to a handler +that filters by selector, parses the status, and registers the resulting +worker URL. A pod that loses readiness, has its selector label removed, or +enters termination is still delivered to the watcher; the handler either +admits it as a fresh worker or silently drops the registration without +ever telling the router. Workers serving to a participant that is no longer +healthy stay selectable, and a same-name replacement at the same address +gets admitted alongside the previous owner -- a later Delete event for the +old UID then withdraws the current occupant from the router. -Behind all three sat a fourth problem that made a partial fix useless rather than -wrong. The tracker was a HashSet and PodInfo derives PartialEq, Eq and Hash -over the whole struct, including is_ready, status and ip. A Ready -> NotReady -transition, a phase change or an address change all produce a PodInfo that no longer -compares equal to the registered one, so any removal looked up by the current PodInfo -would MISS the live entry and the worker would stay in the pool; and because the -address is part of the key, a same-name replacement that reused an address aliased the -previous identity. +THE FIX: the tracker becomes HashMap keyed on Pod UID, the +PodInfo struct gains a `uid` field, and handle_watch_event gains four +withdrawal decisions (readiness loss, selector miss, terminating status, +relist absence) plus a take-over step that withdraws a previous owner at +the same address before admitting its replacement. A late Delete for a +replaced UID finds nothing to withdraw and cannot reach the current +occupant. Pods whose UID is absent or empty are refused at from_pod() so +they cannot be admitted as workers that can never be reliably removed. -THE FIX. The watch stream is consumed as raw Events and the selector decision moved -into the handler, where the Pod itself is in hand: a Pod that is terminating, no longer -matches the selectors, or is present but not healthy is a withdrawal. Termination is -checked FIRST, before any status is parsed and before selector membership, because a -terminating Pod can still report Ready=True and Running, and reading health first would -re-register it and leave the withdrawal to the final Delete event -- the preStop path -this replaces. Apply and Delete are handled directly, and the relist is implemented as -the contract describes: Init opens a set, InitApply records UIDs, InitDone withdraws -every tracked UID the listing did not contain. +Equality and hashing for PodInfo intentionally ignore `uid`: the tracker +keys on UID separately, but value equality inside tests and downstream +consumers compares the worker itself. A manual PartialEq / Eq / Hash pair +on the routing-relevant fields (name, ip, status, is_ready, pod_type, +bootstrap_port) keeps the Eq/Hash contract sound -- a derive that included +uid would compare uids that are deliberately distinct on identical workers. +Eleven new tests cover the withdrawal and take-over paths; they use a +StubRouter (implements WorkerManagement and RouterTrait, routing methods +return 501) instead of a real Router, because Router::add_worker does an +HTTP health check that fails against the synthetic addresses the tests +construct. The 11 tests assert on tracker state, not on router HTTP +behaviour, and the stub is sufficient. -The tracker becomes a HashMap keyed on the Pod UID, and the stored value is the PodInfo -the worker was REGISTERED with rather than as it is now, so a withdrawal issued for an -old identity removes that old URL and role. A UID-keyed map alone does not make that -true at the router, because the router is keyed by URL: if a new UID were admitted at -an address the old UID still held, the map would carry both, and a later Delete for the -old UID would remove the CURRENT occupant from the pool. register_worker therefore -reconciles address ownership before admitting a new UID -- any other UID registered at -the same address is withdrawn first, and a failed withdrawal holds the new add -- so an -address has exactly one owner and a late delete for a replaced UID finds nothing to act -on. A Pod whose UID is absent OR EMPTY is refused rather than admitted, because an -empty string is not an identity: it is the one value every unidentified Pod shares, so -it would alias unrelated Pods into one registration and make a later withdrawal remove -the wrong worker. +AGAINST: vllm-project/router v0.1.16.1 (commit +91b65ee4bdfbe98863a32a2205dced5dbb249ffe) with PR 266 (EPD, commit +516f1cb90eefe565c9fc894d72c8748eae25085b) cherry-picked on top +(0001-epd-routing.patch), plus the lazy-Mooncake patch (0002) and the +read-only observer patch (0003). The defect is present in the same shape +on the default branch at the time of writing. The image build applies +this patch with git apply after 0001, 0002, 0003; a pin whose hunks no +longer match fails loudly and surfaces at build time, not in a cluster. -A Pod that still matches the selectors but whose current status or address cannot be -read is WITHDRAWN, not retained. Eligibility is not health, and a registration that -cannot be shown healthy must not stay selectable. - -A registration is dropped only after the router has actually been told, and a failure -returns false rather than being swallowed. The watcher loop treats that false as a -reason to end the stream, so the existing backoff recreates the watcher and the relist -that comes with it reconciles whatever is still held. That is the only retry added here: -no second timer, no polling loop, and the failed withdrawal is not left waiting for a -Pod event that may never arrive. The backoff is taken alone on that path rather than -followed by the full check-interval sleep, which would push the retry past any sensible -withdrawal budget. The first error is what is logged, on every attempt. - -P/D behaviour is preserved: the add and remove paths keep the existing VllmPDRouter -prefill/decode routing, bootstrap port handling and the regular-mode fallbacks, and are -only reached from the new decision points. - -AGAINST: vllm-project/router v0.1.15 (commit -1fbcde7443d75b36befb61bc081f64c2a1f13a4b), the pin named by -VLLM_ROUTER_VERSION in pack/llm-router/Dockerfile. The image build applies this patch -with git apply from the clone root, which fails loudly on a pin whose hunks no longer -match, so a version bump that absorbs or breaks the patch is detected at build time -rather than in a cluster. - -The added tests drive the discovery decisions through handle_watch_event and assert on -the tracker and the router: readiness loss withdraws while a healthy observation of the -same identity does not; a label deselect withdraws; a terminating Pod withdraws both -while reporting Ready and with no status at all; a relist withdraws an absent UID and -keeps a listed one; a new UID takes an address over from its previous owner, and a late -delete for the replaced UID then cannot reach the current owner; an eligible Pod whose -current status is unusable withdraws; an absent or empty UID is refused; a withdrawal -drops only the exact registration it removed; a delete for an untracked UID is a -delivered no-op. - -The successful-insert leg is not asserted here: Router::add_worker polls the worker's -health endpoint, so a registration that is genuinely added cannot be observed with the -in-crate test router. The full add/withdraw/re-add round trip, and the ordering of a -replacement against a late delete, are measured against the built binary on placement -evidence in the operator's event-matrix harness rather than claimed from a unit test. - -TO REFRESH after a pin bump: clone the new pin, git apply this patch, resolve -any rejects, regenerate with git diff, and keep this header on top. +WHEN TO REMOVE: when upstream restructures service discovery so that a +pod losing readiness, leaving the selector, or entering termination +withdraws the worker from the router. Until then the EPD + Kubernetes +service discovery path this patch fixes is a regression vs the operator's +prior baseline. diff --git a/src/service_discovery.rs b/src/service_discovery.rs -index d765800..798eab7 100644 +index d765800..fc5b063 100644 --- a/src/service_discovery.rs +++ b/src/service_discovery.rs -@@ -1,11 +1,10 @@ +@@ -1,14 +1,14 @@ use crate::routers::RouterTrait; -use futures::{StreamExt, TryStreamExt}; @@ -115,15 +70,61 @@ index d765800..798eab7 100644 Client, }; use std::collections::{HashMap, HashSet}; -@@ -62,6 +61,7 @@ pub enum PodType { - #[derive(Debug, Clone, PartialEq, Eq, Hash)] ++use std::hash::{Hash, Hasher}; + use std::net::{IpAddr, SocketAddr}; + + use rustls; +@@ -59,9 +59,17 @@ pub enum PodType { + } + + /// Represents a Kubernetes pod's information used for worker management +-#[derive(Debug, Clone, PartialEq, Eq, Hash)] ++/// ++/// Equality and hashing intentionally ignore `uid`. The discovery tracker is ++/// keyed on `uid` separately (`HashMap` in ++/// `handle_watch_event`); identity inside the tracker is the UID, while value ++/// equality inside tests and downstream consumers compares the worker itself. ++/// Two `PodInfo`s describe the same worker when their routing-relevant ++/// fields agree, regardless of which Pod incarnation produced them. ++#[derive(Debug, Clone)] pub struct PodInfo { pub name: String, + pub uid: String, pub ip: IpAddr, pub status: String, pub is_ready: bool, -@@ -106,6 +106,12 @@ impl PodInfo { +@@ -69,6 +77,30 @@ pub struct PodInfo { + pub bootstrap_port: Option, + } + ++impl PartialEq for PodInfo { ++ fn eq(&self, other: &Self) -> bool { ++ self.name == other.name ++ && self.ip == other.ip ++ && self.status == other.status ++ && self.is_ready == other.is_ready ++ && self.pod_type == other.pod_type ++ && self.bootstrap_port == other.bootstrap_port ++ } ++} ++ ++impl Eq for PodInfo {} ++ ++impl Hash for PodInfo { ++ fn hash(&self, state: &mut H) { ++ self.name.hash(state); ++ self.ip.hash(state); ++ self.status.hash(state); ++ self.is_ready.hash(state); ++ self.pod_type.hash(state); ++ self.bootstrap_port.hash(state); ++ } ++} ++ + impl PodInfo { + /// Check if a pod matches any of the given selectors + fn matches_selector(pod: &Pod, selector: &HashMap) -> bool { +@@ -106,6 +138,12 @@ impl PodInfo { /// Unified PodInfo creation with optional PD configuration pub fn from_pod(pod: &Pod, config: Option<&ServiceDiscoveryConfig>) -> Option { let name = pod.metadata.name.clone()?; @@ -136,7 +137,7 @@ index d765800..798eab7 100644 let status = pod.status.clone()?; let pod_ip: IpAddr = status.pod_ip?.parse().ok()?; -@@ -155,6 +161,7 @@ impl PodInfo { +@@ -155,6 +193,7 @@ impl PodInfo { Some(PodInfo { name, @@ -144,7 +145,7 @@ index d765800..798eab7 100644 ip: pod_ip, status: pod_status, is_ready, -@@ -230,8 +237,18 @@ pub async fn start_service_discovery( +@@ -230,8 +269,18 @@ pub async fn start_service_discovery( // Create the task that will run in the background let handle = task::spawn(async move { @@ -165,7 +166,7 @@ index d765800..798eab7 100644 // Create a watcher for pods let pods: Api = if let Some(namespace) = &config.namespace { -@@ -251,88 +268,97 @@ pub async fn start_service_discovery( +@@ -251,88 +300,97 @@ pub async fn start_service_discovery( loop { // Create a watcher with the proper parameters according to the kube-rs API @@ -335,7 +336,7 @@ index d765800..798eab7 100644 // If the watcher exits for some reason, wait a bit before restarting warn!( "Kubernetes watcher exited, restarting in {} seconds", -@@ -345,150 +371,438 @@ pub async fn start_service_discovery( +@@ -345,158 +403,615 @@ pub async fn start_service_discovery( Ok(handle) } @@ -864,7 +865,7 @@ index d765800..798eab7 100644 - ); + // Regular mode removal + router.remove_worker(&worker_url); - } ++ } + + // Only now is the registration dropped, because the router has actually been + // told. Guarded on the exact PodInfo so a concurrent update for the same UID, @@ -877,12 +878,189 @@ index d765800..798eab7 100644 + { + tracker.remove(®istered.uid); + } -+ } + } + true } #[cfg(test)] -@@ -509,6 +823,10 @@ mod tests { + mod tests { + use super::*; ++ use axum::{ ++ body::Body, ++ http::Request, ++ response::{IntoResponse, Response}, ++ }; + use k8s_openapi::api::core::v1::{Pod, PodCondition, PodSpec, PodStatus}; + use k8s_openapi::apimachinery::pkg::apis::meta::v1::ObjectMeta; + use k8s_openapi::apimachinery::pkg::apis::meta::v1::Time; ++ use std::sync::Mutex as StdMutex; ++ ++ // A router that records add/remove calls and always succeeds. ++ // ++ // The real Router::add_worker does an HTTP health check against the worker URL, ++ // which fails in unit tests because the URL points at a synthetic address. ++ // The new discovery tests assert on tracker state, not on the router's ++ // HTTP behaviour, so the router here only needs to: ++ // 1. Always succeed on add_worker. ++ // 2. Track which URLs have been added/removed. ++ // 3. Return them from get_worker_urls. ++ // Routing methods are stubbed to 501 since none of the discovery tests ++ // exercise them. ++ struct StubRouter { ++ urls: StdMutex>, ++ } ++ ++ impl StubRouter { ++ fn new() -> Self { ++ Self { ++ urls: StdMutex::new(HashSet::new()), ++ } ++ } ++ } ++ ++ impl std::fmt::Debug for StubRouter { ++ fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { ++ f.debug_struct("StubRouter").finish() ++ } ++ } ++ ++ #[async_trait::async_trait] ++ impl crate::routers::WorkerManagement for StubRouter { ++ async fn add_worker(&self, worker_url: &str) -> Result { ++ self.urls.lock().unwrap().insert(worker_url.to_string()); ++ Ok(worker_url.to_string()) ++ } ++ ++ fn remove_worker(&self, worker_url: &str) { ++ self.urls.lock().unwrap().remove(worker_url); ++ } ++ ++ fn get_worker_urls(&self) -> Vec { ++ self.urls.lock().unwrap().iter().cloned().collect() ++ } ++ } ++ ++ #[async_trait::async_trait] ++ impl crate::routers::RouterTrait for StubRouter { ++ fn as_any(&self) -> &dyn std::any::Any { ++ self ++ } ++ ++ async fn health(&self, _req: Request) -> Response { ++ axum::http::StatusCode::NOT_IMPLEMENTED.into_response() ++ } ++ async fn health_generate(&self, _req: Request) -> Response { ++ axum::http::StatusCode::NOT_IMPLEMENTED.into_response() ++ } ++ async fn get_server_info(&self, _req: Request) -> Response { ++ axum::http::StatusCode::NOT_IMPLEMENTED.into_response() ++ } ++ async fn get_models(&self, _req: Request) -> Response { ++ axum::http::StatusCode::NOT_IMPLEMENTED.into_response() ++ } ++ async fn get_model_info(&self, _req: Request) -> Response { ++ axum::http::StatusCode::NOT_IMPLEMENTED.into_response() ++ } ++ async fn route_generate( ++ &self, ++ _headers: Option<&axum::http::HeaderMap>, ++ _body: &crate::protocols::spec::GenerateRequest, ++ _model_id: Option<&str>, ++ ) -> Response { ++ axum::http::StatusCode::NOT_IMPLEMENTED.into_response() ++ } ++ async fn route_inference_generate( ++ &self, ++ _headers: Option<&axum::http::HeaderMap>, ++ _body: &crate::protocols::spec::InferenceGenerateRequest, ++ _model_id: Option<&str>, ++ ) -> Response { ++ axum::http::StatusCode::NOT_IMPLEMENTED.into_response() ++ } ++ async fn route_chat( ++ &self, ++ _headers: Option<&axum::http::HeaderMap>, ++ _body: &crate::protocols::spec::ChatCompletionRequest, ++ _model_id: Option<&str>, ++ ) -> Response { ++ axum::http::StatusCode::NOT_IMPLEMENTED.into_response() ++ } ++ async fn route_completion( ++ &self, ++ _headers: Option<&axum::http::HeaderMap>, ++ _body: &crate::protocols::spec::CompletionRequest, ++ _model_id: Option<&str>, ++ ) -> Response { ++ axum::http::StatusCode::NOT_IMPLEMENTED.into_response() ++ } ++ async fn route_responses( ++ &self, ++ _headers: Option<&axum::http::HeaderMap>, ++ _body: &crate::protocols::spec::ResponsesRequest, ++ _model_id: Option<&str>, ++ ) -> Response { ++ axum::http::StatusCode::NOT_IMPLEMENTED.into_response() ++ } ++ async fn get_response( ++ &self, ++ _headers: Option<&axum::http::HeaderMap>, ++ _response_id: &str, ++ ) -> Response { ++ axum::http::StatusCode::NOT_IMPLEMENTED.into_response() ++ } ++ async fn cancel_response( ++ &self, ++ _headers: Option<&axum::http::HeaderMap>, ++ _response_id: &str, ++ ) -> Response { ++ axum::http::StatusCode::NOT_IMPLEMENTED.into_response() ++ } ++ async fn route_embeddings( ++ &self, ++ _headers: Option<&axum::http::HeaderMap>, ++ _body: &crate::protocols::spec::EmbeddingRequest, ++ _model_id: Option<&str>, ++ ) -> Response { ++ axum::http::StatusCode::NOT_IMPLEMENTED.into_response() ++ } ++ async fn route_rerank( ++ &self, ++ _headers: Option<&axum::http::HeaderMap>, ++ _body: &crate::protocols::spec::RerankRequest, ++ _model_id: Option<&str>, ++ ) -> Response { ++ axum::http::StatusCode::NOT_IMPLEMENTED.into_response() ++ } ++ async fn flush_cache(&self) -> Response { ++ axum::http::StatusCode::NOT_IMPLEMENTED.into_response() ++ } ++ async fn get_worker_loads(&self) -> Response { ++ axum::http::StatusCode::NOT_IMPLEMENTED.into_response() ++ } ++ fn router_type(&self) -> &'static str { ++ "stub" ++ } ++ fn readiness(&self) -> Response { ++ axum::http::StatusCode::OK.into_response() ++ } ++ async fn route_transparent( ++ &self, ++ _headers: Option<&axum::http::HeaderMap>, ++ _path: &str, ++ _method: &axum::http::Method, ++ _body: serde_json::Value, ++ ) -> Response { ++ axum::http::StatusCode::NOT_IMPLEMENTED.into_response() ++ } ++ } ++ ++ fn stub_router() -> std::sync::Arc { ++ std::sync::Arc::new(StubRouter::new()) ++ } + + // Helper function to create a Pod for testing PodInfo::from_pod + fn create_k8s_pod( +@@ -509,6 +1024,10 @@ mod tests { let mut pod = Pod { metadata: ObjectMeta { name: name.map(String::from), @@ -893,7 +1071,32 @@ index d765800..798eab7 100644 deletion_timestamp, ..Default::default() }, -@@ -832,6 +1150,7 @@ mod tests { +@@ -555,6 +1074,9 @@ mod tests { + Pod { + metadata: ObjectMeta { + name: Some(name.to_string()), ++ // from_pod refuses Pods without a UID, and the discovery tracker ++ // is keyed on UID. Each PD fixture gets its own. ++ uid: Some(format!("uid-{}", name)), + labels: Some(labels), + annotations: Some(annotations), + ..Default::default() +@@ -577,7 +1099,13 @@ mod tests { + } + } + +- // Helper to create a Router instance for testing event handlers ++ // Helper to create a Router instance for testing event handlers. ++ // ++ // Kept alongside `stub_router` for reference: the new event-handler tests ++ // use `stub_router` because the real Router's HTTP health check fails ++ // against the synthetic worker URLs those tests construct. This helper ++ // remains in case a non-stubbed path is added later. ++ #[allow(dead_code)] + async fn create_test_router() -> Arc { + use crate::config::RouterConfig; + use crate::middleware::TokenBucket; +@@ -832,6 +1360,7 @@ mod tests { fn test_pod_info_is_healthy() { let healthy_pod = PodInfo { name: "p1".into(), @@ -901,7 +1104,7 @@ index d765800..798eab7 100644 ip: "1.1.1.1".parse().unwrap(), status: "Running".into(), is_ready: true, -@@ -842,6 +1161,7 @@ mod tests { +@@ -842,6 +1371,7 @@ mod tests { let not_ready_pod = PodInfo { name: "p2".into(), @@ -909,7 +1112,7 @@ index d765800..798eab7 100644 ip: "1.1.1.2".parse().unwrap(), status: "Running".into(), is_ready: false, -@@ -852,6 +1172,7 @@ mod tests { +@@ -852,6 +1382,7 @@ mod tests { let not_running_pod = PodInfo { name: "p3".into(), @@ -917,7 +1120,7 @@ index d765800..798eab7 100644 ip: "1.1.1.3".parse().unwrap(), status: "Pending".into(), is_ready: true, -@@ -865,6 +1186,7 @@ mod tests { +@@ -865,6 +1396,7 @@ mod tests { fn test_pod_info_worker_url_ipv4() { let pod_info = PodInfo { name: "p1".into(), @@ -925,7 +1128,7 @@ index d765800..798eab7 100644 ip: "1.2.3.4".parse().unwrap(), status: "Running".into(), is_ready: true, -@@ -878,6 +1200,7 @@ mod tests { +@@ -878,6 +1410,7 @@ mod tests { fn test_pod_info_worker_url_ipv6() { let pod_info = PodInfo { name: "p1".into(), @@ -933,7 +1136,7 @@ index d765800..798eab7 100644 ip: "2803:6086:5cd3:6062:c68d:e7c7:480:10".parse().unwrap(), status: "Running".into(), is_ready: true, -@@ -894,6 +1217,7 @@ mod tests { +@@ -894,6 +1427,7 @@ mod tests { fn test_pod_info_equality_with_pod_type() { let pod1 = PodInfo { name: "pod1".into(), @@ -941,7 +1144,7 @@ index d765800..798eab7 100644 ip: "1.2.3.4".parse().unwrap(), status: "Running".into(), is_ready: true, -@@ -903,6 +1227,7 @@ mod tests { +@@ -903,6 +1437,7 @@ mod tests { let pod2 = PodInfo { name: "pod1".into(), @@ -949,7 +1152,7 @@ index d765800..798eab7 100644 ip: "1.2.3.4".parse().unwrap(), status: "Running".into(), is_ready: true, -@@ -912,6 +1237,7 @@ mod tests { +@@ -912,6 +1447,7 @@ mod tests { let pod3 = PodInfo { name: "pod1".into(), @@ -957,7 +1160,7 @@ index d765800..798eab7 100644 ip: "1.2.3.4".parse().unwrap(), status: "Running".into(), is_ready: true, -@@ -923,269 +1249,376 @@ mod tests { +@@ -923,269 +1459,376 @@ mod tests { assert_ne!(pod1, pod3); } @@ -970,6 +1173,24 @@ index d765800..798eab7 100644 - ip: "1.2.3.4".parse().unwrap(), - status: "Pending".into(), - is_ready: false, +- pod_type: None, +- bootstrap_port: None, +- }; +- let port = 8080u16; +- +- handle_pod_event( +- &pod_info, +- Arc::clone(&tracked_pods), +- Arc::clone(&router), +- port, +- false, // pd_mode = false +- ) +- .await; +- +- assert!(!tracked_pods.lock().unwrap().contains(&pod_info)); +- assert!(!router +- .get_worker_urls() +- .contains(&pod_info.worker_url(port))); + // The tracker is keyed on Pod UID and holds the PodInfo a worker was registered + // with, so these tests drive the discovery decisions through handle_watch_event + // and assert on the tracker and the router rather than on logs. @@ -979,28 +1200,33 @@ index d765800..798eab7 100644 + map.insert(entry.uid.clone(), entry); + } + Arc::new(Mutex::new(map)) -+ } -+ + } + +- #[tokio::test] +- async fn test_handle_pod_deletion_non_existing_pod() { +- let router = create_test_router().await; +- let tracked_pods = Arc::new(Mutex::new(HashSet::new())); +- let pod_info = PodInfo { +- name: "pod1".into(), +- ip: "1.2.3.4".parse().unwrap(), + fn registered(uid: &str, name: &str, ip: &str) -> PodInfo { + PodInfo { + name: name.to_string(), + uid: uid.to_string(), + ip: ip.parse().unwrap(), -+ status: "Running".into(), -+ is_ready: true, + status: "Running".into(), + is_ready: true, pod_type: None, bootstrap_port: None, - }; - let port = 8080u16; - -- handle_pod_event( +- handle_pod_deletion( - &pod_info, - Arc::clone(&tracked_pods), - Arc::clone(&router), - port, - false, // pd_mode = false -- ) -- .await; + } + } + @@ -1037,11 +1263,7 @@ index d765800..798eab7 100644 + pod.metadata.deletion_timestamp = Some(Time(k8s_openapi::chrono::Utc::now())); + pod + } - -- assert!(!tracked_pods.lock().unwrap().contains(&pod_info)); -- assert!(!router -- .get_worker_urls() -- .contains(&pod_info.worker_url(port))); ++ + async fn drive( + event: Event, + tracked_pods: &Arc>>, @@ -1056,31 +1278,16 @@ index d765800..798eab7 100644 + router, + config.port, + config, -+ ) + ) +- .await; + .await - } ++ } - #[tokio::test] -- async fn test_handle_pod_deletion_non_existing_pod() { +- assert!(tracked_pods.lock().unwrap().is_empty()); +- assert!(router.get_worker_urls().is_empty()); ++ #[tokio::test] + async fn test_readiness_loss_withdraws_worker_and_restore_does_not() { - let router = create_test_router().await; -- let tracked_pods = Arc::new(Mutex::new(HashSet::new())); -- let pod_info = PodInfo { -- name: "pod1".into(), -- ip: "1.2.3.4".parse().unwrap(), -- status: "Running".into(), -- is_ready: true, -- pod_type: None, -- bootstrap_port: None, -- }; -- let port = 8080u16; -- -- handle_pod_deletion( -- &pod_info, -- Arc::clone(&tracked_pods), -- Arc::clone(&router), -- port, -- false, // pd_mode = false ++ let router = stub_router(); + let config = regular_config(); + let entry = registered("uid-a", "pod-a", "10.0.0.1"); + let worker_url = entry.worker_url(config.port); @@ -1093,11 +1300,8 @@ index d765800..798eab7 100644 + &tracked, + &router, + &config, - ) - .await; -- -- assert!(tracked_pods.lock().unwrap().is_empty()); -- assert!(router.get_worker_urls().is_empty()); ++ ) ++ .await; + assert!(delivered); + assert!( + !tracked.lock().unwrap().contains_key("uid-a"), @@ -1124,8 +1328,7 @@ index d765800..798eab7 100644 #[tokio::test] - async fn test_handle_pd_pod_event_prefill_pod() { -+ async fn test_label_deselect_withdraws_worker() { - let router = create_test_router().await; +- let router = create_test_router().await; - let tracked_pods = Arc::new(Mutex::new(HashSet::new())); - let pod_info = PodInfo { - name: "prefill-pod".into(), @@ -1145,6 +1348,8 @@ index d765800..798eab7 100644 - Arc::clone(&router), - port, - false, // pd_mode = false, so it should fallback to regular handling ++ async fn test_label_deselect_withdraws_worker() { ++ let router = stub_router(); + let config = regular_config(); + let entry = registered("uid-b", "pod-b", "10.0.0.2"); + let worker_url = entry.worker_url(config.port); @@ -1172,8 +1377,7 @@ index d765800..798eab7 100644 #[tokio::test] - async fn test_handle_pd_pod_event_decode_pod() { -+ async fn test_terminating_pod_withdraws_even_while_ready() { - let router = create_test_router().await; +- let router = create_test_router().await; - let tracked_pods = Arc::new(Mutex::new(HashSet::new())); - let pod_info = PodInfo { - name: "decode-pod".into(), @@ -1191,6 +1395,8 @@ index d765800..798eab7 100644 - Arc::clone(&router), - port, - false, // pd_mode = false, so it should fallback to regular handling ++ async fn test_terminating_pod_withdraws_even_while_ready() { ++ let router = stub_router(); + let config = regular_config(); + let entry = registered("uid-t", "pod-t", "10.0.0.12"); + @@ -1222,8 +1428,7 @@ index d765800..798eab7 100644 #[tokio::test] - async fn test_handle_pd_pod_deletion_tracked_pod() { -+ async fn test_terminating_pod_without_status_withdraws() { - let router = create_test_router().await; +- let router = create_test_router().await; - let tracked_pods = Arc::new(Mutex::new(HashSet::new())); - let pod_info = PodInfo { - name: "test-pod".into(), @@ -1233,6 +1438,8 @@ index d765800..798eab7 100644 - pod_type: Some(PodType::Prefill), - bootstrap_port: Some(8081), - }; ++ async fn test_terminating_pod_without_status_withdraws() { ++ let router = stub_router(); + let config = regular_config(); + let entry = registered("uid-t2", "pod-t2", "10.0.0.13"); + @@ -1252,7 +1459,7 @@ index d765800..798eab7 100644 - // Add pod to tracked set first + #[tokio::test] + async fn test_relist_withdraws_pod_missing_from_listing() { -+ let router = create_test_router().await; ++ let router = stub_router(); + let config = regular_config(); + let gone = registered("uid-gone", "pod-gone", "10.0.0.3"); + let present = registered("uid-here", "pod-here", "10.0.0.4"); @@ -1310,8 +1517,7 @@ index d765800..798eab7 100644 #[tokio::test] - async fn test_handle_pd_pod_deletion_untracked_pod() { -+ async fn test_new_uid_takes_over_address_from_previous_owner() { - let router = create_test_router().await; +- let router = create_test_router().await; - let tracked_pods = Arc::new(Mutex::new(HashSet::new())); - let pod_info = PodInfo { - name: "untracked-pod".into(), @@ -1331,6 +1537,8 @@ index d765800..798eab7 100644 - Arc::clone(&router), - port, - true, // pd_mode = true ++ async fn test_new_uid_takes_over_address_from_previous_owner() { ++ let router = stub_router(); + let config = regular_config(); + let old = registered("uid-old", "pod-0", "10.0.0.5"); + @@ -1360,8 +1568,7 @@ index d765800..798eab7 100644 #[tokio::test] - async fn test_unified_handler_regular_mode() { -+ async fn test_late_delete_of_replaced_uid_cannot_reach_current_owner() { - let router = create_test_router().await; +- let router = create_test_router().await; - let tracked_pods = Arc::new(Mutex::new(HashSet::new())); - let pod_info = PodInfo { - name: "regular-pod".into(), @@ -1380,6 +1587,8 @@ index d765800..798eab7 100644 - Arc::clone(&router), - port, - false, // pd_mode = false ++ async fn test_late_delete_of_replaced_uid_cannot_reach_current_owner() { ++ let router = stub_router(); + let config = regular_config(); + let new = registered("uid-new2", "pod-0", "10.0.0.6"); + @@ -1402,8 +1611,7 @@ index d765800..798eab7 100644 #[tokio::test] - async fn test_unified_handler_pd_mode_with_prefill() { -+ async fn test_unreadable_pod_status_withdraws_registration() { - let router = create_test_router().await; +- let router = create_test_router().await; - let tracked_pods = Arc::new(Mutex::new(HashSet::new())); - let pod_info = PodInfo { - name: "prefill-pod".into(), @@ -1424,6 +1632,8 @@ index d765800..798eab7 100644 - true, // pd_mode = true - ) - .await; ++ async fn test_unreadable_pod_status_withdraws_registration() { ++ let router = stub_router(); + let config = regular_config(); + let entry = registered("uid-i", "pod-i", "10.0.0.7"); + let tracked = tracker_with(vec![entry]); @@ -1440,10 +1650,12 @@ index d765800..798eab7 100644 + "an eligible pod whose current status is unusable must be withdrawn" + ); + } -+ + +- // Pod should not be tracked since router.add_pd_worker will fail for regular router +- assert!(!tracked_pods.lock().unwrap().contains(&pod_info)); + #[tokio::test] + async fn test_pod_with_empty_uid_is_refused() { -+ let router = create_test_router().await; ++ let router = stub_router(); + let config = regular_config(); + let tracked = tracker_with(vec![]); + @@ -1452,9 +1664,7 @@ index d765800..798eab7 100644 + let mut pod = labelled_pod("uid-x", "pod-x", "10.0.0.7", true, true); + pod.metadata.uid = Some(String::new()); + drive(Event::Apply(pod), &tracked, &router, &config).await; - -- // Pod should not be tracked since router.add_pd_worker will fail for regular router -- assert!(!tracked_pods.lock().unwrap().contains(&pod_info)); ++ + let mut pod = labelled_pod("uid-y", "pod-y", "10.0.0.8", true, true); + pod.metadata.uid = None; + drive(Event::Apply(pod), &tracked, &router, &config).await; @@ -1465,8 +1675,7 @@ index d765800..798eab7 100644 #[tokio::test] - async fn test_unified_handler_deletion_with_pd_mode() { -+ async fn test_withdrawal_keeps_a_concurrently_replaced_registration() { - let router = create_test_router().await; +- let router = create_test_router().await; - let tracked_pods = Arc::new(Mutex::new(HashSet::new())); - let pod_info = PodInfo { - name: "decode-pod".into(), @@ -1475,6 +1684,8 @@ index d765800..798eab7 100644 - is_ready: true, - pod_type: Some(PodType::Decode), - bootstrap_port: None, ++ async fn test_withdrawal_keeps_a_concurrently_replaced_registration() { ++ let router = stub_router(); + let config = regular_config(); + let old = registered("uid-race", "pod-race", "10.0.0.8"); + let new = PodInfo { @@ -1510,7 +1721,7 @@ index d765800..798eab7 100644 - true, // pd_mode = true + #[tokio::test] + async fn test_withdraw_unknown_uid_is_a_no_op() { -+ let router = create_test_router().await; ++ let router = stub_router(); + let config = regular_config(); + let entry = registered("uid-keep", "pod-keep", "10.0.0.10"); + let tracked = tracker_with(vec![entry.clone()]);