diff --git a/pack/llm-router/Dockerfile b/pack/llm-router/Dockerfile index f854e98a9..b0a3880f0 100644 --- a/pack/llm-router/Dockerfile +++ b/pack/llm-router/Dockerfile @@ -29,8 +29,8 @@ 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 VLLM_ROUTER_VERSION=v0.1.15 +ARG LLMD_ROUTER_VERSION=v0.11.0 +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/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. 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()]);