Merge upstream/main into epp-hotpath-log-guards

Signed-off-by: Luke Van Drie <lukevandrie@google.com>

Author: LukeAVanDrie

pkg/common/error/error.go

@@ -26,7 +26,7 @@ import (
 )
 
 // RequestDroppedReasonHeaderKey is the HTTP response header that communicates the specific
-// reason a request was dropped by flow control.
+// reason the EPP dropped a request.
 const RequestDroppedReasonHeaderKey = "x-llm-d-request-dropped-reason"
 
 // RequestDroppedReason is the reason a request was rejected before dispatch or evicted after dispatch.

pkg/epp/framework/plugins/datalayer/extractor/metrics/extractor.go

@@ -87,10 +87,20 @@ func (ext *Extractor) TypedName() fwkplugin.TypedName {
 var _ fwkplugin.ProducerPlugin = &Extractor{}
 
 // Produces declares the custom scalar metric attributes, whose names come from
-// the per-engine mappings in configuration. Core metrics land on the Metrics
-// struct rather than the attribute map and so are not declared here.
+// the per-engine mappings in configuration, plus the per-pod metric fields this
+// extractor writes into the endpoint's Metrics struct that a scorer or filter
+// declares as a Consumes() dependency. Declaring both lets the data-attribute
+// registry and the DAG builder validate the scorers and filters that read them.
+// The Metrics field types mirror the fwkdl.Metrics struct (float64, int, map
+// of model name to count).
 func (ext *Extractor) Produces() map[fwkplugin.DataKey]any {
-	produced := map[fwkplugin.DataKey]any{}
+	produced := map[fwkplugin.DataKey]any{
+		fwkplugin.NewDataKey(KVCacheUsagePercentKey, MetricsExtractorType): float64(0),
+		fwkplugin.NewDataKey(WaitingQueueSizeKey, MetricsExtractorType):    int(0),
+		fwkplugin.NewDataKey(RunningRequestsSizeKey, MetricsExtractorType): int(0),
+		fwkplugin.NewDataKey(ActiveModelsKey, MetricsExtractorType):        map[string]int{},
+		fwkplugin.NewDataKey(WaitingModelsKey, MetricsExtractorType):       map[string]int{},
+	}
 	for _, mapping := range ext.registry.Mappings() {
 		for _, custom := range mapping.CustomMetrics {
 			produced[attrmetrics.ScalarMetricDataKey(custom.AttributeKey)] = attrmetrics.ScalarMetricValue(0)

pkg/epp/framework/plugins/scheduling/scorer/endpointattribute/endpointattribute.go

@@ -36,7 +36,10 @@ const (
 )
 
 // compile-time type assertion
-var _ fwksched.Scorer = &EndpointAttributeScorer{}
+var (
+	_ fwksched.Scorer          = &EndpointAttributeScorer{}
+	_ fwkplugin.ConsumerPlugin = &EndpointAttributeScorer{}
+)
 
 // fixedRangeParameters normalizes the attribute value against a fixed
 // [min, max] range (e.g. kv-cache utilization, which is always in [0, 1]).

pkg/epp/framework/plugins/scheduling/scorer/kvcacheutilization/kvcache_utilization.go

@@ -30,7 +30,10 @@ const (
 )
 
 // compile-time type assertion
-var _ fwksched.Scorer = &KVCacheUtilizationScorer{}
+var (
+	_ fwksched.Scorer          = &KVCacheUtilizationScorer{}
+	_ fwkplugin.ConsumerPlugin = &KVCacheUtilizationScorer{}
+)
 
 // KvCacheUtilizationScorerFactory defines the factory function for KVCacheUtilizationScorer.
 func KvCacheUtilizationScorerFactory(name string, _ *json.Decoder, _ fwkplugin.Handle) (fwkplugin.Plugin, error) {
@@ -59,10 +62,15 @@ func (s *KVCacheUtilizationScorer) Category() fwksched.ScorerCategory {
 	return fwksched.Distribution
 }
 
-// Consumes returns the list of data that is consumed by the plugin.
-func (s *KVCacheUtilizationScorer) Consumes() map[string]any {
-	return map[string]any{
-		metrics.KVCacheUsagePercentKey: float64(0),
+// Consumes declares that the scorer reads the KV-cache utilization field
+// from the endpoint's Metrics struct. The DataKey names the field the
+// core-metrics-extractor publishes; the registry validates the consumer's
+// declared type against the producer's declaration.
+func (s *KVCacheUtilizationScorer) Consumes() fwkplugin.DataDependencies {
+	return fwkplugin.DataDependencies{
+		Required: map[fwkplugin.DataKey]any{
+			fwkplugin.NewDataKey(metrics.KVCacheUsagePercentKey, metrics.MetricsExtractorType): float64(0),
+		},
 	}
 }
 

pkg/epp/framework/plugins/scheduling/scorer/loraaffinity/lora_affinity.go

@@ -30,7 +30,10 @@ const (
 )
 
 // compile-time type assertion
-var _ fwksched.Scorer = &LoraAffinityScorer{}
+var (
+	_ fwksched.Scorer          = &LoraAffinityScorer{}
+	_ fwkplugin.ConsumerPlugin = &LoraAffinityScorer{}
+)
 
 // LoraAffinityScorerFactory defines the factory function for LoraAffinityScorer.
 func LoraAffinityScorerFactory(name string, _ *json.Decoder, _ fwkplugin.Handle) (fwkplugin.Plugin, error) {
@@ -59,11 +62,15 @@ func (s *LoraAffinityScorer) Category() fwksched.ScorerCategory {
 	return fwksched.Affinity
 }
 
-// Consumes returns the list of data that is consumed by the plugin.
-func (s *LoraAffinityScorer) Consumes() map[string]any {
-	return map[string]any{
-		metrics.ActiveModelsKey:  map[string]int{},
-		metrics.WaitingModelsKey: map[string]int{},
+// Consumes declares the scorer reads the per-pod active and waiting model
+// sets from the endpoint's Metrics struct, published by the core-metrics-
+// extractor.
+func (s *LoraAffinityScorer) Consumes() fwkplugin.DataDependencies {
+	return fwkplugin.DataDependencies{
+		Required: map[fwkplugin.DataKey]any{
+			fwkplugin.NewDataKey(metrics.ActiveModelsKey, metrics.MetricsExtractorType):  map[string]int{},
+			fwkplugin.NewDataKey(metrics.WaitingModelsKey, metrics.MetricsExtractorType): map[string]int{},
+		},
 	}
 }
 

pkg/epp/framework/plugins/scheduling/scorer/queuedepth/queue.go

@@ -31,7 +31,10 @@ const (
 )
 
 // compile-time type assertion
-var _ fwksched.Scorer = &QueueScorer{}
+var (
+	_ fwksched.Scorer          = &QueueScorer{}
+	_ fwkplugin.ConsumerPlugin = &QueueScorer{}
+)
 
 // QueueScorerFactory defines the factory function for QueueScorer.
 func QueueScorerFactory(name string, _ *json.Decoder, _ fwkplugin.Handle) (fwkplugin.Plugin, error) {
@@ -61,10 +64,13 @@ func (s *QueueScorer) Category() fwksched.ScorerCategory {
 	return fwksched.Distribution
 }
 
-// Consumes returns the list of data that is consumed by the plugin.
-func (s *QueueScorer) Consumes() map[string]any {
-	return map[string]any{
-		metrics.WaitingQueueSizeKey: int(0),
+// Consumes declares the scorer reads the waiting queue size from the
+// endpoint's Metrics struct, published by the core-metrics-extractor.
+func (s *QueueScorer) Consumes() fwkplugin.DataDependencies {
+	return fwkplugin.DataDependencies{
+		Required: map[fwkplugin.DataKey]any{
+			fwkplugin.NewDataKey(metrics.WaitingQueueSizeKey, metrics.MetricsExtractorType): int(0),
+		},
 	}
 }
 

pkg/epp/framework/plugins/scheduling/scorer/runningrequests/runningrequest.go

@@ -31,7 +31,10 @@ const (
 )
 
 // compile-time type assertion
-var _ fwksched.Scorer = &RunningRequestsSizeScorer{}
+var (
+	_ fwksched.Scorer          = &RunningRequestsSizeScorer{}
+	_ fwkplugin.ConsumerPlugin = &RunningRequestsSizeScorer{}
+)
 
 // RunningRequestsSizeScorerFactory defines the factory function for RunningRequestsSizeScorer.
 func RunningRequestsSizeScorerFactory(name string, _ *json.Decoder, _ fwkplugin.Handle) (fwkplugin.Plugin, error) {
@@ -61,10 +64,13 @@ func (s *RunningRequestsSizeScorer) Category() fwksched.ScorerCategory {
 	return fwksched.Distribution
 }
 
-// Consumes returns the list of data that is consumed by the plugin.
-func (s *RunningRequestsSizeScorer) Consumes() map[string]any {
-	return map[string]any{
-		metrics.RunningRequestsSizeKey: int(0),
+// Consumes declares the scorer reads the running requests size from the
+// endpoint's Metrics struct, published by the core-metrics-extractor.
+func (s *RunningRequestsSizeScorer) Consumes() fwkplugin.DataDependencies {
+	return fwkplugin.DataDependencies{
+		Required: map[fwkplugin.DataKey]any{
+			fwkplugin.NewDataKey(metrics.RunningRequestsSizeKey, metrics.MetricsExtractorType): int(0),
+		},
 	}
 }
 

pkg/epp/requestcontrol/director.go

@@ -304,17 +304,19 @@ func (d *Director) HandleRequest(ctx context.Context, reqCtx *handlers.RequestCo
 	endpointCandidates := d.endpointCandidates.Locate(ctx, reqCtx.Request.Metadata)
 	if len(endpointCandidates) == 0 {
 		return reqCtx, errcommon.Error{
-			Code: errcommon.ServiceUnavailable,
-			Msg:  "failed to find endpoint candidates for serving the request",
+			Code:    errcommon.ServiceUnavailable,
+			Msg:     "failed to find endpoint candidates for serving the request",
+			Headers: map[string]string{errcommon.RequestDroppedReasonHeaderKey: string(errcommon.RequestDroppedReasonNoEndpoints)},
 		}
 	}
 
 	snapshotOfCandidatePods := d.toSchedulerEndpoints(endpointCandidates)
 	snapshotOfCandidatePods = d.runScreeners(ctx, reqCtx.SchedulingRequest, snapshotOfCandidatePods)
 	if len(snapshotOfCandidatePods) == 0 {
 		return reqCtx, errcommon.Error{
-			Code: errcommon.ServiceUnavailable,
-			Msg:  "screeners eliminated all endpoint candidates",
+			Code:    errcommon.ServiceUnavailable,
+			Msg:     "screeners eliminated all endpoint candidates",
+			Headers: map[string]string{errcommon.RequestDroppedReasonHeaderKey: string(errcommon.RequestDroppedReasonNoEndpoints)},
 		}
 	}
 	// Prepare per request data by running DataProducer plugins.

pkg/epp/requestcontrol/director_test.go

@@ -414,11 +414,13 @@ func TestDirector_HandleRequest(t *testing.T) {
 		initialTargetModelName  string // Initial target model in the reqCtx.
 		parser                  fwkrh.Parser
 		wantErrCode             string                   // Expected errcommon code string
+		wantDroppedReason       string                   // If non-empty, expected x-llm-d-request-dropped-reason header on the error
 		wantReqCtx              *handlers.RequestContext // Fields to check in the returned RequestContext
 		targetModelName         string                   // Expected model name after target model resolution
 		admitRequestDenialError error                    // Expected denial error from admission plugin
 		dataProducerPlugin      *mockDataProducerPlugin
 		screener                *mockScreener
+		emptyEndpoints          bool // If true, the director locates no endpoint candidates.
 		preRequestPlugins       []*mockPreRequestPlugin
 		requestHeaderPlugin     *mockRequestHeaderPlugin
 		wantMutatedBody         map[string]any
@@ -995,7 +997,21 @@ func TestDirector_HandleRequest(t *testing.T) {
 			screener: &mockScreener{name: "eliminate-all", screen: func([]fwksched.Endpoint) []fwksched.Endpoint {
 				return nil
 			}},
-			wantErrCode: errcommon.ServiceUnavailable,
+			wantErrCode:       errcommon.ServiceUnavailable,
+			wantDroppedReason: string(errcommon.RequestDroppedReasonNoEndpoints),
+		},
+		{
+			name: "no endpoint candidates located",
+			reqBodyMap: map[string]any{
+				"model":  model,
+				"prompt": "critical prompt",
+			},
+			mockAdmissionController: &mockAdmissionController{admitErr: nil},
+			initialTargetModelName:  model,
+			inferenceObjectiveName:  objectiveName,
+			emptyEndpoints:          true,
+			wantErrCode:             errcommon.ServiceUnavailable,
+			wantDroppedReason:       string(errcommon.RequestDroppedReasonNoEndpoints),
 		},
 		{
 			name: "scheduler returns error",
@@ -1010,6 +1026,30 @@ func TestDirector_HandleRequest(t *testing.T) {
 			wantErrCode:            errcommon.ResourceExhausted,
 			inferenceObjectiveName: objectiveName,
 		},
+		{
+			// The typed error inside a joined scheduler error, including its
+			// drop-reason header, must reach the caller instead of the
+			// untyped-error fallback.
+			name: "scheduler returns joined error with typed capacity rejection",
+			reqBodyMap: map[string]any{
+				"model":  model,
+				"prompt": "prompt that causes scheduling drain",
+			},
+			mockAdmissionController: &mockAdmissionController{admitErr: nil},
+			schedulerMockSetup: func(m *mockScheduler) {
+				m.scheduleErr = errors.Join(
+					errors.New("failed to run scheduler profile 'default'"),
+					fmt.Errorf("profile %q: %w", "default", errcommon.Error{
+						Code:    errcommon.ResourceExhausted,
+						Msg:     "no endpoints available for the given request",
+						Headers: map[string]string{errcommon.RequestDroppedReasonHeaderKey: string(errcommon.RequestDroppedReasonSaturated)},
+					}),
+				)
+			},
+			wantErrCode:            errcommon.ResourceExhausted,
+			wantDroppedReason:      string(errcommon.RequestDroppedReasonSaturated),
+			inferenceObjectiveName: objectiveName,
+		},
 		{
 			name: "scheduler returns nil result and nil error",
 			reqBodyMap: map[string]any{
@@ -1090,6 +1130,9 @@ func TestDirector_HandleRequest(t *testing.T) {
 
 				endpointCandidates := NewCachedEndpointCandidates(context.Background(), NewDatastoreEndpointCandidates(ds), time.Minute)
 				director := NewDirectorWithConfig(ds, mockSched, test.mockAdmissionController, endpointCandidates, config)
+				if test.emptyEndpoints {
+					director.endpointCandidates = &mockEndpointCandidates{}
+				}
 				if len(test.rewrites) > 0 {
 					mockDs := &mockDatastore{
 						pods:     ds.PodList(datastore.AllPodsPredicate),
@@ -1140,6 +1183,9 @@ func TestDirector_HandleRequest(t *testing.T) {
 					var e errcommon.Error
 					if assert.ErrorAs(t, err, &e, "Error should be of type errcommon.Error") {
 						assert.Equal(t, test.wantErrCode, e.Code, "Error code mismatch")
+						if test.wantDroppedReason != "" {
+							assert.Equal(t, test.wantDroppedReason, e.Headers[errcommon.RequestDroppedReasonHeaderKey], "drop-reason header mismatch")
+						}
 					}
 					return
 				}

pkg/epp/scheduling/scheduler.go

@@ -19,7 +19,10 @@ package scheduling
 
 import (
 	"context"
+	"errors"
 	"fmt"
+	"maps"
+	"slices"
 	"time"
 
 	"go.opentelemetry.io/otel/attribute"
@@ -66,6 +69,11 @@ func (s *Scheduler) Schedule(ctx context.Context, request *fwksched.InferenceReq
 	}()
 
 	profileRunResults := map[string]*fwksched.ProfileRunResult{}
+	// Keyed like profileRunResults so a profile that fails in one iteration and
+	// succeeds in a later one leaves no stale error behind. Nil until a profile
+	// fails: delete and len are no-ops on a nil map, and the happy path skips
+	// the allocation.
+	var profileRunErrors map[string]error
 
 	for { // get the next set of profiles to run iteratively based on the request and the previous execution results
 		if verboseEnabled {
@@ -87,12 +95,19 @@ func (s *Scheduler) Schedule(ctx context.Context, request *fwksched.InferenceReq
 			}
 			// run the selected profiles and collect results (current code runs all profiles)
 			profileRunResult, err := runSchedulerProfile(ctx, name, profile, request, candidateEndpoints)
-			if verboseEnabled {
-				if err != nil {
+			if err != nil {
+				if verboseEnabled {
 					loggerVerbose.Info("failed to run scheduler profile", "profile", name, "error", err.Error())
-				} else {
+				}
+				if profileRunErrors == nil {
+					profileRunErrors = map[string]error{}
+				}
+				profileRunErrors[name] = fmt.Errorf("profile %q: %w", name, err)
+			} else {
+				if verboseEnabled {
 					loggerVerbose.Info("Completed running scheduler profile successfully", "profile", name)
 				}
+				delete(profileRunErrors, name)
 			}
 
 			profileRunResults[name] = profileRunResult // if profile failed to run, the run result is nil
@@ -114,6 +129,21 @@ func (s *Scheduler) Schedule(ctx context.Context, request *fwksched.InferenceReq
 		loggerVerbose.Info("Completed running profile handler ProcessResults successfully", "plugin", handlerName)
 	}
 
+	// Profile handlers see failed profiles only as nil results and report them
+	// with fresh untyped errors. Join the retained profile errors so a typed
+	// errcommon.Error raised inside a profile run (e.g. filters draining the
+	// candidate set) stays reachable via errors.As in the caller.
+	if err != nil && len(profileRunErrors) > 0 {
+		errs := make([]error, 0, len(profileRunErrors)+1)
+		errs = append(errs, err)
+		// Sorted so error composition, and therefore errors.As selection when
+		// profiles fail with different typed codes, is deterministic.
+		for _, name := range slices.Sorted(maps.Keys(profileRunErrors)) {
+			errs = append(errs, profileRunErrors[name])
+		}
+		err = errors.Join(errs...)
+	}
+
 	return result, err
 }