From fe94c911ad272a895f4bd4979bd04b23c8f05acd Mon Sep 17 00:00:00 2001 From: Avinash Varma <51354954+avinash-rafay@users.noreply.github.com> Date: Thu, 30 Jul 2026 07:33:30 -0700 Subject: [PATCH 01/15] fix(inference-gateway): forward full request body to EPP KV-router scorer (#11991) Signed-off-by: Avinash Varma Co-authored-by: Avinash Varma (cherry picked from commit ffae596b7b45c8741d3a347be77d5d31ac2df9a0) Signed-off-by: Avinash Varma --- .../epp/pkg/plugins/disagg/shared.go | 12 +- .../pkg/plugins/dynamo_kv_scorer/plugin.go | 93 ++-------- .../plugins/dynamo_kv_scorer/plugin_test.go | 172 ++++++++++-------- 3 files changed, 116 insertions(+), 161 deletions(-) diff --git a/deploy/inference-gateway/epp/pkg/plugins/disagg/shared.go b/deploy/inference-gateway/epp/pkg/plugins/disagg/shared.go index 5a29e8a62993..351ec206b530 100644 --- a/deploy/inference-gateway/epp/pkg/plugins/disagg/shared.go +++ b/deploy/inference-gateway/epp/pkg/plugins/disagg/shared.go @@ -27,8 +27,6 @@ limitations under the License. package disagg import ( - "encoding/json" - "fmt" "os" "strings" "sync" @@ -71,15 +69,7 @@ func readPrefillEnabled(cycleState *schedtypes.CycleState) bool { // 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) - } - data, err := json.Marshal(requestBody) - if err != nil { - return "", fmt.Errorf("failed to marshal request JSON: %w", err) - } - return string(data), nil + return dynscorer.BuildOpenAIRequestJSON(req) } // serializeEndpoints converts endpoints to a JSON string for the FFI filter. 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 f3320f2bd14d..628ffc38bafc 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 @@ -101,6 +101,7 @@ import "C" import ( "encoding/json" "fmt" + "maps" "os" "strings" "sync" @@ -282,92 +283,32 @@ 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's FFI, overriding only the model, so tool-calling and +// reasoning fields survive the router's parse and chat-template render. Errors +// when no payload is available so the scorer falls back to non-KV routing. +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() { - addCompletionPrompt(requestBody, req.Body.Completions.Prompt) - } 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) + // Route on the resolved target model. 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 - } - if cacheSalt := extractTopLevelCacheSalt(req.Body.Payload); cacheSalt != "" { - requestBody["cache_salt"] = cacheSalt - } - - return requestBody, nil -} - -func addCompletionPrompt(requestBody map[string]any, prompt fwkrh.Prompt) { - if len(prompt.TokenIDs) > 0 { - tokenIDs := make([]uint32, len(prompt.TokenIDs)) - copy(tokenIDs, prompt.TokenIDs) - requestBody["prompt"] = tokenIDs - return - } - - // Keep non-token completions on the legacy chat-shaped scorer path. - requestBody["messages"] = []map[string]any{ - { - "role": "user", - "content": prompt.PlainText(), - }, - } -} - -// 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 -} -func extractTopLevelCacheSalt(payload fwkrh.RequestPayload) string { - pm, ok := payload.(fwkrh.PayloadMap) - if !ok { - return "" + data, err := json.Marshal(requestBody) + if err != nil { + return "", fmt.Errorf("failed to marshal request JSON: %w", err) } - cacheSalt, _ := pm["cache_salt"].(string) - return cacheSalt + return string(data), nil } // CallAddRequest registers a request with the router's bookkeeping. 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 887d0a8bb982..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,27 +17,35 @@ limitations under the License. package dynamo_kv_scorer import ( - "reflect" + "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", @@ -46,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"]) @@ -59,20 +63,17 @@ 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) } } -func TestBuildOpenAIRequest_ForwardsLegacyTopLevelCacheSalt(t *testing.T) { +// 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{ - 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", @@ -81,91 +82,114 @@ func TestBuildOpenAIRequest_ForwardsLegacyTopLevelCacheSalt(t *testing.T) { }, } - body, err := BuildOpenAIRequest(req) - if err != nil { - t.Fatalf("BuildOpenAIRequest returned error: %v", err) - } + 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) } } -func TestBuildOpenAIRequest_CompletionsTokenPromptUsesPromptIDs(t *testing.T) { +// 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{ - Completions: &fwkrh.CompletionsRequest{ - Prompt: fwkrh.Prompt{TokenIDs: []uint32{101, 102, 103}}, + 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, err := BuildOpenAIRequest(req) - if err != nil { - t.Fatalf("BuildOpenAIRequest returned error: %v", err) + 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) } - if _, ok := body["messages"]; ok { - t.Fatalf("did not expect token-id completions to synthesize messages: %v", body["messages"]) + msgs, ok := body["messages"].([]any) + if !ok || len(msgs) != 3 { + t.Fatalf("expected 3 messages, got %#v", body["messages"]) } - if got := body["prompt"]; !reflect.DeepEqual(got, []uint32{101, 102, 103}) { - t.Fatalf("expected prompt token IDs, got %#v", got) + + // 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 got := body["model"]; got != "test-model" { - t.Fatalf("expected model=test-model, got %v", got) + 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) } } -func TestBuildOpenAIRequest_CompletionsTextPromptKeepsLegacyMessageShape(t *testing.T) { +// 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{ - Completions: &fwkrh.CompletionsRequest{ - Prompt: fwkrh.Prompt{Raw: "hello"}, + Payload: fwkrh.PayloadMap{ + "model": "alias-model", + "prompt": "hello world", }, }, } - body, err := BuildOpenAIRequest(req) - if err != nil { - t.Fatalf("BuildOpenAIRequest returned error: %v", err) - } - - if _, ok := body["prompt"]; ok { - t.Fatalf("did not expect text completions to change to prompt field: %v", body["prompt"]) + body := ffiBody(t, req) + if got := body["prompt"]; got != "hello world" { + t.Fatalf("expected prompt forwarded, got %v", got) } - messages, ok := body["messages"].([]map[string]any) - if !ok { - t.Fatalf("expected legacy messages shape, got %#v", body["messages"]) - } - if len(messages) != 1 || messages[0]["role"] != "user" || messages[0]["content"] != "hello" { - t.Fatalf("expected single user message with content=hello, got %#v", messages) + if got := body["model"]; got != "test-model" { + t.Fatalf("expected model overridden to test-model, got %v", got) } } -func TestBuildOpenAIRequest_CompletionsStringArrayPromptKeepsLegacyMessageShape(t *testing.T) { +// 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{ - Completions: &fwkrh.CompletionsRequest{ - Prompt: fwkrh.Prompt{Strings: []string{"hello", "world"}}, + ChatCompletions: &fwkrh.ChatCompletionsRequest{ + Messages: []fwkrh.Message{ + {Role: "user", Content: fwkrh.Content{Raw: "hi"}}, + }, }, + // No Payload set — the typed view cannot carry tool-calling fields. }, } - body, err := BuildOpenAIRequest(req) - if err != nil { - t.Fatalf("BuildOpenAIRequest returned error: %v", err) - } - - if _, ok := body["prompt"]; ok { - t.Fatalf("did not expect string-array completions to change to prompt field: %v", body["prompt"]) - } - messages, ok := body["messages"].([]map[string]any) - if !ok { - t.Fatalf("expected legacy messages shape, got %#v", body["messages"]) - } - if len(messages) != 1 || messages[0]["role"] != "user" || messages[0]["content"] != "hello world" { - t.Fatalf("expected single user message with content='hello world', got %#v", messages) + if _, err := BuildOpenAIRequestJSON(req); err == nil { + t.Fatalf("expected an error when the raw payload is unavailable, got nil") } } From 4d5412933202b842c7d9218a24f8e6cbd50ca397 Mon Sep 17 00:00:00 2001 From: Thomas Montfort Date: Tue, 11 Aug 2026 12:24:04 -0700 Subject: [PATCH 02/15] fix(epp): reserve prefill load during routing Signed-off-by: Thomas Montfort (cherry picked from commit be9f00ceed2abc0a07d89221d489054dc8818c38) Signed-off-by: Avinash Varma --- deploy/inference-gateway/epp/go.mod | 2 +- .../epp/pkg/plugins/disagg/decode_scorer.go | 190 ++++++++++----- .../epp/pkg/plugins/disagg/prefill_scorer.go | 71 +++++- .../pkg/plugins/disagg/reservation_test.go | 152 ++++++++++++ .../epp/pkg/plugins/disagg/shared.go | 58 +++++ .../pkg/plugins/dynamo_kv_scorer/plugin.go | 69 ++++++ lib/bindings/c/src/lib.rs | 147 +++++++++++ .../kv_router/prefill_router/activation.rs | 3 + lib/llm/src/kv_router/prefill_router/mod.rs | 2 + .../kv_router/prefill_router/reservations.rs | 230 ++++++++++++++++++ 10 files changed, 856 insertions(+), 68 deletions(-) create mode 100644 deploy/inference-gateway/epp/pkg/plugins/disagg/reservation_test.go create mode 100644 lib/llm/src/kv_router/prefill_router/reservations.rs diff --git a/deploy/inference-gateway/epp/go.mod b/deploy/inference-gateway/epp/go.mod index 692d536ec5d2..dbfbee391900 100644 --- a/deploy/inference-gateway/epp/go.mod +++ b/deploy/inference-gateway/epp/go.mod @@ -45,7 +45,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/decode_scorer.go b/deploy/inference-gateway/epp/pkg/plugins/disagg/decode_scorer.go index 040c7fcb7ebe..3abcc4f8703b 100644 --- a/deploy/inference-gateway/epp/pkg/plugins/disagg/decode_scorer.go +++ b/deploy/inference-gateway/epp/pkg/plugins/disagg/decode_scorer.go @@ -54,6 +54,7 @@ var _ rc.ResponseBodyProcessor = &DynDecodeScorer{} // DecodeRoutingState holds routing information passed from Score() to PreRequest(). type DecodeRoutingState struct { + BookingID string WorkerID string DpRank uint32 PrefillWorkerID string @@ -67,6 +68,7 @@ func (s *DecodeRoutingState) Clone() plugins.StateData { return nil } clone := &DecodeRoutingState{ + BookingID: s.BookingID, WorkerID: s.WorkerID, DpRank: s.DpRank, PrefillWorkerID: s.PrefillWorkerID, @@ -102,16 +104,22 @@ func DynDecodeScorerFactory(name string, rawParameters json.RawMessage, handle p // NewDynDecodeScorer initializes a new DynDecodeScorer. func NewDynDecodeScorer(ctx context.Context) *DynDecodeScorer { return &DynDecodeScorer{ - typedName: plugins.TypedName{Type: DynDecodeScorerType, Name: DynDecodeScorerType}, - pluginState: plugins.NewPluginState(ctx), + typedName: plugins.TypedName{Type: DynDecodeScorerType, Name: DynDecodeScorerType}, + pluginState: plugins.NewPluginState(ctx), + 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 + pluginState *plugins.PluginState + prefillMarkInFlight sync.Map + addRequest func(string, []int64, uint64, uint32, string) error + markPrefillComplete func(string) error + freeBooking func(string) error } // TypedName returns the type and name tuple of this plugin instance. @@ -133,12 +141,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) } @@ -150,46 +174,43 @@ func (s *DynDecodeScorer) Score(ctx context.Context, cycleState *schedtypes.Cycl result, err := dynscorer.CallRouteDecodeRequest(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) 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" + delete(req.Headers, PrefillWorkerIDHeader) + delete(req.Headers, PrefillDpRankHeader) } - // Store routing state for PreRequest bookkeeping - if req.RequestId != "" { - routingState := &DecodeRoutingState{ - WorkerID: workerIDStr, - DpRank: result.DpRank, - TokenData: result.TokenData, - CacheNamespace: result.CacheNamespace, - } - s.pluginState.Write(req.RequestId, plugins.StateKey(decodeStateKey), routingState) + // Store routing state for PreRequest bookkeeping, keyed by booking ID. + routingState := &DecodeRoutingState{ + BookingID: booking.ID, + WorkerID: workerIDStr, + DpRank: result.DpRank, + TokenData: result.TokenData, + CacheNamespace: result.CacheNamespace, } + s.pluginState.Write(booking.ID, plugins.StateKey(decodeStateKey), routingState) // Inject pre-computed tokens into the request body so the frontend // sidecar can skip redundant tokenization. @@ -198,84 +219,135 @@ func (s *DynDecodeScorer) Score(ctx context.Context, cycleState *schedtypes.Cycl return uniformScores(endpoints, 1.0) } +func (s *DynDecodeScorer) cleanupBooking(ctx context.Context, bookingID, reason string) bool { + if bookingID == "" { + return true + } + 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 +} + +func (s *DynDecodeScorer) rollbackPrefillReservation( + ctx context.Context, + cycleState *schedtypes.CycleState, + request *schedtypes.InferenceRequest, + booking *BookingState, + reason string, +) { + if booking != 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, PrefillWorkerIDHeader) + delete(request.Headers, PrefillDpRankHeader) + } +} + // 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") + bookingID := bookingIDFromRequest(request) + if bookingID == "" { + logger.V(logutil.DEBUG).Info("DynDecodeScorer PreRequest: no controller booking ID, skipping") + return + } + if err := ctx.Err(); err != nil { + s.cleanupBooking(ctx, bookingID, "request cancelled before decode booking") return } state, err := plugins.ReadPluginStateKey[*DecodeRoutingState]( - s.pluginState, request.RequestId, plugins.StateKey(decodeStateKey), + s.pluginState, bookingID, plugins.StateKey(decodeStateKey), ) - s.pluginState.Delete(request.RequestId) - - if err != nil { + s.pluginState.Delete(bookingID) + if err != nil || state == nil || state.BookingID != bookingID { logger.V(logutil.DEBUG).Info("DynDecodeScorer PreRequest: no routing state found", - "requestID", request.RequestId) + "bookingID", bookingID) + s.cleanupBooking(ctx, bookingID, "decode routing state missing") return } 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) + "bookingID", bookingID, "workerID", state.WorkerID) + s.cleanupBooking(ctx, bookingID, "decode worker ID invalid") return } - if addErr := dynscorer.CallAddRequest( - request.RequestId, + if addErr := s.addRequest( + bookingID, state.TokenData, workerIDUint, state.DpRank, state.CacheNamespace, ); addErr != nil { logger.V(logutil.DEFAULT).Error(addErr, "DynDecodeScorer PreRequest: failed to add request", - "requestID", request.RequestId) + "bookingID", bookingID) + s.cleanupBooking(ctx, bookingID, "decode booking failed") return } logger.V(logutil.VERBOSE).Info("DynDecodeScorer PreRequest: registered request", - "requestID", request.RequestId, + "bookingID", bookingID, "workerID", state.WorkerID, "dpRank", state.DpRank, "hasCacheNamespace", state.CacheNamespace != "", "tokenCount", len(state.TokenData)) + + go func() { + <-ctx.Done() + s.cleanupBooking(ctx, bookingID, "request context cancelled") + }() } // 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) - } - } - - // 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 { + // Terminal cleanup takes precedence over first-token bookkeeping. This also + // covers empty/error responses where no output token was observed. + if response.EndOfStream { + s.prefillMarkInFlight.Delete(bookingID) + if err := s.freeBooking(bookingID); err != nil { logger.V(logutil.DEFAULT).Error(err, "DynDecodeScorer ResponseBody: failed to free request", - "requestID", request.RequestId) + "bookingID", bookingID, "requestID", request.RequestId) } else { logger.V(logutil.VERBOSE).Info("DynDecodeScorer ResponseBody: freed request", - "requestID", request.RequestId) + "bookingID", bookingID, "requestID", request.RequestId) } + return + } + + // Keep the marker while the FFI call is in flight. On failure, remove it so + // the next response chunk retries instead of leaking the prefill load. + if _, alreadyInFlight := s.prefillMarkInFlight.LoadOrStore(bookingID, struct{}{}); alreadyInFlight { + return + } + if err := s.markPrefillComplete(bookingID); err != nil { + s.prefillMarkInFlight.Delete(bookingID) + logger.V(logutil.DEFAULT).Error(err, "DynDecodeScorer ResponseBody: failed to mark prefill complete", + "bookingID", bookingID, "requestID", request.RequestId) + } else { + logger.V(logutil.VERBOSE).Info("DynDecodeScorer ResponseBody: marked prefill complete", + "bookingID", bookingID, "requestID", 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..1c08c4b1ded3 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,6 +34,8 @@ import ( const ( // DynPrefillScorerType is the plugin type registered in the plugin registry. DynPrefillScorerType = "dyn-prefill-scorer" + + defaultPrefillReservationTimeout = 5 * time.Second ) // compile-time type assertion @@ -60,13 +63,17 @@ func DynPrefillScorerFactory(name string, rawParameters json.RawMessage, _ plugi // NewDynPrefillScorer initializes a new DynPrefillScorer. func NewDynPrefillScorer() *DynPrefillScorer { return &DynPrefillScorer{ - typedName: plugins.TypedName{Type: DynPrefillScorerType, Name: DynPrefillScorerType}, + typedName: plugins.TypedName{Type: DynPrefillScorerType, Name: DynPrefillScorerType}, + reservePrefill: dynscorer.CallRoutePrefillRequestWithReservation, + freeBooking: dynscorer.CallFreeRequest, } } // DynPrefillScorer is a scorer plugin for the prefill scheduling profile. type DynPrefillScorer struct { - typedName plugins.TypedName + typedName plugins.TypedName + reservePrefill func(string, string, string, time.Duration) (*dynscorer.RoutingResult, error) + freeBooking func(string) error } // TypedName returns the type and name tuple of this plugin instance. @@ -88,15 +95,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 +126,36 @@ func (s *DynPrefillScorer) Score(ctx context.Context, cycleState *schedtypes.Cyc "endpointCount", len(endpoints), "endpointsJSON", string(endpointsJSON)) - result, err := dynscorer.CallRoutePrefillRequest(requestJSON, endpointsJSON) + timeout := defaultPrefillReservationTimeout + if deadline, ok := ctx.Deadline(); ok { + remaining := time.Until(deadline) + if remaining < timeout { + timeout = remaining + } + } + if timeout <= 0 { + cycleState.Write(PrefillEnabledStateKey, &PrefillEnabledState{Enabled: false}) + return uniformScores(endpoints, 0) + } + + result, err := s.reservePrefill(booking.ID, requestJSON, endpointsJSON, timeout) 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 +163,25 @@ func (s *DynPrefillScorer) Score(ctx context.Context, cycleState *schedtypes.Cyc delete(req.Headers, PrefillDpRankHeader) } + if err := ctx.Err(); err != nil { + if cleanupErr := s.freeBooking(booking.ID); cleanupErr != nil { + logger.V(logutil.DEFAULT).Error(cleanupErr, "DynPrefillScorer: failed to clean cancelled reservation", + "bookingID", booking.ID) + } + booking.PrefillReserved = false + cycleState.Write(BookingStateKey, booking) + cycleState.Write(PrefillEnabledStateKey, &PrefillEnabledState{Enabled: false}) + return uniformScores(endpoints, 0) + } + + bookingID := booking.ID + go func() { + <-ctx.Done() + if cleanupErr := s.freeBooking(bookingID); cleanupErr != nil { + logger.V(logutil.DEFAULT).Error(cleanupErr, "DynPrefillScorer: cancellation cleanup failed", + "bookingID", bookingID) + } + }() + 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..a3f92ed5b3fe --- /dev/null +++ b/deploy/inference-gateway/epp/pkg/plugins/disagg/reservation_test.go @@ -0,0 +1,152 @@ +/* +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" + + rc "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/framework/interface/requestcontrol" + schedtypes "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/framework/interface/scheduling" +) + +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) + 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 + }, + } + + scorer.ResponseBody(context.Background(), request, &rc.Response{}, nil) + scorer.ResponseBody(context.Background(), request, &rc.Response{}, nil) + scorer.ResponseBody(context.Background(), request, &rc.Response{}, nil) + scorer.ResponseBody(context.Background(), request, &rc.Response{EndOfStream: true}, nil) + + 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 + }, + } + + scorer.ResponseBody( + context.Background(), + requestWithBooking("external-request", bookingID), + &rc.Response{EndOfStream: true}, + nil, + ) + + if markCalls != 0 { + t.Fatalf("mark calls = %d, want 0", markCalls) + } + if freeCalls != 1 { + t.Fatalf("free calls = %d, want 1", freeCalls) + } +} + +func TestPreRequestCancellationCleansBooking(t *testing.T) { + bookingID := ensureBookingState(schedtypes.NewCycleState()).ID + freeCalls := 0 + scorer := &DynDecodeScorer{ + freeBooking: func(got string) error { + if got != bookingID { + t.Fatalf("free booking ID = %q, want %q", got, bookingID) + } + freeCalls++ + return nil + }, + } + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + scorer.PreRequest(ctx, requestWithBooking("external-request", bookingID), nil) + + if freeCalls != 1 { + t.Fatalf("free calls = %d, want 1", freeCalls) + } +} diff --git a/deploy/inference-gateway/epp/pkg/plugins/disagg/shared.go b/deploy/inference-gateway/epp/pkg/plugins/disagg/shared.go index 351ec206b530..807424e92c60 100644 --- a/deploy/inference-gateway/epp/pkg/plugins/disagg/shared.go +++ b/deploy/inference-gateway/epp/pkg/plugins/disagg/shared.go @@ -32,6 +32,7 @@ import ( "sync" "github.com/go-logr/logr" + "github.com/google/uuid" 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" @@ -46,6 +47,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. @@ -67,6 +70,61 @@ func readPrefillEnabled(cycleState *schedtypes.CycleState) bool { return false } +// 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{} + } + 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 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) 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 628ffc38bafc..6ef23acc7d78 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 @@ -71,6 +72,13 @@ query_router_result_t route_prefill_request(RouterHandles *handle, 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, + uint64_t timeout_ms, + CRoutingResult *out_result); + query_router_result_t route_decode_request(RouterHandles *handle, const char *request_json, const char *pods_json, @@ -482,8 +490,69 @@ func CallRoutePrefillRequest(requestJSON string, podsJSON string) (*RoutingResul }, nil } +// CallRoutePrefillRequestWithReservation atomically selects and books a prefill worker. +// The Rust scheduler retracts pending admission when timeout expires. +func CallRoutePrefillRequestWithReservation(reservationID string, requestJSON string, podsJSON string, timeout time.Duration) (*RoutingResult, error) { + if reservationID == "" { + return nil, fmt.Errorf("prefill reservation ID is required") + } + if timeout <= 0 { + return nil, fmt.Errorf("prefill reservation timeout must be positive") + } + if !routerInitialized { + return nil, fmt.Errorf("dynamo router not initialized") + } + + routerHandlesMutex.RLock() + router := routerHandles + routerHandlesMutex.RUnlock() + if router == nil { + 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)) + + var cPodsJSON *C.char + if podsJSON != "" { + cPodsJSON = C.CString(podsJSON) + defer C.free(unsafe.Pointer(cPodsJSON)) + } + + timeoutMS := timeout.Milliseconds() + if timeoutMS < 1 { + timeoutMS = 1 + } + var result C.CRoutingResult + rc := C.route_prefill_request_with_reservation( + router, + cReservationID, + cRequestJSON, + cPodsJSON, + C.uint64_t(timeoutMS), + &result, + ) + if rc != C.QUERY_ROUTER_OK { + return nil, fmt.Errorf("route_prefill_request_with_reservation failed with code %d", rc) + } + + tokens := extractTokenData(&result) + workerID := uint64(result.prefill_worker_id) + dpRank := uint32(result.prefill_dp_rank) + C.free_routing_result(&result) + + return &RoutingResult{ + WorkerID: workerID, + DpRank: dpRank, + TokenData: tokens, + }, 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/lib/bindings/c/src/lib.rs b/lib/bindings/c/src/lib.rs index bee6a5befca7..a09f7c5cdb87 100644 --- a/lib/bindings/c/src/lib.rs +++ b/lib/bindings/c/src/lib.rs @@ -528,6 +528,56 @@ impl RouterHandles { } } + /// Atomically select and reserve a prefill worker for an EPP-owned booking. + #[expect(clippy::too_many_arguments)] + async fn reserve_prefill_worker( + &self, + reservation_id: &str, + tokens: &[u32], + block_mm_infos: Option<&[Option]>, + lora_name: Option, + priority_jump: f64, + strict_priority: u32, + allowed_worker_ids: Option>, + routing_constraints: RoutingConstraints, + ) -> Result<(u64, Option), QueryRouterResult> { + if let Some(ref ids) = allowed_worker_ids { + self.prefill_router.register_workers(ids); + } + + let outcome = self + .prefill_router + .reserve_prefill_worker( + reservation_id, + tokens, + block_mm_infos, + lora_name, + priority_jump, + strict_priority, + allowed_worker_ids, + routing_constraints, + ) + .await + .map_err(|error| { + tracing::error!(%reservation_id, %error, "Prefill reservation failed"); + QueryRouterResult::ErrQueryFailed + })?; + match outcome { + 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 reservation rejected by policy-class queue limit" + ); + Err(QueryRouterResult::ErrBackpressure) + } + } + } + /// Query optimal decode worker for a request. /// For disaggregated mode, set `is_disaggregated` to true to use overlap_score_credit=0 /// (since KV cache is being transferred from prefill, not reused). @@ -984,6 +1034,7 @@ pub unsafe extern "C" fn add_request_with_cache_namespace( Vec::new() }; + let prefill_router = handles.prefill_router.clone(); let decode_router = handles.decode_router.clone(); let result = handles.runtime.secondary().block_on(async { @@ -1073,12 +1124,16 @@ pub unsafe extern "C" fn mark_prefill_complete( Err(_) => return QueryRouterResult::ErrInvalidParam, }; + let prefill_router = handles.prefill_router.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) = prefill_router.release_prefill_reservation(&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, @@ -1130,12 +1185,16 @@ pub unsafe extern "C" fn free_request( Err(_) => return QueryRouterResult::ErrInvalidParam, }; + let prefill_router = handles.prefill_router.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) = prefill_router.release_prefill_reservation(&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, @@ -1484,6 +1543,94 @@ pub unsafe extern "C" fn route_prefill_request( } } +/// Atomically select and reserve the best prefill worker for an EPP-owned booking. +/// +/// Pending scheduler admission is bounded by `timeout_ms`. +/// +/// # 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_with_reservation( + handle: RouterHandlesPtr, + reservation_id: *const c_char, + request_json: *const c_char, + pods_json: *const c_char, + timeout_ms: u64, + out_result: *mut CRoutingResult, +) -> QueryRouterResult { + if handle.is_null() + || reservation_id.is_null() + || request_json.is_null() + || out_result.is_null() + || timeout_ms == 0 + { + 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 }; + let (tokens, priority_jump, strict_priority, routing_constraints) = + match unsafe { preprocess_request(handles, request_json) } { + Ok(values) => values, + Err(code) => return code, + }; + let allowed_worker_ids = unsafe { parse_pods_filter(pods_json) }; + let timeout_duration = Duration::from_millis(timeout_ms); + + let result = handles.runtime.secondary().block_on(async { + tokio::time::timeout(timeout_duration, async { + handles + .reserve_prefill_worker( + &reservation_id, + &tokens, + None, + None, + priority_jump, + strict_priority, + allowed_worker_ids, + routing_constraints, + ) + .await + }) + .await + }); + + match result { + Ok(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; + out.prefill_worker_id = prefill_worker_id; + out.prefill_dp_rank = prefill_dp_rank; + write_tokens_to_result(&tokens, out); + QueryRouterResult::Ok + } + Ok(Err(code)) => code, + Err(_) => { + tracing::warn!(%reservation_id, timeout_ms, "Prefill reservation timed out"); + QueryRouterResult::ErrTimeout + } + } +} + /// 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/activation.rs b/lib/llm/src/kv_router/prefill_router/activation.rs index 85287fb83409..c523c7fd9964 100644 --- a/lib/llm/src/kv_router/prefill_router/activation.rs +++ b/lib/llm/src/kv_router/prefill_router/activation.rs @@ -49,6 +49,7 @@ impl PrefillRouter { ) -> Arc { Arc::new(Self { prefill_router: std::sync::OnceLock::new(), + reservations: Default::default(), model_manager, endpoint_id: std::sync::OnceLock::new(), cancel_token: tokio_util::sync::CancellationToken::new(), @@ -85,6 +86,7 @@ impl PrefillRouter { let router = Arc::new(Self { prefill_router, + reservations: Default::default(), model_manager: model_manager.clone(), endpoint_id: std::sync::OnceLock::new(), cancel_token: cancel_token.clone(), @@ -100,6 +102,7 @@ impl PrefillRouter { activation_task_state: Arc::new(()), }); + Self::spawn_reservation_reaper(&router); // Spawn background task to wait for activation let router_weak = Arc::downgrade(&router); #[cfg(test)] diff --git a/lib/llm/src/kv_router/prefill_router/mod.rs b/lib/llm/src/kv_router/prefill_router/mod.rs index afbba1a0b4e4..9991a84564e5 100644 --- a/lib/llm/src/kv_router/prefill_router/mod.rs +++ b/lib/llm/src/kv_router/prefill_router/mod.rs @@ -35,6 +35,7 @@ use crate::{ mod activation; mod admission; mod query; +mod reservations; use admission::InnerPrefillRouter; @@ -154,6 +155,7 @@ fn strip_terminal_disaggregated_params( /// - Normal: Worker IDs determined by router based on KV cache state pub struct PrefillRouter { prefill_router: OnceLock, + reservations: reservations::PrefillReservationRegistry, model_manager: Arc, endpoint_id: OnceLock, cancel_token: CancellationToken, 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..1ac71f992928 --- /dev/null +++ b/lib/llm/src/kv_router/prefill_router/reservations.rs @@ -0,0 +1,230 @@ +// 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, WorkerWithDpRank}, + sequence::DEFAULT_ACTIVE_REQUEST_EXPIRY_DURATION, +}; +use parking_lot::Mutex; +use tokio::time::{Instant, MissedTickBehavior}; + +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_ENV: &str = "DYN_ROUTER_ACTIVE_REQUEST_EXPIRY_SECS"; +const RESERVATION_EXPIRY_GRACE: Duration = Duration::from_secs(60); +const RESERVATION_REAPER_INTERVAL: Duration = Duration::from_secs(30); + +struct ActivePrefillReservation { + chooser: Arc, + scheduler_id: String, + created_at: Instant, +} + +impl Clone for ActivePrefillReservation { + fn clone(&self) -> Self { + Self { + chooser: self.chooser.clone(), + scheduler_id: self.scheduler_id.clone(), + created_at: self.created_at, + } + } +} + +pub(super) struct PrefillReservationRegistry { + entries: Mutex>, +} + +impl Default for PrefillReservationRegistry { + fn default() -> Self { + Self { + entries: Mutex::new(HashMap::new()), + } + } +} + +impl PrefillReservationRegistry { + fn insert(&self, reservation_id: &str, reservation: ActivePrefillReservation) -> bool { + match self.entries.lock().entry(reservation_id.to_string()) { + Entry::Vacant(entry) => { + entry.insert(reservation); + true + } + Entry::Occupied(_) => false, + } + } + + fn get(&self, reservation_id: &str) -> Option { + self.entries.lock().get(reservation_id).cloned() + } + + 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| entry.scheduler_id == scheduler_id) + { + entries.remove(reservation_id); + } + } + + fn expired_ids(&self, now: Instant, retention: Duration) -> Vec { + self.entries + .lock() + .iter() + .filter(|(_, entry)| now.saturating_duration_since(entry.created_at) >= retention) + .map(|(reservation_id, _)| reservation_id.clone()) + .collect() + } +} + +fn reservation_retention() -> Duration { + let configured = std::env::var(ACTIVE_REQUEST_EXPIRY_ENV) + .ok() + .and_then(|raw| raw.parse::().ok()) + .filter(|seconds| *seconds > 0) + .map(Duration::from_secs) + .unwrap_or(DEFAULT_ACTIVE_REQUEST_EXPIRY_DURATION); + configured + .max(DEFAULT_ACTIVE_REQUEST_EXPIRY_DURATION) + .saturating_add(RESERVATION_EXPIRY_GRACE) +} + +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 PrefillRouter { + /// 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_prefill_worker( + &self, + 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"); + } + if self.lifecycle_state() != PrefillLifecycleState::Active { + return Err(anyhow::anyhow!(PrefillError::NotActivated)); + } + let inner = self + .prefill_router + .get() + .ok_or_else(|| anyhow::anyhow!(PrefillError::NotActivated))?; + let InnerPrefillRouter::KvRouter(router) = inner else { + return Err(anyhow::anyhow!(PrefillError::NotActivated)); + }; + let chooser = router.chooser.clone(); + let scheduler_id = scheduler_id(reservation_id); + let 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, + ) + .await?; + + match outcome { + crate::kv_router::FindBestMatchOutcome::Routed { worker, .. } => { + let reservation = ActivePrefillReservation { + chooser: chooser.clone(), + scheduler_id: scheduler_id.clone(), + created_at: Instant::now(), + }; + if !self.reservations.insert(reservation_id, reservation) { + ignore_missing_request(chooser.free(&scheduler_id).await)?; + anyhow::bail!("prefill reservation {reservation_id:?} already exists"); + } + Ok(PrefillQueryOutcome::Routed { + worker_id: worker.worker_id, + dp_rank: Some(worker.dp_rank), + }) + } + crate::kv_router::FindBestMatchOutcome::QueueRejected { rejection } => { + Ok(PrefillQueryOutcome::QueueRejected { rejection }) + } + } + } + + /// Release a prefill reservation. Missing reservations are idempotent no-ops. + /// + /// The registry entry remains present until scheduler cleanup succeeds, so a timeout or + /// transient router error can be retried safely. + pub async fn release_prefill_reservation(&self, reservation_id: &str) -> Result<()> { + let Some(reservation) = self.reservations.get(reservation_id) else { + return Ok(()); + }; + + ignore_missing_request(reservation.chooser.free(&reservation.scheduler_id).await)?; + self.reservations + .remove_if_scheduler_id(reservation_id, &reservation.scheduler_id); + Ok(()) + } + + pub(super) fn spawn_reservation_reaper(router: &Arc) { + let router: Weak = Arc::downgrade(router); + let cancellation = router + .upgrade() + .expect("router must be alive while starting reservation reaper") + .cancel_token + .child_token(); + + 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 { + tokio::select! { + _ = cancellation.cancelled() => return, + _ = interval.tick() => { + let Some(router) = router.upgrade() else { + return; + }; + let expired = router.reservations.expired_ids(Instant::now(), retention); + for reservation_id in expired { + if let Err(error) = router.release_prefill_reservation(&reservation_id).await { + tracing::warn!( + %reservation_id, + %error, + "Failed to expire stale EPP prefill reservation" From 025b97b3ff579bcced5b6334df37ccef6f1182aa Mon Sep 17 00:00:00 2001 From: Thomas Montfort Date: Tue, 11 Aug 2026 13:01:57 -0700 Subject: [PATCH 03/15] refactor(epp): remove advisory prefill route Signed-off-by: Thomas Montfort (cherry picked from commit 92eadad13b569a004f39d119fbcbd028d60983f5) Signed-off-by: Avinash Varma --- .../pkg/plugins/dynamo_kv_scorer/plugin.go | 43 ------ lib/bindings/c/src/lib.rs | 133 ------------------ 2 files changed, 176 deletions(-) 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 6ef23acc7d78..891f66e31d33 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 @@ -447,49 +447,6 @@ func extractCacheNamespace(result *C.CRoutingResult) string { return "" } -// 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) { - if !routerInitialized { - return nil, fmt.Errorf("dynamo router not initialized") - } - - routerHandlesMutex.RLock() - router := routerHandles - routerHandlesMutex.RUnlock() - if router == nil { - return nil, fmt.Errorf("dynamo router handles not created") - } - - cRequestJSON := C.CString(requestJSON) - defer C.free(unsafe.Pointer(cRequestJSON)) - - var cPodsJSON *C.char - if podsJSON != "" { - cPodsJSON = C.CString(podsJSON) - defer C.free(unsafe.Pointer(cPodsJSON)) - } - - var result C.CRoutingResult - rc := C.route_prefill_request(router, cRequestJSON, cPodsJSON, &result) - if rc != C.QUERY_ROUTER_OK { - return nil, fmt.Errorf("route_prefill_request failed with code %d", rc) - } - - tokens := extractTokenData(&result) - cacheNamespace := extractCacheNamespace(&result) - workerID := uint64(result.prefill_worker_id) - dpRank := uint32(result.prefill_dp_rank) - C.free_routing_result(&result) - - return &RoutingResult{ - WorkerID: workerID, - DpRank: dpRank, - TokenData: tokens, - CacheNamespace: cacheNamespace, - }, nil -} - // CallRoutePrefillRequestWithReservation atomically selects and books a prefill worker. // The Rust scheduler retracts pending admission when timeout expires. func CallRoutePrefillRequestWithReservation(reservationID string, requestJSON string, podsJSON string, timeout time.Duration) (*RoutingResult, error) { diff --git a/lib/bindings/c/src/lib.rs b/lib/bindings/c/src/lib.rs index a09f7c5cdb87..422ec843656b 100644 --- a/lib/bindings/c/src/lib.rs +++ b/lib/bindings/c/src/lib.rs @@ -475,59 +475,6 @@ 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. - #[expect(clippy::too_many_arguments)] - async fn query_prefill_worker( - &self, - tokens: &[u32], - block_mm_infos: Option<&[Option]>, - lora_name: Option, - cache_namespace: Option, - priority_jump: f64, - strict_priority: u32, - allowed_worker_ids: Option>, - routing_constraints: RoutingConstraints, - ) -> Result<(u64, Option), QueryRouterResult> { - if let Some(ref ids) = allowed_worker_ids { - self.prefill_router.register_workers(ids); - } - - let outcome = self - .prefill_router - .query_prefill_worker( - tokens, - block_mm_infos, - lora_name, - cache_namespace, - priority_jump, - strict_priority, - allowed_worker_ids, - routing_constraints, - ) - .await - .map_err(|e| { - tracing::error!(error = ?e, "Prefill query 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!( - policy_class = %rejection.policy_class, - limit_kind = %rejection.limit_kind, - current = rejection.current, - limit = rejection.limit, - "Prefill query rejected by policy-class queue limit" - ); - Err(QueryRouterResult::ErrBackpressure) - } - } - } - /// Atomically select and reserve a prefill worker for an EPP-owned booking. #[expect(clippy::too_many_arguments)] async fn reserve_prefill_worker( @@ -1463,86 +1410,6 @@ fn write_cache_namespace_to_result(cache_namespace: Option<&str>, out: &mut CRou std::mem::forget(namespace_boxed); } -/// Route a request to select the best **prefill** worker only. -/// -/// 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. -/// -/// 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 -/// - `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( - handle: RouterHandlesPtr, - 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() { - return QueryRouterResult::ErrInvalidParam; - } - - let handles = unsafe { &*handle }; - - let (tokens, cache_namespace, priority_jump, strict_priority, routing_constraints) = - match unsafe { preprocess_request(handles, request_json) } { - Ok(t) => t, - Err(code) => 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, - cache_namespace.clone(), - 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(), - priority_jump, - strict_priority, - "Routed prefill request" - ); - - Ok((prefill_worker_id, prefill_dp_rank)) - }); - - match result { - Ok((prefill_worker_id, prefill_dp_rank)) => { - let out = unsafe { &mut *out_result }; - *out = CRoutingResult::default(); - out.is_disaggregated = true; - out.prefill_worker_id = prefill_worker_id; - out.prefill_dp_rank = prefill_dp_rank; - write_tokens_to_result(&tokens, out); - write_cache_namespace_to_result(cache_namespace.as_deref(), out); - QueryRouterResult::Ok - } - Err(code) => code, - } -} - /// Atomically select and reserve the best prefill worker for an EPP-owned booking. /// /// Pending scheduler admission is bounded by `timeout_ms`. From 84e8f290dbe33a40a620cbef134992b3a03868c7 Mon Sep 17 00:00:00 2001 From: Thomas Montfort Date: Tue, 11 Aug 2026 14:05:03 -0700 Subject: [PATCH 04/15] fix(epp): cancel pending prefill reservations Signed-off-by: Thomas Montfort (cherry picked from commit 583cdb2888a200b24600d4b6c9194876fad4110d) Signed-off-by: Avinash Varma --- .../epp/pkg/plugins/disagg/prefill_scorer.go | 52 +++-- .../pkg/plugins/disagg/reservation_test.go | 65 ++++++ .../pkg/plugins/dynamo_kv_scorer/plugin.go | 43 +++- lib/bindings/c/src/lib.rs | 69 +++--- .../kv_router/prefill_router/reservations.rs | 205 +++++++++++++----- 5 files changed, 324 insertions(+), 110 deletions(-) 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 1c08c4b1ded3..43b4e409b1c3 100644 --- a/deploy/inference-gateway/epp/pkg/plugins/disagg/prefill_scorer.go +++ b/deploy/inference-gateway/epp/pkg/plugins/disagg/prefill_scorer.go @@ -21,7 +21,6 @@ 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" @@ -34,8 +33,6 @@ import ( const ( // DynPrefillScorerType is the plugin type registered in the plugin registry. DynPrefillScorerType = "dyn-prefill-scorer" - - defaultPrefillReservationTimeout = 5 * time.Second ) // compile-time type assertion @@ -65,6 +62,7 @@ func NewDynPrefillScorer() *DynPrefillScorer { return &DynPrefillScorer{ typedName: plugins.TypedName{Type: DynPrefillScorerType, Name: DynPrefillScorerType}, reservePrefill: dynscorer.CallRoutePrefillRequestWithReservation, + cancelPrefill: dynscorer.CallCancelPrefillReservation, freeBooking: dynscorer.CallFreeRequest, } } @@ -72,10 +70,16 @@ func NewDynPrefillScorer() *DynPrefillScorer { // DynPrefillScorer is a scorer plugin for the prefill scheduling profile. type DynPrefillScorer struct { typedName plugins.TypedName - reservePrefill func(string, string, string, time.Duration) (*dynscorer.RoutingResult, error) + reservePrefill func(string, string, string) (*dynscorer.RoutingResult, error) + cancelPrefill func(string) error freeBooking func(string) error } +type prefillReservationResult struct { + result *dynscorer.RoutingResult + err error +} + // TypedName returns the type and name tuple of this plugin instance. func (s *DynPrefillScorer) TypedName() plugins.TypedName { return s.typedName @@ -126,19 +130,40 @@ func (s *DynPrefillScorer) Score(ctx context.Context, cycleState *schedtypes.Cyc "endpointCount", len(endpoints), "endpointsJSON", string(endpointsJSON)) - timeout := defaultPrefillReservationTimeout - if deadline, ok := ctx.Deadline(); ok { - remaining := time.Until(deadline) - if remaining < timeout { - timeout = remaining + resultCh := make(chan prefillReservationResult, 1) + bookingID := booking.ID + go func() { + result, err := s.reservePrefill(bookingID, requestJSON, endpointsJSON) + resultCh <- prefillReservationResult{result: result, err: err} + }() + + var result *dynscorer.RoutingResult + select { + case <-ctx.Done(): + logger.V(logutil.VERBOSE).Info("DynPrefillScorer: scheduling cancelled during prefill reservation", + "error", ctx.Err().Error()) + if cancelErr := s.cancelPrefill(bookingID); cancelErr != nil { + logger.V(logutil.DEFAULT).Error(cancelErr, "DynPrefillScorer: failed to cancel pending prefill reservation", + "bookingID", bookingID) } - } - if timeout <= 0 { + go func(bookingID string) { + outcome := <-resultCh + if outcome.err != nil { + return + } + if cleanupErr := s.freeBooking(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 } - - result, err := s.reservePrefill(booking.ID, requestJSON, endpointsJSON, timeout) if err != nil { logger.V(logutil.DEFAULT).Error(err, "DynPrefillScorer: FFI prefill reservation failed") booking.PrefillReserved = false @@ -174,7 +199,6 @@ func (s *DynPrefillScorer) Score(ctx context.Context, cycleState *schedtypes.Cyc return uniformScores(endpoints, 0) } - bookingID := booking.ID go func() { <-ctx.Done() if cleanupErr := s.freeBooking(bookingID); cleanupErr != nil { diff --git a/deploy/inference-gateway/epp/pkg/plugins/disagg/reservation_test.go b/deploy/inference-gateway/epp/pkg/plugins/disagg/reservation_test.go index a3f92ed5b3fe..d23f08085178 100644 --- a/deploy/inference-gateway/epp/pkg/plugins/disagg/reservation_test.go +++ b/deploy/inference-gateway/epp/pkg/plugins/disagg/reservation_test.go @@ -20,9 +20,13 @@ 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 { @@ -150,3 +154,64 @@ func TestPreRequestCancellationCleansBooking(t *testing.T) { 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) + freeCalls := make(chan string, 1) + 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 + }, + freeBooking: func(bookingID string) error { + freeCalls <- 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() + + 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 := <-freeCalls: + 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 readPrefillEnabled(cycleState) { + t.Fatal("prefill remained enabled after reservation cancellation") + } +} 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 891f66e31d33..d77ef82b4156 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 @@ -76,9 +76,11 @@ query_router_result_t route_prefill_request_with_reservation(RouterHandles *hand const char *reservation_id, const char *request_json, const char *pods_json, - uint64_t timeout_ms, CRoutingResult *out_result); +query_router_result_t cancel_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, @@ -448,14 +450,11 @@ func extractCacheNamespace(result *C.CRoutingResult) string { } // CallRoutePrefillRequestWithReservation atomically selects and books a prefill worker. -// The Rust scheduler retracts pending admission when timeout expires. -func CallRoutePrefillRequestWithReservation(reservationID string, requestJSON string, podsJSON string, timeout time.Duration) (*RoutingResult, error) { +// 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 timeout <= 0 { - return nil, fmt.Errorf("prefill reservation timeout must be positive") - } if !routerInitialized { return nil, fmt.Errorf("dynamo router not initialized") } @@ -478,17 +477,12 @@ func CallRoutePrefillRequestWithReservation(reservationID string, requestJSON st defer C.free(unsafe.Pointer(cPodsJSON)) } - timeoutMS := timeout.Milliseconds() - if timeoutMS < 1 { - timeoutMS = 1 - } var result C.CRoutingResult rc := C.route_prefill_request_with_reservation( router, cReservationID, cRequestJSON, cPodsJSON, - C.uint64_t(timeoutMS), &result, ) if rc != C.QUERY_ROUTER_OK { @@ -507,6 +501,33 @@ func CallRoutePrefillRequestWithReservation(reservationID string, requestJSON st }, 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 +} + // CallRouteDecodeRequest routes a request to the best decode worker. // When isDisaggregated is true, overlap_score_credit=0 is used (KV cache transferred from prefill). diff --git a/lib/bindings/c/src/lib.rs b/lib/bindings/c/src/lib.rs index 422ec843656b..73cb84fb76c4 100644 --- a/lib/bindings/c/src/lib.rs +++ b/lib/bindings/c/src/lib.rs @@ -1412,7 +1412,8 @@ fn write_cache_namespace_to_result(cache_namespace: Option<&str>, out: &mut CRou /// Atomically select and reserve the best prefill worker for an EPP-owned booking. /// -/// Pending scheduler admission is bounded by `timeout_ms`. +/// 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 @@ -1426,14 +1427,12 @@ pub unsafe extern "C" fn route_prefill_request_with_reservation( reservation_id: *const c_char, request_json: *const c_char, pods_json: *const c_char, - timeout_ms: u64, out_result: *mut CRoutingResult, ) -> QueryRouterResult { if handle.is_null() || reservation_id.is_null() || request_json.is_null() || out_result.is_null() - || timeout_ms == 0 { return QueryRouterResult::ErrInvalidParam; } @@ -1449,28 +1448,20 @@ pub unsafe extern "C" fn route_prefill_request_with_reservation( Err(code) => return code, }; let allowed_worker_ids = unsafe { parse_pods_filter(pods_json) }; - let timeout_duration = Duration::from_millis(timeout_ms); - let result = handles.runtime.secondary().block_on(async { - tokio::time::timeout(timeout_duration, async { - handles - .reserve_prefill_worker( - &reservation_id, - &tokens, - None, - None, - priority_jump, - strict_priority, - allowed_worker_ids, - routing_constraints, - ) - .await - }) - .await - }); + let result = handles.runtime.secondary().block_on(handles.reserve_prefill_worker( + &reservation_id, + &tokens, + None, + None, + priority_jump, + strict_priority, + allowed_worker_ids, + routing_constraints, + )); match result { - Ok(Ok((prefill_worker_id, prefill_dp_rank))) => { + Ok((prefill_worker_id, prefill_dp_rank)) => { let prefill_dp_rank = prefill_dp_rank.unwrap_or(u32::MAX); tracing::info!( %reservation_id, @@ -1490,14 +1481,38 @@ pub unsafe extern "C" fn route_prefill_request_with_reservation( write_tokens_to_result(&tokens, out); QueryRouterResult::Ok } - Ok(Err(code)) => code, - Err(_) => { - tracing::warn!(%reservation_id, timeout_ms, "Prefill reservation timed out"); - QueryRouterResult::ErrTimeout - } + Err(code) => code, } } +/// 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 + .prefill_router + .cancel_prefill_reservation(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/reservations.rs b/lib/llm/src/kv_router/prefill_router/reservations.rs index 1ac71f992928..e43b1410d399 100644 --- a/lib/llm/src/kv_router/prefill_router/reservations.rs +++ b/lib/llm/src/kv_router/prefill_router/reservations.rs @@ -9,11 +9,12 @@ use std::{ use anyhow::Result; use dynamo_kv_router::{ - protocols::{BlockExtraInfo, RoutingConstraints, WorkerId, WorkerWithDpRank}, + protocols::{BlockExtraInfo, RoutingConstraints, WorkerId}, sequence::DEFAULT_ACTIVE_REQUEST_EXPIRY_DURATION, }; use parking_lot::Mutex; use tokio::time::{Instant, MissedTickBehavior}; +use tokio_util::sync::CancellationToken; use super::{ InnerPrefillRouter, PrefillError, PrefillLifecycleState, PrefillQueryOutcome, PrefillRouter, @@ -23,6 +24,7 @@ use crate::kv_router::{KvRouter, sequence::SequenceError}; const PREFILL_SCHEDULER_ID_PREFIX: &str = "epp-prefill/"; const ACTIVE_REQUEST_EXPIRY_ENV: &str = "DYN_ROUTER_ACTIVE_REQUEST_EXPIRY_SECS"; 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); struct ActivePrefillReservation { @@ -41,8 +43,25 @@ impl Clone for ActivePrefillReservation { } } +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 }, + 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, +} + pub(super) struct PrefillReservationRegistry { - entries: Mutex>, + entries: Mutex>, } impl Default for PrefillReservationRegistry { @@ -54,38 +73,111 @@ impl Default for PrefillReservationRegistry { } impl PrefillReservationRegistry { - fn insert(&self, reservation_id: &str, reservation: ActivePrefillReservation) -> bool { - match self.entries.lock().entry(reservation_id.to_string()) { + fn begin(&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(), + }); + BeginReservation::Pending(cancellation) + } + Entry::Occupied(entry) => match entry.get() { + PrefillReservationEntry::Cancelled { .. } => { + entry.remove(); + BeginReservation::Cancelled + } + PrefillReservationEntry::Pending { .. } | PrefillReservationEntry::Active(_) => { + BeginReservation::AlreadyExists + } + }, + } + } + + fn activate(&self, reservation_id: &str, reservation: ActivePrefillReservation) -> bool { + let mut entries = self.entries.lock(); + let Some(entry) = entries.get(reservation_id) else { + return false; + }; + + let PrefillReservationEntry::Pending { cancellation } = entry else { + return false; + }; + if cancellation.is_cancelled() { + entries.remove(reservation_id); + return false; + } + + entries.insert( + reservation_id.to_string(), + PrefillReservationEntry::Active(reservation), + ); + true + } + + fn remove_pending(&self, reservation_id: &str) { + let mut entries = self.entries.lock(); + if matches!( + entries.get(reservation_id), + Some(PrefillReservationEntry::Pending { .. }) + ) { + entries.remove(reservation_id); + } + } + + fn cancel(&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(reservation); - true + entry.insert(PrefillReservationEntry::Cancelled { + created_at: Instant::now(), + }); } - Entry::Occupied(_) => false, } } - fn get(&self, reservation_id: &str) -> Option { - self.entries.lock().get(reservation_id).cloned() + 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| entry.scheduler_id == scheduler_id) - { + if entries.get(reservation_id).is_some_and(|entry| { + matches!(entry, PrefillReservationEntry::Active(active) if active.scheduler_id == scheduler_id) + }) { entries.remove(reservation_id); } } - fn expired_ids(&self, now: Instant, retention: Duration) -> Vec { + fn expired_active_ids(&self, now: Instant, retention: Duration) -> Vec { self.entries .lock() .iter() - .filter(|(_, entry)| now.saturating_duration_since(entry.created_at) >= retention) + .filter(|(_, entry)| { + matches!(entry, PrefillReservationEntry::Active(active) + if now.saturating_duration_since(active.created_at) >= retention) + }) .map(|(reservation_id, _)| reservation_id.clone()) .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() -> Duration { @@ -143,8 +235,23 @@ impl PrefillRouter { }; let chooser = router.chooser.clone(); let scheduler_id = scheduler_id(reservation_id); - let outcome = chooser - .find_best_match_details( + let cancellation = match self.reservations.begin(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") + } + }; + + let outcome = tokio::select! { + biased; + _ = cancellation.cancelled() => { + self.reservations.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, @@ -158,8 +265,15 @@ impl PrefillRouter { None, allowed_worker_ids, routing_constraints, - ) - .await?; + ) => outcome, + }; + let outcome = match outcome { + Ok(outcome) => outcome, + Err(error) => { + self.reservations.remove_pending(reservation_id); + return Err(error); + } + }; match outcome { crate::kv_router::FindBestMatchOutcome::Routed { worker, .. } => { @@ -168,9 +282,9 @@ impl PrefillRouter { scheduler_id: scheduler_id.clone(), created_at: Instant::now(), }; - if !self.reservations.insert(reservation_id, reservation) { + if !self.reservations.activate(reservation_id, reservation) { ignore_missing_request(chooser.free(&scheduler_id).await)?; - anyhow::bail!("prefill reservation {reservation_id:?} already exists"); + anyhow::bail!("prefill reservation {reservation_id:?} was cancelled"); } Ok(PrefillQueryOutcome::Routed { worker_id: worker.worker_id, @@ -178,53 +292,28 @@ impl PrefillRouter { }) } crate::kv_router::FindBestMatchOutcome::QueueRejected { rejection } => { + self.reservations.remove_pending(reservation_id); Ok(PrefillQueryOutcome::QueueRejected { rejection }) } } } - /// Release a prefill reservation. Missing reservations are idempotent no-ops. + /// Cancel a pending prefill reservation without waiting for scheduler cleanup. /// - /// The registry entry remains present until scheduler cleanup succeeds, so a timeout or - /// transient router error can be retried safely. + /// 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_prefill_reservation(&self, reservation_id: &str) { + if !reservation_id.is_empty() { + self.reservations.cancel(reservation_id); + } + } + + /// Release a prefill reservation. Missing reservations are idempotent no-ops. pub async fn release_prefill_reservation(&self, reservation_id: &str) -> Result<()> { - let Some(reservation) = self.reservations.get(reservation_id) else { + let Some(reservation) = self.reservations.get_active(reservation_id) else { return Ok(()); }; ignore_missing_request(reservation.chooser.free(&reservation.scheduler_id).await)?; self.reservations .remove_if_scheduler_id(reservation_id, &reservation.scheduler_id); - Ok(()) - } - - pub(super) fn spawn_reservation_reaper(router: &Arc) { - let router: Weak = Arc::downgrade(router); - let cancellation = router - .upgrade() - .expect("router must be alive while starting reservation reaper") - .cancel_token - .child_token(); - - 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 { - tokio::select! { - _ = cancellation.cancelled() => return, - _ = interval.tick() => { - let Some(router) = router.upgrade() else { - return; - }; - let expired = router.reservations.expired_ids(Instant::now(), retention); - for reservation_id in expired { - if let Err(error) = router.release_prefill_reservation(&reservation_id).await { - tracing::warn!( - %reservation_id, - %error, - "Failed to expire stale EPP prefill reservation" From ea7e01d21d61e6657be09f877509721a269ab58c Mon Sep 17 00:00:00 2001 From: Thomas Montfort Date: Tue, 11 Aug 2026 14:40:02 -0700 Subject: [PATCH 05/15] fix(epp): bound prefill reservation lifecycle Signed-off-by: Thomas Montfort (cherry picked from commit 2ac4c67df0049a0908659c2ace04571cfeb2d4e0) Signed-off-by: Avinash Varma --- .../epp/pkg/plugins/disagg/decode_scorer.go | 41 +--- .../epp/pkg/plugins/disagg/prefill_scorer.go | 116 +++++++--- .../pkg/plugins/disagg/reservation_test.go | 185 ++++++++++++++++ .../epp/pkg/plugins/disagg/shared.go | 160 ++++++++++++++ .../pkg/plugins/dynamo_kv_scorer/plugin.go | 37 +++- lib/bindings/c/src/lib.rs | 72 ++++++- .../kv_router/prefill_router/reservations.rs | 201 +++++++++++++++--- 7 files changed, 711 insertions(+), 101 deletions(-) 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 3abcc4f8703b..80d730468931 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" @@ -116,7 +115,6 @@ func NewDynDecodeScorer(ctx context.Context) *DynDecodeScorer { type DynDecodeScorer struct { typedName plugins.TypedName pluginState *plugins.PluginState - prefillMarkInFlight sync.Map addRequest func(string, []int64, uint64, uint32, string) error markPrefillComplete func(string) error freeBooking func(string) error @@ -223,6 +221,9 @@ func (s *DynDecodeScorer) cleanupBooking(ctx context.Context, bookingID, reason if bookingID == "" { return true } + if lifecycle := findBookingLifecycle(bookingID); lifecycle != nil { + return lifecycle.cleanup(ctx, reason) + } if err := s.freeBooking(bookingID); err != nil { log.FromContext(ctx).V(logutil.DEFAULT).Error(err, "DynDecodeScorer: booking cleanup failed", "bookingID", bookingID, "reason", reason) @@ -306,11 +307,6 @@ func (s *DynDecodeScorer) PreRequest(ctx context.Context, request *schedtypes.In "dpRank", state.DpRank, "hasCacheNamespace", state.CacheNamespace != "", "tokenCount", len(state.TokenData)) - - go func() { - <-ctx.Done() - s.cleanupBooking(ctx, bookingID, "request context cancelled") - }() } // ResponseBody handles streaming chunks and end-of-stream cleanup. @@ -321,33 +317,14 @@ func (s *DynDecodeScorer) ResponseBody(ctx context.Context, request *schedtypes. return } - logger := log.FromContext(ctx) - - // Terminal cleanup takes precedence over first-token bookkeeping. This also - // covers empty/error responses where no output token was observed. + lifecycle := findBookingLifecycle(bookingID) + if lifecycle == nil { + lifecycle = registerBookingLifecycle(bookingID, s.freeBooking) + } if response.EndOfStream { - s.prefillMarkInFlight.Delete(bookingID) - if err := s.freeBooking(bookingID); err != nil { - logger.V(logutil.DEFAULT).Error(err, "DynDecodeScorer ResponseBody: failed to free request", - "bookingID", bookingID, "requestID", request.RequestId) - } else { - logger.V(logutil.VERBOSE).Info("DynDecodeScorer ResponseBody: freed request", - "bookingID", bookingID, "requestID", request.RequestId) - } + lifecycle.cleanup(ctx, "response end of stream") return } - // Keep the marker while the FFI call is in flight. On failure, remove it so - // the next response chunk retries instead of leaking the prefill load. - if _, alreadyInFlight := s.prefillMarkInFlight.LoadOrStore(bookingID, struct{}{}); alreadyInFlight { - return - } - if err := s.markPrefillComplete(bookingID); err != nil { - s.prefillMarkInFlight.Delete(bookingID) - logger.V(logutil.DEFAULT).Error(err, "DynDecodeScorer ResponseBody: failed to mark prefill complete", - "bookingID", bookingID, "requestID", request.RequestId) - } else { - logger.V(logutil.VERBOSE).Info("DynDecodeScorer ResponseBody: marked prefill complete", - "bookingID", bookingID, "requestID", request.RequestId) - } + 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 43b4e409b1c3..3f3a58e8e520 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 = 5 * time.Second + defaultMaxPrefillReservations = 32 ) // 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,30 +59,51 @@ 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}, - reservePrefill: dynscorer.CallRoutePrefillRequestWithReservation, - cancelPrefill: dynscorer.CallCancelPrefillReservation, - freeBooking: dynscorer.CallFreeRequest, + typedName: plugins.TypedName{Type: DynPrefillScorerType, Name: DynPrefillScorerType}, + beginPrefill: dynscorer.CallBeginPrefillReservation, + reservePrefill: dynscorer.CallRoutePrefillRequestWithReservation, + cancelPrefill: dynscorer.CallCancelPrefillReservation, + 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 - reservePrefill func(string, string, string) (*dynscorer.RoutingResult, error) - cancelPrefill func(string) error - freeBooking func(string) error + typedName plugins.TypedName + beginPrefill func(string) error + reservePrefill func(string, string, string) (*dynscorer.RoutingResult, error) + cancelPrefill func(string) error + freeBooking func(string) error + reservationAdmissionTimeout time.Duration + reservationSlots chan struct{} } type prefillReservationResult struct { @@ -80,6 +111,32 @@ type prefillReservationResult struct { 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) +} + // TypedName returns the type and name tuple of this plugin instance. func (s *DynPrefillScorer) TypedName() plugins.TypedName { return s.typedName @@ -130,18 +187,36 @@ func (s *DynPrefillScorer) Score(ctx context.Context, cycleState *schedtypes.Cyc "endpointCount", len(endpoints), "endpointsJSON", string(endpointsJSON)) - resultCh := make(chan prefillReservationResult, 1) 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 <-ctx.Done(): + case <-admissionCtx.Done(): logger.V(logutil.VERBOSE).Info("DynPrefillScorer: scheduling cancelled during prefill reservation", - "error", ctx.Err().Error()) + "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) @@ -188,24 +263,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 { - if cleanupErr := s.freeBooking(booking.ID); cleanupErr != nil { - logger.V(logutil.DEFAULT).Error(cleanupErr, "DynPrefillScorer: failed to clean cancelled reservation", - "bookingID", booking.ID) - } + 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) } - go func() { - <-ctx.Done() - if cleanupErr := s.freeBooking(bookingID); cleanupErr != nil { - logger.V(logutil.DEFAULT).Error(cleanupErr, "DynPrefillScorer: cancellation cleanup failed", - "bookingID", bookingID) - } - }() - 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 index d23f08085178..8ca6e4470daa 100644 --- a/deploy/inference-gateway/epp/pkg/plugins/disagg/reservation_test.go +++ b/deploy/inference-gateway/epp/pkg/plugins/disagg/reservation_test.go @@ -93,7 +93,22 @@ func TestResponseBodyRetriesPrefillMarkAndFreesTerminalResponse(t *testing.T) { 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) @@ -124,6 +139,15 @@ func TestResponseBodyFreesWithoutMarkingEmptyTerminalResponse(t *testing.T) { &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) @@ -215,3 +239,164 @@ func TestPrefillScoreCancelsPendingReservation(t *testing.T) { t.Fatal("prefill remained enabled after reservation cancellation") } } + +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) + + 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) + scorer.ResponseBody(context.Background(), request, &rc.Response{}, nil) + <-markStarted + lifecycle := findBookingLifecycle(bookingID) + if lifecycle == nil { + t.Fatal("expected booking lifecycle after first response chunk") + } + + started := time.Now() + scorer.ResponseBody(context.Background(), request, &rc.Response{EndOfStream: true}, nil) + if elapsed := time.Since(started); elapsed > 100*time.Millisecond { + t.Fatalf("EOS callback blocked for %s waiting on prefill mark", elapsed) + } + 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") + } +} diff --git a/deploy/inference-gateway/epp/pkg/plugins/disagg/shared.go b/deploy/inference-gateway/epp/pkg/plugins/disagg/shared.go index 807424e92c60..4d33eddb50eb 100644 --- a/deploy/inference-gateway/epp/pkg/plugins/disagg/shared.go +++ b/deploy/inference-gateway/epp/pkg/plugins/disagg/shared.go @@ -27,12 +27,15 @@ limitations under the License. package disagg import ( + "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" @@ -213,6 +216,163 @@ func getEnvBoolOrDefault(key string, def bool) bool { var enforceDisaggDeprecationOnce sync.Once +const ( + prefillMarkMaxAttempts = 3 + prefillMarkRetryBackoff = 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 free_request +// caller even when EOS and context cancellation race. +type bookingLifecycle struct { + bookingID string + freeBooking func(string) error + + cleanupOnce sync.Once + mu sync.Mutex + cleaned bool + + stopCancellation func() bool + stopMarker context.CancelFunc + markerDone chan struct{} + cleanupDone chan struct{} +} + +func registerBookingLifecycle(bookingID string, freeBooking func(string) error) *bookingLifecycle { + lifecycle := &bookingLifecycle{ + bookingID: bookingID, + freeBooking: freeBooking, + } + 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) +} + +func (l *bookingLifecycle) armCancellation(ctx context.Context) { + l.mu.Lock() + defer l.mu.Unlock() + if l.cleaned || l.stopCancellation != nil { + return + } + l.stopCancellation = context.AfterFunc(ctx, func() { + l.cleanup(ctx, "request context cancelled") + }) +} + +// startPrefillMarker makes first-token bookkeeping bounded and independent of +// the response callback. EOS or request cancellation stops any pending retry. +func (l *bookingLifecycle) startPrefillMarker(markPrefillComplete func(string) error, logger logr.Logger, requestID string) { + l.mu.Lock() + if l.cleaned || 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() + + go func() { + defer close(markerDone) + for attempt := 1; attempt <= prefillMarkMaxAttempts; attempt++ { + if markerCtx.Err() != nil { + return + } + if err := markPrefillComplete(l.bookingID); err == nil { + logger.V(logutil.VERBOSE).Info("DynDecodeScorer ResponseBody: marked prefill complete", + "bookingID", l.bookingID, "requestID", requestID, "attempt", attempt) + return + } else { + logger.V(logutil.DEFAULT).Error(err, "DynDecodeScorer ResponseBody: failed to mark prefill complete", + "bookingID", l.bookingID, "requestID", requestID, "attempt", attempt) + } + if attempt == prefillMarkMaxAttempts { + return + } + + timer := time.NewTimer(prefillMarkRetryBackoff * time.Duration(attempt)) + select { + case <-markerCtx.Done(): + if !timer.Stop() { + select { + case <-timer.C: + default: + } + } + return + case <-timer.C: + } + } + }() +} + +// cleanup stops cancellation and retry ownership, then releases the booking +// off the response callback. If a mark call is already in flight, cleanup +// waits for that one bounded call before issuing free_request, preventing +// concurrent mark/free FFI operations for the same booking. +func (l *bookingLifecycle) cleanup(ctx context.Context, reason string) bool { + started := false + l.cleanupOnce.Do(func() { + started = true + l.mu.Lock() + l.cleaned = true + stopCancellation := l.stopCancellation + stopMarker := l.stopMarker + markerDone := l.markerDone + cleanupDone := make(chan struct{}) + l.cleanupDone = cleanupDone + l.mu.Unlock() + + if stopCancellation != nil { + stopCancellation() + } + if stopMarker != nil { + stopMarker() + } + + go func() { + if markerDone != nil { + <-markerDone + } + if err := l.freeBooking(l.bookingID); err != nil { + log.FromContext(ctx).V(logutil.DEFAULT).Error(err, "Dynamo EPP booking cleanup failed", + "bookingID", l.bookingID, "reason", reason) + } else { + log.FromContext(ctx).V(logutil.VERBOSE).Info("Dynamo EPP booking cleaned up", + "bookingID", l.bookingID, "reason", reason) + } + close(cleanupDone) + bookingLifecycles.Delete(l.bookingID) + }() + }) + return started +} + +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 d77ef82b4156..4fa12d0acaac 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 @@ -67,6 +67,9 @@ 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, @@ -449,6 +452,34 @@ func extractCacheNamespace(result *C.CRoutingResult) string { return "" } +// 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) { @@ -495,9 +526,9 @@ func CallRoutePrefillRequestWithReservation(reservationID string, requestJSON st C.free_routing_result(&result) return &RoutingResult{ - WorkerID: workerID, - DpRank: dpRank, - TokenData: tokens, + WorkerID: workerID, + DpRank: dpRank, + TokenData: tokens, }, nil } diff --git a/lib/bindings/c/src/lib.rs b/lib/bindings/c/src/lib.rs index 73cb84fb76c4..104f71a46798 100644 --- a/lib/bindings/c/src/lib.rs +++ b/lib/bindings/c/src/lib.rs @@ -1410,6 +1410,41 @@ fn write_cache_namespace_to_result(cache_namespace: Option<&str>, out: &mut CRou std::mem::forget(namespace_boxed); } +/// Register a prefill booking before potentially slow request preprocessing. +/// +/// 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. +/// +/// # 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 + .prefill_router + .begin_prefill_reservation(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 @@ -1442,23 +1477,38 @@ pub unsafe extern "C" fn route_prefill_request_with_reservation( _ => return QueryRouterResult::ErrInvalidParam, }; let handles = unsafe { &*handle }; + if let Err(error) = handles + .prefill_router + .begin_prefill_reservation(&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(values) => values, - Err(code) => return code, + Err(code) => { + handles + .prefill_router + .abort_prefill_reservation(&reservation_id); + return code; + } }; let allowed_worker_ids = unsafe { parse_pods_filter(pods_json) }; - let result = handles.runtime.secondary().block_on(handles.reserve_prefill_worker( - &reservation_id, - &tokens, - None, - None, - priority_jump, - strict_priority, - allowed_worker_ids, - routing_constraints, - )); + let result = handles + .runtime + .secondary() + .block_on(handles.reserve_prefill_worker( + &reservation_id, + &tokens, + None, + None, + priority_jump, + strict_priority, + allowed_worker_ids, + routing_constraints, + )); match result { Ok((prefill_worker_id, prefill_dp_rank)) => { diff --git a/lib/llm/src/kv_router/prefill_router/reservations.rs b/lib/llm/src/kv_router/prefill_router/reservations.rs index e43b1410d399..fd1bdde3b2d6 100644 --- a/lib/llm/src/kv_router/prefill_router/reservations.rs +++ b/lib/llm/src/kv_router/prefill_router/reservations.rs @@ -8,10 +8,7 @@ use std::{ }; use anyhow::Result; -use dynamo_kv_router::{ - protocols::{BlockExtraInfo, RoutingConstraints, WorkerId}, - sequence::DEFAULT_ACTIVE_REQUEST_EXPIRY_DURATION, -}; +use dynamo_kv_router::protocols::{BlockExtraInfo, RoutingConstraints, WorkerId}; use parking_lot::Mutex; use tokio::time::{Instant, MissedTickBehavior}; use tokio_util::sync::CancellationToken; @@ -22,7 +19,7 @@ use super::{ use crate::kv_router::{KvRouter, sequence::SequenceError}; const PREFILL_SCHEDULER_ID_PREFIX: &str = "epp-prefill/"; -const ACTIVE_REQUEST_EXPIRY_ENV: &str = "DYN_ROUTER_ACTIVE_REQUEST_EXPIRY_SECS"; +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); @@ -47,11 +44,15 @@ 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 }, + Pending { + cancellation: CancellationToken, + }, Active(ActivePrefillReservation), /// Preserve a cancellation that reaches Rust before the blocking reserve /// call creates the pending entry. - Cancelled { created_at: Instant }, + Cancelled { + created_at: Instant, + }, } enum BeginReservation { @@ -88,9 +89,10 @@ impl PrefillReservationRegistry { entry.remove(); BeginReservation::Cancelled } - PrefillReservationEntry::Pending { .. } | PrefillReservationEntry::Active(_) => { - BeginReservation::AlreadyExists + PrefillReservationEntry::Pending { cancellation } => { + BeginReservation::Pending(cancellation.clone()) } + PrefillReservationEntry::Active(_) => BeginReservation::AlreadyExists, }, } } @@ -180,16 +182,12 @@ impl PrefillReservationRegistry { } } +fn reservation_retention_from_expiry(active_request_expiry: Duration) -> Duration { + active_request_expiry.saturating_add(RESERVATION_EXPIRY_GRACE) +} + fn reservation_retention() -> Duration { - let configured = std::env::var(ACTIVE_REQUEST_EXPIRY_ENV) - .ok() - .and_then(|raw| raw.parse::().ok()) - .filter(|seconds| *seconds > 0) - .map(Duration::from_secs) - .unwrap_or(DEFAULT_ACTIVE_REQUEST_EXPIRY_DURATION); - configured - .max(DEFAULT_ACTIVE_REQUEST_EXPIRY_DURATION) - .saturating_add(RESERVATION_EXPIRY_GRACE) + reservation_retention_from_expiry(ACTIVE_REQUEST_EXPIRY_DURATION) } fn scheduler_id(reservation_id: &str) -> String { @@ -204,6 +202,37 @@ fn ignore_missing_request(result: std::result::Result<(), SequenceError>) -> Res } impl PrefillRouter { + /// 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_prefill_reservation(&self, reservation_id: &str) -> Result<()> { + if reservation_id.is_empty() { + anyhow::bail!("prefill reservation ID must not be empty"); + } + if self.lifecycle_state() != PrefillLifecycleState::Active { + return Err(anyhow::anyhow!(PrefillError::NotActivated)); + } + if self.prefill_router.get().is_none() { + return Err(anyhow::anyhow!(PrefillError::NotActivated)); + } + + match self.reservations.begin(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 a pending reservation when preprocessing fails before scheduler admission. + pub fn abort_prefill_reservation(&self, reservation_id: &str) { + self.reservations.remove_pending(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 @@ -223,18 +252,6 @@ impl PrefillRouter { if reservation_id.is_empty() { anyhow::bail!("prefill reservation ID must not be empty"); } - if self.lifecycle_state() != PrefillLifecycleState::Active { - return Err(anyhow::anyhow!(PrefillError::NotActivated)); - } - let inner = self - .prefill_router - .get() - .ok_or_else(|| anyhow::anyhow!(PrefillError::NotActivated))?; - let InnerPrefillRouter::KvRouter(router) = inner else { - return Err(anyhow::anyhow!(PrefillError::NotActivated)); - }; - let chooser = router.chooser.clone(); - let scheduler_id = scheduler_id(reservation_id); let cancellation = match self.reservations.begin(reservation_id) { BeginReservation::Pending(cancellation) => cancellation, BeginReservation::Cancelled => { @@ -244,6 +261,20 @@ impl PrefillRouter { anyhow::bail!("prefill reservation {reservation_id:?} already exists") } }; + if self.lifecycle_state() != PrefillLifecycleState::Active { + self.reservations.remove_pending(reservation_id); + return Err(anyhow::anyhow!(PrefillError::NotActivated)); + } + let Some(inner) = self.prefill_router.get() else { + self.reservations.remove_pending(reservation_id); + return Err(anyhow::anyhow!(PrefillError::NotActivated)); + }; + let InnerPrefillRouter::KvRouter(router) = inner else { + self.reservations.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; @@ -317,3 +348,113 @@ impl PrefillRouter { ignore_missing_request(reservation.chooser.free(&reservation.scheduler_id).await)?; self.reservations .remove_if_scheduler_id(reservation_id, &reservation.scheduler_id); + Ok(()) + } + + pub(super) fn spawn_reservation_reaper(router: &Arc) { + let router: Weak = Arc::downgrade(router); + let cancellation = router + .upgrade() + .expect("router must be alive while starting reservation reaper") + .cancel_token + .child_token(); + + 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 { + tokio::select! { + _ = cancellation.cancelled() => return, + _ = interval.tick() => { + let Some(router) = router.upgrade() else { + return; + }; + let now = Instant::now(); + router + .reservations + .remove_expired_cancellations(now, CANCELLED_RESERVATION_RETENTION); + let expired = router.reservations.expired_active_ids(now, retention); + for reservation_id in expired { + if let Err(error) = router.release_prefill_reservation(&reservation_id).await { + tracing::warn!( + %reservation_id, + %error, + "Failed to expire stale EPP prefill reservation" + ); + } + } + } + } + } + }); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn cancellation_before_reservation_is_observed() { + let registry = PrefillReservationRegistry::default(); + registry.cancel("reservation-1"); + + assert!(matches!( + registry.begin("reservation-1"), + BeginReservation::Cancelled + )); + assert!(matches!( + registry.begin("reservation-1"), + BeginReservation::Pending(_) + )); + } + + #[test] + fn pre_registered_pending_reservation_survives_tombstone_reaping() { + let registry = PrefillReservationRegistry::default(); + let BeginReservation::Pending(cancellation) = registry.begin("reservation-1") else { + panic!("expected pending reservation"); + }; + + registry.cancel("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("reservation-1") else { + panic!("expected pending reservation to remain registered"); + }; + assert!(existing.is_cancelled()); + } + + #[test] + fn cancellation_signals_pending_reservation() { + let registry = PrefillReservationRegistry::default(); + let BeginReservation::Pending(cancellation) = registry.begin("reservation-1") else { + panic!("expected pending reservation"); + }; + + registry.cancel("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) + ); + } +} From 260b305c689c72a46f10ede731c5a58c2ac443d8 Mon Sep 17 00:00:00 2001 From: Thomas Montfort Date: Tue, 11 Aug 2026 16:15:14 -0700 Subject: [PATCH 06/15] refactor(epp): scope reservations to C handle Signed-off-by: Thomas Montfort (cherry picked from commit 2ead078d43e686349a9e545ed58b1fc48de0d538) (cherry picked from commit 1b79a16846e02c230770d117a322e48711acc59b) Signed-off-by: Avinash Varma --- lib/bindings/c/src/lib.rs | 46 ++++--- .../kv_router/prefill_router/activation.rs | 3 - lib/llm/src/kv_router/prefill_router/mod.rs | 3 +- .../kv_router/prefill_router/reservations.rs | 126 ++++++++---------- 4 files changed, 89 insertions(+), 89 deletions(-) diff --git a/lib/bindings/c/src/lib.rs b/lib/bindings/c/src/lib.rs index 104f71a46798..80945eb344c4 100644 --- a/lib/bindings/c/src/lib.rs +++ b/lib/bindings/c/src/lib.rs @@ -29,7 +29,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; @@ -463,6 +463,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, @@ -493,8 +494,9 @@ impl RouterHandles { } let outcome = self - .prefill_router - .reserve_prefill_worker( + .epp_reservations + .reserve( + &self.prefill_router, reservation_id, tokens, block_mm_infos, @@ -871,7 +873,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, @@ -881,9 +887,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, @@ -1071,14 +1085,14 @@ pub unsafe extern "C" fn mark_prefill_complete( Err(_) => return QueryRouterResult::ErrInvalidParam, }; - let prefill_router = handles.prefill_router.clone(); + 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) = prefill_router.release_prefill_reservation(&request_id_str).await { + 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 { @@ -1132,14 +1146,14 @@ pub unsafe extern "C" fn free_request( Err(_) => return QueryRouterResult::ErrInvalidParam, }; - let prefill_router = handles.prefill_router.clone(); + 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) = prefill_router.release_prefill_reservation(&request_id_str).await { + 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 { @@ -1434,8 +1448,8 @@ pub unsafe extern "C" fn begin_prefill_reservation( }; let handles = unsafe { &*handle }; match handles - .prefill_router - .begin_prefill_reservation(reservation_id) + .epp_reservations + .begin(&handles.prefill_router, reservation_id) { Ok(()) => QueryRouterResult::Ok, Err(error) => { @@ -1478,8 +1492,8 @@ pub unsafe extern "C" fn route_prefill_request_with_reservation( }; let handles = unsafe { &*handle }; if let Err(error) = handles - .prefill_router - .begin_prefill_reservation(&reservation_id) + .epp_reservations + .begin(&handles.prefill_router, &reservation_id) { tracing::warn!(%reservation_id, %error, "Failed to begin EPP prefill reservation"); return QueryRouterResult::ErrQueryFailed; @@ -1488,9 +1502,7 @@ pub unsafe extern "C" fn route_prefill_request_with_reservation( match unsafe { preprocess_request(handles, request_json) } { Ok(values) => values, Err(code) => { - handles - .prefill_router - .abort_prefill_reservation(&reservation_id); + handles.epp_reservations.abort(&reservation_id); return code; } }; @@ -1557,9 +1569,7 @@ pub unsafe extern "C" fn cancel_prefill_reservation( _ => return QueryRouterResult::ErrInvalidParam, }; let handles = unsafe { &*handle }; - handles - .prefill_router - .cancel_prefill_reservation(reservation_id); + handles.epp_reservations.cancel(reservation_id); QueryRouterResult::Ok } diff --git a/lib/llm/src/kv_router/prefill_router/activation.rs b/lib/llm/src/kv_router/prefill_router/activation.rs index c523c7fd9964..85287fb83409 100644 --- a/lib/llm/src/kv_router/prefill_router/activation.rs +++ b/lib/llm/src/kv_router/prefill_router/activation.rs @@ -49,7 +49,6 @@ impl PrefillRouter { ) -> Arc { Arc::new(Self { prefill_router: std::sync::OnceLock::new(), - reservations: Default::default(), model_manager, endpoint_id: std::sync::OnceLock::new(), cancel_token: tokio_util::sync::CancellationToken::new(), @@ -86,7 +85,6 @@ impl PrefillRouter { let router = Arc::new(Self { prefill_router, - reservations: Default::default(), model_manager: model_manager.clone(), endpoint_id: std::sync::OnceLock::new(), cancel_token: cancel_token.clone(), @@ -102,7 +100,6 @@ impl PrefillRouter { activation_task_state: Arc::new(()), }); - Self::spawn_reservation_reaper(&router); // Spawn background task to wait for activation let router_weak = Arc::downgrade(&router); #[cfg(test)] diff --git a/lib/llm/src/kv_router/prefill_router/mod.rs b/lib/llm/src/kv_router/prefill_router/mod.rs index 9991a84564e5..fbdfd95410c4 100644 --- a/lib/llm/src/kv_router/prefill_router/mod.rs +++ b/lib/llm/src/kv_router/prefill_router/mod.rs @@ -37,6 +37,8 @@ mod admission; mod query; mod reservations; +pub use reservations::EppReservationManager; + use admission::InnerPrefillRouter; #[derive(Clone, Copy, Debug, PartialEq, Eq)] @@ -155,7 +157,6 @@ fn strip_terminal_disaggregated_params( /// - Normal: Worker IDs determined by router based on KV cache state pub struct PrefillRouter { prefill_router: OnceLock, - reservations: reservations::PrefillReservationRegistry, model_manager: Arc, endpoint_id: OnceLock, cancel_token: CancellationToken, diff --git a/lib/llm/src/kv_router/prefill_router/reservations.rs b/lib/llm/src/kv_router/prefill_router/reservations.rs index fd1bdde3b2d6..4723408565ab 100644 --- a/lib/llm/src/kv_router/prefill_router/reservations.rs +++ b/lib/llm/src/kv_router/prefill_router/reservations.rs @@ -61,11 +61,11 @@ enum BeginReservation { AlreadyExists, } -pub(super) struct PrefillReservationRegistry { +pub struct EppReservationManager { entries: Mutex>, } -impl Default for PrefillReservationRegistry { +impl Default for EppReservationManager { fn default() -> Self { Self { entries: Mutex::new(HashMap::new()), @@ -73,8 +73,8 @@ impl Default for PrefillReservationRegistry { } } -impl PrefillReservationRegistry { - fn begin(&self, reservation_id: &str) -> BeginReservation { +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) => { @@ -128,7 +128,7 @@ impl PrefillReservationRegistry { } } - fn cancel(&self, reservation_id: &str) { + fn cancel_entry(&self, reservation_id: &str) { let mut entries = self.entries.lock(); match entries.entry(reservation_id.to_string()) { Entry::Occupied(entry) => { @@ -201,23 +201,23 @@ fn ignore_missing_request(result: std::result::Result<(), SequenceError>) -> Res } } -impl PrefillRouter { +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_prefill_reservation(&self, reservation_id: &str) -> Result<()> { + 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 self.lifecycle_state() != PrefillLifecycleState::Active { + if router.lifecycle_state() != PrefillLifecycleState::Active { return Err(anyhow::anyhow!(PrefillError::NotActivated)); } - if self.prefill_router.get().is_none() { + if router.prefill_router.get().is_none() { return Err(anyhow::anyhow!(PrefillError::NotActivated)); } - match self.reservations.begin(reservation_id) { + match self.begin_entry(reservation_id) { BeginReservation::Pending(_) => Ok(()), BeginReservation::Cancelled => { anyhow::bail!("prefill reservation {reservation_id:?} was cancelled") @@ -229,8 +229,8 @@ impl PrefillRouter { } /// Drop a pending reservation when preprocessing fails before scheduler admission. - pub fn abort_prefill_reservation(&self, reservation_id: &str) { - self.reservations.remove_pending(reservation_id); + pub fn abort(&self, reservation_id: &str) { + self.remove_pending(reservation_id); } /// Atomically select and reserve a prefill worker for an externally dispatched request. @@ -238,8 +238,9 @@ impl PrefillRouter { /// 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_prefill_worker( + pub async fn reserve( &self, + router: &PrefillRouter, reservation_id: &str, token_ids: &[u32], block_mm_infos: Option<&[Option]>, @@ -252,7 +253,7 @@ impl PrefillRouter { if reservation_id.is_empty() { anyhow::bail!("prefill reservation ID must not be empty"); } - let cancellation = match self.reservations.begin(reservation_id) { + let cancellation = match self.begin_entry(reservation_id) { BeginReservation::Pending(cancellation) => cancellation, BeginReservation::Cancelled => { anyhow::bail!("prefill reservation {reservation_id:?} was cancelled") @@ -261,16 +262,16 @@ impl PrefillRouter { anyhow::bail!("prefill reservation {reservation_id:?} already exists") } }; - if self.lifecycle_state() != PrefillLifecycleState::Active { - self.reservations.remove_pending(reservation_id); + if router.lifecycle_state() != PrefillLifecycleState::Active { + self.remove_pending(reservation_id); return Err(anyhow::anyhow!(PrefillError::NotActivated)); } - let Some(inner) = self.prefill_router.get() else { - self.reservations.remove_pending(reservation_id); + 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.reservations.remove_pending(reservation_id); + self.remove_pending(reservation_id); return Err(anyhow::anyhow!(PrefillError::NotActivated)); }; let chooser = router.chooser.clone(); @@ -279,7 +280,7 @@ impl PrefillRouter { let outcome = tokio::select! { biased; _ = cancellation.cancelled() => { - self.reservations.remove_pending(reservation_id); + self.remove_pending(reservation_id); anyhow::bail!("prefill reservation {reservation_id:?} was cancelled") } outcome = chooser.find_best_match_details( @@ -301,7 +302,7 @@ impl PrefillRouter { let outcome = match outcome { Ok(outcome) => outcome, Err(error) => { - self.reservations.remove_pending(reservation_id); + self.remove_pending(reservation_id); return Err(error); } }; @@ -313,7 +314,7 @@ impl PrefillRouter { scheduler_id: scheduler_id.clone(), created_at: Instant::now(), }; - if !self.reservations.activate(reservation_id, reservation) { + if !self.activate(reservation_id, reservation) { ignore_missing_request(chooser.free(&scheduler_id).await)?; anyhow::bail!("prefill reservation {reservation_id:?} was cancelled"); } @@ -323,7 +324,7 @@ impl PrefillRouter { }) } crate::kv_router::FindBestMatchOutcome::QueueRejected { rejection } => { - self.reservations.remove_pending(reservation_id); + self.remove_pending(reservation_id); Ok(PrefillQueryOutcome::QueueRejected { rejection }) } } @@ -333,31 +334,28 @@ impl PrefillRouter { /// /// 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_prefill_reservation(&self, reservation_id: &str) { + pub fn cancel(&self, reservation_id: &str) { if !reservation_id.is_empty() { - self.reservations.cancel(reservation_id); + self.cancel_entry(reservation_id); } } /// Release a prefill reservation. Missing reservations are idempotent no-ops. - pub async fn release_prefill_reservation(&self, reservation_id: &str) -> Result<()> { - let Some(reservation) = self.reservations.get_active(reservation_id) else { + 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.reservations - .remove_if_scheduler_id(reservation_id, &reservation.scheduler_id); + self.remove_if_scheduler_id(reservation_id, &reservation.scheduler_id); Ok(()) } - pub(super) fn spawn_reservation_reaper(router: &Arc) { - let router: Weak = Arc::downgrade(router); - let cancellation = router - .upgrade() - .expect("router must be alive while starting reservation reaper") - .cancel_token - .child_token(); + /// 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(); @@ -368,26 +366,20 @@ impl PrefillRouter { interval.set_missed_tick_behavior(MissedTickBehavior::Skip); loop { - tokio::select! { - _ = cancellation.cancelled() => return, - _ = interval.tick() => { - let Some(router) = router.upgrade() else { - return; - }; - let now = Instant::now(); - router - .reservations - .remove_expired_cancellations(now, CANCELLED_RESERVATION_RETENTION); - let expired = router.reservations.expired_active_ids(now, retention); - for reservation_id in expired { - if let Err(error) = router.release_prefill_reservation(&reservation_id).await { - tracing::warn!( - %reservation_id, - %error, - "Failed to expire stale EPP prefill reservation" - ); - } - } + interval.tick().await; + let Some(manager) = manager.upgrade() else { + return; + }; + let now = Instant::now(); + manager.remove_expired_cancellations(now, CANCELLED_RESERVATION_RETENTION); + let expired = manager.expired_active_ids(now, retention); + for reservation_id in expired { + if let Err(error) = manager.release(&reservation_id).await { + tracing::warn!( + %reservation_id, + %error, + "Failed to expire stale EPP prefill reservation" + ); } } } @@ -401,34 +393,34 @@ mod tests { #[test] fn cancellation_before_reservation_is_observed() { - let registry = PrefillReservationRegistry::default(); - registry.cancel("reservation-1"); + let registry = EppReservationManager::default(); + registry.cancel_entry("reservation-1"); assert!(matches!( - registry.begin("reservation-1"), + registry.begin_entry("reservation-1"), BeginReservation::Cancelled )); assert!(matches!( - registry.begin("reservation-1"), + registry.begin_entry("reservation-1"), BeginReservation::Pending(_) )); } #[test] fn pre_registered_pending_reservation_survives_tombstone_reaping() { - let registry = PrefillReservationRegistry::default(); - let BeginReservation::Pending(cancellation) = registry.begin("reservation-1") else { + let registry = EppReservationManager::default(); + let BeginReservation::Pending(cancellation) = registry.begin_entry("reservation-1") else { panic!("expected pending reservation"); }; - registry.cancel("reservation-1"); + 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("reservation-1") else { + let BeginReservation::Pending(existing) = registry.begin_entry("reservation-1") else { panic!("expected pending reservation to remain registered"); }; assert!(existing.is_cancelled()); @@ -436,12 +428,12 @@ mod tests { #[test] fn cancellation_signals_pending_reservation() { - let registry = PrefillReservationRegistry::default(); - let BeginReservation::Pending(cancellation) = registry.begin("reservation-1") else { + let registry = EppReservationManager::default(); + let BeginReservation::Pending(cancellation) = registry.begin_entry("reservation-1") else { panic!("expected pending reservation"); }; - registry.cancel("reservation-1"); + registry.cancel_entry("reservation-1"); assert!(cancellation.is_cancelled()); registry.remove_pending("reservation-1"); } From 07a59fce4ec04d0c65a061249dba05109c7fe230 Mon Sep 17 00:00:00 2001 From: Thomas Montfort Date: Tue, 11 Aug 2026 17:57:22 -0700 Subject: [PATCH 07/15] fix(epp): harden booking cleanup Signed-off-by: Thomas Montfort (cherry picked from commit cab8cacca56debc048831709a227ddcec2888756) (cherry picked from commit 17dabaa8a93ae041b24ab07e3f557ce5e872b2bb) Signed-off-by: Avinash Varma --- .../epp/pkg/plugins/disagg/decode_scorer.go | 132 +++--------- .../epp/pkg/plugins/disagg/prefill_scorer.go | 11 +- .../pkg/plugins/disagg/reservation_test.go | 202 +++++++++++++++--- .../epp/pkg/plugins/disagg/shared.go | 143 +++++++++---- .../pkg/plugins/dynamo_kv_scorer/plugin.go | 30 +++ lib/bindings/c/src/lib.rs | 48 +++++ 6 files changed, 397 insertions(+), 169 deletions(-) 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 80d730468931..98ffdddf0390 100644 --- a/deploy/inference-gateway/epp/pkg/plugins/disagg/decode_scorer.go +++ b/deploy/inference-gateway/epp/pkg/plugins/disagg/decode_scorer.go @@ -41,45 +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 { - BookingID string - WorkerID string - DpRank uint32 - PrefillWorkerID string - TokenData []int64 - CacheNamespace string -} - -// Clone implements plugins.StateData. -func (s *DecodeRoutingState) Clone() plugins.StateData { - if s == nil { - return nil - } - clone := &DecodeRoutingState{ - BookingID: s.BookingID, - WorkerID: s.WorkerID, - DpRank: s.DpRank, - PrefillWorkerID: s.PrefillWorkerID, - CacheNamespace: s.CacheNamespace, - } - 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{} @@ -101,10 +69,10 @@ 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), + routeDecode: dynscorer.CallRouteDecodeRequest, addRequest: dynscorer.CallAddRequest, markPrefillComplete: dynscorer.CallMarkPrefillComplete, freeBooking: dynscorer.CallFreeRequest, @@ -114,7 +82,7 @@ func NewDynDecodeScorer(ctx context.Context) *DynDecodeScorer { // DynDecodeScorer is a scorer plugin for the decode scheduling profile. type DynDecodeScorer struct { typedName plugins.TypedName - pluginState *plugins.PluginState + routeDecode func(string, string, bool) (*dynscorer.RoutingResult, error) addRequest func(string, []int64, uint64, uint32, string) error markPrefillComplete func(string) error freeBooking func(string) error @@ -169,7 +137,7 @@ 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") @@ -183,6 +151,28 @@ func (s *DynDecodeScorer) Score(ctx context.Context, cycleState *schedtypes.Cycl 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, result.CacheNamespace) + 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, @@ -200,16 +190,6 @@ func (s *DynDecodeScorer) Score(ctx context.Context, cycleState *schedtypes.Cycl delete(req.Headers, PrefillDpRankHeader) } - // Store routing state for PreRequest bookkeeping, keyed by booking ID. - routingState := &DecodeRoutingState{ - BookingID: booking.ID, - WorkerID: workerIDStr, - DpRank: result.DpRank, - TokenData: result.TokenData, - CacheNamespace: result.CacheNamespace, - } - s.pluginState.Write(booking.ID, plugins.StateKey(decodeStateKey), routingState) - // Inject pre-computed tokens into the request body so the frontend // sidecar can skip redundant tokenization. setTokenizedPrompt(req, result.TokenData, logger) @@ -241,7 +221,10 @@ func (s *DynDecodeScorer) rollbackPrefillReservation( booking *BookingState, reason string, ) { - if booking != nil && booking.PrefillReserved && s.cleanupBooking(ctx, booking.ID, reason) { + if booking != nil { + if findBookingLifecycle(booking.ID) != nil || booking.PrefillReserved { + s.cleanupBooking(ctx, booking.ID, reason) + } booking.PrefillReserved = false cycleState.Write(BookingStateKey, booking) } @@ -251,64 +234,13 @@ func (s *DynDecodeScorer) rollbackPrefillReservation( 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) } } -// 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) - bookingID := bookingIDFromRequest(request) - if bookingID == "" { - logger.V(logutil.DEBUG).Info("DynDecodeScorer PreRequest: no controller booking ID, skipping") - return - } - if err := ctx.Err(); err != nil { - s.cleanupBooking(ctx, bookingID, "request cancelled before decode booking") - return - } - - state, err := plugins.ReadPluginStateKey[*DecodeRoutingState]( - s.pluginState, bookingID, plugins.StateKey(decodeStateKey), - ) - s.pluginState.Delete(bookingID) - if err != nil || state == nil || state.BookingID != bookingID { - logger.V(logutil.DEBUG).Info("DynDecodeScorer PreRequest: no routing state found", - "bookingID", bookingID) - s.cleanupBooking(ctx, bookingID, "decode routing state missing") - return - } - - var workerIDUint uint64 - if _, parseErr := fmt.Sscanf(state.WorkerID, "%d", &workerIDUint); parseErr != nil { - logger.V(logutil.DEFAULT).Error(parseErr, "DynDecodeScorer PreRequest: invalid worker ID", - "bookingID", bookingID, "workerID", state.WorkerID) - s.cleanupBooking(ctx, bookingID, "decode worker ID invalid") - return - } - - if addErr := s.addRequest( - bookingID, - state.TokenData, - workerIDUint, - state.DpRank, - state.CacheNamespace, - ); addErr != nil { - logger.V(logutil.DEFAULT).Error(addErr, "DynDecodeScorer PreRequest: failed to add request", - "bookingID", bookingID) - s.cleanupBooking(ctx, bookingID, "decode booking failed") - return - } - - logger.V(logutil.VERBOSE).Info("DynDecodeScorer PreRequest: registered request", - "bookingID", bookingID, - "workerID", state.WorkerID, - "dpRank", state.DpRank, - "hasCacheNamespace", state.CacheNamespace != "", - "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) { 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 3f3a58e8e520..845cd574ee80 100644 --- a/deploy/inference-gateway/epp/pkg/plugins/disagg/prefill_scorer.go +++ b/deploy/inference-gateway/epp/pkg/plugins/disagg/prefill_scorer.go @@ -89,6 +89,7 @@ func newDynPrefillScorer(cfg DynPrefillScorerConfig) *DynPrefillScorer { beginPrefill: dynscorer.CallBeginPrefillReservation, reservePrefill: dynscorer.CallRoutePrefillRequestWithReservation, cancelPrefill: dynscorer.CallCancelPrefillReservation, + releasePrefill: dynscorer.CallReleasePrefillReservation, freeBooking: dynscorer.CallFreeRequest, reservationAdmissionTimeout: reservationTimeout, reservationSlots: make(chan struct{}, maxReservations), @@ -101,6 +102,7 @@ type DynPrefillScorer struct { 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{} @@ -137,6 +139,13 @@ func (s *DynPrefillScorer) beginReservation(bookingID string) error { 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. func (s *DynPrefillScorer) TypedName() plugins.TypedName { return s.typedName @@ -226,7 +235,7 @@ func (s *DynPrefillScorer) Score(ctx context.Context, cycleState *schedtypes.Cyc if outcome.err != nil { return } - if cleanupErr := s.freeBooking(bookingID); cleanupErr != nil { + if cleanupErr := s.releaseLatePrefillReservation(bookingID); cleanupErr != nil { logger.V(logutil.DEFAULT).Error(cleanupErr, "DynPrefillScorer: failed to release late prefill reservation", "bookingID", bookingID) } diff --git a/deploy/inference-gateway/epp/pkg/plugins/disagg/reservation_test.go b/deploy/inference-gateway/epp/pkg/plugins/disagg/reservation_test.go index 8ca6e4470daa..f7dbcfc96abf 100644 --- a/deploy/inference-gateway/epp/pkg/plugins/disagg/reservation_test.go +++ b/deploy/inference-gateway/epp/pkg/plugins/disagg/reservation_test.go @@ -157,28 +157,6 @@ func TestResponseBodyFreesWithoutMarkingEmptyTerminalResponse(t *testing.T) { } } -func TestPreRequestCancellationCleansBooking(t *testing.T) { - bookingID := ensureBookingState(schedtypes.NewCycleState()).ID - freeCalls := 0 - scorer := &DynDecodeScorer{ - freeBooking: func(got string) error { - if got != bookingID { - t.Fatalf("free booking ID = %q, want %q", got, bookingID) - } - freeCalls++ - return nil - }, - } - ctx, cancel := context.WithCancel(context.Background()) - cancel() - - scorer.PreRequest(ctx, requestWithBooking("external-request", bookingID), nil) - - if freeCalls != 1 { - t.Fatalf("free calls = %d, want 1", freeCalls) - } -} - func TestPrefillScoreCancelsPendingReservation(t *testing.T) { ctx, cancel := context.WithCancel(context.Background()) defer cancel() @@ -186,7 +164,8 @@ func TestPrefillScoreCancelsPendingReservation(t *testing.T) { reserveStarted := make(chan struct{}) allowReserveReturn := make(chan struct{}) cancelCalls := make(chan string, 1) - freeCalls := make(chan string, 1) + releaseCalls := make(chan string, 1) + combinedFreeCalls := 0 scorer := &DynPrefillScorer{ reservePrefill: func(string, string, string) (*dynscorer.RoutingResult, error) { close(reserveStarted) @@ -198,8 +177,12 @@ func TestPrefillScoreCancelsPendingReservation(t *testing.T) { close(allowReserveReturn) return nil }, - freeBooking: func(bookingID string) error { - freeCalls <- bookingID + releasePrefill: func(bookingID string) error { + releaseCalls <- bookingID + return nil + }, + freeBooking: func(string) error { + combinedFreeCalls++ return nil }, } @@ -228,13 +211,16 @@ func TestPrefillScoreCancelsPendingReservation(t *testing.T) { t.Fatalf("cancel booking ID = %q, want %q", got, bookingID) } select { - case got := <-freeCalls: + 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") } @@ -400,3 +386,167 @@ func TestPrefillScoreBoundsConcurrentReservations(t *testing.T) { 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") + } +} diff --git a/deploy/inference-gateway/epp/pkg/plugins/disagg/shared.go b/deploy/inference-gateway/epp/pkg/plugins/disagg/shared.go index 4d33eddb50eb..87fa1ecc84f7 100644 --- a/deploy/inference-gateway/epp/pkg/plugins/disagg/shared.go +++ b/deploy/inference-gateway/epp/pkg/plugins/disagg/shared.go @@ -217,22 +217,31 @@ func getEnvBoolOrDefault(key string, def bool) bool { var enforceDisaggDeprecationOnce sync.Once const ( - prefillMarkMaxAttempts = 3 - prefillMarkRetryBackoff = 100 * time.Millisecond + prefillMarkMaxAttempts = 3 + prefillMarkRetryBackoff = 100 * time.Millisecond + cleanupMaxAttempts = 3 + cleanupRetryBackoff = 100 * time.Millisecond + maxConcurrentBookingCleanups = 32 ) -var bookingLifecycles sync.Map +var ( + bookingLifecycles sync.Map + bookingCleanupSlots = make(chan struct{}, maxConcurrentBookingCleanups) +) // bookingLifecycle owns cleanup for one EPP booking across the prefill scorer, -// decode scorer, and response callbacks. A booking has exactly one free_request -// caller even when EOS and context cancellation race. +// 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 - cleanupOnce sync.Once - mu sync.Mutex - cleaned bool + mu sync.Mutex + cleanupStarted bool + cleanupSucceeded bool + cleanupExhausted bool + decodeRegistrationDone chan struct{} + decodeRegistrationOpen bool stopCancellation func() bool stopMarker context.CancelFunc @@ -260,10 +269,36 @@ func findBookingLifecycle(bookingID string) *bookingLifecycle { 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.cleaned || l.stopCancellation != nil { + if l.cleanupStarted || l.stopCancellation != nil { return } l.stopCancellation = context.AfterFunc(ctx, func() { @@ -275,7 +310,7 @@ func (l *bookingLifecycle) armCancellation(ctx context.Context) { // the response callback. EOS or request cancellation stops any pending retry. func (l *bookingLifecycle) startPrefillMarker(markPrefillComplete func(string) error, logger logr.Logger, requestID string) { l.mu.Lock() - if l.cleaned || l.markerDone != nil { + if l.cleanupStarted || l.markerDone != nil { l.mu.Unlock() return } @@ -319,46 +354,70 @@ func (l *bookingLifecycle) startPrefillMarker(markPrefillComplete func(string) e }() } -// cleanup stops cancellation and retry ownership, then releases the booking -// off the response callback. If a mark call is already in flight, cleanup -// waits for that one bounded call before issuing free_request, preventing -// concurrent mark/free FFI operations for the same booking. +// cleanup stops cancellation and marker ownership, waits for in-flight work, +// and retries free_request on a bounded executor. A terminal failure remains in +// bookingLifecycles as a tombstone so a duplicate cleanup cannot hide it. func (l *bookingLifecycle) cleanup(ctx context.Context, reason string) bool { - started := false - l.cleanupOnce.Do(func() { - started = true - l.mu.Lock() - l.cleaned = true - stopCancellation := l.stopCancellation - stopMarker := l.stopMarker - markerDone := l.markerDone - cleanupDone := make(chan struct{}) - l.cleanupDone = cleanupDone + l.mu.Lock() + if l.cleanupStarted { l.mu.Unlock() + return false + } + l.cleanupStarted = true + stopCancellation := l.stopCancellation + stopMarker := l.stopMarker + markerDone := l.markerDone + decodeRegistrationDone := l.decodeRegistrationDone + cleanupDone := make(chan struct{}) + l.cleanupDone = cleanupDone + l.mu.Unlock() - if stopCancellation != nil { - stopCancellation() + if stopCancellation != nil { + stopCancellation() + } + if stopMarker != nil { + stopMarker() + } + + go func() { + if markerDone != nil { + <-markerDone } - if stopMarker != nil { - stopMarker() + if decodeRegistrationDone != nil { + <-decodeRegistrationDone } - go func() { - if markerDone != nil { - <-markerDone + logger := log.FromContext(ctx) + for attempt := 1; attempt <= cleanupMaxAttempts; attempt++ { + bookingCleanupSlots <- struct{}{} + err := l.freeBooking(l.bookingID) + <-bookingCleanupSlots + if err == nil { + l.mu.Lock() + l.cleanupSucceeded = true + l.mu.Unlock() + logger.V(logutil.VERBOSE).Info("Dynamo EPP booking cleaned up", + "bookingID", l.bookingID, "reason", reason, "attempt", attempt) + close(cleanupDone) + bookingLifecycles.Delete(l.bookingID) + return } - if err := l.freeBooking(l.bookingID); err != nil { - log.FromContext(ctx).V(logutil.DEFAULT).Error(err, "Dynamo EPP booking cleanup failed", - "bookingID", l.bookingID, "reason", reason) - } else { - log.FromContext(ctx).V(logutil.VERBOSE).Info("Dynamo EPP booking cleaned up", - "bookingID", l.bookingID, "reason", reason) + + logger.V(logutil.DEFAULT).Error(err, "Dynamo EPP booking cleanup failed", + "bookingID", l.bookingID, "reason", reason, "attempt", attempt) + if attempt == cleanupMaxAttempts { + l.mu.Lock() + l.cleanupExhausted = true + l.mu.Unlock() + logger.V(logutil.DEFAULT).Error(err, "Dynamo EPP booking cleanup exhausted retries; retaining tombstone", + "bookingID", l.bookingID, "reason", reason, "attempts", cleanupMaxAttempts) + close(cleanupDone) + return } - close(cleanupDone) - bookingLifecycles.Delete(l.bookingID) - }() - }) - return started + time.Sleep(cleanupRetryBackoff * time.Duration(attempt)) + } + }() + return true } func (l *bookingLifecycle) markerComplete() <-chan struct{} { 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 4fa12d0acaac..5384823850b5 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 @@ -84,6 +84,9 @@ query_router_result_t route_prefill_request_with_reservation(RouterHandles *hand 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, @@ -559,6 +562,33 @@ func CallCancelPrefillReservation(reservationID string) error { 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). diff --git a/lib/bindings/c/src/lib.rs b/lib/bindings/c/src/lib.rs index 80945eb344c4..977c6e9cfdf7 100644 --- a/lib/bindings/c/src/lib.rs +++ b/lib/bindings/c/src/lib.rs @@ -1124,6 +1124,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. From f0e647e7500a4249574597baaae857cf6821a4d7 Mon Sep 17 00:00:00 2001 From: Thomas Montfort Date: Tue, 11 Aug 2026 18:20:19 -0700 Subject: [PATCH 08/15] fix(epp): harden reservation ownership Signed-off-by: Thomas Montfort (cherry picked from commit c32ce78b15f03085c5d7dd6d8e1cd171f436e9ba) (cherry picked from commit 0ce6273372cf0490615d6b79ab1cb983880fef97) Signed-off-by: Avinash Varma --- .../epp/pkg/plugins/disagg/decode_scorer.go | 7 +- .../pkg/plugins/disagg/reservation_test.go | 84 +++++- .../epp/pkg/plugins/disagg/shared.go | 2 +- .../kv_router/prefill_router/reservations.rs | 244 ++++++++++++++++-- 4 files changed, 303 insertions(+), 34 deletions(-) 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 98ffdddf0390..fde9bff5935b 100644 --- a/deploy/inference-gateway/epp/pkg/plugins/disagg/decode_scorer.go +++ b/deploy/inference-gateway/epp/pkg/plugins/disagg/decode_scorer.go @@ -251,12 +251,17 @@ func (s *DynDecodeScorer) ResponseBody(ctx context.Context, request *schedtypes. lifecycle := findBookingLifecycle(bookingID) if lifecycle == nil { - lifecycle = registerBookingLifecycle(bookingID, s.freeBooking) + // Only Score creates controller-owned booking lifecycles. Do not let an + // inbound header create router bookkeeping on a response-only path. + return } 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/reservation_test.go b/deploy/inference-gateway/epp/pkg/plugins/disagg/reservation_test.go index f7dbcfc96abf..6c748eb0df90 100644 --- a/deploy/inference-gateway/epp/pkg/plugins/disagg/reservation_test.go +++ b/deploy/inference-gateway/epp/pkg/plugins/disagg/reservation_test.go @@ -68,6 +68,7 @@ func TestBookingStateDoesNotTrustExternalRequestID(t *testing.T) { 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{ @@ -89,6 +90,7 @@ func TestResponseBodyRetriesPrefillMarkAndFreesTerminalResponse(t *testing.T) { return nil }, } + registerBookingLifecycle(bookingID, scorer.freeBooking) scorer.ResponseBody(context.Background(), request, &rc.Response{}, nil) scorer.ResponseBody(context.Background(), request, &rc.Response{}, nil) @@ -132,6 +134,7 @@ func TestResponseBodyFreesWithoutMarkingEmptyTerminalResponse(t *testing.T) { return nil }, } + registerBookingLifecycle(bookingID, scorer.freeBooking) scorer.ResponseBody( context.Background(), @@ -241,6 +244,8 @@ func TestResponseBodyBoundsPersistentPrefillMarkRetries(t *testing.T) { }, } request := requestWithBooking("external-request", bookingID) + request.Headers[RoutingModeHeader] = "disaggregated" + registerBookingLifecycle(bookingID, scorer.freeBooking) for range 10 { scorer.ResponseBody(context.Background(), request, &rc.Response{}, nil) @@ -286,6 +291,8 @@ func TestResponseBodyEOSDoesNotWaitForInFlightPrefillMark(t *testing.T) { }, } 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) @@ -293,10 +300,16 @@ func TestResponseBodyEOSDoesNotWaitForInFlightPrefillMark(t *testing.T) { t.Fatal("expected booking lifecycle after first response chunk") } - started := time.Now() - scorer.ResponseBody(context.Background(), request, &rc.Response{EndOfStream: true}, nil) - if elapsed := time.Since(started); elapsed > 100*time.Millisecond { - t.Fatalf("EOS callback blocked for %s waiting on prefill mark", elapsed) + 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 { @@ -550,3 +563,66 @@ func TestBookingLifecycleRetainsTombstoneAfterCleanupExhaustion(t *testing.T) { 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 87fa1ecc84f7..0014482abdef 100644 --- a/deploy/inference-gateway/epp/pkg/plugins/disagg/shared.go +++ b/deploy/inference-gateway/epp/pkg/plugins/disagg/shared.go @@ -398,8 +398,8 @@ func (l *bookingLifecycle) cleanup(ctx context.Context, reason string) bool { l.mu.Unlock() logger.V(logutil.VERBOSE).Info("Dynamo EPP booking cleaned up", "bookingID", l.bookingID, "reason", reason, "attempt", attempt) - close(cleanupDone) bookingLifecycles.Delete(l.bookingID) + close(cleanupDone) return } diff --git a/lib/llm/src/kv_router/prefill_router/reservations.rs b/lib/llm/src/kv_router/prefill_router/reservations.rs index 4723408565ab..5519350bbef5 100644 --- a/lib/llm/src/kv_router/prefill_router/reservations.rs +++ b/lib/llm/src/kv_router/prefill_router/reservations.rs @@ -23,6 +23,7 @@ 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); struct ActivePrefillReservation { chooser: Arc, @@ -46,6 +47,8 @@ enum PrefillReservationEntry { /// prevents scheduler booking. Pending { cancellation: CancellationToken, + created_at: Instant, + claimed: bool, }, Active(ActivePrefillReservation), /// Preserve a cancellation that reaches Rust before the blocking reserve @@ -61,6 +64,12 @@ enum BeginReservation { AlreadyExists, } +enum Activation { + Active, + Cancelled, + Lost, +} + pub struct EppReservationManager { entries: Mutex>, } @@ -81,6 +90,8 @@ impl EppReservationManager { let cancellation = CancellationToken::new(); entry.insert(PrefillReservationEntry::Pending { cancellation: cancellation.clone(), + created_at: Instant::now(), + claimed: false, }); BeginReservation::Pending(cancellation) } @@ -89,7 +100,7 @@ impl EppReservationManager { entry.remove(); BeginReservation::Cancelled } - PrefillReservationEntry::Pending { cancellation } => { + PrefillReservationEntry::Pending { cancellation, .. } => { BeginReservation::Pending(cancellation.clone()) } PrefillReservationEntry::Active(_) => BeginReservation::AlreadyExists, @@ -97,42 +108,123 @@ impl EppReservationManager { } } - fn activate(&self, reservation_id: &str, reservation: ActivePrefillReservation) -> bool { + // 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 false; + return Activation::Lost; }; - let PrefillReservationEntry::Pending { cancellation } = entry else { - return false; + 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; + } }; - if cancellation.is_cancelled() { - entries.remove(reservation_id); - return false; - } entries.insert( reservation_id.to_string(), PrefillReservationEntry::Active(reservation), ); - true + 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 { .. }) + 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() { + if let PrefillReservationEntry::Pending { cancellation, .. } = entry.get() { cancellation.cancel(); } } @@ -144,6 +236,26 @@ impl EppReservationManager { } } + 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()), @@ -228,9 +340,9 @@ impl EppReservationManager { } } - /// Drop a pending reservation when preprocessing fails before scheduler admission. + /// Drop an unclaimed pending reservation when preprocessing fails before scheduler admission. pub fn abort(&self, reservation_id: &str) { - self.remove_pending(reservation_id); + self.abort_unclaimed(reservation_id); } /// Atomically select and reserve a prefill worker for an externally dispatched request. @@ -253,7 +365,7 @@ impl EppReservationManager { if reservation_id.is_empty() { anyhow::bail!("prefill reservation ID must not be empty"); } - let cancellation = match self.begin_entry(reservation_id) { + let cancellation = match self.claim_pending(reservation_id) { BeginReservation::Pending(cancellation) => cancellation, BeginReservation::Cancelled => { anyhow::bail!("prefill reservation {reservation_id:?} was cancelled") @@ -314,14 +426,34 @@ impl EppReservationManager { scheduler_id: scheduler_id.clone(), created_at: Instant::now(), }; - if !self.activate(reservation_id, reservation) { - ignore_missing_request(chooser.free(&scheduler_id).await)?; - anyhow::bail!("prefill reservation {reservation_id:?} was cancelled"); + 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"); + } } - Ok(PrefillQueryOutcome::Routed { - worker_id: worker.worker_id, - dp_rank: Some(worker.dp_rank), - }) } crate::kv_router::FindBestMatchOutcome::QueueRejected { rejection } => { self.remove_pending(reservation_id); @@ -371,15 +503,31 @@ impl EppReservationManager { return; }; let now = Instant::now(); + manager.expire_pending(now, retention); manager.remove_expired_cancellations(now, CANCELLED_RESERVATION_RETENTION); let expired = manager.expired_active_ids(now, retention); for reservation_id in expired { - if let Err(error) = manager.release(&reservation_id).await { - tracing::warn!( - %reservation_id, - %error, - "Failed to expire stale EPP prefill reservation" - ); + 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" + ); + } } } } @@ -426,6 +574,46 @@ mod tests { 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(); From 054e05302c613df0ab20dc8f154c5ad26af2f762 Mon Sep 17 00:00:00 2001 From: Thomas Montfort Date: Wed, 12 Aug 2026 10:39:41 -0700 Subject: [PATCH 09/15] fix(epp): bound booking cleanup executor Signed-off-by: Thomas Montfort (cherry picked from commit fa3e7bcbc05a8a3cb4c0c5d0c6969a98bcd0550c) (cherry picked from commit f891b8214e5fbdd49efebb81e318614d62c17db6) Signed-off-by: Avinash Varma --- .../pkg/plugins/disagg/booking_executor.go | 390 ++++++++++++++++++ .../plugins/disagg/booking_executor_test.go | 271 ++++++++++++ .../epp/pkg/plugins/disagg/shared.go | 139 +++---- 3 files changed, 714 insertions(+), 86 deletions(-) create mode 100644 deploy/inference-gateway/epp/pkg/plugins/disagg/booking_executor.go create mode 100644 deploy/inference-gateway/epp/pkg/plugins/disagg/booking_executor_test.go 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/shared.go b/deploy/inference-gateway/epp/pkg/plugins/disagg/shared.go index 0014482abdef..2a5bf1f8d973 100644 --- a/deploy/inference-gateway/epp/pkg/plugins/disagg/shared.go +++ b/deploy/inference-gateway/epp/pkg/plugins/disagg/shared.go @@ -217,17 +217,13 @@ func getEnvBoolOrDefault(key string, def bool) bool { var enforceDisaggDeprecationOnce sync.Once const ( - prefillMarkMaxAttempts = 3 - prefillMarkRetryBackoff = 100 * time.Millisecond - cleanupMaxAttempts = 3 - cleanupRetryBackoff = 100 * time.Millisecond - maxConcurrentBookingCleanups = 32 + prefillMarkMaxAttempts = 3 + prefillMarkRetryBackoff = 100 * time.Millisecond + cleanupMaxAttempts = 3 + cleanupRetryBackoff = 100 * time.Millisecond ) -var ( - bookingLifecycles sync.Map - bookingCleanupSlots = make(chan struct{}, maxConcurrentBookingCleanups) -) +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 @@ -235,11 +231,20 @@ var ( 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 @@ -250,9 +255,14 @@ type bookingLifecycle 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 { @@ -306,8 +316,9 @@ func (l *bookingLifecycle) armCancellation(ctx context.Context) { }) } -// startPrefillMarker makes first-token bookkeeping bounded and independent of -// the response callback. EOS or request cancellation stops any pending retry. +// 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 { @@ -320,43 +331,24 @@ func (l *bookingLifecycle) startPrefillMarker(markPrefillComplete func(string) e l.markerDone = markerDone l.mu.Unlock() - go func() { - defer close(markerDone) - for attempt := 1; attempt <= prefillMarkMaxAttempts; attempt++ { - if markerCtx.Err() != nil { - return - } - if err := markPrefillComplete(l.bookingID); err == nil { - logger.V(logutil.VERBOSE).Info("DynDecodeScorer ResponseBody: marked prefill complete", - "bookingID", l.bookingID, "requestID", requestID, "attempt", attempt) - return - } else { - logger.V(logutil.DEFAULT).Error(err, "DynDecodeScorer ResponseBody: failed to mark prefill complete", - "bookingID", l.bookingID, "requestID", requestID, "attempt", attempt) - } - if attempt == prefillMarkMaxAttempts { - return - } - - timer := time.NewTimer(prefillMarkRetryBackoff * time.Duration(attempt)) - select { - case <-markerCtx.Done(): - if !timer.Stop() { - select { - case <-timer.C: - default: - } - } - return - case <-timer.C: - } - } - }() + 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 stops cancellation and marker ownership, waits for in-flight work, -// and retries free_request on a bounded executor. A terminal failure remains in -// bookingLifecycles as a tombstone so a duplicate cleanup cannot hide it. +// 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 { @@ -364,10 +356,11 @@ func (l *bookingLifecycle) cleanup(ctx context.Context, reason string) bool { return false } l.cleanupStarted = true + l.cleanupStartedAt = time.Now() + l.cleanupReason = reason + l.cleanupLogger = log.FromContext(ctx) stopCancellation := l.stopCancellation stopMarker := l.stopMarker - markerDone := l.markerDone - decodeRegistrationDone := l.decodeRegistrationDone cleanupDone := make(chan struct{}) l.cleanupDone = cleanupDone l.mu.Unlock() @@ -379,47 +372,21 @@ func (l *bookingLifecycle) cleanup(ctx context.Context, reason string) bool { stopMarker() } - go func() { - if markerDone != nil { - <-markerDone - } - if decodeRegistrationDone != nil { - <-decodeRegistrationDone - } - - logger := log.FromContext(ctx) - for attempt := 1; attempt <= cleanupMaxAttempts; attempt++ { - bookingCleanupSlots <- struct{}{} - err := l.freeBooking(l.bookingID) - <-bookingCleanupSlots - if err == nil { - l.mu.Lock() - l.cleanupSucceeded = true - l.mu.Unlock() - logger.V(logutil.VERBOSE).Info("Dynamo EPP booking cleaned up", - "bookingID", l.bookingID, "reason", reason, "attempt", attempt) - bookingLifecycles.Delete(l.bookingID) - close(cleanupDone) - return - } - - logger.V(logutil.DEFAULT).Error(err, "Dynamo EPP booking cleanup failed", - "bookingID", l.bookingID, "reason", reason, "attempt", attempt) - if attempt == cleanupMaxAttempts { - l.mu.Lock() - l.cleanupExhausted = true - l.mu.Unlock() - logger.V(logutil.DEFAULT).Error(err, "Dynamo EPP booking cleanup exhausted retries; retaining tombstone", - "bookingID", l.bookingID, "reason", reason, "attempts", cleanupMaxAttempts) - close(cleanupDone) - return - } - time.Sleep(cleanupRetryBackoff * time.Duration(attempt)) - } - }() + 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() From f2b062dd0ff0899fd37bf4cc0b40962d75c7ff22 Mon Sep 17 00:00:00 2001 From: Thomas Montfort Date: Wed, 12 Aug 2026 11:02:22 -0700 Subject: [PATCH 10/15] fix(epp): release late reservation errors Signed-off-by: Thomas Montfort (cherry picked from commit 334324294cd807a05dcf62fb58538f755d483af9) (cherry picked from commit 001d66692efe6be4dcf06d7541cdecb18f788f52) Signed-off-by: Avinash Varma --- .../epp/pkg/plugins/disagg/prefill_scorer.go | 8 +-- .../pkg/plugins/disagg/reservation_test.go | 51 +++++++++++++++++++ 2 files changed, 55 insertions(+), 4 deletions(-) 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 845cd574ee80..0629086746f7 100644 --- a/deploy/inference-gateway/epp/pkg/plugins/disagg/prefill_scorer.go +++ b/deploy/inference-gateway/epp/pkg/plugins/disagg/prefill_scorer.go @@ -231,10 +231,10 @@ func (s *DynPrefillScorer) Score(ctx context.Context, cycleState *schedtypes.Cyc "bookingID", bookingID) } go func(bookingID string) { - outcome := <-resultCh - if outcome.err != nil { - return - } + // 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) diff --git a/deploy/inference-gateway/epp/pkg/plugins/disagg/reservation_test.go b/deploy/inference-gateway/epp/pkg/plugins/disagg/reservation_test.go index 6c748eb0df90..81f41ac0fde1 100644 --- a/deploy/inference-gateway/epp/pkg/plugins/disagg/reservation_test.go +++ b/deploy/inference-gateway/epp/pkg/plugins/disagg/reservation_test.go @@ -229,6 +229,57 @@ func TestPrefillScoreCancelsPendingReservation(t *testing.T) { } } +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 From 08e3fc967e2cff3c721d7502578066bd949535a0 Mon Sep 17 00:00:00 2001 From: Thomas Montfort Date: Wed, 12 Aug 2026 11:02:29 -0700 Subject: [PATCH 11/15] fix(router): bound stale reservation reaping Signed-off-by: Thomas Montfort (cherry picked from commit 7b90075122d1ef48f9b678e477c9f311d94e4e85) (cherry picked from commit 530aec1d100f50e99f03933292852cec6eccb115) Signed-off-by: Avinash Varma --- .../kv_router/prefill_router/reservations.rs | 112 +++++++++++++----- 1 file changed, 80 insertions(+), 32 deletions(-) diff --git a/lib/llm/src/kv_router/prefill_router/reservations.rs b/lib/llm/src/kv_router/prefill_router/reservations.rs index 5519350bbef5..c1999c8bd9e8 100644 --- a/lib/llm/src/kv_router/prefill_router/reservations.rs +++ b/lib/llm/src/kv_router/prefill_router/reservations.rs @@ -9,6 +9,7 @@ use std::{ 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; @@ -24,11 +25,15 @@ 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 { @@ -37,6 +42,7 @@ impl Clone for ActivePrefillReservation { chooser: self.chooser.clone(), scheduler_id: self.scheduler_id.clone(), created_at: self.created_at, + reap_attempts: self.reap_attempts, } } } @@ -274,15 +280,46 @@ impl EppReservationManager { } } - fn expired_active_ids(&self, now: Instant, retention: Duration) -> Vec { - self.entries - .lock() - .iter() - .filter(|(_, entry)| { - matches!(entry, PrefillReservationEntry::Active(active) - if now.saturating_duration_since(active.created_at) >= retention) + // 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, }) - .map(|(reservation_id, _)| reservation_id.clone()) + .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() } @@ -425,6 +462,7 @@ impl EppReservationManager { 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 { @@ -483,30 +521,11 @@ impl EppReservationManager { Ok(()) } - /// 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.expired_active_ids(now, retention); - for reservation_id in expired { + 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), @@ -530,6 +549,35 @@ impl EppReservationManager { } } } + }) + .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; } }); } From 3e8b2f97988a7116ebe4a1e090e56bf4b24a1cbf Mon Sep 17 00:00:00 2001 From: Thomas Montfort Date: Wed, 12 Aug 2026 11:36:16 -0700 Subject: [PATCH 12/15] fix(epp): raise prefill admission defaults Signed-off-by: Thomas Montfort (cherry picked from commit ece553d712e294ee3b2786c1e3d0b66b4a526842) (cherry picked from commit 94901a1581e3194c9a4077b9a1fd0379518b0d76) Signed-off-by: Avinash Varma --- .../epp/pkg/plugins/disagg/prefill_scorer.go | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) 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 0629086746f7..b6c00383e885 100644 --- a/deploy/inference-gateway/epp/pkg/plugins/disagg/prefill_scorer.go +++ b/deploy/inference-gateway/epp/pkg/plugins/disagg/prefill_scorer.go @@ -35,8 +35,8 @@ const ( // DynPrefillScorerType is the plugin type registered in the plugin registry. DynPrefillScorerType = "dyn-prefill-scorer" - defaultPrefillReservationAdmissionTimeout = 5 * time.Second - defaultMaxPrefillReservations = 32 + defaultPrefillReservationAdmissionTimeout = 60 * time.Second + defaultMaxPrefillReservations = 64 ) // compile-time type assertion From 0284de3381787a358ac6c03359d90b024d20200b Mon Sep 17 00:00:00 2001 From: Thomas Montfort <61255722+tmonty12@users.noreply.github.com> Date: Wed, 12 Aug 2026 12:03:19 -0700 Subject: [PATCH 13/15] Potential fix for pull request finding 'CodeQL / Incorrect conversion between integer types' Co-authored-by: Copilot Autofix powered by AI <62310815+github-advanced-security[bot]@users.noreply.github.com> Signed-off-by: Thomas Montfort <61255722+tmonty12@users.noreply.github.com> (cherry picked from commit 858b76d53a6a0f3d9093741deb785437f2c23fe2) Signed-off-by: Avinash Varma --- .../epp/pkg/plugins/disagg/booking_executor.go | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/deploy/inference-gateway/epp/pkg/plugins/disagg/booking_executor.go b/deploy/inference-gateway/epp/pkg/plugins/disagg/booking_executor.go index 30af2c703d12..765120530d9c 100644 --- a/deploy/inference-gateway/epp/pkg/plugins/disagg/booking_executor.go +++ b/deploy/inference-gateway/epp/pkg/plugins/disagg/booking_executor.go @@ -57,7 +57,8 @@ func bookingCleanupRetention() time.Duration { 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) + const maxDurationSeconds = uint64((1<<63 - 1) / int64(time.Second)) + maxSeconds := maxDurationSeconds - uint64(bookingCleanupRetentionGrace/time.Second) if err != nil || seconds == 0 || seconds > maxSeconds { return minimumBookingCleanupRetention } From 65c3194d4c9d8dc04a2febab9ebc77b1e42d718f Mon Sep 17 00:00:00 2001 From: Avinash Varma Date: Tue, 1 Sep 2026 01:37:21 -0700 Subject: [PATCH 14/15] fix(inference-gateway): adapt prefill reservation port to 1.4.2 APIs The upstream change was developed against release/1.3.0, which predates the cache-namespace plumbing on this release line. Thread cache_namespace through the reservation path so it matches the v1.4.2 signatures: - reservations::reserve() and RouterHandles::reserve_prefill_worker() take cache_namespace and forward it to find_best_match_details(), which requires it on 1.4.2. - route_prefill_request_with_reservation() destructures the 5-tuple PreprocessedRequest and writes the namespace back via write_cache_namespace_to_result(), as the removed advisory route did. - CallRoutePrefillRequestWithReservation() surfaces CacheNamespace on RoutingResult, matching the decode path. - DynDecodeScorer keeps the 5-arg addRequest and passes state.CacheNamespace. Also drops an unused prefill_router binding in add_request_with_cache_namespace that is not present in the upstream PR head. Signed-off-by: Avinash Varma --- .../epp/pkg/plugins/dynamo_kv_scorer/plugin.go | 8 +++++--- lib/bindings/c/src/lib.rs | 7 +++++-- lib/llm/src/kv_router/prefill_router/reservations.rs | 2 ++ 3 files changed, 12 insertions(+), 5 deletions(-) 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 5384823850b5..cc40d4448cb2 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 @@ -524,14 +524,16 @@ func CallRoutePrefillRequestWithReservation(reservationID string, requestJSON st } tokens := extractTokenData(&result) + cacheNamespace := extractCacheNamespace(&result) workerID := uint64(result.prefill_worker_id) dpRank := uint32(result.prefill_dp_rank) C.free_routing_result(&result) return &RoutingResult{ - WorkerID: workerID, - DpRank: dpRank, - TokenData: tokens, + WorkerID: workerID, + DpRank: dpRank, + TokenData: tokens, + CacheNamespace: cacheNamespace, }, nil } diff --git a/lib/bindings/c/src/lib.rs b/lib/bindings/c/src/lib.rs index 977c6e9cfdf7..6c11aafbd81f 100644 --- a/lib/bindings/c/src/lib.rs +++ b/lib/bindings/c/src/lib.rs @@ -484,6 +484,7 @@ impl RouterHandles { tokens: &[u32], block_mm_infos: Option<&[Option]>, lora_name: Option, + cache_namespace: Option, priority_jump: f64, strict_priority: u32, allowed_worker_ids: Option>, @@ -501,6 +502,7 @@ impl RouterHandles { tokens, block_mm_infos, lora_name, + cache_namespace, priority_jump, strict_priority, allowed_worker_ids, @@ -995,7 +997,6 @@ pub unsafe extern "C" fn add_request_with_cache_namespace( Vec::new() }; - let prefill_router = handles.prefill_router.clone(); let decode_router = handles.decode_router.clone(); let result = handles.runtime.secondary().block_on(async { @@ -1546,7 +1547,7 @@ pub unsafe extern "C" fn route_prefill_request_with_reservation( tracing::warn!(%reservation_id, %error, "Failed to begin EPP prefill reservation"); return QueryRouterResult::ErrQueryFailed; } - let (tokens, priority_jump, strict_priority, routing_constraints) = + let (tokens, cache_namespace, priority_jump, strict_priority, routing_constraints) = match unsafe { preprocess_request(handles, request_json) } { Ok(values) => values, Err(code) => { @@ -1564,6 +1565,7 @@ pub unsafe extern "C" fn route_prefill_request_with_reservation( &tokens, None, None, + cache_namespace.clone(), priority_jump, strict_priority, allowed_worker_ids, @@ -1589,6 +1591,7 @@ pub unsafe extern "C" fn route_prefill_request_with_reservation( out.prefill_worker_id = prefill_worker_id; out.prefill_dp_rank = prefill_dp_rank; write_tokens_to_result(&tokens, out); + write_cache_namespace_to_result(cache_namespace.as_deref(), out); QueryRouterResult::Ok } Err(code) => code, diff --git a/lib/llm/src/kv_router/prefill_router/reservations.rs b/lib/llm/src/kv_router/prefill_router/reservations.rs index c1999c8bd9e8..68f420e51966 100644 --- a/lib/llm/src/kv_router/prefill_router/reservations.rs +++ b/lib/llm/src/kv_router/prefill_router/reservations.rs @@ -394,6 +394,7 @@ impl EppReservationManager { token_ids: &[u32], block_mm_infos: Option<&[Option]>, lora_name: Option, + cache_namespace: Option, priority_jump: f64, strict_priority: u32, allowed_worker_ids: Option>, @@ -440,6 +441,7 @@ impl EppReservationManager { true, false, lora_name, + cache_namespace, priority_jump, strict_priority, None, From 0a12bfba5d4cd57cf6213bc33b8691d4988c27b7 Mon Sep 17 00:00:00 2001 From: Avinash Varma Date: Tue, 1 Sep 2026 01:49:57 -0700 Subject: [PATCH 15/15] fix(inference-gateway): guard nil prefill cancellation hook DynPrefillScorer.Score() called the cancelPrefill func field directly on the admission-timeout path, so any scorer built without that hook panics with a nil pointer dereference instead of reporting a configuration error. beginPrefill and releasePrefill are already accessed through nil-guarded wrappers; this adds the matching wrapper for cancelPrefill. Reproduced by TestPrefillScoreBoundsConcurrentReservations, which constructs a DynPrefillScorer without a cancelPrefill hook. The upstream PR head carries the same unguarded call and the same test, so it panics there too. Signed-off-by: Avinash Varma --- .../epp/pkg/plugins/disagg/prefill_scorer.go | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) 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 b6c00383e885..06684c180360 100644 --- a/deploy/inference-gateway/epp/pkg/plugins/disagg/prefill_scorer.go +++ b/deploy/inference-gateway/epp/pkg/plugins/disagg/prefill_scorer.go @@ -139,6 +139,13 @@ func (s *DynPrefillScorer) beginReservation(bookingID string) error { return s.beginPrefill(bookingID) } +func (s *DynPrefillScorer) cancelReservation(bookingID string) error { + if s.cancelPrefill == nil { + return fmt.Errorf("prefill reservation cancellation is not configured") + } + return s.cancelPrefill(bookingID) +} + func (s *DynPrefillScorer) releaseLatePrefillReservation(bookingID string) error { if s.releasePrefill == nil { return fmt.Errorf("prefill reservation release is not configured") @@ -226,7 +233,7 @@ func (s *DynPrefillScorer) Score(ctx context.Context, cycleState *schedtypes.Cyc 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 { + if cancelErr := s.cancelReservation(bookingID); cancelErr != nil { logger.V(logutil.DEFAULT).Error(cancelErr, "DynPrefillScorer: failed to cancel pending prefill reservation", "bookingID", bookingID) }