diff --git a/docs/architecture.md b/docs/architecture.md index 57ed5d15d6..39898fed8f 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -26,7 +26,7 @@ The design enables: -- Support for **multiple base models** within a shared cluster (see [serving multiple inference pools](https://gateway-api-inference-extension.sigs.k8s.io/guides/serving-multiple-inference-pools-latest/)) +- Support for **multiple base models** within a shared cluster (see [InferencePool & InferenceModel Design](#inferencepool--inferencemodel-design)) - Efficient routing based on **KV cache locality**, **session affinity**, **load**, and **model metadata** - Disaggregated **Prefill/Decode (P/D)** execution diff --git a/pkg/epp/scheduling/scheduler_profile.go b/pkg/epp/scheduling/scheduler_profile.go index 76e9a39d42..8a6255bae4 100644 --- a/pkg/epp/scheduling/scheduler_profile.go +++ b/pkg/epp/scheduling/scheduler_profile.go @@ -35,6 +35,14 @@ import ( "github.com/llm-d/llm-d-router/pkg/epp/metrics" ) +// tracerScope is the OTel instrumentation scope for scheduler-level scoring +// spans emitted from this package. +const tracerScope = "llm-d-router/pkg/epp/scheduling" + +// internalSpanKind is hoisted to avoid allocating a span-start option on every +// scoring call. +var internalSpanKind = trace.WithSpanKind(trace.SpanKindInternal) + // NewSchedulerProfile creates a new SchedulerProfile object and returns its pointer. func NewSchedulerProfile() *SchedulerProfile { return &SchedulerProfile{ @@ -171,6 +179,26 @@ func (p *SchedulerProfile) runScorerPlugins(ctx context.Context, request *fwksch logger := log.FromContext(ctx) logger.V(logutil.DEBUG).Info("Before running scorer plugins", "endpoints", endpoints) + // Parent span over the whole scorer chain. Per-scorer child spans (and any + // plugin-internal spans) nest under it. Attributes are request- and + // chain-level only; no per-endpoint keys, to keep span cardinality bounded. + // The tracer is resolved once and threaded into runScorer so the per-scorer + // spans reuse it rather than rebuilding instrumentation options per scorer. + tracer := tracing.Tracer(tracerScope) + ctx, span := tracer.Start(ctx, "llm_d.epp.scoring", internalSpanKind) + defer span.End() + // On the default (tracing-disabled) path Start returns a non-recording span; + // skip all attribute and child-span construction so the scoring hot path + // stays allocation-free, matching the rest of this package. + tracingActive := span.IsRecording() + if tracingActive { + span.SetAttributes( + attribute.Int("llm_d.epp.scorer.count", len(p.scorers)), + attribute.Int("llm_d.epp.scoring.candidate_endpoints", len(endpoints)), + ) + setRequestSpanAttributes(span, request) + } + weightedScorePerEndpoint := make(map[fwksched.Endpoint]float64, len(endpoints)) for _, endpoint := range endpoints { weightedScorePerEndpoint[endpoint] = float64(0) // initialize weighted score per endpoint with 0 value @@ -187,9 +215,7 @@ func (p *SchedulerProfile) runScorerPlugins(ctx context.Context, request *fwksch // Iterate through each scorer in the chain and accumulate the weighted scores. for _, scorer := range p.scorers { logger.V(logutil.VERBOSE).Info("Running scorer plugin", "plugin", scorer.TypedName()) - before := time.Now() - scores := scorer.Score(ctx, request, endpoints) - metrics.RecordPluginProcessingLatency(scorerExtensionPoint, scorer.TypedName().Type, scorer.TypedName().Name, time.Since(before)) + scores := runScorer(ctx, tracer, tracingActive, scorer, request, endpoints) for endpoint, score := range scores { // weight is relative to the sum of weights if debugEnabled { debug.Info("Calculated score", "plugin", scorer.TypedName(), "endpoint", endpoint.GetMetadata().NamespacedName, "score", score) @@ -203,6 +229,69 @@ func (p *SchedulerProfile) runScorerPlugins(ctx context.Context, request *fwksch return weightedScorePerEndpoint } +// runScorer invokes a single weighted scorer and records its latency metric. +// When tracing is active it wraps the call in an llm_d.epp.scorer. span +// annotated with the scorer's identity, weight, candidate count, and aggregate +// score signals; aggregates are derived from the returned score map only, with +// no per-endpoint attribute keys, to keep span cardinality bounded. When +// tracing is inactive no span or attribute work is performed. +func runScorer(ctx context.Context, tracer trace.Tracer, tracingActive bool, scorer *WeightedScorer, request *fwksched.InferenceRequest, endpoints []fwksched.Endpoint) map[fwksched.Endpoint]float64 { + typedName := scorer.TypedName() + + if !tracingActive { + before := time.Now() + scores := scorer.Score(ctx, request, endpoints) + metrics.RecordPluginProcessingLatency(scorerExtensionPoint, typedName.Type, typedName.Name, time.Since(before)) + return scores + } + + ctx, span := tracer.Start(ctx, "llm_d.epp.scorer."+typedName.Type, internalSpanKind) + defer span.End() + span.SetAttributes( + attribute.String("llm_d.epp.scorer.type", typedName.Type), + attribute.String("llm_d.epp.scorer.name", typedName.Name), + attribute.Float64("llm_d.epp.scorer.weight", scorer.Weight()), + attribute.Int("llm_d.epp.scorer.candidate_endpoints", len(endpoints)), + ) + + before := time.Now() + scores := scorer.Score(ctx, request, endpoints) + metrics.RecordPluginProcessingLatency(scorerExtensionPoint, typedName.Type, typedName.Name, time.Since(before)) + + if len(scores) > 0 { + var maxScore, totalScore float64 + first := true + for _, s := range scores { + if first || s > maxScore { + maxScore = s + } + first = false + totalScore += s + } + span.SetAttributes( + attribute.Float64("llm_d.epp.scorer.score.max", maxScore), + attribute.Float64("llm_d.epp.scorer.score.avg", totalScore/float64(len(scores))), + attribute.Int("llm_d.epp.scorer.endpoints_scored", len(scores)), + ) + } + + return scores +} + +// setRequestSpanAttributes mirrors the request-identity attribute schema used +// by the scorer plugins so scheduler-level spans correlate with the same trace. +func setRequestSpanAttributes(span trace.Span, request *fwksched.InferenceRequest) { + if request == nil { + return + } + if request.TargetModel != "" { + span.SetAttributes(attribute.String("gen_ai.request.model", request.TargetModel)) + } + if request.RequestID != "" { + span.SetAttributes(attribute.String("gen_ai.request.id", request.RequestID)) + } +} + func (p *SchedulerProfile) runPickerPlugin(ctx context.Context, request *fwksched.InferenceRequest, weightedScorePerEndpoint map[fwksched.Endpoint]float64) *fwksched.ProfileRunResult { logger := log.FromContext(ctx) diff --git a/pkg/epp/scheduling/scheduler_profile_test.go b/pkg/epp/scheduling/scheduler_profile_test.go index 7fe6b8084b..9d40d0fc70 100644 --- a/pkg/epp/scheduling/scheduler_profile_test.go +++ b/pkg/epp/scheduling/scheduler_profile_test.go @@ -23,6 +23,10 @@ import ( "github.com/google/go-cmp/cmp" "github.com/google/uuid" + "go.opentelemetry.io/otel" + sdktrace "go.opentelemetry.io/otel/sdk/trace" + "go.opentelemetry.io/otel/sdk/trace/tracetest" + tracenoop "go.opentelemetry.io/otel/trace/noop" k8stypes "k8s.io/apimachinery/pkg/types" fwkdl "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/datalayer" @@ -603,6 +607,225 @@ func (p *filterOnlyPlugin) Filter(_ context.Context, _ *fwksched.InferenceReques return endpoints } +// fixedScoresScorer returns a caller-supplied score per endpoint name, letting +// tests assert aggregate span attributes (max/avg) over a non-uniform map. +type fixedScoresScorer struct { + typedName fwkplugin.TypedName + scores map[string]float64 +} + +func (s *fixedScoresScorer) TypedName() fwkplugin.TypedName { return s.typedName } + +func (s *fixedScoresScorer) Category() fwksched.ScorerCategory { return fwksched.Distribution } + +func (s *fixedScoresScorer) Score(_ context.Context, _ *fwksched.InferenceRequest, endpoints []fwksched.Endpoint) map[fwksched.Endpoint]float64 { + // A nil score table models a scorer that declines to score any endpoint, + // exercising the empty-map aggregate guard in runScorer. + if s.scores == nil { + return map[fwksched.Endpoint]float64{} + } + out := make(map[fwksched.Endpoint]float64, len(endpoints)) + for _, e := range endpoints { + out[e] = s.scores[e.GetMetadata().NamespacedName.Name] + } + return out +} + +// installSpanRecorder routes spans to an in-memory recorder for the duration of +// the test and restores an explicit no-op provider afterward. Restoring the +// no-op provider (rather than the global proxy returned by GetTracerProvider) +// ensures later tests do not inherit this test's recording SDK provider. +func installSpanRecorder(t *testing.T) *tracetest.SpanRecorder { + t.Helper() + recorder := tracetest.NewSpanRecorder() + tp := sdktrace.NewTracerProvider(sdktrace.WithSpanProcessor(recorder)) + otel.SetTracerProvider(tp) + t.Cleanup(func() { otel.SetTracerProvider(tracenoop.NewTracerProvider()) }) + return recorder +} + +// findSpan returns the first recorded span with the given name, or nil. +func findSpan(spans tracetest.SpanStubs, name string) *tracetest.SpanStub { + for i := range spans { + if spans[i].Name == name { + return &spans[i] + } + } + return nil +} + +func spanFloat(t *testing.T, span *tracetest.SpanStub, key string) float64 { + t.Helper() + for _, a := range span.Attributes { + if string(a.Key) == key { + return a.Value.AsFloat64() + } + } + t.Fatalf("span %q missing attribute %q", span.Name, key) + return 0 +} + +func spanInt(t *testing.T, span *tracetest.SpanStub, key string) int64 { + t.Helper() + for _, a := range span.Attributes { + if string(a.Key) == key { + return a.Value.AsInt64() + } + } + t.Fatalf("span %q missing attribute %q", span.Name, key) + return 0 +} + +func spanHasAttr(span *tracetest.SpanStub, key string) bool { + for _, a := range span.Attributes { + if string(a.Key) == key { + return true + } + } + return false +} + +// TestRunScorerPluginsTracing verifies the scheduler scoring path emits a parent +// llm_d.epp.scoring span with one llm_d.epp.scorer. child per scorer, +// carrying the documented identity, weight, candidate-count, and aggregate +// score attributes, and no per-endpoint attribute keys. +func TestRunScorerPluginsTracing(t *testing.T) { + recorder := installSpanRecorder(t) + + scorerA := &fixedScoresScorer{ + typedName: fwkplugin.TypedName{Type: "scorer-a", Name: "a"}, + scores: map[string]float64{"pod1": 0.2, "pod2": 0.8}, + } + scorerB := &fixedScoresScorer{ + typedName: fwkplugin.TypedName{Type: "scorer-b", Name: "b"}, + scores: map[string]float64{"pod1": 0.5, "pod2": 0.5}, + } + picker := &testPlugin{TypeRes: "picker", PickRes: k8stypes.NamespacedName{Name: "pod1"}} + + profile := NewSchedulerProfile(). + WithScorers(NewWeightedScorer(scorerA, 0.25), NewWeightedScorer(scorerB, 0.75)). + WithPicker(picker) + + input := []fwksched.Endpoint{ + fwksched.NewEndpoint(&fwkdl.EndpointMetadata{NamespacedName: k8stypes.NamespacedName{Name: "pod1"}}, nil, nil), + fwksched.NewEndpoint(&fwkdl.EndpointMetadata{NamespacedName: k8stypes.NamespacedName{Name: "pod2"}}, nil, nil), + } + request := &fwksched.InferenceRequest{TargetModel: "test-model", RequestID: "req-123"} + + if _, err := profile.Run(context.Background(), request, input); err != nil { + t.Fatalf("unexpected error: %v", err) + } + + spans := tracetest.SpanStubsFromReadOnlySpans(recorder.Ended()) + parent := findSpan(spans, "llm_d.epp.scoring") + if parent == nil { + t.Fatalf("missing parent span llm_d.epp.scoring; got %d spans", len(spans)) + } + if got := spanInt(t, parent, "llm_d.epp.scorer.count"); got != 2 { + t.Errorf("parent llm_d.epp.scorer.count = %d, want 2", got) + } + if got := spanInt(t, parent, "llm_d.epp.scoring.candidate_endpoints"); got != 2 { + t.Errorf("parent candidate_endpoints = %d, want 2", got) + } + + childA := findSpan(spans, "llm_d.epp.scorer.scorer-a") + childB := findSpan(spans, "llm_d.epp.scorer.scorer-b") + if childA == nil || childB == nil { + t.Fatalf("missing per-scorer child spans (a=%v b=%v)", childA != nil, childB != nil) + } + + // Both children must nest under the parent scoring span. + if childA.Parent.SpanID() != parent.SpanContext.SpanID() { + t.Errorf("scorer-a span is not a child of the scoring span") + } + + if got := spanFloat(t, childA, "llm_d.epp.scorer.weight"); got != 0.25 { + t.Errorf("scorer-a weight = %v, want 0.25", got) + } + if got := spanInt(t, childA, "llm_d.epp.scorer.candidate_endpoints"); got != 2 { + t.Errorf("scorer-a candidate_endpoints = %d, want 2", got) + } + // scorer-a scores {0.2, 0.8}: max 0.8, avg 0.5. + if got := spanFloat(t, childA, "llm_d.epp.scorer.score.max"); got != 0.8 { + t.Errorf("scorer-a score.max = %v, want 0.8", got) + } + if got := spanFloat(t, childA, "llm_d.epp.scorer.score.avg"); got != 0.5 { + t.Errorf("scorer-a score.avg = %v, want 0.5", got) + } + if got := spanInt(t, childA, "llm_d.epp.scorer.endpoints_scored"); got != 2 { + t.Errorf("scorer-a endpoints_scored = %d, want 2", got) + } + + // Cardinality guard: no per-pod/per-endpoint identifier attribute keys. + for _, span := range []*tracetest.SpanStub{parent, childA, childB} { + for _, a := range span.Attributes { + key := string(a.Key) + if strings.Contains(key, "pod") || strings.Contains(key, "endpoint.") || strings.Contains(key, "namespacedname") { + t.Errorf("span %q carries per-endpoint attribute key %q", span.Name, key) + } + } + } +} + +// TestRunScorerPluginsTracingDisabled verifies scoring works and records no +// spans when the global provider is the default no-op, and that an empty +// candidate set still emits the parent span without dividing by zero. +func TestRunScorerPluginsTracingDisabled(t *testing.T) { + // Explicitly install a no-op provider so this test exercises the + // tracing-disabled path regardless of provider state left by other tests. + otel.SetTracerProvider(tracenoop.NewTracerProvider()) + t.Cleanup(func() { otel.SetTracerProvider(tracenoop.NewTracerProvider()) }) + + scorer := &testPlugin{TypeRes: "noop", ScoreRes: 0.5} + picker := &testPlugin{TypeRes: "picker", PickRes: k8stypes.NamespacedName{Name: "pod1"}} + + profile := NewSchedulerProfile(). + WithScorers(NewWeightedScorer(scorer, 1)). + WithPicker(picker) + + input := []fwksched.Endpoint{ + fwksched.NewEndpoint(&fwkdl.EndpointMetadata{NamespacedName: k8stypes.NamespacedName{Name: "pod1"}}, nil, nil), + } + request := &fwksched.InferenceRequest{TargetModel: "test-model", RequestID: uuid.NewString()} + + if _, err := profile.Run(context.Background(), request, input); err != nil { + t.Fatalf("unexpected error with no-op tracer: %v", err) + } + if scorer.ScoreCallCount != 1 { + t.Errorf("scorer called %d times, want 1", scorer.ScoreCallCount) + } +} + +// TestRunScorerEmptyCandidateAvg verifies a scorer that returns an empty score +// map does not trigger a divide-by-zero when computing the average attribute. +func TestRunScorerEmptyCandidateAvg(t *testing.T) { + recorder := installSpanRecorder(t) + + empty := &fixedScoresScorer{typedName: fwkplugin.TypedName{Type: "empty", Name: "e"}} + picker := &testPlugin{TypeRes: "picker", PickRes: k8stypes.NamespacedName{Name: "pod1"}} + profile := NewSchedulerProfile(). + WithScorers(NewWeightedScorer(empty, 1)). + WithPicker(picker) + + input := []fwksched.Endpoint{ + fwksched.NewEndpoint(&fwkdl.EndpointMetadata{NamespacedName: k8stypes.NamespacedName{Name: "pod1"}}, nil, nil), + } + request := &fwksched.InferenceRequest{TargetModel: "test-model", RequestID: "req-1"} + + if _, err := profile.Run(context.Background(), request, input); err != nil { + t.Fatalf("unexpected error: %v", err) + } + + child := findSpan(tracetest.SpanStubsFromReadOnlySpans(recorder.Ended()), "llm_d.epp.scorer.empty") + if child == nil { + t.Fatal("missing scorer span for empty scorer") + } + // With no scored endpoints, aggregate attributes are omitted entirely. + if spanHasAttr(child, "llm_d.epp.scorer.score.avg") { + t.Error("empty scorer span should not carry a score.avg attribute") + } +} + func findEndpoints(endpoints []fwksched.Endpoint, names ...k8stypes.NamespacedName) []fwksched.Endpoint { res := []fwksched.Endpoint{} for _, endpoint := range endpoints {