diff --git a/deploy/inference-gateway/epp/go.mod b/deploy/inference-gateway/epp/go.mod index f059c321e6fb..e6ab7da45d72 100644 --- a/deploy/inference-gateway/epp/go.mod +++ b/deploy/inference-gateway/epp/go.mod @@ -46,7 +46,7 @@ require ( github.com/google/cel-go v0.28.0 // indirect github.com/google/gnostic-models v0.7.0 // indirect github.com/google/go-cmp v0.7.0 // indirect - github.com/google/uuid v1.6.0 // indirect + github.com/google/uuid v1.6.0 github.com/grafana/regexp v0.0.0-20250905093917-f7b3be9d1853 // indirect github.com/grpc-ecosystem/grpc-gateway/v2 v2.28.0 // indirect github.com/hashicorp/golang-lru/v2 v2.0.7 // indirect diff --git a/deploy/inference-gateway/epp/pkg/plugins/disagg/booking_executor.go b/deploy/inference-gateway/epp/pkg/plugins/disagg/booking_executor.go new file mode 100644 index 000000000000..30af2c703d12 --- /dev/null +++ b/deploy/inference-gateway/epp/pkg/plugins/disagg/booking_executor.go @@ -0,0 +1,390 @@ +/* +Copyright 2026 NVIDIA Corporation. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package disagg + +import ( + "context" + "os" + "strconv" + "sync" + "time" + + "github.com/go-logr/logr" + logutil "sigs.k8s.io/gateway-api-inference-extension/pkg/common/observability/logging" +) + +const ( + maxConcurrentPrefillMarkers = 32 + maxConcurrentBookingCleanups = 32 + prefillMarkerQueueCapacity = 1024 + bookingCleanupQueueCapacity = 4096 + bookingReconcileInterval = time.Second + bookingCleanupRetryInterval = 30 * time.Second + minimumBookingCleanupRetention = 10 * time.Minute + bookingCleanupRetentionGrace = 2 * time.Minute + routerActiveRequestExpirySeconds = "DYN_ROUTER_ACTIVE_REQUEST_EXPIRY_SECS" +) + +var defaultBookingExecutor = newBookingExecutor(bookingExecutorConfig{ + markerWorkers: maxConcurrentPrefillMarkers, + cleanupWorkers: maxConcurrentBookingCleanups, + markerQueueSize: prefillMarkerQueueCapacity, + cleanupQueueSize: bookingCleanupQueueCapacity, + reconcileInterval: bookingReconcileInterval, + cleanupRetryDelay: bookingCleanupRetryInterval, + cleanupRetention: bookingCleanupRetention(), + cleanupBackoff: cleanupRetryBackoff, +}) + +func bookingCleanupRetention() time.Duration { + return bookingCleanupRetentionFromLookup(os.Getenv) +} + +func bookingCleanupRetentionFromLookup(getEnv func(string) string) time.Duration { + raw := getEnv(routerActiveRequestExpirySeconds) + seconds, err := strconv.ParseUint(raw, 10, 64) + maxSeconds := uint64((time.Duration(1<<63-1) - bookingCleanupRetentionGrace) / time.Second) + if err != nil || seconds == 0 || seconds > maxSeconds { + return minimumBookingCleanupRetention + } + + retention := time.Duration(seconds)*time.Second + bookingCleanupRetentionGrace + if retention < minimumBookingCleanupRetention { + return minimumBookingCleanupRetention + } + return retention +} + +type bookingExecutorConfig struct { + markerWorkers int + cleanupWorkers int + markerQueueSize int + cleanupQueueSize int + reconcileInterval time.Duration + cleanupRetryDelay time.Duration + cleanupRetention time.Duration + cleanupBackoff time.Duration +} + +// bookingExecutor bounds both queued and active CGO bookkeeping calls. Cleanup +// has a dedicated pool so first-token marker pressure cannot delay terminal +// request cleanup. +type bookingExecutor struct { + markerQueue chan prefillMarkerWork + cleanupQueue chan *bookingLifecycle + reconcileInterval time.Duration + cleanupRetryDelay time.Duration + cleanupRetention time.Duration + cleanupBackoff time.Duration + stopCh chan struct{} + stopOnce sync.Once + wg sync.WaitGroup +} + +type prefillMarkerWork struct { + lifecycle *bookingLifecycle + ctx context.Context + done chan struct{} + markPrefillComplete func(string) error + logger logr.Logger + requestID string +} + +func newBookingExecutor(cfg bookingExecutorConfig) *bookingExecutor { + if cfg.markerWorkers <= 0 || cfg.cleanupWorkers <= 0 || cfg.markerQueueSize <= 0 || cfg.cleanupQueueSize <= 0 { + panic("booking executor worker and queue limits must be positive") + } + if cfg.reconcileInterval <= 0 || cfg.cleanupRetryDelay <= 0 || cfg.cleanupRetention <= 0 || cfg.cleanupBackoff <= 0 { + panic("booking executor durations must be positive") + } + + executor := &bookingExecutor{ + markerQueue: make(chan prefillMarkerWork, cfg.markerQueueSize), + cleanupQueue: make(chan *bookingLifecycle, cfg.cleanupQueueSize), + reconcileInterval: cfg.reconcileInterval, + cleanupRetryDelay: cfg.cleanupRetryDelay, + cleanupRetention: cfg.cleanupRetention, + cleanupBackoff: cfg.cleanupBackoff, + stopCh: make(chan struct{}), + } + for range cfg.markerWorkers { + executor.wg.Add(1) + go executor.runMarkerWorker() + } + for range cfg.cleanupWorkers { + executor.wg.Add(1) + go executor.runCleanupWorker() + } + executor.wg.Add(1) + go executor.runReconciler() + return executor +} + +func (e *bookingExecutor) stop() { + e.stopOnce.Do(func() { close(e.stopCh) }) + e.wg.Wait() +} + +func (e *bookingExecutor) enqueueMarker(work prefillMarkerWork) bool { + select { + case <-e.stopCh: + return false + default: + } + select { + case e.markerQueue <- work: + return true + case <-e.stopCh: + return false + default: + return false + } +} + +func (e *bookingExecutor) runMarkerWorker() { + defer e.wg.Done() + for { + select { + case work := <-e.markerQueue: + e.runPrefillMarker(work) + case <-e.stopCh: + return + } + } +} + +func (e *bookingExecutor) runPrefillMarker(work prefillMarkerWork) { + defer close(work.done) + for attempt := 1; attempt <= prefillMarkMaxAttempts; attempt++ { + select { + case <-e.stopCh: + return + default: + } + if work.ctx.Err() != nil { + return + } + if err := work.markPrefillComplete(work.lifecycle.bookingID); err == nil { + work.logger.V(logutil.VERBOSE).Info("DynDecodeScorer ResponseBody: marked prefill complete", + "bookingID", work.lifecycle.bookingID, "requestID", work.requestID, "attempt", attempt) + return + } else { + work.logger.V(logutil.DEFAULT).Error(err, "DynDecodeScorer ResponseBody: failed to mark prefill complete", + "bookingID", work.lifecycle.bookingID, "requestID", work.requestID, "attempt", attempt) + } + if attempt == prefillMarkMaxAttempts { + return + } + + timer := time.NewTimer(prefillMarkRetryBackoff * time.Duration(attempt)) + select { + case <-work.ctx.Done(): + stopTimer(timer) + return + case <-e.stopCh: + stopTimer(timer) + return + case <-timer.C: + } + } +} + +func (e *bookingExecutor) enqueueCleanup(lifecycle *bookingLifecycle) bool { + now := time.Now() + lifecycle.mu.Lock() + if !lifecycle.cleanupStarted || lifecycle.cleanupSucceeded || lifecycle.cleanupExpired || + lifecycle.cleanupQueued || lifecycle.cleanupRunning || + (lifecycle.cleanupExhausted && now.Before(lifecycle.cleanupRetryAt)) { + lifecycle.mu.Unlock() + return false + } + lifecycle.cleanupQueued = true + lifecycle.mu.Unlock() + + select { + case e.cleanupQueue <- lifecycle: + return true + case <-e.stopCh: + default: + } + + lifecycle.mu.Lock() + lifecycle.cleanupQueued = false + lifecycle.mu.Unlock() + return false +} + +func (e *bookingExecutor) runCleanupWorker() { + defer e.wg.Done() + for { + select { + case lifecycle := <-e.cleanupQueue: + e.runCleanup(lifecycle) + case <-e.stopCh: + return + } + } +} + +func (e *bookingExecutor) runCleanup(lifecycle *bookingLifecycle) { + lifecycle.mu.Lock() + lifecycle.cleanupQueued = false + if lifecycle.cleanupSucceeded || lifecycle.cleanupExpired { + lifecycle.mu.Unlock() + return + } + lifecycle.cleanupRunning = true + decodeRegistrationDone := lifecycle.decodeRegistrationDone + logger := lifecycle.cleanupLogger + reason := lifecycle.cleanupReason + lifecycle.mu.Unlock() + + if !e.waitForLifecycleWork(decodeRegistrationDone) { + lifecycle.mu.Lock() + lifecycle.cleanupRunning = false + lifecycle.mu.Unlock() + return + } + + for attempt := 1; attempt <= cleanupMaxAttempts; attempt++ { + select { + case <-e.stopCh: + lifecycle.mu.Lock() + lifecycle.cleanupRunning = false + lifecycle.mu.Unlock() + return + default: + } + + err := lifecycle.freeBooking(lifecycle.bookingID) + if err == nil { + lifecycle.mu.Lock() + lifecycle.cleanupRunning = false + lifecycle.cleanupSucceeded = true + done := lifecycle.closeCleanupDoneLocked() + lifecycle.mu.Unlock() + bookingLifecycles.CompareAndDelete(lifecycle.bookingID, lifecycle) + logger.V(logutil.VERBOSE).Info("Dynamo EPP booking cleaned up", + "bookingID", lifecycle.bookingID, "reason", reason, "attempt", attempt) + if done != nil { + close(done) + } + return + } + + logger.V(logutil.DEFAULT).Error(err, "Dynamo EPP booking cleanup failed", + "bookingID", lifecycle.bookingID, "reason", reason, "attempt", attempt) + if attempt == cleanupMaxAttempts { + lifecycle.mu.Lock() + lifecycle.cleanupRunning = false + lifecycle.cleanupExhausted = true + lifecycle.cleanupRetryAt = time.Now().Add(e.cleanupRetryDelay) + done := lifecycle.closeCleanupDoneLocked() + lifecycle.mu.Unlock() + logger.V(logutil.DEFAULT).Error(err, "Dynamo EPP booking cleanup exhausted retries; retaining finite tombstone for background retry", + "bookingID", lifecycle.bookingID, "reason", reason, "attempts", cleanupMaxAttempts, + "retention", e.cleanupRetention) + if done != nil { + close(done) + } + return + } + + timer := time.NewTimer(e.cleanupBackoff * time.Duration(attempt)) + select { + case <-timer.C: + case <-e.stopCh: + stopTimer(timer) + lifecycle.mu.Lock() + lifecycle.cleanupRunning = false + lifecycle.mu.Unlock() + return + } + } +} + +func (e *bookingExecutor) waitForLifecycleWork(done <-chan struct{}) bool { + if done == nil { + return true + } + select { + case <-done: + return true + case <-e.stopCh: + return false + } +} + +func (e *bookingExecutor) runReconciler() { + defer e.wg.Done() + ticker := time.NewTicker(e.reconcileInterval) + defer ticker.Stop() + for { + select { + case now := <-ticker.C: + e.reconcile(now) + case <-e.stopCh: + return + } + } +} + +func (e *bookingExecutor) reconcile(now time.Time) { + bookingLifecycles.Range(func(key, value any) bool { + lifecycle := value.(*bookingLifecycle) + if lifecycle.executor != e { + return true + } + + lifecycle.mu.Lock() + if !lifecycle.cleanupStarted || lifecycle.cleanupSucceeded || lifecycle.cleanupExpired { + lifecycle.mu.Unlock() + return true + } + if now.Sub(lifecycle.cleanupStartedAt) >= e.cleanupRetention { + lifecycle.cleanupExpired = true + done := lifecycle.closeCleanupDoneLocked() + logger := lifecycle.cleanupLogger + reason := lifecycle.cleanupReason + lifecycle.mu.Unlock() + bookingLifecycles.CompareAndDelete(key, lifecycle) + logger.V(logutil.DEFAULT).Info("Dynamo EPP booking cleanup ownership expired; relying on router stale-booking reaper", + "bookingID", lifecycle.bookingID, "reason", reason, "retention", e.cleanupRetention) + if done != nil { + close(done) + } + return true + } + eligible := !lifecycle.cleanupQueued && !lifecycle.cleanupRunning && + (!lifecycle.cleanupExhausted || !now.Before(lifecycle.cleanupRetryAt)) + lifecycle.mu.Unlock() + if eligible { + e.enqueueCleanup(lifecycle) + } + return true + }) +} + +func stopTimer(timer *time.Timer) { + if timer.Stop() { + return + } + select { + case <-timer.C: + default: + } +} diff --git a/deploy/inference-gateway/epp/pkg/plugins/disagg/booking_executor_test.go b/deploy/inference-gateway/epp/pkg/plugins/disagg/booking_executor_test.go new file mode 100644 index 000000000000..82178922aace --- /dev/null +++ b/deploy/inference-gateway/epp/pkg/plugins/disagg/booking_executor_test.go @@ -0,0 +1,271 @@ +/* +Copyright 2026 NVIDIA Corporation. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package disagg + +import ( + "context" + "errors" + "sync/atomic" + "testing" + "time" + + "github.com/go-logr/logr" + "github.com/google/uuid" +) + +func newTestBookingExecutor(t *testing.T, mutate func(*bookingExecutorConfig)) *bookingExecutor { + t.Helper() + cfg := bookingExecutorConfig{ + markerWorkers: 1, + cleanupWorkers: 1, + markerQueueSize: 1, + cleanupQueueSize: 1, + reconcileInterval: 5 * time.Millisecond, + cleanupRetryDelay: 10 * time.Millisecond, + cleanupRetention: time.Second, + cleanupBackoff: time.Millisecond, + } + if mutate != nil { + mutate(&cfg) + } + executor := newBookingExecutor(cfg) + t.Cleanup(executor.stop) + return executor +} + +func registerTestLifecycle(t *testing.T, executor *bookingExecutor, freeBooking func(string) error) *bookingLifecycle { + t.Helper() + bookingID := uuid.NewString() + lifecycle := registerBookingLifecycleWithExecutor(bookingID, freeBooking, executor) + t.Cleanup(func() { bookingLifecycles.Delete(bookingID) }) + return lifecycle +} + +func waitForLifecycleRemoval(t *testing.T, bookingID string) { + t.Helper() + deadline := time.Now().Add(time.Second) + for findBookingLifecycle(bookingID) != nil { + if time.Now().After(deadline) { + t.Fatalf("booking lifecycle %q was not removed", bookingID) + } + time.Sleep(time.Millisecond) + } +} + +func TestBookingCleanupRetentionFollowsRouterExpiryOverride(t *testing.T) { + tests := []struct { + name string + raw string + want time.Duration + }{ + {name: "unset", want: minimumBookingCleanupRetention}, + {name: "invalid", raw: "invalid", want: minimumBookingCleanupRetention}, + {name: "zero", raw: "0", want: minimumBookingCleanupRetention}, + {name: "shorter than minimum", raw: "60", want: minimumBookingCleanupRetention}, + {name: "long router expiry", raw: "3600", want: 62 * time.Minute}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + got := bookingCleanupRetentionFromLookup(func(string) string { return test.raw }) + if got != test.want { + t.Fatalf("cleanup retention = %s, want %s", got, test.want) + } + }) + } +} + +func TestBookingExecutorBoundsPrefillMarkerQueue(t *testing.T) { + executor := newTestBookingExecutor(t, nil) + firstStarted := make(chan struct{}) + releaseFirst := make(chan struct{}) + var calls atomic.Int32 + marker := func(string) error { + call := calls.Add(1) + if call == 1 { + close(firstStarted) + select { + case <-releaseFirst: + case <-executor.stopCh: + } + } + return nil + } + + first := registerTestLifecycle(t, executor, func(string) error { return nil }) + second := registerTestLifecycle(t, executor, func(string) error { return nil }) + overflow := registerTestLifecycle(t, executor, func(string) error { return nil }) + + first.startPrefillMarker(marker, logr.Discard(), "first") + select { + case <-firstStarted: + case <-time.After(time.Second): + t.Fatal("first marker did not start") + } + second.startPrefillMarker(marker, logr.Discard(), "second") + if got := len(executor.markerQueue); got != 1 { + t.Fatalf("marker queue length = %d, want 1", got) + } + overflow.startPrefillMarker(marker, logr.Discard(), "overflow") + select { + case <-overflow.markerComplete(): + case <-time.After(50 * time.Millisecond): + t.Fatal("overflow marker was not rejected immediately") + } + if got := calls.Load(); got != 1 { + t.Fatalf("active marker calls = %d, want 1", got) + } + + close(releaseFirst) + for name, lifecycle := range map[string]*bookingLifecycle{"first": first, "second": second} { + select { + case <-lifecycle.markerComplete(): + case <-time.After(time.Second): + t.Fatalf("%s marker did not finish", name) + } + } + if got := calls.Load(); got != 2 { + t.Fatalf("marker calls = %d, want 2 accepted jobs", got) + } +} + +func TestBookingExecutorReconcilesCleanupQueueOverflow(t *testing.T) { + executor := newTestBookingExecutor(t, nil) + firstStarted := make(chan struct{}) + releaseFirst := make(chan struct{}) + first := registerTestLifecycle(t, executor, func(string) error { + close(firstStarted) + select { + case <-releaseFirst: + case <-executor.stopCh: + } + return nil + }) + second := registerTestLifecycle(t, executor, func(string) error { return nil }) + overflow := registerTestLifecycle(t, executor, func(string) error { return nil }) + + if !first.cleanup(context.Background(), "first") { + t.Fatal("first cleanup did not start") + } + select { + case <-firstStarted: + case <-time.After(time.Second): + t.Fatal("first cleanup did not enter the worker") + } + if !second.cleanup(context.Background(), "second") { + t.Fatal("second cleanup did not start") + } + if got := len(executor.cleanupQueue); got != 1 { + t.Fatalf("cleanup queue length = %d, want 1", got) + } + if !overflow.cleanup(context.Background(), "overflow") { + t.Fatal("overflow cleanup did not retain ownership") + } + overflow.mu.Lock() + overflowQueued := overflow.cleanupQueued + overflowRunning := overflow.cleanupRunning + overflow.mu.Unlock() + if overflowQueued || overflowRunning { + t.Fatal("overflow cleanup unexpectedly entered the full executor") + } + + close(releaseFirst) + for name, lifecycle := range map[string]*bookingLifecycle{ + "first": first, + "second": second, + "overflow": overflow, + } { + select { + case <-lifecycle.cleanupComplete(): + case <-time.After(time.Second): + t.Fatalf("%s cleanup did not finish", name) + } + if findBookingLifecycle(lifecycle.bookingID) != nil { + t.Fatalf("%s cleanup retained its lifecycle", name) + } + } +} + +func TestBookingExecutorRetriesExhaustedCleanup(t *testing.T) { + executor := newTestBookingExecutor(t, func(cfg *bookingExecutorConfig) { + cfg.cleanupRetryDelay = 100 * time.Millisecond + }) + var calls atomic.Int32 + lifecycle := registerTestLifecycle(t, executor, func(string) error { + if calls.Add(1) <= cleanupMaxAttempts { + return errors.New("transient cleanup outage") + } + return nil + }) + + if !lifecycle.cleanup(context.Background(), "retry exhausted cleanup") { + t.Fatal("cleanup did not start") + } + select { + case <-lifecycle.cleanupComplete(): + case <-time.After(time.Second): + t.Fatal("initial cleanup attempts did not finish") + } + if got := calls.Load(); got != cleanupMaxAttempts { + t.Fatalf("initial cleanup calls = %d, want %d", got, cleanupMaxAttempts) + } + if findBookingLifecycle(lifecycle.bookingID) == nil { + t.Fatal("exhausted cleanup lost retry ownership") + } + + waitForLifecycleRemoval(t, lifecycle.bookingID) + if got := calls.Load(); got != cleanupMaxAttempts+1 { + t.Fatalf("cleanup calls after recovery = %d, want %d", got, cleanupMaxAttempts+1) + } +} + +func TestBookingExecutorExpiresCleanupTombstone(t *testing.T) { + executor := newTestBookingExecutor(t, func(cfg *bookingExecutorConfig) { + cfg.cleanupRetryDelay = time.Hour + cfg.cleanupRetention = 250 * time.Millisecond + }) + var calls atomic.Int32 + lifecycle := registerTestLifecycle(t, executor, func(string) error { + calls.Add(1) + return errors.New("persistent cleanup outage") + }) + + if !lifecycle.cleanup(context.Background(), "expire cleanup ownership") { + t.Fatal("cleanup did not start") + } + select { + case <-lifecycle.cleanupComplete(): + case <-time.After(time.Second): + t.Fatal("cleanup attempts did not finish") + } + if got := calls.Load(); got != cleanupMaxAttempts { + t.Fatalf("cleanup calls = %d, want %d", got, cleanupMaxAttempts) + } + if findBookingLifecycle(lifecycle.bookingID) == nil { + t.Fatal("cleanup tombstone expired before its retention deadline") + } + + waitForLifecycleRemoval(t, lifecycle.bookingID) + lifecycle.mu.Lock() + expired := lifecycle.cleanupExpired + lifecycle.mu.Unlock() + if !expired { + t.Fatal("cleanup lifecycle was removed without recording expiry") + } + if got := calls.Load(); got != cleanupMaxAttempts { + t.Fatalf("cleanup calls after expiry = %d, want %d", got, cleanupMaxAttempts) + } +} diff --git a/deploy/inference-gateway/epp/pkg/plugins/disagg/decode_scorer.go b/deploy/inference-gateway/epp/pkg/plugins/disagg/decode_scorer.go index d867a48da9fe..adfff843a8b0 100644 --- a/deploy/inference-gateway/epp/pkg/plugins/disagg/decode_scorer.go +++ b/deploy/inference-gateway/epp/pkg/plugins/disagg/decode_scorer.go @@ -21,7 +21,6 @@ import ( "encoding/json" "fmt" "strconv" - "sync" log "sigs.k8s.io/controller-runtime/pkg/log" logutil "sigs.k8s.io/gateway-api-inference-extension/pkg/common/observability/logging" @@ -42,41 +41,13 @@ const ( DpRankHeader = "x-dynamo-dp-rank" PrefillDpRankHeader = "x-dynamo-prefill-dp-rank" RoutingModeHeader = "x-dynamo-routing-mode" - - decodeStateKey = "dynamo-decode-routing-state" ) // compile-time type assertions var _ schedtypes.Scorer = &DynDecodeScorer{} var _ plugins.Plugin = &DynDecodeScorer{} -var _ rc.PreRequest = &DynDecodeScorer{} var _ rc.ResponseBodyProcessor = &DynDecodeScorer{} -// DecodeRoutingState holds routing information passed from Score() to PreRequest(). -type DecodeRoutingState struct { - WorkerID string - DpRank uint32 - PrefillWorkerID string - TokenData []int64 -} - -// Clone implements plugins.StateData. -func (s *DecodeRoutingState) Clone() plugins.StateData { - if s == nil { - return nil - } - clone := &DecodeRoutingState{ - WorkerID: s.WorkerID, - DpRank: s.DpRank, - PrefillWorkerID: s.PrefillWorkerID, - } - if s.TokenData != nil { - clone.TokenData = make([]int64, len(s.TokenData)) - copy(clone.TokenData, s.TokenData) - } - return clone -} - // DynDecodeScorerConfig holds the configuration for the DynDecodeScorer plugin. type DynDecodeScorerConfig struct{} @@ -98,18 +69,23 @@ func DynDecodeScorerFactory(name string, rawParameters json.RawMessage, handle p } // NewDynDecodeScorer initializes a new DynDecodeScorer. -func NewDynDecodeScorer(ctx context.Context) *DynDecodeScorer { +func NewDynDecodeScorer(_ context.Context) *DynDecodeScorer { return &DynDecodeScorer{ - typedName: plugins.TypedName{Type: DynDecodeScorerType, Name: DynDecodeScorerType}, - pluginState: plugins.NewPluginState(ctx), + typedName: plugins.TypedName{Type: DynDecodeScorerType, Name: DynDecodeScorerType}, + routeDecode: dynscorer.CallRouteDecodeRequest, + addRequest: dynscorer.CallAddRequest, + markPrefillComplete: dynscorer.CallMarkPrefillComplete, + freeBooking: dynscorer.CallFreeRequest, } } // DynDecodeScorer is a scorer plugin for the decode scheduling profile. type DynDecodeScorer struct { - typedName plugins.TypedName - pluginState *plugins.PluginState - firstTokenSeen sync.Map + typedName plugins.TypedName + routeDecode func(string, string, bool) (*dynscorer.RoutingResult, error) + addRequest func(string, []int64, uint64, uint32) error + markPrefillComplete func(string) error + freeBooking func(string) error } // TypedName returns the type and name tuple of this plugin instance. @@ -131,12 +107,28 @@ func (s *DynDecodeScorer) Category() schedtypes.ScorerCategory { // Score scores endpoints for decode suitability. func (s *DynDecodeScorer) Score(ctx context.Context, cycleState *schedtypes.CycleState, req *schedtypes.InferenceRequest, endpoints []schedtypes.Endpoint) map[schedtypes.Endpoint]float64 { logger := log.FromContext(ctx) + if req == nil { + cycleState.Write(PrefillEnabledStateKey, &PrefillEnabledState{Enabled: false}) + return uniformScores(endpoints, 1.0) + } + booking := ensureBookingState(cycleState) + attachBookingID(req, booking.ID) - isDisaggregated := readPrefillEnabled(cycleState) + if err := ctx.Err(); err != nil { + logger.V(logutil.VERBOSE).Info("DynDecodeScorer: scheduling already cancelled", "error", err.Error()) + s.cleanupBooking(ctx, booking.ID, "decode scheduling cancelled before routing") + return uniformScores(endpoints, 1.0) + } + prefillEnabled := readPrefillEnabled(cycleState) + if booking.PrefillReserved && !prefillEnabled { + s.rollbackPrefillReservation(ctx, cycleState, req, booking, "prefill scheduling did not complete") + } + isDisaggregated := prefillEnabled && booking.PrefillReserved requestJSON, err := buildRequestJSON(req) if err != nil { logger.V(logutil.DEFAULT).Error(err, "DynDecodeScorer: failed to build request") + s.rollbackPrefillReservation(ctx, cycleState, req, booking, "decode request serialization failed") return uniformScores(endpoints, 1.0) } @@ -145,47 +137,57 @@ func (s *DynDecodeScorer) Score(ctx context.Context, cycleState *schedtypes.Cycl "endpointCount", len(endpoints), "endpointsJSON", string(endpointsJSON)) - result, err := dynscorer.CallRouteDecodeRequest(requestJSON, endpointsJSON, isDisaggregated) + result, err := s.routeDecode(requestJSON, endpointsJSON, isDisaggregated) if err != nil { logger.V(logutil.DEFAULT).Error(err, "DynDecodeScorer: FFI decode routing failed") + s.rollbackPrefillReservation(ctx, cycleState, req, booking, "decode routing failed") + return uniformScores(endpoints, 1.0) + } + if err := ctx.Err(); err != nil { + logger.V(logutil.VERBOSE).Info("DynDecodeScorer: scheduling cancelled after decode routing", "error", err.Error()) + s.cleanupBooking(ctx, booking.ID, "decode scheduling cancelled after routing") return uniformScores(endpoints, 1.0) } workerIDStr := fmt.Sprintf("%d", result.WorkerID) dpRankStr := strconv.FormatUint(uint64(result.DpRank), 10) + lifecycle := registerBookingLifecycle(booking.ID, s.freeBooking) + lifecycle.armCancellation(ctx) + if !lifecycle.startDecodeRegistration() { + logger.V(logutil.VERBOSE).Info("DynDecodeScorer: request cancelled before decode booking registration", + "bookingID", booking.ID) + s.rollbackPrefillReservation(ctx, cycleState, req, booking, "decode booking registration cancelled") + return uniformScores(endpoints, 1.0) + } + addErr := s.addRequest(booking.ID, result.TokenData, result.WorkerID, result.DpRank) + lifecycle.finishDecodeRegistration() + if addErr != nil { + logger.V(logutil.DEFAULT).Error(addErr, "DynDecodeScorer: failed to add decode booking", + "bookingID", booking.ID) + s.rollbackPrefillReservation(ctx, cycleState, req, booking, "decode booking failed") + return uniformScores(endpoints, 1.0) + } + if err := ctx.Err(); err != nil { + logger.V(logutil.VERBOSE).Info("DynDecodeScorer: scheduling cancelled during decode booking registration", "error", err.Error()) + s.rollbackPrefillReservation(ctx, cycleState, req, booking, "decode scheduling cancelled during booking registration") + return uniformScores(endpoints, 1.0) + } + logger.V(logutil.DEFAULT).Info("[EPP-SCORER] FFI returned tokens from C bindings tokenization", + "bookingID", booking.ID, "decodeWorkerID", workerIDStr, "decodeDpRank", result.DpRank, "isDisaggregated", isDisaggregated, "tokenCount", len(result.TokenData)) - if req.Headers == nil { - req.Headers = map[string]string{} - } req.Headers[WorkerIDHeader] = workerIDStr req.Headers[DpRankHeader] = dpRankStr - if isDisaggregated { req.Headers[RoutingModeHeader] = "disaggregated" - if prefillID, ok := req.Headers[PrefillWorkerIDHeader]; ok { - logger.V(logutil.DEFAULT).Info("DynDecodeScorer: prefill worker header present", - "prefillWorkerID", prefillID) - } else { - logger.V(logutil.DEFAULT).Error(nil, - "DynDecodeScorer: x-dynamo-prefill-instance-id header missing — DynPrefillScorer did not set it") - } } else { req.Headers[RoutingModeHeader] = "aggregated" - } - - // Store routing state for PreRequest bookkeeping - if req.RequestId != "" { - routingState := &DecodeRoutingState{ - WorkerID: workerIDStr, - DpRank: result.DpRank, - TokenData: result.TokenData, - } - s.pluginState.Write(req.RequestId, plugins.StateKey(decodeStateKey), routingState) + delete(req.Headers, PrefillWorkerIDHeader) + delete(req.Headers, PrefillDpRankHeader) } // Inject pre-computed tokens into the request body so the frontend @@ -195,77 +197,71 @@ func (s *DynDecodeScorer) Score(ctx context.Context, cycleState *schedtypes.Cycl return uniformScores(endpoints, 1.0) } -// PreRequest registers the request with the Dynamo router's bookkeeping. -func (s *DynDecodeScorer) PreRequest(ctx context.Context, request *schedtypes.InferenceRequest, _ *schedtypes.SchedulingResult) { - logger := log.FromContext(ctx) - - if request == nil || request.RequestId == "" { - logger.V(logutil.DEBUG).Info("DynDecodeScorer PreRequest: no request ID, skipping") - return +func (s *DynDecodeScorer) cleanupBooking(ctx context.Context, bookingID, reason string) bool { + if bookingID == "" { + return true } - - state, err := plugins.ReadPluginStateKey[*DecodeRoutingState]( - s.pluginState, request.RequestId, plugins.StateKey(decodeStateKey), - ) - s.pluginState.Delete(request.RequestId) - - if err != nil { - logger.V(logutil.DEBUG).Info("DynDecodeScorer PreRequest: no routing state found", - "requestID", request.RequestId) - return + if lifecycle := findBookingLifecycle(bookingID); lifecycle != nil { + return lifecycle.cleanup(ctx, reason) } - - var workerIDUint uint64 - if _, parseErr := fmt.Sscanf(state.WorkerID, "%d", &workerIDUint); parseErr != nil { - logger.V(logutil.DEFAULT).Error(parseErr, "DynDecodeScorer PreRequest: invalid worker ID", - "requestID", request.RequestId, "workerID", state.WorkerID) - return + if err := s.freeBooking(bookingID); err != nil { + log.FromContext(ctx).V(logutil.DEFAULT).Error(err, "DynDecodeScorer: booking cleanup failed", + "bookingID", bookingID, "reason", reason) + return false } + log.FromContext(ctx).V(logutil.VERBOSE).Info("DynDecodeScorer: booking cleaned up", + "bookingID", bookingID, "reason", reason) + return true +} - if addErr := dynscorer.CallAddRequest(request.RequestId, state.TokenData, workerIDUint, state.DpRank); addErr != nil { - logger.V(logutil.DEFAULT).Error(addErr, "DynDecodeScorer PreRequest: failed to add request", - "requestID", request.RequestId) - return +func (s *DynDecodeScorer) rollbackPrefillReservation( + ctx context.Context, + cycleState *schedtypes.CycleState, + request *schedtypes.InferenceRequest, + booking *BookingState, + reason string, +) { + if booking != nil { + if findBookingLifecycle(booking.ID) != nil || booking.PrefillReserved { + s.cleanupBooking(ctx, booking.ID, reason) + } + booking.PrefillReserved = false + cycleState.Write(BookingStateKey, booking) + } + cycleState.Write(PrefillEnabledStateKey, &PrefillEnabledState{Enabled: false}) + if request != nil { + if request.Headers == nil { + request.Headers = map[string]string{} + } + request.Headers[RoutingModeHeader] = "aggregated" + delete(request.Headers, WorkerIDHeader) + delete(request.Headers, DpRankHeader) + delete(request.Headers, PrefillWorkerIDHeader) + delete(request.Headers, PrefillDpRankHeader) } - - logger.V(logutil.VERBOSE).Info("DynDecodeScorer PreRequest: registered request", - "requestID", request.RequestId, - "workerID", state.WorkerID, - "dpRank", state.DpRank, - "tokenCount", len(state.TokenData)) } // ResponseBody handles streaming chunks and end-of-stream cleanup. // On the first token it marks prefill as complete; on EndOfStream it frees the request. func (s *DynDecodeScorer) ResponseBody(ctx context.Context, request *schedtypes.InferenceRequest, response *rc.Response, _ *fwkdl.EndpointMetadata) { - if request == nil || request.RequestId == "" { + bookingID := bookingIDFromRequest(request) + if bookingID == "" || response == nil { return } - logger := log.FromContext(ctx) - - // Mark prefill complete on first token - if _, alreadySeen := s.firstTokenSeen.LoadOrStore(request.RequestId, true); !alreadySeen { - if err := dynscorer.CallMarkPrefillComplete(request.RequestId); err != nil { - logger.V(logutil.DEFAULT).Error(err, "DynDecodeScorer ResponseBody: failed to mark prefill complete", - "requestID", request.RequestId) - } else { - logger.V(logutil.VERBOSE).Info("DynDecodeScorer ResponseBody: marked prefill complete", - "requestID", request.RequestId) - } + lifecycle := findBookingLifecycle(bookingID) + if lifecycle == nil { + // Only Score creates controller-owned booking lifecycles. Do not let an + // inbound header create router bookkeeping on a response-only path. + return } - - // Free request on end of stream — must always run regardless of - // earlier errors to avoid leaking router bookkeeping state. - if response != nil && response.EndOfStream { - s.firstTokenSeen.Delete(request.RequestId) - - if err := dynscorer.CallFreeRequest(request.RequestId); err != nil { - logger.V(logutil.DEFAULT).Error(err, "DynDecodeScorer ResponseBody: failed to free request", - "requestID", request.RequestId) - } else { - logger.V(logutil.VERBOSE).Info("DynDecodeScorer ResponseBody: freed request", - "requestID", request.RequestId) - } + if response.EndOfStream { + lifecycle.cleanup(ctx, "response end of stream") + return } + if request.Headers[RoutingModeHeader] != "disaggregated" { + return + } + + lifecycle.startPrefillMarker(s.markPrefillComplete, log.FromContext(ctx), request.RequestId) } diff --git a/deploy/inference-gateway/epp/pkg/plugins/disagg/prefill_scorer.go b/deploy/inference-gateway/epp/pkg/plugins/disagg/prefill_scorer.go index 3572abc3e6f8..b6c00383e885 100644 --- a/deploy/inference-gateway/epp/pkg/plugins/disagg/prefill_scorer.go +++ b/deploy/inference-gateway/epp/pkg/plugins/disagg/prefill_scorer.go @@ -21,6 +21,7 @@ import ( "encoding/json" "fmt" "strconv" + "time" log "sigs.k8s.io/controller-runtime/pkg/log" logutil "sigs.k8s.io/gateway-api-inference-extension/pkg/common/observability/logging" @@ -33,13 +34,22 @@ import ( const ( // DynPrefillScorerType is the plugin type registered in the plugin registry. DynPrefillScorerType = "dyn-prefill-scorer" + + defaultPrefillReservationAdmissionTimeout = 60 * time.Second + defaultMaxPrefillReservations = 64 ) // compile-time type assertion var _ schedtypes.Scorer = &DynPrefillScorer{} // DynPrefillScorerConfig holds the configuration for the DynPrefillScorer plugin. -type DynPrefillScorerConfig struct{} +type DynPrefillScorerConfig struct { + // ReservationTimeoutSeconds bounds waiting for scheduler admission. The + // request context deadline still wins when it is shorter. + ReservationTimeoutSeconds int `json:"reservationTimeoutSeconds"` + // MaxConcurrentReservations bounds the number of blocking CGO admissions. + MaxConcurrentReservations int `json:"maxConcurrentReservations"` +} // DynPrefillScorerFactory defines the factory function for DynPrefillScorer. func DynPrefillScorerFactory(name string, rawParameters json.RawMessage, _ plugins.Handle) (plugins.Plugin, error) { @@ -49,24 +59,91 @@ func DynPrefillScorerFactory(name string, rawParameters json.RawMessage, _ plugi return nil, fmt.Errorf("failed to parse %s plugin parameters: %w", DynPrefillScorerType, err) } } + if cfg.ReservationTimeoutSeconds < 0 || cfg.MaxConcurrentReservations < 0 { + return nil, fmt.Errorf("%s reservation timeout and concurrency must not be negative", DynPrefillScorerType) + } if err := dynscorer.InitFFI(); err != nil { return nil, fmt.Errorf("Dynamo FFI init for prefill scorer failed: %w", err) } - return NewDynPrefillScorer().WithName(name), nil + return newDynPrefillScorer(cfg).WithName(name), nil } // NewDynPrefillScorer initializes a new DynPrefillScorer. func NewDynPrefillScorer() *DynPrefillScorer { + return newDynPrefillScorer(DynPrefillScorerConfig{}) +} + +func newDynPrefillScorer(cfg DynPrefillScorerConfig) *DynPrefillScorer { + reservationTimeout := defaultPrefillReservationAdmissionTimeout + if cfg.ReservationTimeoutSeconds > 0 { + reservationTimeout = time.Duration(cfg.ReservationTimeoutSeconds) * time.Second + } + maxReservations := defaultMaxPrefillReservations + if cfg.MaxConcurrentReservations > 0 { + maxReservations = cfg.MaxConcurrentReservations + } return &DynPrefillScorer{ - typedName: plugins.TypedName{Type: DynPrefillScorerType, Name: DynPrefillScorerType}, + typedName: plugins.TypedName{Type: DynPrefillScorerType, Name: DynPrefillScorerType}, + beginPrefill: dynscorer.CallBeginPrefillReservation, + reservePrefill: dynscorer.CallRoutePrefillRequestWithReservation, + cancelPrefill: dynscorer.CallCancelPrefillReservation, + releasePrefill: dynscorer.CallReleasePrefillReservation, + freeBooking: dynscorer.CallFreeRequest, + reservationAdmissionTimeout: reservationTimeout, + reservationSlots: make(chan struct{}, maxReservations), } } // DynPrefillScorer is a scorer plugin for the prefill scheduling profile. type DynPrefillScorer struct { - typedName plugins.TypedName + typedName plugins.TypedName + beginPrefill func(string) error + reservePrefill func(string, string, string) (*dynscorer.RoutingResult, error) + cancelPrefill func(string) error + releasePrefill func(string) error + freeBooking func(string) error + reservationAdmissionTimeout time.Duration + reservationSlots chan struct{} +} + +type prefillReservationResult struct { + result *dynscorer.RoutingResult + err error +} + +func (s *DynPrefillScorer) admissionTimeout() time.Duration { + if s.reservationAdmissionTimeout <= 0 { + return defaultPrefillReservationAdmissionTimeout + } + return s.reservationAdmissionTimeout +} + +func (s *DynPrefillScorer) acquireReservationSlot(ctx context.Context) (func(), bool) { + if s.reservationSlots == nil { + return func() {}, true + } + select { + case s.reservationSlots <- struct{}{}: + return func() { <-s.reservationSlots }, true + case <-ctx.Done(): + return nil, false + } +} + +func (s *DynPrefillScorer) beginReservation(bookingID string) error { + if s.beginPrefill == nil { + return nil + } + return s.beginPrefill(bookingID) +} + +func (s *DynPrefillScorer) releaseLatePrefillReservation(bookingID string) error { + if s.releasePrefill == nil { + return fmt.Errorf("prefill reservation release is not configured") + } + return s.releasePrefill(bookingID) } // TypedName returns the type and name tuple of this plugin instance. @@ -88,15 +165,29 @@ func (s *DynPrefillScorer) Category() schedtypes.ScorerCategory { // Score scores endpoints for prefill suitability. func (s *DynPrefillScorer) Score(ctx context.Context, cycleState *schedtypes.CycleState, req *schedtypes.InferenceRequest, endpoints []schedtypes.Endpoint) map[schedtypes.Endpoint]float64 { logger := log.FromContext(ctx) + if req == nil { + cycleState.Write(PrefillEnabledStateKey, &PrefillEnabledState{Enabled: false}) + return uniformScores(endpoints, 0) + } + if err := ctx.Err(); err != nil { + logger.V(logutil.VERBOSE).Info("DynPrefillScorer: scheduling already cancelled", "error", err.Error()) + return uniformScores(endpoints, 0) + } if !readPrefillEnabled(cycleState) { logger.V(logutil.VERBOSE).Info("DynPrefillScorer: prefill not enabled, returning zero scores") return uniformScores(endpoints, 0) } + booking := ensureBookingState(cycleState) + attachBookingID(req, booking.ID) + delete(req.Headers, PrefillWorkerIDHeader) + delete(req.Headers, PrefillDpRankHeader) + requestJSON, err := buildRequestJSON(req) if err != nil { logger.V(logutil.DEFAULT).Error(err, "DynPrefillScorer: failed to build request") + cycleState.Write(PrefillEnabledStateKey, &PrefillEnabledState{Enabled: false}) return uniformScores(endpoints, 0) } @@ -105,22 +196,75 @@ func (s *DynPrefillScorer) Score(ctx context.Context, cycleState *schedtypes.Cyc "endpointCount", len(endpoints), "endpointsJSON", string(endpointsJSON)) - result, err := dynscorer.CallRoutePrefillRequest(requestJSON, endpointsJSON) + bookingID := booking.ID + admissionCtx, cancelAdmission := context.WithTimeout(ctx, s.admissionTimeout()) + defer cancelAdmission() + releaseSlot, acquired := s.acquireReservationSlot(admissionCtx) + if !acquired { + logger.V(logutil.DEFAULT).Info("DynPrefillScorer: prefill reservation admission budget exhausted", + "bookingID", bookingID, "error", admissionCtx.Err().Error()) + cycleState.Write(PrefillEnabledStateKey, &PrefillEnabledState{Enabled: false}) + return uniformScores(endpoints, 0) + } + if err := s.beginReservation(bookingID); err != nil { + releaseSlot() + logger.V(logutil.DEFAULT).Error(err, "DynPrefillScorer: failed to begin prefill reservation", + "bookingID", bookingID) + cycleState.Write(PrefillEnabledStateKey, &PrefillEnabledState{Enabled: false}) + return uniformScores(endpoints, 0) + } + + resultCh := make(chan prefillReservationResult, 1) + go func() { + defer releaseSlot() + result, err := s.reservePrefill(bookingID, requestJSON, endpointsJSON) + resultCh <- prefillReservationResult{result: result, err: err} + }() + + var result *dynscorer.RoutingResult + select { + case <-admissionCtx.Done(): + logger.V(logutil.VERBOSE).Info("DynPrefillScorer: scheduling cancelled during prefill reservation", + "error", admissionCtx.Err().Error()) + if cancelErr := s.cancelPrefill(bookingID); cancelErr != nil { + logger.V(logutil.DEFAULT).Error(cancelErr, "DynPrefillScorer: failed to cancel pending prefill reservation", + "bookingID", bookingID) + } + go func(bookingID string) { + // Rust can retain an active reservation after returning an error when + // its first cleanup attempt fails. Release is idempotent for every + // other late outcome. + <-resultCh + if cleanupErr := s.releaseLatePrefillReservation(bookingID); cleanupErr != nil { + logger.V(logutil.DEFAULT).Error(cleanupErr, "DynPrefillScorer: failed to release late prefill reservation", + "bookingID", bookingID) + } + }(bookingID) + booking.PrefillReserved = false + cycleState.Write(BookingStateKey, booking) + cycleState.Write(PrefillEnabledStateKey, &PrefillEnabledState{Enabled: false}) + return uniformScores(endpoints, 0) + case outcome := <-resultCh: + result = outcome.result + err = outcome.err + } if err != nil { - logger.V(logutil.DEFAULT).Error(err, "DynPrefillScorer: FFI prefill routing failed") + logger.V(logutil.DEFAULT).Error(err, "DynPrefillScorer: FFI prefill reservation failed") + booking.PrefillReserved = false + cycleState.Write(BookingStateKey, booking) cycleState.Write(PrefillEnabledStateKey, &PrefillEnabledState{Enabled: false}) return uniformScores(endpoints, 0) } + booking.PrefillReserved = true + cycleState.Write(BookingStateKey, booking) prefillWorkerID := strconv.FormatUint(result.WorkerID, 10) - logger.V(logutil.DEFAULT).Info("DynPrefillScorer: prefill worker selected", + logger.V(logutil.DEFAULT).Info("DynPrefillScorer: prefill worker reserved", + "bookingID", booking.ID, "prefillWorkerID", prefillWorkerID, "prefillDpRank", result.DpRank, "tokenCount", len(result.TokenData)) - if req.Headers == nil { - req.Headers = map[string]string{} - } req.Headers[PrefillWorkerIDHeader] = prefillWorkerID if result.DpRank != dynscorer.UnsetDpRank { req.Headers[PrefillDpRankHeader] = strconv.FormatUint(uint64(result.DpRank), 10) @@ -128,5 +272,15 @@ func (s *DynPrefillScorer) Score(ctx context.Context, cycleState *schedtypes.Cyc delete(req.Headers, PrefillDpRankHeader) } + lifecycle := registerBookingLifecycle(booking.ID, s.freeBooking) + lifecycle.armCancellation(ctx) + if err := ctx.Err(); err != nil { + lifecycle.cleanup(ctx, "prefill scheduling cancelled after reservation") + booking.PrefillReserved = false + cycleState.Write(BookingStateKey, booking) + cycleState.Write(PrefillEnabledStateKey, &PrefillEnabledState{Enabled: false}) + return uniformScores(endpoints, 0) + } + return uniformScores(endpoints, 1.0) } diff --git a/deploy/inference-gateway/epp/pkg/plugins/disagg/reservation_test.go b/deploy/inference-gateway/epp/pkg/plugins/disagg/reservation_test.go new file mode 100644 index 000000000000..81f41ac0fde1 --- /dev/null +++ b/deploy/inference-gateway/epp/pkg/plugins/disagg/reservation_test.go @@ -0,0 +1,679 @@ +/* +Copyright 2026 NVIDIA Corporation. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package disagg + +import ( + "context" + "errors" + "testing" + "time" + + rc "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/framework/interface/requestcontrol" + fwkrh "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/framework/interface/requesthandling" + schedtypes "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/framework/interface/scheduling" + + dynscorer "github.com/nvidia/dynamo/deploy/inference-gateway/pkg/plugins/dynamo_kv_scorer" +) + +func requestWithBooking(externalRequestID, bookingID string) *schedtypes.InferenceRequest { + return &schedtypes.InferenceRequest{ + RequestId: externalRequestID, + Headers: map[string]string{BookingIDHeader: bookingID}, + } +} + +func TestBookingStateDoesNotTrustExternalRequestID(t *testing.T) { + externalRequestID := "shared-client-request-id" + requestA := &schedtypes.InferenceRequest{ + RequestId: externalRequestID, + Headers: map[string]string{BookingIDHeader: "caller-controlled"}, + } + requestB := &schedtypes.InferenceRequest{RequestId: externalRequestID} + stateA := schedtypes.NewCycleState() + stateB := schedtypes.NewCycleState() + + bookingA := ensureBookingState(stateA) + bookingB := ensureBookingState(stateB) + attachBookingID(requestA, bookingA.ID) + attachBookingID(requestB, bookingB.ID) + + if bookingA.ID == externalRequestID || bookingB.ID == externalRequestID { + t.Fatal("controller booking ID reused an external request ID") + } + if bookingA.ID == bookingB.ID { + t.Fatal("independent scheduling cycles received the same booking ID") + } + if got := bookingIDFromRequest(requestA); got != bookingA.ID { + t.Fatalf("request A booking header = %q, want %q", got, bookingA.ID) + } + if got := bookingIDFromRequest(requestB); got != bookingB.ID { + t.Fatalf("request B booking header = %q, want %q", got, bookingB.ID) + } +} + +func TestResponseBodyRetriesPrefillMarkAndFreesTerminalResponse(t *testing.T) { + bookingID := ensureBookingState(schedtypes.NewCycleState()).ID + request := requestWithBooking("external-request", bookingID) + request.Headers[RoutingModeHeader] = "disaggregated" + markCalls := 0 + freeCalls := 0 + scorer := &DynDecodeScorer{ + markPrefillComplete: func(got string) error { + if got != bookingID { + t.Fatalf("mark booking ID = %q, want %q", got, bookingID) + } + markCalls++ + if markCalls == 1 { + return errors.New("transient mark failure") + } + return nil + }, + freeBooking: func(got string) error { + if got != bookingID { + t.Fatalf("free booking ID = %q, want %q", got, bookingID) + } + freeCalls++ + return nil + }, + } + registerBookingLifecycle(bookingID, scorer.freeBooking) + + scorer.ResponseBody(context.Background(), request, &rc.Response{}, nil) + scorer.ResponseBody(context.Background(), request, &rc.Response{}, nil) + scorer.ResponseBody(context.Background(), request, &rc.Response{}, nil) + + lifecycle := findBookingLifecycle(bookingID) + if lifecycle == nil { + t.Fatal("expected booking lifecycle after first response chunk") + } + select { + case <-lifecycle.markerComplete(): + case <-time.After(time.Second): + t.Fatal("prefill marker did not finish") + } + scorer.ResponseBody(context.Background(), request, &rc.Response{EndOfStream: true}, nil) + select { + case <-lifecycle.cleanupComplete(): + case <-time.After(time.Second): + t.Fatal("terminal booking cleanup did not finish") + } + + if markCalls != 2 { + t.Fatalf("mark calls = %d, want one retry then success", markCalls) + } + if freeCalls != 1 { + t.Fatalf("free calls = %d, want 1", freeCalls) + } +} + +func TestResponseBodyFreesWithoutMarkingEmptyTerminalResponse(t *testing.T) { + bookingID := ensureBookingState(schedtypes.NewCycleState()).ID + markCalls := 0 + freeCalls := 0 + scorer := &DynDecodeScorer{ + markPrefillComplete: func(string) error { + markCalls++ + return nil + }, + freeBooking: func(string) error { + freeCalls++ + return nil + }, + } + registerBookingLifecycle(bookingID, scorer.freeBooking) + + scorer.ResponseBody( + context.Background(), + requestWithBooking("external-request", bookingID), + &rc.Response{EndOfStream: true}, + nil, + ) + lifecycle := findBookingLifecycle(bookingID) + if lifecycle == nil { + t.Fatal("expected booking lifecycle for terminal response") + } + select { + case <-lifecycle.cleanupComplete(): + case <-time.After(time.Second): + t.Fatal("terminal booking cleanup did not finish") + } + + if markCalls != 0 { + t.Fatalf("mark calls = %d, want 0", markCalls) + } + if freeCalls != 1 { + t.Fatalf("free calls = %d, want 1", freeCalls) + } +} + +func TestPrefillScoreCancelsPendingReservation(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + reserveStarted := make(chan struct{}) + allowReserveReturn := make(chan struct{}) + cancelCalls := make(chan string, 1) + releaseCalls := make(chan string, 1) + combinedFreeCalls := 0 + scorer := &DynPrefillScorer{ + reservePrefill: func(string, string, string) (*dynscorer.RoutingResult, error) { + close(reserveStarted) + <-allowReserveReturn + return &dynscorer.RoutingResult{WorkerID: 7}, nil + }, + cancelPrefill: func(bookingID string) error { + cancelCalls <- bookingID + close(allowReserveReturn) + return nil + }, + releasePrefill: func(bookingID string) error { + releaseCalls <- bookingID + return nil + }, + freeBooking: func(string) error { + combinedFreeCalls++ + return nil + }, + } + cycleState := schedtypes.NewCycleState() + cycleState.Write(PrefillEnabledStateKey, &PrefillEnabledState{Enabled: true}) + req := &schedtypes.InferenceRequest{ + TargetModel: "model", + Headers: map[string]string{}, + Body: &fwkrh.InferenceRequestBody{ + Payload: fwkrh.PayloadMap{"model": "model", "prompt": "hello"}, + }, + } + + scoresCh := make(chan map[schedtypes.Endpoint]float64, 1) + go func() { + scoresCh <- scorer.Score(ctx, cycleState, req, nil) + }() + <-reserveStarted + cancel() + + if scores := <-scoresCh; len(scores) != 0 { + t.Fatalf("scores = %v, want aggregate fallback with no prefill endpoints", scores) + } + bookingID := bookingIDFromRequest(req) + if got := <-cancelCalls; got != bookingID { + t.Fatalf("cancel booking ID = %q, want %q", got, bookingID) + } + select { + case got := <-releaseCalls: + if got != bookingID { + t.Fatalf("late release booking ID = %q, want %q", got, bookingID) + } + case <-time.After(time.Second): + t.Fatal("late successful reservation was not released") + } + if combinedFreeCalls != 0 { + t.Fatalf("combined free_request calls = %d, want 0", combinedFreeCalls) + } + if readPrefillEnabled(cycleState) { + t.Fatal("prefill remained enabled after reservation cancellation") + } +} + +func TestPrefillScoreReleasesLateErrorResult(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + reserveStarted := make(chan struct{}) + allowReserveReturn := make(chan struct{}) + releaseCalls := make(chan string, 1) + scorer := &DynPrefillScorer{ + reservePrefill: func(string, string, string) (*dynscorer.RoutingResult, error) { + close(reserveStarted) + <-allowReserveReturn + return nil, errors.New("reservation failed after activation") + }, + cancelPrefill: func(string) error { + close(allowReserveReturn) + return nil + }, + releasePrefill: func(bookingID string) error { + releaseCalls <- bookingID + return nil + }, + } + cycleState := schedtypes.NewCycleState() + cycleState.Write(PrefillEnabledStateKey, &PrefillEnabledState{Enabled: true}) + req := &schedtypes.InferenceRequest{ + TargetModel: "model", + Headers: map[string]string{}, + Body: &fwkrh.InferenceRequestBody{ + Payload: fwkrh.PayloadMap{"model": "model", "prompt": "hello"}, + }, + } + + scoresCh := make(chan map[schedtypes.Endpoint]float64, 1) + go func() { + scoresCh <- scorer.Score(ctx, cycleState, req, nil) + }() + <-reserveStarted + cancel() + <-scoresCh + + bookingID := bookingIDFromRequest(req) + select { + case got := <-releaseCalls: + if got != bookingID { + t.Fatalf("late release booking ID = %q, want %q", got, bookingID) + } + case <-time.After(time.Second): + t.Fatal("late error-valued reservation result was not released") + } +} + +func TestResponseBodyBoundsPersistentPrefillMarkRetries(t *testing.T) { + bookingID := ensureBookingState(schedtypes.NewCycleState()).ID + markCalls := 0 + freeCalls := 0 + scorer := &DynDecodeScorer{ + markPrefillComplete: func(string) error { + markCalls++ + return errors.New("persistent mark failure") + }, + freeBooking: func(string) error { + freeCalls++ + return nil + }, + } + request := requestWithBooking("external-request", bookingID) + request.Headers[RoutingModeHeader] = "disaggregated" + registerBookingLifecycle(bookingID, scorer.freeBooking) + + for range 10 { + scorer.ResponseBody(context.Background(), request, &rc.Response{}, nil) + } + lifecycle := findBookingLifecycle(bookingID) + if lifecycle == nil { + t.Fatal("expected booking lifecycle after first response chunk") + } + select { + case <-lifecycle.markerComplete(): + case <-time.After(2 * time.Second): + t.Fatal("bounded marker retries did not finish") + } + if markCalls != prefillMarkMaxAttempts { + t.Fatalf("mark calls = %d, want %d", markCalls, prefillMarkMaxAttempts) + } + + scorer.ResponseBody(context.Background(), request, &rc.Response{EndOfStream: true}, nil) + select { + case <-lifecycle.cleanupComplete(): + case <-time.After(time.Second): + t.Fatal("terminal booking cleanup did not finish") + } + if freeCalls != 1 { + t.Fatalf("free calls = %d, want 1", freeCalls) + } +} + +func TestResponseBodyEOSDoesNotWaitForInFlightPrefillMark(t *testing.T) { + bookingID := ensureBookingState(schedtypes.NewCycleState()).ID + markStarted := make(chan struct{}) + allowMarkReturn := make(chan struct{}) + freeCalls := 0 + scorer := &DynDecodeScorer{ + markPrefillComplete: func(string) error { + close(markStarted) + <-allowMarkReturn + return nil + }, + freeBooking: func(string) error { + freeCalls++ + return nil + }, + } + request := requestWithBooking("external-request", bookingID) + request.Headers[RoutingModeHeader] = "disaggregated" + registerBookingLifecycle(bookingID, scorer.freeBooking) + scorer.ResponseBody(context.Background(), request, &rc.Response{}, nil) + <-markStarted + lifecycle := findBookingLifecycle(bookingID) + if lifecycle == nil { + t.Fatal("expected booking lifecycle after first response chunk") + } + + responseDone := make(chan struct{}) + go func() { + scorer.ResponseBody(context.Background(), request, &rc.Response{EndOfStream: true}, nil) + close(responseDone) + }() + select { + case <-responseDone: + case <-time.After(time.Second): + close(allowMarkReturn) + t.Fatal("EOS callback blocked waiting on prefill mark") + } + close(allowMarkReturn) + select { + case <-lifecycle.cleanupComplete(): + case <-time.After(time.Second): + t.Fatal("terminal booking cleanup did not finish") + } + if freeCalls != 1 { + t.Fatalf("free calls = %d, want 1", freeCalls) + } +} + +func TestBookingLifecycleCleansUpOnceForEOSAndCancellation(t *testing.T) { + bookingID := ensureBookingState(schedtypes.NewCycleState()).ID + freeCalls := 0 + lifecycle := registerBookingLifecycle(bookingID, func(string) error { + freeCalls++ + return nil + }) + ctx, cancel := context.WithCancel(context.Background()) + lifecycle.armCancellation(ctx) + lifecycle.cleanup(ctx, "response end of stream") + cancel() + select { + case <-lifecycle.cleanupComplete(): + case <-time.After(time.Second): + t.Fatal("booking cleanup did not finish") + } + if freeCalls != 1 { + t.Fatalf("free calls = %d, want 1", freeCalls) + } +} + +func TestPrefillScoreBoundsConcurrentReservations(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + reserveStarted := make(chan struct{}) + allowReserveReturn := make(chan struct{}) + reserveCalls := 0 + scorer := &DynPrefillScorer{ + reservePrefill: func(string, string, string) (*dynscorer.RoutingResult, error) { + reserveCalls++ + close(reserveStarted) + <-allowReserveReturn + return &dynscorer.RoutingResult{WorkerID: 7}, nil + }, + freeBooking: func(string) error { return nil }, + reservationAdmissionTimeout: 20 * time.Millisecond, + reservationSlots: make(chan struct{}, 1), + } + newRequest := func() (*schedtypes.CycleState, *schedtypes.InferenceRequest) { + cycleState := schedtypes.NewCycleState() + cycleState.Write(PrefillEnabledStateKey, &PrefillEnabledState{Enabled: true}) + return cycleState, &schedtypes.InferenceRequest{ + TargetModel: "model", + Headers: map[string]string{}, + Body: &fwkrh.InferenceRequestBody{ + Payload: fwkrh.PayloadMap{"model": "model", "prompt": "hello"}, + }, + } + } + firstState, firstRequest := newRequest() + firstScores := make(chan map[schedtypes.Endpoint]float64, 1) + go func() { + firstScores <- scorer.Score(ctx, firstState, firstRequest, nil) + }() + <-reserveStarted + + secondState, secondRequest := newRequest() + if scores := scorer.Score(context.Background(), secondState, secondRequest, nil); len(scores) != 0 { + t.Fatalf("scores = %v, want aggregate fallback with no prefill endpoints", scores) + } + if reserveCalls != 1 { + t.Fatalf("reserve calls = %d, want 1 bounded in-flight admission", reserveCalls) + } + + close(allowReserveReturn) + <-firstScores + cancel() + lifecycle := findBookingLifecycle(bookingIDFromRequest(firstRequest)) + if lifecycle == nil { + t.Fatal("expected lifecycle for successful first reservation") + } + select { + case <-lifecycle.cleanupComplete(): + case <-time.After(time.Second): + t.Fatal("first reservation cleanup did not finish") + } +} + +func newDecodeRequest() *schedtypes.InferenceRequest { + return &schedtypes.InferenceRequest{ + TargetModel: "model", + Headers: map[string]string{}, + Body: &fwkrh.InferenceRequestBody{ + Payload: fwkrh.PayloadMap{"model": "model", "prompt": "hello"}, + }, + } +} + +func TestDecodeScoreCancellationWaitsForRegistration(t *testing.T) { + cycleState := schedtypes.NewCycleState() + booking := ensureBookingState(cycleState) + request := newDecodeRequest() + addStarted := make(chan struct{}) + allowAddReturn := make(chan struct{}) + freeCalls := make(chan string, 1) + scorer := &DynDecodeScorer{ + routeDecode: func(string, string, bool) (*dynscorer.RoutingResult, error) { + return &dynscorer.RoutingResult{WorkerID: 7, DpRank: 1, TokenData: []int64{1}}, nil + }, + addRequest: func(string, []int64, uint64, uint32, string) error { + close(addStarted) + <-allowAddReturn + return nil + }, + freeBooking: func(bookingID string) error { + freeCalls <- bookingID + return nil + }, + } + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + scoresDone := make(chan map[schedtypes.Endpoint]float64, 1) + go func() { + scoresDone <- scorer.Score(ctx, cycleState, request, nil) + }() + <-addStarted + lifecycle := findBookingLifecycle(booking.ID) + if lifecycle == nil { + t.Fatal("expected lifecycle while decode registration is in flight") + } + cancel() + select { + case got := <-freeCalls: + t.Fatalf("booking %q was freed before decode registration returned", got) + case <-time.After(20 * time.Millisecond): + } + close(allowAddReturn) + <-scoresDone + select { + case <-lifecycle.cleanupComplete(): + case <-time.After(time.Second): + t.Fatal("booking cleanup did not finish after registration returned") + } + select { + case got := <-freeCalls: + if got != booking.ID { + t.Fatalf("freed booking ID = %q, want %q", got, booking.ID) + } + case <-time.After(time.Second): + t.Fatal("booking was not cleaned up after cancellation") + } +} + +func TestDecodeScoreRegistrationFailureRedirectsAggregate(t *testing.T) { + cycleState := schedtypes.NewCycleState() + booking := ensureBookingState(cycleState) + request := newDecodeRequest() + freeStarted := make(chan struct{}) + allowFreeReturn := make(chan struct{}) + scorer := &DynDecodeScorer{ + routeDecode: func(string, string, bool) (*dynscorer.RoutingResult, error) { + return &dynscorer.RoutingResult{WorkerID: 7, DpRank: 1, TokenData: []int64{1}}, nil + }, + addRequest: func(string, []int64, uint64, uint32, string) error { + return errors.New("decode booking failed") + }, + freeBooking: func(string) error { + close(freeStarted) + <-allowFreeReturn + return nil + }, + } + + scorer.Score(context.Background(), cycleState, request, nil) + lifecycle := findBookingLifecycle(booking.ID) + if lifecycle == nil { + t.Fatal("expected lifecycle after decode booking failure") + } + select { + case <-freeStarted: + case <-time.After(time.Second): + t.Fatal("failed decode booking was not cleaned up") + } + if got := request.Headers[RoutingModeHeader]; got != "aggregated" { + t.Fatalf("routing mode = %q, want aggregated fallback", got) + } + for _, header := range []string{WorkerIDHeader, DpRankHeader, PrefillWorkerIDHeader, PrefillDpRankHeader} { + if _, ok := request.Headers[header]; ok { + t.Fatalf("redirect retained %s header", header) + } + } + close(allowFreeReturn) + select { + case <-lifecycle.cleanupComplete(): + case <-time.After(time.Second): + t.Fatal("failed decode booking cleanup did not finish") + } +} + +func TestBookingLifecycleRetriesCleanup(t *testing.T) { + bookingID := ensureBookingState(schedtypes.NewCycleState()).ID + calls := 0 + lifecycle := registerBookingLifecycle(bookingID, func(string) error { + calls++ + if calls == 1 { + return errors.New("transient cleanup failure") + } + return nil + }) + if !lifecycle.cleanup(context.Background(), "test retry") { + t.Fatal("initial cleanup did not start") + } + select { + case <-lifecycle.cleanupComplete(): + case <-time.After(time.Second): + t.Fatal("cleanup retry did not finish") + } + if calls != 2 { + t.Fatalf("cleanup calls = %d, want 2", calls) + } + if findBookingLifecycle(bookingID) != nil { + t.Fatal("successful cleanup retained a lifecycle") + } +} + +func TestBookingLifecycleRetainsTombstoneAfterCleanupExhaustion(t *testing.T) { + bookingID := ensureBookingState(schedtypes.NewCycleState()).ID + defer bookingLifecycles.Delete(bookingID) + calls := 0 + lifecycle := registerBookingLifecycle(bookingID, func(string) error { + calls++ + return errors.New("persistent cleanup failure") + }) + if !lifecycle.cleanup(context.Background(), "test exhaustion") { + t.Fatal("initial cleanup did not start") + } + select { + case <-lifecycle.cleanupComplete(): + case <-time.After(time.Second): + t.Fatal("cleanup retries did not finish") + } + if calls != cleanupMaxAttempts { + t.Fatalf("cleanup calls = %d, want %d", calls, cleanupMaxAttempts) + } + if got := findBookingLifecycle(bookingID); got != lifecycle { + t.Fatal("exhausted cleanup did not retain its lifecycle tombstone") + } + if lifecycle.cleanup(context.Background(), "duplicate cleanup") { + t.Fatal("exhausted cleanup started a second owner") + } +} + +func TestResponseBodySkipsMarkForAggregatedRequest(t *testing.T) { + bookingID := ensureBookingState(schedtypes.NewCycleState()).ID + markCalled := make(chan struct{}, 1) + freeCalls := 0 + scorer := &DynDecodeScorer{ + markPrefillComplete: func(string) error { + markCalled <- struct{}{} + return nil + }, + freeBooking: func(string) error { + freeCalls++ + return nil + }, + } + request := requestWithBooking("external-request", bookingID) + request.Headers[RoutingModeHeader] = "aggregated" + lifecycle := registerBookingLifecycle(bookingID, scorer.freeBooking) + + scorer.ResponseBody(context.Background(), request, &rc.Response{}, nil) + select { + case <-markCalled: + t.Fatal("aggregated response marked prefill complete") + case <-time.After(100 * time.Millisecond): + } + scorer.ResponseBody(context.Background(), request, &rc.Response{EndOfStream: true}, nil) + select { + case <-lifecycle.cleanupComplete(): + case <-time.After(time.Second): + t.Fatal("aggregated response cleanup did not finish") + } + if freeCalls != 1 { + t.Fatalf("free calls = %d, want 1", freeCalls) + } +} + +func TestResponseBodyIgnoresUntrackedBookingHeader(t *testing.T) { + bookingID := ensureBookingState(schedtypes.NewCycleState()).ID + markCalled := make(chan struct{}, 1) + freeCalled := make(chan struct{}, 1) + scorer := &DynDecodeScorer{ + markPrefillComplete: func(string) error { + markCalled <- struct{}{} + return nil + }, + freeBooking: func(string) error { + freeCalled <- struct{}{} + return nil + }, + } + request := requestWithBooking("client-request", bookingID) + request.Headers[RoutingModeHeader] = "disaggregated" + + scorer.ResponseBody(context.Background(), request, &rc.Response{}, nil) + scorer.ResponseBody(context.Background(), request, &rc.Response{EndOfStream: true}, nil) + select { + case <-markCalled: + t.Fatal("untracked booking header marked prefill complete") + case <-freeCalled: + t.Fatal("untracked booking header freed router bookkeeping") + case <-time.After(100 * time.Millisecond): + } +} diff --git a/deploy/inference-gateway/epp/pkg/plugins/disagg/shared.go b/deploy/inference-gateway/epp/pkg/plugins/disagg/shared.go index 5d381b03a103..a8ef3a49a3eb 100644 --- a/deploy/inference-gateway/epp/pkg/plugins/disagg/shared.go +++ b/deploy/inference-gateway/epp/pkg/plugins/disagg/shared.go @@ -27,13 +27,15 @@ limitations under the License. package disagg import ( - "encoding/json" - "fmt" + "context" "os" "strings" "sync" + "time" "github.com/go-logr/logr" + "github.com/google/uuid" + log "sigs.k8s.io/controller-runtime/pkg/log" logutil "sigs.k8s.io/gateway-api-inference-extension/pkg/common/observability/logging" plugins "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/framework/interface/plugin" fwkrh "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/framework/interface/requesthandling" @@ -48,6 +50,8 @@ const ( // PrefillEnabledStateKey tracks whether this request should use disaggregated routing. PrefillEnabledStateKey = plugins.StateKey("disagg-prefill-enabled") + BookingStateKey = plugins.StateKey("dynamo-epp-booking") + BookingIDHeader = "x-dynamo-epp-booking-id" ) // PrefillEnabledState stores whether prefill is enabled for the current scheduling cycle. @@ -69,17 +73,64 @@ func readPrefillEnabled(cycleState *schedtypes.CycleState) bool { return false } -// buildRequestJSON builds an OpenAI-compatible JSON string from a GAIE LLMRequest. -func buildRequestJSON(req *schedtypes.InferenceRequest) (string, error) { - requestBody, err := dynscorer.BuildOpenAIRequest(req) - if err != nil { - return "", fmt.Errorf("failed to build OpenAI request: %w", err) +// BookingState owns the controller-generated ID used for router bookkeeping. +type BookingState struct { + ID string + PrefillReserved bool +} + +// Clone implements plugins.StateData. +func (s *BookingState) Clone() plugins.StateData { + if s == nil { + return &BookingState{} } - data, err := json.Marshal(requestBody) - if err != nil { - return "", fmt.Errorf("failed to marshal request JSON: %w", err) + return &BookingState{ID: s.ID, PrefillReserved: s.PrefillReserved} +} + +func readBookingState(cycleState *schedtypes.CycleState) (*BookingState, bool) { + state, err := schedtypes.ReadCycleStateKey[*BookingState](cycleState, BookingStateKey) + if err != nil || state == nil { + return nil, false + } + if _, err := uuid.Parse(state.ID); err != nil { + return nil, false } - return string(data), nil + return state, true +} + +func ensureBookingState(cycleState *schedtypes.CycleState) *BookingState { + if state, ok := readBookingState(cycleState); ok { + return state + } + state := &BookingState{ID: uuid.NewString()} + cycleState.Write(BookingStateKey, state) + return state +} + +func attachBookingID(request *schedtypes.InferenceRequest, bookingID string) { + if request == nil { + return + } + if request.Headers == nil { + request.Headers = map[string]string{} + } + request.Headers[BookingIDHeader] = bookingID +} + +func bookingIDFromRequest(request *schedtypes.InferenceRequest) string { + if request == nil || request.Headers == nil { + return "" + } + bookingID := request.Headers[BookingIDHeader] + if _, err := uuid.Parse(bookingID); err != nil { + return "" + } + return bookingID +} + +// buildRequestJSON builds an OpenAI-compatible JSON string from a GAIE LLMRequest. +func buildRequestJSON(req *schedtypes.InferenceRequest) (string, error) { + return dynscorer.BuildOpenAIRequestJSON(req) } // serializeEndpoints converts endpoints to a JSON string for the FFI filter. @@ -163,6 +214,189 @@ func getEnvBoolOrDefault(key string, def bool) bool { var enforceDisaggDeprecationOnce sync.Once +const ( + prefillMarkMaxAttempts = 3 + prefillMarkRetryBackoff = 100 * time.Millisecond + cleanupMaxAttempts = 3 + cleanupRetryBackoff = 100 * time.Millisecond +) + +var bookingLifecycles sync.Map + +// bookingLifecycle owns cleanup for one EPP booking across the prefill scorer, +// decode scorer, and response callbacks. A booking has exactly one cleanup owner +// even when EOS, decode registration, and context cancellation race. +type bookingLifecycle struct { + bookingID string + freeBooking func(string) error + executor *bookingExecutor + + mu sync.Mutex + cleanupStarted bool + cleanupSucceeded bool + cleanupExhausted bool + cleanupExpired bool + cleanupQueued bool + cleanupRunning bool + cleanupDoneClosed bool + cleanupStartedAt time.Time + cleanupRetryAt time.Time + cleanupReason string + cleanupLogger logr.Logger + decodeRegistrationDone chan struct{} + decodeRegistrationOpen bool + + stopCancellation func() bool + stopMarker context.CancelFunc + markerDone chan struct{} + cleanupDone chan struct{} +} + +func registerBookingLifecycle(bookingID string, freeBooking func(string) error) *bookingLifecycle { + return registerBookingLifecycleWithExecutor(bookingID, freeBooking, defaultBookingExecutor) +} + +func registerBookingLifecycleWithExecutor(bookingID string, freeBooking func(string) error, executor *bookingExecutor) *bookingLifecycle { + lifecycle := &bookingLifecycle{ + bookingID: bookingID, + freeBooking: freeBooking, + executor: executor, + } + actual, loaded := bookingLifecycles.LoadOrStore(bookingID, lifecycle) + if loaded { + return actual.(*bookingLifecycle) + } + return lifecycle +} + +func findBookingLifecycle(bookingID string) *bookingLifecycle { + lifecycle, ok := bookingLifecycles.Load(bookingID) + if !ok { + return nil + } + return lifecycle.(*bookingLifecycle) +} + +// startDecodeRegistration installs a barrier before adding the decode booking. +// Cleanup waits on this barrier so cancellation cannot free the booking before +// add_request has either installed it or reported failure. +func (l *bookingLifecycle) startDecodeRegistration() bool { + l.mu.Lock() + defer l.mu.Unlock() + if l.cleanupStarted || l.decodeRegistrationDone != nil { + return false + } + l.decodeRegistrationDone = make(chan struct{}) + l.decodeRegistrationOpen = true + return true +} + +func (l *bookingLifecycle) finishDecodeRegistration() { + l.mu.Lock() + if !l.decodeRegistrationOpen { + l.mu.Unlock() + return + } + l.decodeRegistrationOpen = false + done := l.decodeRegistrationDone + l.mu.Unlock() + close(done) +} + +func (l *bookingLifecycle) armCancellation(ctx context.Context) { + l.mu.Lock() + defer l.mu.Unlock() + if l.cleanupStarted || l.stopCancellation != nil { + return + } + l.stopCancellation = context.AfterFunc(ctx, func() { + l.cleanup(ctx, "request context cancelled") + }) +} + +// startPrefillMarker submits first-token bookkeeping to a bounded executor and +// returns without blocking the response callback. If the marker queue is full, +// terminal cleanup still releases both the prefill and decode bookings. +func (l *bookingLifecycle) startPrefillMarker(markPrefillComplete func(string) error, logger logr.Logger, requestID string) { + l.mu.Lock() + if l.cleanupStarted || l.markerDone != nil { + l.mu.Unlock() + return + } + markerCtx, stopMarker := context.WithCancel(context.Background()) + markerDone := make(chan struct{}) + l.stopMarker = stopMarker + l.markerDone = markerDone + l.mu.Unlock() + + if !l.executor.enqueueMarker(prefillMarkerWork{ + lifecycle: l, + ctx: markerCtx, + done: markerDone, + markPrefillComplete: markPrefillComplete, + logger: logger, + requestID: requestID, + }) { + stopMarker() + close(markerDone) + logger.V(logutil.DEFAULT).Info("DynDecodeScorer ResponseBody: prefill marker queue full; deferring release to terminal cleanup", + "bookingID", l.bookingID, "requestID", requestID) + } +} + +// cleanup transfers ownership to the bounded cleanup executor. Queue overflow +// retains the lifecycle for reconciler submission, and exhausted work remains +// retryable only for a finite retention window. +func (l *bookingLifecycle) cleanup(ctx context.Context, reason string) bool { + l.mu.Lock() + if l.cleanupStarted { + l.mu.Unlock() + return false + } + l.cleanupStarted = true + l.cleanupStartedAt = time.Now() + l.cleanupReason = reason + l.cleanupLogger = log.FromContext(ctx) + stopCancellation := l.stopCancellation + stopMarker := l.stopMarker + cleanupDone := make(chan struct{}) + l.cleanupDone = cleanupDone + l.mu.Unlock() + + if stopCancellation != nil { + stopCancellation() + } + if stopMarker != nil { + stopMarker() + } + + if !l.executor.enqueueCleanup(l) { + l.cleanupLogger.V(logutil.DEFAULT).Info("Dynamo EPP booking cleanup queue full; retained for reconciler submission", + "bookingID", l.bookingID, "reason", reason) + } + return true +} + +func (l *bookingLifecycle) closeCleanupDoneLocked() chan struct{} { + if l.cleanupDone == nil || l.cleanupDoneClosed { + return nil + } + l.cleanupDoneClosed = true + return l.cleanupDone +} + +func (l *bookingLifecycle) markerComplete() <-chan struct{} { + l.mu.Lock() + defer l.mu.Unlock() + return l.markerDone +} + +func (l *bookingLifecycle) cleanupComplete() <-chan struct{} { + l.mu.Lock() + defer l.mu.Unlock() + return l.cleanupDone +} + func warnDeprecatedEnforceDisagg(logger logr.Logger) { if getEnvBoolOrDefault("DYN_ENFORCE_DISAGG", false) { enforceDisaggDeprecationOnce.Do(func() { diff --git a/deploy/inference-gateway/epp/pkg/plugins/dynamo_kv_scorer/plugin.go b/deploy/inference-gateway/epp/pkg/plugins/dynamo_kv_scorer/plugin.go index 8630970e662f..7f2505f675e4 100644 --- a/deploy/inference-gateway/epp/pkg/plugins/dynamo_kv_scorer/plugin.go +++ b/deploy/inference-gateway/epp/pkg/plugins/dynamo_kv_scorer/plugin.go @@ -41,6 +41,7 @@ enum { QUERY_ROUTER_ERR_QUERY_FAILED = 4, QUERY_ROUTER_ERR_DISAGG_ENFORCED = 5, QUERY_ROUTER_ERR_TIMEOUT = 6, + QUERY_ROUTER_ERR_BACKPRESSURE = 7, }; // opaque handle forward-decl for Router bindings @@ -64,11 +65,26 @@ query_router_result_t create_routers(const char *namespace_c_str, bool enforce_disagg, RouterHandles **out_handle); +query_router_result_t begin_prefill_reservation(RouterHandles *handle, + const char *reservation_id); + query_router_result_t route_prefill_request(RouterHandles *handle, const char *request_json, const char *pods_json, CRoutingResult *out_result); +query_router_result_t route_prefill_request_with_reservation(RouterHandles *handle, + const char *reservation_id, + const char *request_json, + const char *pods_json, + CRoutingResult *out_result); + +query_router_result_t cancel_prefill_reservation(RouterHandles *handle, + const char *reservation_id); + +query_router_result_t release_prefill_reservation(RouterHandles *handle, + const char *reservation_id); + query_router_result_t route_decode_request(RouterHandles *handle, const char *request_json, const char *pods_json, @@ -97,6 +113,7 @@ import "C" import ( "encoding/json" "fmt" + "maps" "os" "strings" "sync" @@ -278,65 +295,31 @@ func SerializeEndpointsToJSON(endpoints []schedtypes.Endpoint) (string, error) { return string(data), nil } -func BuildOpenAIRequest(req *schedtypes.InferenceRequest) (map[string]any, error) { - requestBody := make(map[string]any) - +// BuildOpenAIRequestJSON forwards the full request body (req.Body.Payload) to +// the Rust router FFI, overriding only the resolved model. This preserves +// tool calls, reasoning fields, and other request data needed to render and +// tokenize the same prompt as the worker. +func BuildOpenAIRequestJSON(req *schedtypes.InferenceRequest) (string, error) { if req == nil || req.Body == nil { - return nil, fmt.Errorf("missing request body") + return "", fmt.Errorf("missing request body") } - if req.Body.ChatCompletions != nil && len(req.Body.ChatCompletions.Messages) > 0 { - messages := make([]map[string]any, 0, len(req.Body.ChatCompletions.Messages)) - anyNonEmpty := false - for _, msg := range req.Body.ChatCompletions.Messages { - content := msg.Content.PlainText() - if strings.TrimSpace(content) != "" { - anyNonEmpty = true - } - messages = append(messages, map[string]any{ - "role": msg.Role, - "content": content, - }) - } - if !anyNonEmpty { - return nil, fmt.Errorf("empty chat messages") - } - requestBody["messages"] = messages - } else if req.Body.Completions != nil && !req.Body.Completions.Prompt.IsEmpty() { - requestBody["messages"] = []map[string]any{ - {"role": "user", "content": req.Body.Completions.Prompt.PlainText()}, - } - } else { - return nil, fmt.Errorf("no messages or prompt provided") + pm, ok := req.Body.Payload.(fwkrh.PayloadMap) + if !ok || len(pm) == 0 { + return "", fmt.Errorf("request payload unavailable; cannot build KV-routing request") } + requestBody := make(map[string]any, len(pm)) + maps.Copy(requestBody, pm) if strings.TrimSpace(req.TargetModel) != "" { requestBody["model"] = req.TargetModel - } else { - requestBody["model"] = "default" } - // Forward the caller's nvext block so the Rust router can lift - // nvext.agent_hints.priority into priority_jump. - if nvext := extractNvext(req.Body.Payload); nvext != nil { - requestBody["nvext"] = nvext + data, err := json.Marshal(requestBody) + if err != nil { + return "", fmt.Errorf("failed to marshal request JSON: %w", err) } - - return requestBody, nil -} - -// extractNvext returns the caller-supplied nvext object from the PayloadMap, -// or nil when the payload is not a map or does not contain an nvext object. -// -// This is how routing hints — most notably nvext.agent_hints.priority — reach -// the Rust router via the FFI JSON. -func extractNvext(payload fwkrh.RequestPayload) map[string]any { - pm, ok := payload.(fwkrh.PayloadMap) - if !ok { - return nil - } - nvext, _ := pm["nvext"].(map[string]any) - return nvext + return string(data), nil } // CallAddRequest registers a request with the router's bookkeeping. @@ -450,9 +433,40 @@ func extractTokenData(result *C.CRoutingResult) []int64 { return nil } -// CallRoutePrefillRequest routes a request to the best prefill worker. -// It tokenizes the request and queries only the prefill router. -func CallRoutePrefillRequest(requestJSON string, podsJSON string) (*RoutingResult, error) { +// CallBeginPrefillReservation records a pending booking before the blocking +// route call performs request preprocessing. It is fast and safe to call before +// starting the reservation goroutine. +func CallBeginPrefillReservation(reservationID string) error { + if reservationID == "" { + return fmt.Errorf("prefill reservation ID is required") + } + if !routerInitialized { + return fmt.Errorf("dynamo router not initialized") + } + + routerHandlesMutex.RLock() + router := routerHandles + routerHandlesMutex.RUnlock() + if router == nil { + return fmt.Errorf("dynamo router handles not created") + } + + cReservationID := C.CString(reservationID) + defer C.free(unsafe.Pointer(cReservationID)) + + rc := C.begin_prefill_reservation(router, cReservationID) + if rc != C.QUERY_ROUTER_OK { + return fmt.Errorf("begin_prefill_reservation failed with code %d", rc) + } + return nil +} + +// CallRoutePrefillRequestWithReservation atomically selects and books a prefill worker. +// The caller cancels pending admission through CallCancelPrefillReservation. +func CallRoutePrefillRequestWithReservation(reservationID string, requestJSON string, podsJSON string) (*RoutingResult, error) { + if reservationID == "" { + return nil, fmt.Errorf("prefill reservation ID is required") + } if !routerInitialized { return nil, fmt.Errorf("dynamo router not initialized") } @@ -464,6 +478,8 @@ func CallRoutePrefillRequest(requestJSON string, podsJSON string) (*RoutingResul return nil, fmt.Errorf("dynamo router handles not created") } + cReservationID := C.CString(reservationID) + defer C.free(unsafe.Pointer(cReservationID)) cRequestJSON := C.CString(requestJSON) defer C.free(unsafe.Pointer(cRequestJSON)) @@ -474,9 +490,15 @@ func CallRoutePrefillRequest(requestJSON string, podsJSON string) (*RoutingResul } var result C.CRoutingResult - rc := C.route_prefill_request(router, cRequestJSON, cPodsJSON, &result) + rc := C.route_prefill_request_with_reservation( + router, + cReservationID, + cRequestJSON, + cPodsJSON, + &result, + ) if rc != C.QUERY_ROUTER_OK { - return nil, fmt.Errorf("route_prefill_request failed with code %d", rc) + return nil, fmt.Errorf("route_prefill_request_with_reservation failed with code %d", rc) } tokens := extractTokenData(&result) @@ -484,11 +506,70 @@ func CallRoutePrefillRequest(requestJSON string, podsJSON string) (*RoutingResul dpRank := uint32(result.prefill_dp_rank) C.free_routing_result(&result) - return &RoutingResult{WorkerID: workerID, DpRank: dpRank, TokenData: tokens}, nil + return &RoutingResult{ + WorkerID: workerID, + DpRank: dpRank, + TokenData: tokens, + }, nil +} + +// CallCancelPrefillReservation cancels a pending prefill reservation without waiting for +// scheduler cleanup. It is safe to call concurrently with the blocking reservation call. +func CallCancelPrefillReservation(reservationID string) error { + if reservationID == "" { + return fmt.Errorf("prefill reservation ID is required") + } + if !routerInitialized { + return fmt.Errorf("dynamo router not initialized") + } + + routerHandlesMutex.RLock() + router := routerHandles + routerHandlesMutex.RUnlock() + if router == nil { + return fmt.Errorf("dynamo router handles not created") + } + + cReservationID := C.CString(reservationID) + defer C.free(unsafe.Pointer(cReservationID)) + + rc := C.cancel_prefill_reservation(router, cReservationID) + if rc != C.QUERY_ROUTER_OK { + return fmt.Errorf("cancel_prefill_reservation failed with code %d", rc) + } + return nil +} + +// CallReleasePrefillReservation releases only the EPP prefill reservation. +// It does not remove a decode booking that aggregate fallback may have installed. +func CallReleasePrefillReservation(reservationID string) error { + if reservationID == "" { + return fmt.Errorf("prefill reservation ID is required") + } + if !routerInitialized { + return fmt.Errorf("dynamo router not initialized") + } + + routerHandlesMutex.RLock() + router := routerHandles + routerHandlesMutex.RUnlock() + if router == nil { + return fmt.Errorf("dynamo router handles not created") + } + + cReservationID := C.CString(reservationID) + defer C.free(unsafe.Pointer(cReservationID)) + + rc := C.release_prefill_reservation(router, cReservationID) + if rc != C.QUERY_ROUTER_OK { + return fmt.Errorf("release_prefill_reservation failed with code %d", rc) + } + return nil } // CallRouteDecodeRequest routes a request to the best decode worker. // When isDisaggregated is true, overlap_score_credit=0 is used (KV cache transferred from prefill). + func CallRouteDecodeRequest(requestJSON string, podsJSON string, isDisaggregated bool) (*RoutingResult, error) { if !routerInitialized { return nil, fmt.Errorf("dynamo router not initialized") diff --git a/deploy/inference-gateway/epp/pkg/plugins/dynamo_kv_scorer/plugin_test.go b/deploy/inference-gateway/epp/pkg/plugins/dynamo_kv_scorer/plugin_test.go index df92b02284ab..f91b110c6a4e 100644 --- a/deploy/inference-gateway/epp/pkg/plugins/dynamo_kv_scorer/plugin_test.go +++ b/deploy/inference-gateway/epp/pkg/plugins/dynamo_kv_scorer/plugin_test.go @@ -17,26 +17,35 @@ limitations under the License. package dynamo_kv_scorer import ( + "encoding/json" "testing" fwkrh "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/framework/interface/requesthandling" schedtypes "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/framework/interface/scheduling" ) -// TestBuildOpenAIRequest_ForwardsAgentHintsPriority pins the contract that -// nvext.agent_hints.priority arriving on the original request body is -// preserved in the JSON sent across FFI to the Rust router. Without this, -// the router falls back to priority_jump=0.0 for every request and queue -// ordering silently regresses. -func TestBuildOpenAIRequest_ForwardsAgentHintsPriority(t *testing.T) { +// ffiBody builds the FFI JSON and parses it back into a map for assertions. +func ffiBody(t *testing.T, req *schedtypes.InferenceRequest) map[string]any { + t.Helper() + s, err := BuildOpenAIRequestJSON(req) + if err != nil { + t.Fatalf("BuildOpenAIRequestJSON returned error: %v", err) + } + var body map[string]any + if err := json.Unmarshal([]byte(s), &body); err != nil { + t.Fatalf("failed to unmarshal FFI JSON: %v (json=%s)", err, s) + } + return body +} + +// TestBuildOpenAIRequestJSON_ForwardsAgentHintsPriority pins the contract that +// nvext.agent_hints.priority on the original request body is preserved in the +// JSON sent across FFI to the Rust router. Without it, the router falls back to +// priority_jump=0.0 for every request and queue ordering silently regresses. +func TestBuildOpenAIRequestJSON_ForwardsAgentHintsPriority(t *testing.T) { req := &schedtypes.InferenceRequest{ TargetModel: "test-model", Body: &fwkrh.InferenceRequestBody{ - ChatCompletions: &fwkrh.ChatCompletionsRequest{ - Messages: []fwkrh.Message{ - {Role: "user", Content: fwkrh.Content{Raw: "hi"}}, - }, - }, Payload: fwkrh.PayloadMap{ "messages": []any{map[string]any{"role": "user", "content": "hi"}}, "model": "test-model", @@ -45,11 +54,7 @@ func TestBuildOpenAIRequest_ForwardsAgentHintsPriority(t *testing.T) { }, } - body, err := BuildOpenAIRequest(req) - if err != nil { - t.Fatalf("BuildOpenAIRequest returned error: %v", err) - } - + body := ffiBody(t, req) nvext, ok := body["nvext"].(map[string]any) if !ok { t.Fatalf("expected nvext to be a map, got %T", body["nvext"]) @@ -58,7 +63,133 @@ func TestBuildOpenAIRequest_ForwardsAgentHintsPriority(t *testing.T) { if !ok { t.Fatalf("expected agent_hints to be a map, got %T", nvext["agent_hints"]) } - if got := hints["priority"]; got != 7 { - t.Fatalf("expected priority=7 forwarded to FFI body, got %v", got) + if got := hints["priority"]; got != float64(7) { // JSON numbers decode to float64 + t.Fatalf("expected priority=7 forwarded to FFI body, got %v (%T)", got, got) + } +} + +// TestBuildOpenAIRequestJSON_ForwardsLegacyTopLevelCacheSalt verifies a +// top-level cache_salt on the request body is forwarded to the router. +func TestBuildOpenAIRequestJSON_ForwardsLegacyTopLevelCacheSalt(t *testing.T) { + req := &schedtypes.InferenceRequest{ + TargetModel: "test-model", + Body: &fwkrh.InferenceRequestBody{ + Payload: fwkrh.PayloadMap{ + "messages": []any{map[string]any{"role": "user", "content": "hi"}}, + "model": "test-model", + "cache_salt": "tenant-legacy", + }, + }, + } + + body := ffiBody(t, req) + if got := body["cache_salt"]; got != "tenant-legacy" { + t.Fatalf("expected legacy cache_salt forwarded to FFI body, got %v", got) + } +} + +// TestBuildOpenAIRequestJSON_PreservesToolCallFields pins the contract that a +// multi-turn tool conversation survives intact in the JSON sent across FFI. The +// router parses this body and re-renders the model's chat template to tokenize, +// so it needs the full message structure. If tool_calls / tool_call_id are +// dropped, the router's strict parse fails ("missing field tool_call_id") and +// the request is unroutable (503 no healthy upstream); reasoning/tool parsing is +// also lost when the template renders an incomplete prompt. +func TestBuildOpenAIRequestJSON_PreservesToolCallFields(t *testing.T) { + toolCall := map[string]any{ + "id": "call-abc", + "type": "function", + "function": map[string]any{ + "name": "get_current_weather", + "arguments": `{"location":"Tokyo"}`, + }, + } + req := &schedtypes.InferenceRequest{ + TargetModel: "test-model", + Body: &fwkrh.InferenceRequestBody{ + Payload: fwkrh.PayloadMap{ + "model": "alias-model", + "messages": []any{ + map[string]any{"role": "user", "content": "weather in Tokyo?"}, + map[string]any{"role": "assistant", "content": nil, "tool_calls": []any{toolCall}}, + map[string]any{"role": "tool", "tool_call_id": "call-abc", "content": "18C rain"}, + }, + }, + }, + } + + body := ffiBody(t, req) + + // Target model must override the caller's alias. + if got := body["model"]; got != "test-model" { + t.Fatalf("expected model=test-model, got %v", got) + } + + msgs, ok := body["messages"].([]any) + if !ok || len(msgs) != 3 { + t.Fatalf("expected 3 messages, got %#v", body["messages"]) + } + + // Assistant turn must retain tool_calls. + assistant, ok := msgs[1].(map[string]any) + if !ok { + t.Fatalf("expected assistant message map, got %T", msgs[1]) + } + if _, ok := assistant["tool_calls"].([]any); !ok { + t.Fatalf("assistant message lost tool_calls: %#v", assistant) + } + + // Tool turn must retain tool_call_id (the field whose loss caused the 503). + tool, ok := msgs[2].(map[string]any) + if !ok { + t.Fatalf("expected tool message map, got %T", msgs[2]) + } + if got := tool["tool_call_id"]; got != "call-abc" { + t.Fatalf("tool message lost tool_call_id: got %v in %#v", got, tool) + } +} + +// TestBuildOpenAIRequestJSON_ForwardsCompletionsPayload verifies a /v1/completions +// body is forwarded verbatim (the Rust preprocessor handles it via the prompt +// field), with only the model overridden to the resolved target. +func TestBuildOpenAIRequestJSON_ForwardsCompletionsPayload(t *testing.T) { + req := &schedtypes.InferenceRequest{ + TargetModel: "test-model", + Body: &fwkrh.InferenceRequestBody{ + Payload: fwkrh.PayloadMap{ + "model": "alias-model", + "prompt": "hello world", + }, + }, + } + + body := ffiBody(t, req) + if got := body["prompt"]; got != "hello world" { + t.Fatalf("expected prompt forwarded, got %v", got) + } + if got := body["model"]; got != "test-model" { + t.Fatalf("expected model overridden to test-model, got %v", got) + } +} + +// TestBuildOpenAIRequestJSON_MissingPayloadReturnsError verifies that when the +// raw payload is unavailable the request is not KV-routable: an error is +// returned so the scorer falls back to non-KV routing rather than a lossy +// role/content reconstruction (which would drop tool-calling fields). +func TestBuildOpenAIRequestJSON_MissingPayloadReturnsError(t *testing.T) { + req := &schedtypes.InferenceRequest{ + TargetModel: "test-model", + Body: &fwkrh.InferenceRequestBody{ + ChatCompletions: &fwkrh.ChatCompletionsRequest{ + Messages: []fwkrh.Message{ + {Role: "user", Content: fwkrh.Content{Raw: "hi"}}, + }, + }, + // No Payload set — the typed view cannot carry tool-calling fields. + }, + } + + if _, err := BuildOpenAIRequestJSON(req); err == nil { + t.Fatalf("expected an error when the raw payload is unavailable, got nil") } } diff --git a/hack/Dockerfile.epp b/hack/Dockerfile.epp new file mode 100644 index 000000000000..66ded1c3c164 --- /dev/null +++ b/hack/Dockerfile.epp @@ -0,0 +1,66 @@ +# Builds the Dynamo EPP (endpoint-picker) binary from a single build context +# (the repo root), so it works with a plain `docker build` — no buildx or named +# build contexts required. Mirrors deploy/inference-gateway/epp/Dockerfile, but +# without the sccache / --build-context machinery. +# +# Intended for a native linux/amd64 host (the Rust + Go compiles run natively). +# +# Build from the repo root: +# DOCKER_BUILDKIT=1 docker build -f hack/Dockerfile.epp \ +# --build-arg COMMIT_SHA=$(git rev-parse HEAD) \ +# --build-arg BUILD_REF=$(git rev-parse --abbrev-ref HEAD) \ +# -t dynamo-epp:1.3.0-tool-fix . + +ARG RUST_IMAGE=rust:1.93.1 +ARG BUILDER_IMAGE=golang:1.26.3 +ARG BASE_IMAGE=ubuntu:24.04 + +# ============================================================================= +# Stage 1: Build Dynamo FFI static library (Rust) +# ============================================================================= +FROM ${RUST_IMAGE} AS rust-builder +RUN apt-get update && apt-get install -y --no-install-recommends \ + protobuf-compiler libclang-dev && rm -rf /var/lib/apt/lists/* +WORKDIR /dynamo +COPY .cargo/ .cargo/ +COPY Cargo.toml Cargo.lock README.md ./ +COPY lib/ lib/ +# ext-proc is a workspace member in the root Cargo.toml, so cargo needs its +# manifest on disk to resolve the workspace even though only libdynamo_llm is +# built here. Without it: "failed to load manifest for workspace member". +COPY deploy/inference-gateway/ext-proc/ deploy/inference-gateway/ext-proc/ +RUN cargo build --release -p libdynamo_llm && \ + mkdir -p /out && cp target/release/libdynamo_llm_capi.a /out/ && \ + HEADER=$(find target/release/build -name llm_engine.h -path "*/out/*" | head -1) && \ + [ -n "$HEADER" ] && cp "$HEADER" /out/ || { echo "ERROR: llm_engine.h not found"; exit 1; } + +# ============================================================================= +# Stage 2: Build Go EPP binary (CGO, links the Rust FFI static library) +# ============================================================================= +FROM ${BUILDER_IMAGE} AS go-builder +ARG COMMIT_SHA +ARG BUILD_REF +RUN apt-get update && apt-get install -y --no-install-recommends \ + gcc g++ libc6-dev && rm -rf /var/lib/apt/lists/* +WORKDIR /workspace +COPY deploy/inference-gateway/epp/go.mod deploy/inference-gateway/epp/go.sum ./ +RUN go mod download +COPY deploy/inference-gateway/epp/ . +COPY --from=rust-builder /out/libdynamo_llm_capi.a pkg/plugins/dynamo_kv_scorer/lib/ +COPY --from=rust-builder /out/llm_engine.h pkg/plugins/dynamo_kv_scorer/include/ +RUN CGO_ENABLED=1 GOOS=linux GOARCH=amd64 go build \ + -ldflags="-X sigs.k8s.io/gateway-api-inference-extension/version.GitVersion=${BUILD_REF} \ + -X sigs.k8s.io/gateway-api-inference-extension/version.GitCommit=${COMMIT_SHA}" \ + -o epp ./cmd/epp + +# ============================================================================= +# Stage 3: Runtime +# ============================================================================= +FROM ${BASE_IMAGE} +RUN apt-get update && apt-get install -y --no-install-recommends \ + ca-certificates libstdc++6 && rm -rf /var/lib/apt/lists/* +WORKDIR / +COPY --from=go-builder /workspace/epp . +RUN useradd -r -u 65532 -g nogroup nonroot +USER 65532:65534 +ENTRYPOINT ["/epp"] diff --git a/hack/Dockerfile.frontend-overlay b/hack/Dockerfile.frontend-overlay new file mode 100644 index 000000000000..1cefdd3cd410 --- /dev/null +++ b/hack/Dockerfile.frontend-overlay @@ -0,0 +1,22 @@ +# Overlays the tool-fix EPP binary onto the upstream monolithic frontend image. +# +# The official dynamo-frontend image serves multiple roles (frontend, workers, +# and the EPP via entrypoint /epp). Only /epp changes for this fix, so instead +# of rebuilding the whole frontend we copy the freshly-built /epp over the +# stock image — keeping every other role byte-identical to the upstream release. +# +# EPP_IMAGE must be built first (see hack/Dockerfile.epp / hack/Makefile). +# +# DOCKER_BUILDKIT=1 docker build -f hack/Dockerfile.frontend-overlay \ +# --build-arg BASE_IMAGE=nvcr.io/nvidia/ai-dynamo/dynamo-frontend:1.3.0 \ +# --build-arg EPP_IMAGE=dynamo-epp:1.3.0-tool-fix \ +# -t registry.dev.rafay-edge.net/tf/dynamo-frontend:1.3.0-tool-fix . + +ARG BASE_IMAGE=nvcr.io/nvidia/ai-dynamo/dynamo-frontend:1.3.0 +ARG EPP_IMAGE=dynamo-epp:1.3.0-tool-fix + +FROM ${EPP_IMAGE} AS epp + +FROM ${BASE_IMAGE} +# Stock frontend runs as the "dynamo" user; keep ownership consistent. +COPY --chown=dynamo: --from=epp /epp /epp diff --git a/hack/Makefile b/hack/Makefile new file mode 100644 index 000000000000..3a601f8a4070 --- /dev/null +++ b/hack/Makefile @@ -0,0 +1,67 @@ +# Build the tool-fix frontend image: the fixed EPP binary overlaid on the +# upstream monolithic dynamo-frontend image. +# +# Two steps: +# 1) epp-image — compile the EPP from source (hack/Dockerfile.epp) +# 2) frontend-image — overlay /epp onto the stock frontend (hack/Dockerfile.frontend-overlay) +# +# Usage (native linux/amd64 host, plain docker — no buildx needed): +# make -C hack push +# Override anything on the command line, e.g.: +# make -C hack push VERSION=1.3.0 REGISTRY=registry.dev.rafay-edge.net/tf + +REPO_ROOT := $(abspath $(dir $(lastword $(MAKEFILE_LIST)))/..) + +PLATFORM ?= linux/amd64 +VERSION ?= 1.3.0 +REGISTRY ?= registry.dev.rafay-edge.net/tf + +# Upstream frontend image the fix is overlaid onto. +BASE_IMAGE ?= nvcr.io/nvidia/ai-dynamo/dynamo-frontend:$(VERSION) +# Local intermediate EPP image (not pushed). +EPP_IMAGE ?= dynamo-epp:$(VERSION)-tool-fix +# Final pushed image (drop-in replacement for the frontend/EPP pod). +IMAGE ?= $(REGISTRY)/dynamo-frontend:$(VERSION)-tool-fix + +COMMIT_SHA := $(shell git -C $(REPO_ROOT) rev-parse HEAD 2>/dev/null) +BUILD_REF := $(shell git -C $(REPO_ROOT) rev-parse --abbrev-ref HEAD 2>/dev/null) + +export DOCKER_BUILDKIT := 1 + +.PHONY: all push frontend-image epp-image print + +## all: build the EPP image, overlay it, and push (default) +all: push + +## epp-image: compile the fixed EPP binary into a local image +epp-image: + docker build \ + --platform=$(PLATFORM) \ + -f $(REPO_ROOT)/hack/Dockerfile.epp \ + --build-arg COMMIT_SHA=$(COMMIT_SHA) \ + --build-arg BUILD_REF=$(BUILD_REF) \ + -t $(EPP_IMAGE) \ + $(REPO_ROOT) + +## frontend-image: overlay the fixed /epp onto the upstream frontend image +frontend-image: epp-image + docker build \ + --platform=$(PLATFORM) \ + -f $(REPO_ROOT)/hack/Dockerfile.frontend-overlay \ + --build-arg BASE_IMAGE=$(BASE_IMAGE) \ + --build-arg EPP_IMAGE=$(EPP_IMAGE) \ + -t $(IMAGE) \ + $(REPO_ROOT)/hack + +## push: build and push the final image +push: frontend-image + docker push $(IMAGE) + +## print: show the resolved image names +print: + @echo "PLATFORM = $(PLATFORM)" + @echo "BASE_IMAGE = $(BASE_IMAGE)" + @echo "EPP_IMAGE = $(EPP_IMAGE)" + @echo "IMAGE = $(IMAGE)" + @echo "COMMIT_SHA = $(COMMIT_SHA)" + @echo "BUILD_REF = $(BUILD_REF)" diff --git a/hack/README.md b/hack/README.md new file mode 100644 index 000000000000..5ffcfd835388 --- /dev/null +++ b/hack/README.md @@ -0,0 +1,49 @@ +# hack/ — tool-fix frontend image build + +Builds a drop-in `dynamo-frontend` image carrying the EPP KV-router tool-call +fix (see the change in `deploy/inference-gateway/epp/pkg/plugins/dynamo_kv_scorer`). + +The official `dynamo-frontend` image is monolithic — it serves the frontend, +workers, and the EPP (entrypoint `/epp`) from one image. Only `/epp` changes for +this fix, so we compile the fixed EPP and **overlay** it onto the stock frontend +image rather than rebuilding the whole thing. Every other role stays identical to +the upstream release. + +## Files + +- `Dockerfile.epp` — compiles the EPP binary from a single build context (repo + root), so a plain `docker build` works (no buildx / named contexts). Native + `linux/amd64` build. +- `Dockerfile.frontend-overlay` — copies the built `/epp` onto the upstream + frontend image (`BASE_IMAGE`). +- `Makefile` — orchestrates: `epp-image` → `frontend-image` → `push`. + +## Usage + +On a native `linux/amd64` host with Docker (BuildKit) and registry access: + +```bash +docker login registry.dev.rafay-edge.net +make -C hack push +``` + +Override defaults as needed: + +```bash +make -C hack push \ + VERSION=1.3.0 \ + REGISTRY=registry.dev.rafay-edge.net/tf \ + BASE_IMAGE=nvcr.io/nvidia/ai-dynamo/dynamo-frontend:1.3.0 +``` + +`make -C hack print` shows the resolved image names without building. + +Then point the EPP pod's `extension.image` (or the shared frontend tag) at the +pushed image, e.g. `registry.dev.rafay-edge.net/tf/dynamo-frontend:1.3.0-tool-fix`. + +## Verify + +```bash +docker run --rm --entrypoint /epp \ + registry.dev.rafay-edge.net/tf/dynamo-frontend:1.3.0-tool-fix --help # prints EPP flags +``` diff --git a/lib/bindings/c/src/lib.rs b/lib/bindings/c/src/lib.rs index 21d8c9c60645..1f7e1d677f19 100644 --- a/lib/bindings/c/src/lib.rs +++ b/lib/bindings/c/src/lib.rs @@ -25,7 +25,7 @@ use dynamo_runtime::{DistributedRuntime, Worker}; use dynamo_runtime::Runtime; use dynamo_llm::discovery::{ModelManager, WORKER_TYPE_DECODE}; -use dynamo_llm::kv_router::prefill_router::PrefillQueryOutcome; +use dynamo_llm::kv_router::prefill_router::{EppReservationManager, PrefillQueryOutcome}; use dynamo_llm::kv_router::{KvRouter, PrefillRouter}; use dynamo_runtime::pipeline::RouterMode; @@ -434,6 +434,7 @@ impl Default for CRoutingResult { /// Container holding routers and preprocessor for query routing pub struct RouterHandles { prefill_router: Arc, + epp_reservations: Arc, decode_router: Arc, #[allow(dead_code)] model_manager: Arc, @@ -446,13 +447,11 @@ pub struct RouterHandles { } impl RouterHandles { - /// Query optimal prefill worker for a request. - /// - /// When `allowed_worker_ids` is Some, only workers in that set are considered. - /// Returns worker_id on success. + /// Atomically select and reserve a prefill worker for an EPP-owned booking. #[expect(clippy::too_many_arguments)] - async fn query_prefill_worker( + async fn reserve_prefill_worker( &self, + reservation_id: &str, tokens: &[u32], block_mm_infos: Option<&[Option]>, lora_name: Option, @@ -466,8 +465,10 @@ impl RouterHandles { } let outcome = self - .prefill_router - .query_prefill_worker( + .epp_reservations + .reserve( + &self.prefill_router, + reservation_id, tokens, block_mm_infos, lora_name, @@ -477,20 +478,20 @@ impl RouterHandles { routing_constraints, ) .await - .map_err(|e| { - tracing::error!(error = ?e, "Prefill query failed"); + .map_err(|error| { + tracing::error!(%reservation_id, %error, "Prefill reservation failed"); QueryRouterResult::ErrQueryFailed })?; match outcome { - // Advisory only: the external caller owns dispatch and lifecycle state. PrefillQueryOutcome::Routed { worker_id, dp_rank } => Ok((worker_id, dp_rank)), PrefillQueryOutcome::QueueRejected { rejection } => { tracing::warn!( + %reservation_id, policy_class = %rejection.policy_class, limit_kind = %rejection.limit_kind, current = rejection.current, limit = rejection.limit, - "Prefill query rejected by policy-class queue limit" + "Prefill reservation rejected by policy-class queue limit" ); Err(QueryRouterResult::ErrBackpressure) } @@ -833,7 +834,11 @@ pub unsafe extern "C" fn create_routers( // to activate the PrefillRouter. spawn_prefill_discovery_watcher(drt.clone(), actual_namespace.clone(), prefill_tx); + let epp_reservations = Arc::new(EppReservationManager::default()); + EppReservationManager::spawn_reaper(&epp_reservations); + Ok(( + epp_reservations, prefill_router, decode_router, model_manager, @@ -843,9 +848,17 @@ pub unsafe extern "C" fn create_routers( }); match result { - Ok((prefill_router, decode_router, model_manager, namespace_str, preprocessor)) => { + Ok(( + epp_reservations, + prefill_router, + decode_router, + model_manager, + namespace_str, + preprocessor, + )) => { let handles = RouterHandles { prefill_router, + epp_reservations, decode_router, model_manager, namespace: namespace_str, @@ -894,6 +907,7 @@ pub unsafe extern "C" fn add_request( Vec::new() }; + let prefill_router = handles.prefill_router.clone(); let decode_router = handles.decode_router.clone(); let result = handles.runtime.secondary().block_on(async { @@ -981,12 +995,16 @@ pub unsafe extern "C" fn mark_prefill_complete( Err(_) => return QueryRouterResult::ErrInvalidParam, }; + let epp_reservations = handles.epp_reservations.clone(); let decode_router = handles.decode_router.clone(); let result = handles.runtime.secondary().block_on(async { let timeout_duration = Duration::from_secs(BOOKKEEPING_TIMEOUT_SEC); tokio::time::timeout(timeout_duration, async { + if let Err(e) = epp_reservations.release(&request_id_str).await { + tracing::warn!(request_id = %request_id_str, error = %e, "Failed to release prefill reservation"); + } if let Err(e) = decode_router.mark_prefill_completed(&request_id_str).await { tracing::warn!( request_id = %request_id_str, @@ -1016,6 +1034,54 @@ pub unsafe extern "C" fn mark_prefill_complete( } } +/// Release only the prefill reservation for an EPP-owned booking. +/// +/// This is used when a cancelled prefill admission completes after the request +/// has already fallen back to aggregate decode routing. It must not free any +/// decode booking that the fallback path may have installed. +/// +/// # Safety +/// - `handle` must be a valid RouterHandles handle +/// - `request_id` must be a valid non-empty null-terminated UTF-8 string +#[unsafe(no_mangle)] +pub unsafe extern "C" fn release_prefill_reservation( + handle: RouterHandlesPtr, + request_id: *const c_char, +) -> QueryRouterResult { + if handle.is_null() || request_id.is_null() { + return QueryRouterResult::ErrInvalidParam; + } + + let handles = unsafe { &*handle }; + let request_id = match unsafe { CStr::from_ptr(request_id) }.to_str() { + Ok(value) if !value.is_empty() => value.to_owned(), + _ => return QueryRouterResult::ErrInvalidParam, + }; + let epp_reservations = handles.epp_reservations.clone(); + let result = handles.runtime.secondary().block_on(async { + tokio::time::timeout(Duration::from_secs(BOOKKEEPING_TIMEOUT_SEC), async { + epp_reservations.release(&request_id).await + }) + .await + }); + + match result { + Ok(Ok(())) => QueryRouterResult::Ok, + Ok(Err(error)) => { + tracing::warn!(%request_id, %error, "Failed to release EPP prefill reservation"); + QueryRouterResult::ErrQueryFailed + } + Err(_) => { + tracing::warn!( + %request_id, + timeout_secs = BOOKKEEPING_TIMEOUT_SEC, + "release_prefill_reservation timed out" + ); + QueryRouterResult::ErrTimeout + } + } +} + /// Free a request from the router's bookkeeping. /// /// Call this when the stream is closed (completed or cancelled) to release all resources. @@ -1038,12 +1104,16 @@ pub unsafe extern "C" fn free_request( Err(_) => return QueryRouterResult::ErrInvalidParam, }; + let epp_reservations = handles.epp_reservations.clone(); let decode_router = handles.decode_router.clone(); let result = handles.runtime.secondary().block_on(async { let timeout_duration = Duration::from_secs(BOOKKEEPING_TIMEOUT_SEC); tokio::time::timeout(timeout_duration, async { + if let Err(e) = epp_reservations.release(&request_id_str).await { + tracing::warn!(request_id = %request_id_str, error = %e, "Failed to release prefill reservation"); + } if let Err(e) = decode_router.free(&request_id_str).await { tracing::warn!( request_id = %request_id_str, @@ -1241,72 +1311,117 @@ fn write_tokens_to_result(tokens: &[u32], out: &mut CRoutingResult) { std::mem::forget(tokens_boxed); } -/// Route a request to select the best **prefill** worker only. +/// Register a prefill booking before potentially slow request preprocessing. /// -/// This is used in disaggregated mode where the EPP runs separate prefill and decode -/// scoring profiles. It tokenizes the request and queries only the prefill router. +/// This is intentionally synchronous and lock-only: Go calls it before it starts +/// the blocking route operation, so cancellation always observes an actual Pending +/// entry rather than relying on a timed tombstone. /// -/// The returned `CRoutingResult` contains: -/// - `prefill_worker_id`: the selected prefill worker -/// - `decode_worker_id`: 0 (unused — decode is handled by `route_decode_request`) -/// - `is_disaggregated`: always true (this function is only called in disagg mode) -/// - `token_ids` / `token_count`: the tokenized request (caller must free via `free_routing_result`) +/// # Safety +/// - `handle` must be a valid RouterHandles handle +/// - `reservation_id` must be a non-empty null-terminated UTF-8 string +#[unsafe(no_mangle)] +pub unsafe extern "C" fn begin_prefill_reservation( + handle: RouterHandlesPtr, + reservation_id: *const c_char, +) -> QueryRouterResult { + if handle.is_null() || reservation_id.is_null() { + return QueryRouterResult::ErrInvalidParam; + } + + let reservation_id = match unsafe { CStr::from_ptr(reservation_id) }.to_str() { + Ok(value) if !value.is_empty() => value, + _ => return QueryRouterResult::ErrInvalidParam, + }; + let handles = unsafe { &*handle }; + match handles + .epp_reservations + .begin(&handles.prefill_router, reservation_id) + { + Ok(()) => QueryRouterResult::Ok, + Err(error) => { + tracing::warn!(%reservation_id, %error, "Failed to begin EPP prefill reservation"); + QueryRouterResult::ErrQueryFailed + } + } +} + +/// Atomically select and reserve the best prefill worker for an EPP-owned booking. +/// +/// Cancellation drops the scheduling future. Release/1.3.0 skips a cancelled queued +/// admission when it is subsequently dequeued. /// /// # Safety /// - `handle` must be a valid RouterHandles handle +/// - `reservation_id` must be a non-empty null-terminated UTF-8 string /// - `request_json` must be a valid null-terminated C string containing JSON /// - `pods_json` must be a valid null-terminated C string containing JSON, or null /// - `out_result` must be a valid pointer #[unsafe(no_mangle)] -pub unsafe extern "C" fn route_prefill_request( +pub unsafe extern "C" fn route_prefill_request_with_reservation( handle: RouterHandlesPtr, + reservation_id: *const c_char, request_json: *const c_char, pods_json: *const c_char, out_result: *mut CRoutingResult, ) -> QueryRouterResult { - if handle.is_null() || request_json.is_null() || out_result.is_null() { + if handle.is_null() + || reservation_id.is_null() + || request_json.is_null() + || out_result.is_null() + { return QueryRouterResult::ErrInvalidParam; } + let reservation_id = match unsafe { CStr::from_ptr(reservation_id) }.to_str() { + Ok(value) if !value.is_empty() => value.to_owned(), + _ => return QueryRouterResult::ErrInvalidParam, + }; let handles = unsafe { &*handle }; - + if let Err(error) = handles + .epp_reservations + .begin(&handles.prefill_router, &reservation_id) + { + tracing::warn!(%reservation_id, %error, "Failed to begin EPP prefill reservation"); + return QueryRouterResult::ErrQueryFailed; + } let (tokens, priority_jump, strict_priority, routing_constraints) = match unsafe { preprocess_request(handles, request_json) } { - Ok(t) => t, - Err(code) => return code, + Ok(values) => values, + Err(code) => { + handles.epp_reservations.abort(&reservation_id); + return code; + } }; - let allowed_worker_ids = unsafe { parse_pods_filter(pods_json) }; - let result = handles.runtime.secondary().block_on(async { - let (prefill_worker_id, prefill_dp_rank) = handles - .query_prefill_worker( - &tokens, - None, - None, - priority_jump, - strict_priority, - allowed_worker_ids, - routing_constraints, - ) - .await?; - - let prefill_dp_rank = prefill_dp_rank.unwrap_or(u32::MAX); - - tracing::info!( - prefill_worker_id = prefill_worker_id, - prefill_dp_rank = prefill_dp_rank, - token_count = tokens.len(), + let result = handles + .runtime + .secondary() + .block_on(handles.reserve_prefill_worker( + &reservation_id, + &tokens, + None, + None, priority_jump, strict_priority, - "Routed prefill request" - ); - - Ok((prefill_worker_id, prefill_dp_rank)) - }); + allowed_worker_ids, + routing_constraints, + )); match result { Ok((prefill_worker_id, prefill_dp_rank)) => { + let prefill_dp_rank = prefill_dp_rank.unwrap_or(u32::MAX); + tracing::info!( + %reservation_id, + prefill_worker_id, + prefill_dp_rank, + token_count = tokens.len(), + priority_jump, + strict_priority, + "Reserved prefill request" + ); + let out = unsafe { &mut *out_result }; *out = CRoutingResult::default(); out.is_disaggregated = true; @@ -1319,6 +1434,32 @@ pub unsafe extern "C" fn route_prefill_request( } } +/// Signal cancellation for a pending EPP prefill reservation. +/// +/// This function never waits for scheduler admission or booking cleanup. If admission has +/// already succeeded, normal request cleanup (or the Go late-result drain) owns release. +/// +/// # Safety +/// - `handle` must be a valid RouterHandles handle +/// - `reservation_id` must be a valid non-empty null-terminated UTF-8 string +#[unsafe(no_mangle)] +pub unsafe extern "C" fn cancel_prefill_reservation( + handle: RouterHandlesPtr, + reservation_id: *const c_char, +) -> QueryRouterResult { + if handle.is_null() || reservation_id.is_null() { + return QueryRouterResult::ErrInvalidParam; + } + + let reservation_id = match unsafe { CStr::from_ptr(reservation_id) }.to_str() { + Ok(value) if !value.is_empty() => value, + _ => return QueryRouterResult::ErrInvalidParam, + }; + let handles = unsafe { &*handle }; + handles.epp_reservations.cancel(reservation_id); + QueryRouterResult::Ok +} + /// Route a request to select the best **decode** worker only. /// /// This is used in both aggregated and disaggregated modes. diff --git a/lib/llm/src/kv_router/prefill_router/mod.rs b/lib/llm/src/kv_router/prefill_router/mod.rs index d018ff7511c0..860ef0311808 100644 --- a/lib/llm/src/kv_router/prefill_router/mod.rs +++ b/lib/llm/src/kv_router/prefill_router/mod.rs @@ -33,6 +33,9 @@ use crate::{ mod activation; mod admission; mod query; +mod reservations; + +pub use reservations::EppReservationManager; use admission::InnerPrefillRouter; diff --git a/lib/llm/src/kv_router/prefill_router/reservations.rs b/lib/llm/src/kv_router/prefill_router/reservations.rs new file mode 100644 index 000000000000..c1999c8bd9e8 --- /dev/null +++ b/lib/llm/src/kv_router/prefill_router/reservations.rs @@ -0,0 +1,688 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use std::{ + collections::{HashMap, hash_map::Entry}, + sync::{Arc, Weak}, + time::Duration, +}; + +use anyhow::Result; +use dynamo_kv_router::protocols::{BlockExtraInfo, RoutingConstraints, WorkerId}; +use futures::{StreamExt, stream}; +use parking_lot::Mutex; +use tokio::time::{Instant, MissedTickBehavior}; +use tokio_util::sync::CancellationToken; + +use super::{ + InnerPrefillRouter, PrefillError, PrefillLifecycleState, PrefillQueryOutcome, PrefillRouter, +}; +use crate::kv_router::{KvRouter, sequence::SequenceError}; + +const PREFILL_SCHEDULER_ID_PREFIX: &str = "epp-prefill/"; +const ACTIVE_REQUEST_EXPIRY_DURATION: Duration = Duration::from_secs(300); +const RESERVATION_EXPIRY_GRACE: Duration = Duration::from_secs(60); +const CANCELLED_RESERVATION_RETENTION: Duration = Duration::from_secs(60); +const RESERVATION_REAPER_INTERVAL: Duration = Duration::from_secs(30); +const RESERVATION_RELEASE_TIMEOUT: Duration = Duration::from_secs(5); +// Four timeout waves cap each batch at 20 seconds, below the 30-second interval. +const RESERVATION_REAPER_BATCH_SIZE: usize = 128; +const RESERVATION_REAPER_CONCURRENCY: usize = 32; + +struct ActivePrefillReservation { + chooser: Arc, + scheduler_id: String, + created_at: Instant, + reap_attempts: u64, +} + +impl Clone for ActivePrefillReservation { + fn clone(&self) -> Self { + Self { + chooser: self.chooser.clone(), + scheduler_id: self.scheduler_id.clone(), + created_at: self.created_at, + reap_attempts: self.reap_attempts, + } + } +} + +enum PrefillReservationEntry { + /// Cancelling this token drops the selection future. Release/1.3.0 keeps a + /// queued entry until its next dequeue, where the closed response receiver + /// prevents scheduler booking. + Pending { + cancellation: CancellationToken, + created_at: Instant, + claimed: bool, + }, + Active(ActivePrefillReservation), + /// Preserve a cancellation that reaches Rust before the blocking reserve + /// call creates the pending entry. + Cancelled { + created_at: Instant, + }, +} + +enum BeginReservation { + Pending(CancellationToken), + Cancelled, + AlreadyExists, +} + +enum Activation { + Active, + Cancelled, + Lost, +} + +pub struct EppReservationManager { + entries: Mutex>, +} + +impl Default for EppReservationManager { + fn default() -> Self { + Self { + entries: Mutex::new(HashMap::new()), + } + } +} + +impl EppReservationManager { + fn begin_entry(&self, reservation_id: &str) -> BeginReservation { + let mut entries = self.entries.lock(); + match entries.entry(reservation_id.to_string()) { + Entry::Vacant(entry) => { + let cancellation = CancellationToken::new(); + entry.insert(PrefillReservationEntry::Pending { + cancellation: cancellation.clone(), + created_at: Instant::now(), + claimed: false, + }); + BeginReservation::Pending(cancellation) + } + Entry::Occupied(entry) => match entry.get() { + PrefillReservationEntry::Cancelled { .. } => { + entry.remove(); + BeginReservation::Cancelled + } + PrefillReservationEntry::Pending { cancellation, .. } => { + BeginReservation::Pending(cancellation.clone()) + } + PrefillReservationEntry::Active(_) => BeginReservation::AlreadyExists, + }, + } + } + + // Claim the one scheduler-admission attempt for this reservation. `begin` remains + // idempotent while preprocessing runs, but only one `reserve` may turn its Pending + // state into an admission future. + fn claim_pending(&self, reservation_id: &str) -> BeginReservation { + let mut entries = self.entries.lock(); + match entries.entry(reservation_id.to_string()) { + Entry::Vacant(entry) => { + let cancellation = CancellationToken::new(); + entry.insert(PrefillReservationEntry::Pending { + cancellation: cancellation.clone(), + created_at: Instant::now(), + claimed: true, + }); + BeginReservation::Pending(cancellation) + } + Entry::Occupied(mut entry) => { + if matches!(entry.get(), PrefillReservationEntry::Cancelled { .. }) { + entry.remove(); + return BeginReservation::Cancelled; + } + + match entry.get_mut() { + PrefillReservationEntry::Pending { + cancellation, + claimed, + .. + } => { + if *claimed { + BeginReservation::AlreadyExists + } else { + *claimed = true; + BeginReservation::Pending(cancellation.clone()) + } + } + PrefillReservationEntry::Active(_) => BeginReservation::AlreadyExists, + PrefillReservationEntry::Cancelled { .. } => unreachable!("handled above"), + } + } + } + } + + fn activate(&self, reservation_id: &str, reservation: ActivePrefillReservation) -> Activation { + let mut entries = self.entries.lock(); + let Some(entry) = entries.get(reservation_id) else { + return Activation::Lost; + }; + + let cancelled = match entry { + PrefillReservationEntry::Pending { + cancellation, + claimed: true, + .. + } => cancellation.is_cancelled(), + // A reaped pending entry retains a cancellation tombstone. If its admission + // won just before cancellation, retain the active booking so release/reaper + // can retry cleanup rather than leaking it until scheduler expiry. + PrefillReservationEntry::Cancelled { .. } => true, + PrefillReservationEntry::Pending { .. } | PrefillReservationEntry::Active(_) => { + return Activation::Lost; + } + }; + + entries.insert( + reservation_id.to_string(), + PrefillReservationEntry::Active(reservation), + ); + if cancelled { + Activation::Cancelled + } else { + Activation::Active + } + } + + fn remove_pending(&self, reservation_id: &str) { + let mut entries = self.entries.lock(); + if matches!( + entries.get(reservation_id), + Some(PrefillReservationEntry::Pending { claimed: true, .. }) + ) { + entries.remove(reservation_id); + } + } + + // Preprocessing happens before reserve claims the entry. Never let a duplicate + // caller's preprocessing failure remove an admission another caller already owns. + fn abort_unclaimed(&self, reservation_id: &str) { + let mut entries = self.entries.lock(); + let state = match entries.get(reservation_id) { + Some(PrefillReservationEntry::Pending { + cancellation, + claimed: false, + .. + }) if cancellation.is_cancelled() => Some(true), + Some(PrefillReservationEntry::Pending { claimed: false, .. }) => Some(false), + _ => None, + }; + match state { + Some(true) => { + entries.insert( + reservation_id.to_string(), + PrefillReservationEntry::Cancelled { + created_at: Instant::now(), + }, + ); + } + Some(false) => { + entries.remove(reservation_id); + } + None => {} + } + } + + fn cancel_entry(&self, reservation_id: &str) { + let mut entries = self.entries.lock(); + match entries.entry(reservation_id.to_string()) { + Entry::Occupied(entry) => { + if let PrefillReservationEntry::Pending { cancellation, .. } = entry.get() { + cancellation.cancel(); + } + } + Entry::Vacant(entry) => { + entry.insert(PrefillReservationEntry::Cancelled { + created_at: Instant::now(), + }); + } + } + } + + fn expire_pending(&self, now: Instant, retention: Duration) { + let mut entries = self.entries.lock(); + for entry in entries.values_mut() { + let cancellation = match entry { + PrefillReservationEntry::Pending { + cancellation, + created_at, + .. + } if now.saturating_duration_since(*created_at) >= retention => { + Some(cancellation.clone()) + } + _ => None, + }; + if let Some(cancellation) = cancellation { + cancellation.cancel(); + *entry = PrefillReservationEntry::Cancelled { created_at: now }; + } + } + } + + fn get_active(&self, reservation_id: &str) -> Option { + match self.entries.lock().get(reservation_id) { + Some(PrefillReservationEntry::Active(reservation)) => Some(reservation.clone()), + Some(PrefillReservationEntry::Pending { .. }) + | Some(PrefillReservationEntry::Cancelled { .. }) + | None => None, + } + } + + fn remove_if_scheduler_id(&self, reservation_id: &str, scheduler_id: &str) { + let mut entries = self.entries.lock(); + if entries.get(reservation_id).is_some_and(|entry| { + matches!(entry, PrefillReservationEntry::Active(active) if active.scheduler_id == scheduler_id) + }) { + entries.remove(reservation_id); + } + } + + // Claim the least-attempted generation so retained failures cannot starve the backlog. + fn claim_expired_active_ids( + &self, + now: Instant, + retention: Duration, + limit: usize, + ) -> Vec { + if limit == 0 { + return Vec::new(); + } + + let mut entries = self.entries.lock(); + let minimum_attempts = entries + .values() + .filter_map(|entry| match entry { + PrefillReservationEntry::Active(active) + if now.saturating_duration_since(active.created_at) >= retention => + { + Some(active.reap_attempts) + } + _ => None, + }) + .min(); + let Some(minimum_attempts) = minimum_attempts else { + return Vec::new(); + }; + + entries + .iter_mut() + .filter_map(|(reservation_id, entry)| match entry { + PrefillReservationEntry::Active(active) + if now.saturating_duration_since(active.created_at) >= retention + && active.reap_attempts == minimum_attempts => + { + active.reap_attempts = active.reap_attempts.saturating_add(1); + Some(reservation_id.clone()) + } + _ => None, + }) + .take(limit) + .collect() + } + + fn remove_expired_cancellations(&self, now: Instant, retention: Duration) { + self.entries.lock().retain(|_, entry| { + !matches!(entry, PrefillReservationEntry::Cancelled { created_at } + if now.saturating_duration_since(*created_at) >= retention) + }); + } +} + +fn reservation_retention_from_expiry(active_request_expiry: Duration) -> Duration { + active_request_expiry.saturating_add(RESERVATION_EXPIRY_GRACE) +} + +fn reservation_retention() -> Duration { + reservation_retention_from_expiry(ACTIVE_REQUEST_EXPIRY_DURATION) +} + +fn scheduler_id(reservation_id: &str) -> String { + format!("{PREFILL_SCHEDULER_ID_PREFIX}{reservation_id}") +} + +fn ignore_missing_request(result: std::result::Result<(), SequenceError>) -> Result<()> { + match result { + Ok(()) | Err(SequenceError::RequestNotFound { .. }) => Ok(()), + Err(error) => Err(error.into()), + } +} + +impl EppReservationManager { + /// Register the pending state before potentially slow request preprocessing. + /// + /// Calling this repeatedly for the same in-flight booking is idempotent. A + /// cancellation that arrived first is observed here and never queues work. + pub fn begin(&self, router: &PrefillRouter, reservation_id: &str) -> Result<()> { + if reservation_id.is_empty() { + anyhow::bail!("prefill reservation ID must not be empty"); + } + if router.lifecycle_state() != PrefillLifecycleState::Active { + return Err(anyhow::anyhow!(PrefillError::NotActivated)); + } + if router.prefill_router.get().is_none() { + return Err(anyhow::anyhow!(PrefillError::NotActivated)); + } + + match self.begin_entry(reservation_id) { + BeginReservation::Pending(_) => Ok(()), + BeginReservation::Cancelled => { + anyhow::bail!("prefill reservation {reservation_id:?} was cancelled") + } + BeginReservation::AlreadyExists => { + anyhow::bail!("prefill reservation {reservation_id:?} already exists") + } + } + } + + /// Drop an unclaimed pending reservation when preprocessing fails before scheduler admission. + pub fn abort(&self, reservation_id: &str) { + self.abort_unclaimed(reservation_id); + } + + /// Atomically select and reserve a prefill worker for an externally dispatched request. + /// + /// The scheduler observes the booking before selecting the next request. The caller owns + /// the returned reservation until first output, terminal completion, or cancellation. + #[expect(clippy::too_many_arguments)] + pub async fn reserve( + &self, + router: &PrefillRouter, + reservation_id: &str, + token_ids: &[u32], + block_mm_infos: Option<&[Option]>, + lora_name: Option, + priority_jump: f64, + strict_priority: u32, + allowed_worker_ids: Option>, + routing_constraints: RoutingConstraints, + ) -> Result { + if reservation_id.is_empty() { + anyhow::bail!("prefill reservation ID must not be empty"); + } + let cancellation = match self.claim_pending(reservation_id) { + BeginReservation::Pending(cancellation) => cancellation, + BeginReservation::Cancelled => { + anyhow::bail!("prefill reservation {reservation_id:?} was cancelled") + } + BeginReservation::AlreadyExists => { + anyhow::bail!("prefill reservation {reservation_id:?} already exists") + } + }; + if router.lifecycle_state() != PrefillLifecycleState::Active { + self.remove_pending(reservation_id); + return Err(anyhow::anyhow!(PrefillError::NotActivated)); + } + let Some(inner) = router.prefill_router.get() else { + self.remove_pending(reservation_id); + return Err(anyhow::anyhow!(PrefillError::NotActivated)); + }; + let InnerPrefillRouter::KvRouter(router) = inner else { + self.remove_pending(reservation_id); + return Err(anyhow::anyhow!(PrefillError::NotActivated)); + }; + let chooser = router.chooser.clone(); + let scheduler_id = scheduler_id(reservation_id); + + let outcome = tokio::select! { + biased; + _ = cancellation.cancelled() => { + self.remove_pending(reservation_id); + anyhow::bail!("prefill reservation {reservation_id:?} was cancelled") + } + outcome = chooser.find_best_match_details( + Some(&scheduler_id), + token_ids, + block_mm_infos, + None, + true, + false, + lora_name, + priority_jump, + strict_priority, + None, + None, + allowed_worker_ids, + routing_constraints, + ) => outcome, + }; + let outcome = match outcome { + Ok(outcome) => outcome, + Err(error) => { + self.remove_pending(reservation_id); + return Err(error); + } + }; + + match outcome { + crate::kv_router::FindBestMatchOutcome::Routed { worker, .. } => { + let reservation = ActivePrefillReservation { + chooser: chooser.clone(), + scheduler_id: scheduler_id.clone(), + created_at: Instant::now(), + reap_attempts: 0, + }; + match self.activate(reservation_id, reservation) { + Activation::Active => Ok(PrefillQueryOutcome::Routed { + worker_id: worker.worker_id, + dp_rank: Some(worker.dp_rank), + }), + Activation::Cancelled => { + if let Err(error) = self.release(reservation_id).await { + tracing::warn!( + %reservation_id, + %error, + "Failed to release cancelled EPP prefill reservation; retaining it for reaper retry" + ); + } + anyhow::bail!("prefill reservation {reservation_id:?} was cancelled"); + } + Activation::Lost => { + if let Err(error) = + ignore_missing_request(chooser.free(&scheduler_id).await) + { + tracing::warn!( + %reservation_id, + %error, + "Failed to release EPP prefill reservation with lost ownership" + ); + } + anyhow::bail!("prefill reservation {reservation_id:?} lost ownership"); + } + } + } + crate::kv_router::FindBestMatchOutcome::QueueRejected { rejection } => { + self.remove_pending(reservation_id); + Ok(PrefillQueryOutcome::QueueRejected { rejection }) + } + } + } + + /// Cancel a pending prefill reservation without waiting for scheduler cleanup. + /// + /// An active reservation remains owned by the normal booking lifecycle. If cancellation + /// races with admission, the Go caller drains the reserve result and releases that booking. + pub fn cancel(&self, reservation_id: &str) { + if !reservation_id.is_empty() { + self.cancel_entry(reservation_id); + } + } + + /// Release a prefill reservation. Missing reservations are idempotent no-ops. + pub async fn release(&self, reservation_id: &str) -> Result<()> { + let Some(reservation) = self.get_active(reservation_id) else { + return Ok(()); + }; + + ignore_missing_request(reservation.chooser.free(&reservation.scheduler_id).await)?; + self.remove_if_scheduler_id(reservation_id, &reservation.scheduler_id); + Ok(()) + } + + async fn release_expired_batch(self: &Arc, reservation_ids: Vec) { + stream::iter(reservation_ids) + .for_each_concurrent(RESERVATION_REAPER_CONCURRENCY, |reservation_id| { + let manager = Arc::clone(self); + async move { + match tokio::time::timeout( + RESERVATION_RELEASE_TIMEOUT, + manager.release(&reservation_id), + ) + .await + { + Ok(Ok(())) => {} + Ok(Err(error)) => { + tracing::warn!( + %reservation_id, + %error, + "Failed to expire stale EPP prefill reservation" + ); + } + Err(_) => { + tracing::warn!( + %reservation_id, + timeout_secs = RESERVATION_RELEASE_TIMEOUT.as_secs(), + "Timed out expiring stale EPP prefill reservation" + ); + } + } + } + }) + .await; + } + + /// Reap stale EPP-owned bookings. The task only weakly holds the manager, + /// so destroying the C router handle stops it without affecting the + /// prefill router lifecycle. + pub fn spawn_reaper(manager: &Arc) { + let manager: Weak = Arc::downgrade(manager); + + tokio::spawn(async move { + let retention = reservation_retention(); + let mut interval = tokio::time::interval_at( + Instant::now() + RESERVATION_REAPER_INTERVAL, + RESERVATION_REAPER_INTERVAL, + ); + interval.set_missed_tick_behavior(MissedTickBehavior::Skip); + + loop { + interval.tick().await; + let Some(manager) = manager.upgrade() else { + return; + }; + let now = Instant::now(); + manager.expire_pending(now, retention); + manager.remove_expired_cancellations(now, CANCELLED_RESERVATION_RETENTION); + let expired = + manager.claim_expired_active_ids(now, retention, RESERVATION_REAPER_BATCH_SIZE); + manager.release_expired_batch(expired).await; + } + }); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn cancellation_before_reservation_is_observed() { + let registry = EppReservationManager::default(); + registry.cancel_entry("reservation-1"); + + assert!(matches!( + registry.begin_entry("reservation-1"), + BeginReservation::Cancelled + )); + assert!(matches!( + registry.begin_entry("reservation-1"), + BeginReservation::Pending(_) + )); + } + + #[test] + fn pre_registered_pending_reservation_survives_tombstone_reaping() { + let registry = EppReservationManager::default(); + let BeginReservation::Pending(cancellation) = registry.begin_entry("reservation-1") else { + panic!("expected pending reservation"); + }; + + registry.cancel_entry("reservation-1"); + registry.remove_expired_cancellations( + Instant::now() + CANCELLED_RESERVATION_RETENTION + Duration::from_secs(1), + CANCELLED_RESERVATION_RETENTION, + ); + + assert!(cancellation.is_cancelled()); + let BeginReservation::Pending(existing) = registry.begin_entry("reservation-1") else { + panic!("expected pending reservation to remain registered"); + }; + assert!(existing.is_cancelled()); + } + + #[test] + fn only_one_reserve_claims_pending_reservation() { + let registry = EppReservationManager::default(); + assert!(matches!( + registry.begin_entry("reservation-1"), + BeginReservation::Pending(_) + )); + assert!(matches!( + registry.claim_pending("reservation-1"), + BeginReservation::Pending(_) + )); + assert!(matches!( + registry.claim_pending("reservation-1"), + BeginReservation::AlreadyExists + )); + } + + #[test] + fn expired_pending_reservation_becomes_cancelled_tombstone() { + let registry = EppReservationManager::default(); + let BeginReservation::Pending(cancellation) = registry.begin_entry("reservation-1") else { + panic!("expected pending reservation"); + }; + + registry.expire_pending( + Instant::now() + Duration::from_secs(61), + Duration::from_secs(60), + ); + + assert!(cancellation.is_cancelled()); + assert!(matches!( + registry.claim_pending("reservation-1"), + BeginReservation::Cancelled + )); + assert!(matches!( + registry.begin_entry("reservation-1"), + BeginReservation::Pending(_) + )); + } + + #[test] + fn cancellation_signals_pending_reservation() { + let registry = EppReservationManager::default(); + let BeginReservation::Pending(cancellation) = registry.begin_entry("reservation-1") else { + panic!("expected pending reservation"); + }; + + registry.cancel_entry("reservation-1"); + assert!(cancellation.is_cancelled()); + registry.remove_pending("reservation-1"); + } + + #[test] + fn reservation_retention_tracks_active_request_expiry() { + assert_eq!( + reservation_retention_from_expiry(Duration::from_secs(30)), + Duration::from_secs(90) + ); + assert_eq!( + reservation_retention_from_expiry(Duration::from_secs(3600)), + Duration::from_secs(3660) + ); + } +}