Add picker tracing decorator to scheduling pipeline (llm-d#1708)

chethanuk · web-flow · commit 8bef5f0d7977 · 2026-07-14T20:48:07.000Z
* Trace the picker stage with an inline pick_endpoints span Wrap the picker call in runPickerPlugin in a single inline "pick_endpoints" span via tracing.Tracer(schedplugins.TracerScope) (SpanKindInternal), recording the candidate and selected endpoint counts (llm_d.epp.picker.candidate_endpoints, llm_d.epp.picker.selected_endpoints = len(result.TargetEndpoints), nil-guarded) plus the conditional gen_ai.request.{model,id} keys. request is threaded into runPickerPlugin (sole caller Run) so the span carries the request keys, matching the filter and scorer spans. Follows the single-step span convention from llm-d#1565 and llm-d#1693; picking behavior and the per-plugin latency metric are unchanged. Refs: llm-d#1694, llm-d#1483 Signed-off-by: ChethanUK <chethanuk@outlook.com> * feat(scheduling): record top endpoint scores on picker span The picker almost always returns a single target (with an occasional fallback), so the count of selected endpoints carries little diagnostic signal. What explains a routing decision is the score distribution across the strongest candidates: how close the runner-up was, and how the top scores were spread. Replace the selected_endpoints count attribute with the names and weighted scores of the highest-scoring candidates, ordered by descending score and capped at five so the payload stays bounded on production fleet sizes (~100 pods). The top set is captured before Pick runs because pickers reorder the candidate slice in place. Signed-off-by: ChethanUK <chethanuk@outlook.com> --------- Signed-off-by: ChethanUK <chethanuk@outlook.com> Signed-off-by: chethanuk <chethanuk@outlook.com>
diff --git a/pkg/epp/scheduling/constants.go b/pkg/epp/scheduling/constants.go
@@ -19,3 +19,8 @@ package scheduling
 // TracerScope is the OTel instrumentation scope for spans emitted by the
 // scheduling engine.
 const TracerScope = "llm-d-router/pkg/epp/scheduling"
+
+// maxTracedEndpointScores bounds how many of the highest-scoring endpoints are
+// recorded on the picker span. Candidate sets can reach production fleet sizes
+// (~100 pods); the strongest few carry the signal for why an endpoint was picked.
+const maxTracedEndpointScores = 5
diff --git a/pkg/epp/scheduling/scheduler_profile.go b/pkg/epp/scheduling/scheduler_profile.go
@@ -19,6 +19,7 @@ package scheduling
 import (
 	"context"
 	"fmt"
+	"sort"
 	"strings"
 	"time"
 
@@ -125,7 +126,7 @@ func (p *SchedulerProfile) Run(ctx context.Context, request *fwksched.InferenceR
 	// if we got here, there is at least one endpoint to score
 	weightedScorePerEndpoint := p.runScorerPlugins(ctx, request, endpoints)
 
-	result := p.runPickerPlugin(ctx, weightedScorePerEndpoint)
+	result := p.runPickerPlugin(ctx, request, weightedScorePerEndpoint)
 
 	return result, nil
 }
@@ -202,7 +203,7 @@ func (p *SchedulerProfile) runScorerPlugins(ctx context.Context, request *fwksch
 	return weightedScorePerEndpoint
 }
 
-func (p *SchedulerProfile) runPickerPlugin(ctx context.Context, weightedScorePerEndpoint map[fwksched.Endpoint]float64) *fwksched.ProfileRunResult {
+func (p *SchedulerProfile) runPickerPlugin(ctx context.Context, request *fwksched.InferenceRequest, weightedScorePerEndpoint map[fwksched.Endpoint]float64) *fwksched.ProfileRunResult {
 	logger := log.FromContext(ctx)
 
 	// Allocate the ScoredEndpoint values as a single contiguous backing array
@@ -222,6 +223,33 @@ func (p *SchedulerProfile) runPickerPlugin(ctx context.Context, weightedScorePer
 	}
 	logger.V(logutil.VERBOSE).Info("Running picker plugin", "plugin", p.picker.TypedName())
 	logger.V(logutil.DEBUG).Info("Candidate pods for picking", "endpoints-weighted-score", scoredEndpoints)
+
+	ctx, span := tracing.Tracer(TracerScope).Start(ctx, "pick_endpoints",
+		trace.WithSpanKind(trace.SpanKindInternal),
+	)
+	defer span.End()
+
+	span.SetAttributes(attribute.Int("llm_d.epp.picker.candidate_endpoints", len(scoredEndpoints)))
+	// The picker almost always returns a single target, so its count carries
+	// little signal. The score distribution across the strongest candidates is
+	// what explains why an endpoint was chosen, so record the highest-scoring
+	// few (names with their weighted scores). Captured before Pick because
+	// pickers reorder scoredEndpoints in place.
+	if names, scores := topScoredEndpoints(scoredEndpoints, maxTracedEndpointScores); len(names) > 0 {
+		span.SetAttributes(
+			attribute.StringSlice("llm_d.epp.picker.top_endpoints", names),
+			attribute.Float64Slice("llm_d.epp.picker.top_scores", scores),
+		)
+	}
+	if request != nil {
+		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))
+		}
+	}
+
 	before := time.Now()
 	result := p.picker.Pick(ctx, scoredEndpoints)
 	metrics.RecordPluginProcessingLatency(pickerExtensionPoint, p.picker.TypedName().Type, p.picker.TypedName().Name, time.Since(before))
@@ -230,6 +258,31 @@ func (p *SchedulerProfile) runPickerPlugin(ctx context.Context, weightedScorePer
 	return result
 }
 
+// topScoredEndpoints returns the names and weighted scores of the highest
+// scoring candidates, ordered by descending score with the endpoint name as a
+// stable tiebreaker and capped at limit. The returned slices are index-aligned.
+func topScoredEndpoints(scored []*fwksched.ScoredEndpoint, limit int) ([]string, []float64) {
+	ranked := make([]*fwksched.ScoredEndpoint, len(scored))
+	copy(ranked, scored)
+	sort.Slice(ranked, func(i, j int) bool {
+		if ranked[i].Score != ranked[j].Score {
+			return ranked[i].Score > ranked[j].Score
+		}
+		return ranked[i].GetMetadata().NamespacedName.String() <
+			ranked[j].GetMetadata().NamespacedName.String()
+	})
+	if limit < len(ranked) {
+		ranked = ranked[:limit]
+	}
+	names := make([]string, len(ranked))
+	scores := make([]float64, len(ranked))
+	for i, se := range ranked {
+		names[i] = se.GetMetadata().NamespacedName.String()
+		scores[i] = se.Score
+	}
+	return names, scores
+}
+
 func enforceScoreRange(score float64) float64 {
 	if score < 0 {
 		return 0
diff --git a/pkg/epp/scheduling/scheduler_profile_picker_tracing_test.go b/pkg/epp/scheduling/scheduler_profile_picker_tracing_test.go
@@ -0,0 +1,250 @@
+/*
+Copyright 2025 The Kubernetes Authors.
+
+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 scheduling
+
+import (
+	"context"
+	"reflect"
+	"testing"
+
+	"go.opentelemetry.io/otel"
+	"go.opentelemetry.io/otel/attribute"
+	sdktrace "go.opentelemetry.io/otel/sdk/trace"
+	"go.opentelemetry.io/otel/sdk/trace/tracetest"
+	"go.opentelemetry.io/otel/trace"
+	k8stypes "k8s.io/apimachinery/pkg/types"
+
+	fwkdl "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/datalayer"
+	fwkplugin "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/plugin"
+	fwksched "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/scheduling"
+)
+
+// setupPickerSpanRecorder installs an in-memory span recorder as the global
+// tracer provider and returns it, restoring the previous provider on cleanup.
+func setupPickerSpanRecorder(t *testing.T) *tracetest.SpanRecorder {
+	t.Helper()
+	recorder := tracetest.NewSpanRecorder()
+	tp := sdktrace.NewTracerProvider(sdktrace.WithSpanProcessor(recorder))
+	origTP := otel.GetTracerProvider()
+	otel.SetTracerProvider(tp)
+	t.Cleanup(func() { otel.SetTracerProvider(origTP) })
+	return recorder
+}
+
+func findPickerSpans(spans []sdktrace.ReadOnlySpan, name string) []sdktrace.ReadOnlySpan {
+	var out []sdktrace.ReadOnlySpan
+	for _, s := range spans {
+		if s.Name() == name {
+			out = append(out, s)
+		}
+	}
+	return out
+}
+
+func pickerSpanAttributes(span sdktrace.ReadOnlySpan) map[attribute.Key]attribute.Value {
+	attrs := make(map[attribute.Key]attribute.Value)
+	for _, kv := range span.Attributes() {
+		attrs[kv.Key] = kv.Value
+	}
+	return attrs
+}
+
+func newWeightedScores(names ...string) map[fwksched.Endpoint]float64 {
+	scores := make(map[fwksched.Endpoint]float64, len(names))
+	for i, name := range names {
+		endpoint := fwksched.NewEndpoint(
+			&fwkdl.EndpointMetadata{NamespacedName: k8stypes.NamespacedName{Name: name}}, nil, nil)
+		scores[endpoint] = float64(i)
+	}
+	return scores
+}
+
+// fakePicker returns the first selected candidates as targets, or a nil result
+// when nilResult is set, and optionally emits a child span from its context.
+type fakePicker struct {
+	typedName fwkplugin.TypedName
+	selected  int
+	nilResult bool
+	emitChild bool
+}
+
+var _ fwksched.Picker = &fakePicker{}
+
+func (f *fakePicker) TypedName() fwkplugin.TypedName { return f.typedName }
+
+func (f *fakePicker) Pick(ctx context.Context, scoredPods []*fwksched.ScoredEndpoint) *fwksched.ProfileRunResult {
+	if f.emitChild {
+		_, span := otel.Tracer("test").Start(ctx, "inner_picker_span")
+		span.End()
+	}
+	if f.nilResult {
+		return nil
+	}
+	targets := make([]fwksched.Endpoint, 0, f.selected)
+	for i := 0; i < f.selected && i < len(scoredPods); i++ {
+		targets = append(targets, scoredPods[i].Endpoint)
+	}
+	return &fwksched.ProfileRunResult{TargetEndpoints: targets}
+}
+
+func TestRunPickerPluginSingleSpan(t *testing.T) {
+	recorder := setupPickerSpanRecorder(t)
+
+	picker := &fakePicker{typedName: fwkplugin.TypedName{Type: "max-score", Name: "instance-a"}, selected: 1}
+	profile := NewSchedulerProfile().WithPicker(picker)
+	scores := newWeightedScores("pod1", "pod2", "pod3")
+
+	ctx, root := otel.Tracer("test").Start(context.Background(), "root")
+	result := profile.runPickerPlugin(ctx, &fwksched.InferenceRequest{TargetModel: "m1", RequestID: "r1"}, scores)
+	root.End()
+
+	if result == nil || len(result.TargetEndpoints) != 1 {
+		t.Fatalf("runPickerPlugin returned %v, want 1 target", result)
+	}
+
+	spans := findPickerSpans(recorder.Ended(), "pick_endpoints")
+	if len(spans) != 1 {
+		t.Fatalf("got %d pick_endpoints spans, want 1", len(spans))
+	}
+	span := spans[0]
+	if span.SpanKind() != trace.SpanKindInternal {
+		t.Errorf("span kind = %v, want Internal", span.SpanKind())
+	}
+	if span.Parent().SpanID() != root.SpanContext().SpanID() {
+		t.Errorf("parent span ID = %v, want root %v", span.Parent().SpanID(), root.SpanContext().SpanID())
+	}
+
+	attrs := pickerSpanAttributes(span)
+	for _, tc := range []struct {
+		key  attribute.Key
+		want string
+	}{
+		{"gen_ai.request.model", "m1"},
+		{"gen_ai.request.id", "r1"},
+	} {
+		if got := attrs[tc.key].AsString(); got != tc.want {
+			t.Errorf("%s = %q, want %q", tc.key, got, tc.want)
+		}
+	}
+	if got := attrs["llm_d.epp.picker.candidate_endpoints"].AsInt64(); got != 3 {
+		t.Errorf("candidate_endpoints = %d, want 3", got)
+	}
+	// newWeightedScores assigns score i to the i-th name, so the span records
+	// candidates highest score first.
+	if got, want := attrs["llm_d.epp.picker.top_endpoints"].AsStringSlice(), []string{"/pod3", "/pod2", "/pod1"}; !reflect.DeepEqual(got, want) {
+		t.Errorf("top_endpoints = %v, want %v", got, want)
+	}
+	if got, want := attrs["llm_d.epp.picker.top_scores"].AsFloat64Slice(), []float64{2, 1, 0}; !reflect.DeepEqual(got, want) {
+		t.Errorf("top_scores = %v, want %v", got, want)
+	}
+}
+
+func TestRunPickerPluginTopScoresCapped(t *testing.T) {
+	recorder := setupPickerSpanRecorder(t)
+
+	picker := &fakePicker{typedName: fwkplugin.TypedName{Type: "max-score", Name: "multi"}, selected: 2}
+	profile := NewSchedulerProfile().WithPicker(picker)
+	// Seven candidates (scores 0..6) exceed the cap of five.
+	scores := newWeightedScores("pod1", "pod2", "pod3", "pod4", "pod5", "pod6", "pod7")
+
+	ctx, root := otel.Tracer("test").Start(context.Background(), "root")
+	profile.runPickerPlugin(ctx, &fwksched.InferenceRequest{}, scores)
+	root.End()
+
+	spans := findPickerSpans(recorder.Ended(), "pick_endpoints")
+	if len(spans) != 1 {
+		t.Fatalf("got %d pick_endpoints spans, want 1", len(spans))
+	}
+	attrs := pickerSpanAttributes(spans[0])
+	if got, want := attrs["llm_d.epp.picker.top_endpoints"].AsStringSlice(), []string{"/pod7", "/pod6", "/pod5", "/pod4", "/pod3"}; !reflect.DeepEqual(got, want) {
+		t.Errorf("top_endpoints = %v, want the 5 highest-scoring %v", got, want)
+	}
+	if got, want := attrs["llm_d.epp.picker.top_scores"].AsFloat64Slice(), []float64{6, 5, 4, 3, 2}; !reflect.DeepEqual(got, want) {
+		t.Errorf("top_scores = %v, want %v", got, want)
+	}
+}
+
+func TestRunPickerPluginNilResult(t *testing.T) {
+	recorder := setupPickerSpanRecorder(t)
+
+	picker := &fakePicker{typedName: fwkplugin.TypedName{Type: "noop-picker", Name: "noop"}, nilResult: true}
+	profile := NewSchedulerProfile().WithPicker(picker)
+	scores := newWeightedScores("pod1", "pod2")
+
+	ctx, root := otel.Tracer("test").Start(context.Background(), "root")
+	result := profile.runPickerPlugin(ctx, &fwksched.InferenceRequest{}, scores)
+	root.End()
+
+	if result != nil {
+		t.Fatalf("runPickerPlugin returned %v, want nil", result)
+	}
+	spans := findPickerSpans(recorder.Ended(), "pick_endpoints")
+	if len(spans) != 1 {
+		t.Fatalf("got %d pick_endpoints spans, want 1 (span must end on nil result)", len(spans))
+	}
+	// Scores are captured from the candidates before Pick, so they are recorded
+	// even when the picker selects nothing.
+	if got, want := pickerSpanAttributes(spans[0])["llm_d.epp.picker.top_scores"].AsFloat64Slice(), []float64{1, 0}; !reflect.DeepEqual(got, want) {
+		t.Errorf("top_scores = %v, want %v", got, want)
+	}
+}
+
+func TestRunPickerPluginNestsDelegateSpan(t *testing.T) {
+	recorder := setupPickerSpanRecorder(t)
+
+	picker := &fakePicker{typedName: fwkplugin.TypedName{Type: "child", Name: "c"}, selected: 1, emitChild: true}
+	profile := NewSchedulerProfile().WithPicker(picker)
+	scores := newWeightedScores("pod1")
+
+	ctx, root := otel.Tracer("test").Start(context.Background(), "root")
+	profile.runPickerPlugin(ctx, &fwksched.InferenceRequest{}, scores)
+	root.End()
+
+	outer := findPickerSpans(recorder.Ended(), "pick_endpoints")
+	inner := findPickerSpans(recorder.Ended(), "inner_picker_span")
+	if len(outer) != 1 || len(inner) != 1 {
+		t.Fatalf("got %d pick_endpoints and %d inner spans, want 1 each", len(outer), len(inner))
+	}
+	if inner[0].Parent().SpanID() != outer[0].SpanContext().SpanID() {
+		t.Errorf("inner span parent = %v, want pick_endpoints span %v",
+			inner[0].Parent().SpanID(), outer[0].SpanContext().SpanID())
+	}
+}
+
+func TestRunPickerPluginOmitsEmptyGenAI(t *testing.T) {
+	recorder := setupPickerSpanRecorder(t)
+
+	picker := &fakePicker{typedName: fwkplugin.TypedName{Type: "p", Name: "n"}, selected: 1}
+	profile := NewSchedulerProfile().WithPicker(picker)
+	scores := newWeightedScores("pod1")
+
+	ctx, root := otel.Tracer("test").Start(context.Background(), "root")
+	profile.runPickerPlugin(ctx, &fwksched.InferenceRequest{}, scores)
+	root.End()
+
+	spans := findPickerSpans(recorder.Ended(), "pick_endpoints")
+	if len(spans) != 1 {
+		t.Fatalf("got %d pick_endpoints spans, want 1", len(spans))
+	}
+	attrs := pickerSpanAttributes(spans[0])
+	if _, ok := attrs["gen_ai.request.model"]; ok {
+		t.Error("gen_ai.request.model set for empty TargetModel")
+	}
+	if _, ok := attrs["gen_ai.request.id"]; ok {
+		t.Error("gen_ai.request.id set for empty RequestID")
+	}
+}
