address comments from shaneutt

Signed-off-by: Guangya Liu <gyliu513@gmail.com>

Author: gyliu513

pkg/framework/interface/modelselector/plugins.go

@@ -42,5 +42,8 @@ type Scorer interface {
 // Picker picks the final model(s) to send the request to.
 type Picker interface {
 	plugin.Plugin
+	// Pick selects a target model from the scored candidates. Implementations
+	// must return a non-nil *PipelineRunResult with a non-nil TargetModel;
+	// callers rely on this guarantee and do not nil-check the result.
 	Pick(ctx context.Context, cycleState *plugin.CycleState, scoredModels []*ScoredModel) *PipelineRunResult
 }

pkg/handlers/request_test.go

@@ -19,13 +19,18 @@ package handlers
 import (
 	"context"
 	"encoding/json"
+	"errors"
 	"strconv"
 	"strings"
 	"testing"
 
 	basepb "github.com/envoyproxy/go-control-plane/envoy/config/core/v3"
 	extProcPb "github.com/envoyproxy/go-control-plane/envoy/service/ext_proc/v3"
 	"github.com/google/go-cmp/cmp"
+	"go.opentelemetry.io/otel"
+	otelcodes "go.opentelemetry.io/otel/codes"
+	sdktrace "go.opentelemetry.io/otel/sdk/trace"
+	"go.opentelemetry.io/otel/sdk/trace/tracetest"
 	"google.golang.org/protobuf/testing/protocmp"
 	metricsutils "k8s.io/component-base/metrics/testutil"
 	crmetrics "sigs.k8s.io/controller-runtime/pkg/metrics"
@@ -589,6 +594,82 @@ func (p *fakeRequestPlugin) ProcessRequest(ctx context.Context, _ *plugin.CycleS
 
 var _ requesthandling.RequestProcessor = &fakeRequestPlugin{}
 
+// TestRunRequestPlugins_Spans verifies that runRequestPlugins emits a
+// "request_plugins" stage span with a child "plugin.*" span per plugin, and
+// that a failing plugin's span is marked with an error status.
+func TestRunRequestPlugins_Spans(t *testing.T) {
+	recorder := tracetest.NewSpanRecorder()
+	tp := sdktrace.NewTracerProvider(sdktrace.WithSpanProcessor(recorder))
+	origTP := otel.GetTracerProvider()
+	otel.SetTracerProvider(tp)
+	t.Cleanup(func() { otel.SetTracerProvider(origTP) })
+
+	wantErr := errors.New("boom")
+	plugins := []requesthandling.RequestProcessor{
+		&fakeRequestPlugin{name: "ok", mutateFn: func(context.Context, *requesthandling.InferenceRequest) error { return nil }},
+		&fakeRequestPlugin{name: "boom", mutateFn: func(context.Context, *requesthandling.InferenceRequest) error { return wantErr }},
+	}
+
+	s := &Server{}
+	err := s.runRequestPlugins(context.Background(), plugin.NewCycleState(), requesthandling.NewInferenceRequest(), plugins)
+	if !errors.Is(err, wantErr) {
+		t.Fatalf("runRequestPlugins error = %v, want %v", err, wantErr)
+	}
+
+	ended := recorder.Ended()
+
+	// Stage span grouping the per-plugin spans.
+	stage := findSpan(t, ended, "request_plugins")
+
+	// Both plugins ran (the first succeeded, the second failed and aborted the
+	// loop), so two plugin spans should exist, each parented to the stage span.
+	pluginSpans := findSpans(ended, "plugin.fake")
+	if len(pluginSpans) != 2 {
+		t.Fatalf("expected 2 plugin.fake spans, got %d", len(pluginSpans))
+	}
+	for _, ps := range pluginSpans {
+		if ps.Parent().SpanID() != stage.SpanContext().SpanID() {
+			t.Errorf("plugin span parent = %v, want stage span %v", ps.Parent().SpanID(), stage.SpanContext().SpanID())
+		}
+	}
+
+	// Exactly one plugin span (the failing one) carries an error status.
+	var errored int
+	for _, ps := range pluginSpans {
+		if ps.Status().Code == otelcodes.Error {
+			errored++
+			if ps.Status().Description != wantErr.Error() {
+				t.Errorf("error span description = %q, want %q", ps.Status().Description, wantErr.Error())
+			}
+		}
+	}
+	if errored != 1 {
+		t.Errorf("expected exactly 1 plugin span with error status, got %d", errored)
+	}
+}
+
+// findSpan returns the single ended span with the given name, failing the test
+// if there is not exactly one.
+func findSpan(t *testing.T, spans []sdktrace.ReadOnlySpan, name string) sdktrace.ReadOnlySpan {
+	t.Helper()
+	matches := findSpans(spans, name)
+	if len(matches) != 1 {
+		t.Fatalf("expected exactly 1 span named %q, got %d", name, len(matches))
+	}
+	return matches[0]
+}
+
+// findSpans returns all ended spans with the given name.
+func findSpans(spans []sdktrace.ReadOnlySpan, name string) []sdktrace.ReadOnlySpan {
+	var matches []sdktrace.ReadOnlySpan
+	for _, s := range spans {
+		if s.Name() == name {
+			matches = append(matches, s)
+		}
+	}
+	return matches
+}
+
 // TestHandleRequestBody_MultiPluginHeaderMutations tests the end-to-end behavior of
 // HandleRequestBody when multiple request plugins set and/or remove headers.
 // Each sub-test verifies the HeaderMutation in the resulting ProcessingResponse.

pkg/handlers/server.go

@@ -23,8 +23,6 @@ import (
 	"time"
 
 	extProcPb "github.com/envoyproxy/go-control-plane/envoy/service/ext_proc/v3"
-	"go.opentelemetry.io/otel"
-	"go.opentelemetry.io/otel/attribute"
 	"go.opentelemetry.io/otel/trace"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
@@ -33,11 +31,11 @@ import (
 	envoy "github.com/llm-d/llm-d-inference-payload-processor/pkg/common/envoy"
 	errcommon "github.com/llm-d/llm-d-inference-payload-processor/pkg/common/error"
 	logutil "github.com/llm-d/llm-d-inference-payload-processor/pkg/common/observability/logging"
+	"github.com/llm-d/llm-d-inference-payload-processor/pkg/common/observability/tracing"
 	datasource "github.com/llm-d/llm-d-inference-payload-processor/pkg/framework/interface/datalayer/datasource"
 	"github.com/llm-d/llm-d-inference-payload-processor/pkg/framework/interface/plugin"
 	"github.com/llm-d/llm-d-inference-payload-processor/pkg/framework/interface/requesthandling"
 	"github.com/llm-d/llm-d-inference-payload-processor/pkg/metrics"
-	"github.com/llm-d/llm-d-inference-payload-processor/version"
 )
 
 const (
@@ -89,14 +87,7 @@ func (s *Server) Process(srv extProcPb.ExternalProcessor_ProcessServer) error {
 	ctx := srv.Context()
 
 	// Start tracing span for the request
-	tracer := otel.Tracer(
-		"llm-d-inference-payload-processor/pkg/handlers",
-		trace.WithInstrumentationVersion(version.BuildRef),
-		trace.WithInstrumentationAttributes(
-			attribute.String("commit-sha", version.CommitSHA),
-		),
-	)
-	ctx, span := tracer.Start(ctx, "gateway.request", trace.WithSpanKind(trace.SpanKindServer))
+	ctx, span := tracing.Tracer(handlersTracerScope).Start(ctx, "gateway.request", trace.WithSpanKind(trace.SpanKindServer))
 	defer span.End()
 
 	logger := log.FromContext(ctx)

pkg/modelselector/model_selector_pipeline_test.go

@@ -20,6 +20,11 @@ import (
 	"context"
 	"testing"
 
+	"go.opentelemetry.io/otel"
+	"go.opentelemetry.io/otel/attribute"
+	sdktrace "go.opentelemetry.io/otel/sdk/trace"
+	"go.opentelemetry.io/otel/sdk/trace/tracetest"
+
 	"github.com/llm-d/llm-d-inference-payload-processor/pkg/framework/interface/datalayer"
 	"github.com/llm-d/llm-d-inference-payload-processor/pkg/framework/interface/modelselector"
 	"github.com/llm-d/llm-d-inference-payload-processor/pkg/framework/interface/plugin"
@@ -80,6 +85,78 @@ func (p *testPicker) Pick(_ context.Context, _ *plugin.CycleState, scoredModels
 	return &modelselector.PipelineRunResult{TargetModel: best.Model}
 }
 
+// TestRunPickerPlugin_Span verifies that runPickerPlugin emits a
+// "plugin.<type>" span carrying the plugin and picker attributes, including the
+// selected model.
+func TestRunPickerPlugin_Span(t *testing.T) {
+	recorder := tracetest.NewSpanRecorder()
+	tp := sdktrace.NewTracerProvider(sdktrace.WithSpanProcessor(recorder))
+	origTP := otel.GetTracerProvider()
+	otel.SetTracerProvider(tp)
+	t.Cleanup(func() { otel.SetTracerProvider(origTP) })
+
+	picker := &testPicker{typedName: plugin.TypedName{Type: "test-picker", Name: "max-score"}}
+	pipeline := NewModelSelectorPipeline().WithPicker(picker)
+
+	scoredModels := map[string]*modelselector.ScoredModel{
+		"model-a": {Model: datalayer.NewModel("model-a"), Score: 0.2},
+		"model-b": {Model: datalayer.NewModel("model-b"), Score: 0.9},
+	}
+
+	result := pipeline.runPickerPlugin(context.Background(), plugin.NewCycleState(), scoredModels)
+	if result == nil || result.TargetModel == nil {
+		t.Fatalf("runPickerPlugin returned no target model")
+	}
+	if got := result.TargetModel.GetName(); got != "model-b" {
+		t.Fatalf("selected model = %q, want %q", got, "model-b")
+	}
+
+	span := findSpanByName(t, recorder.Ended(), "plugin.test-picker")
+	attrs := attrMap(span.Attributes())
+
+	wantStrings := map[attribute.Key]string{
+		"llm_d.plugin.extension_point": pickerExtensionPoint,
+		"llm_d.plugin.type":            "test-picker",
+		"llm_d.plugin.name":            "max-score",
+		"llm_d.picker.selected_model":  "model-b",
+	}
+	for k, want := range wantStrings {
+		if got := attrs[k].AsString(); got != want {
+			t.Errorf("attribute %q = %q, want %q", k, got, want)
+		}
+	}
+	if got := attrs["llm_d.picker.candidate_count"].AsInt64(); got != 2 {
+		t.Errorf("candidate_count = %d, want 2", got)
+	}
+}
+
+// findSpanByName returns the single ended span with the given name, failing the
+// test if there is not exactly one.
+func findSpanByName(t *testing.T, spans []sdktrace.ReadOnlySpan, name string) sdktrace.ReadOnlySpan {
+	t.Helper()
+	var match sdktrace.ReadOnlySpan
+	count := 0
+	for _, s := range spans {
+		if s.Name() == name {
+			match = s
+			count++
+		}
+	}
+	if count != 1 {
+		t.Fatalf("expected exactly 1 span named %q, got %d", name, count)
+	}
+	return match
+}
+
+// attrMap indexes a span's attributes by key for easy lookup.
+func attrMap(kvs []attribute.KeyValue) map[attribute.Key]attribute.Value {
+	m := make(map[attribute.Key]attribute.Value, len(kvs))
+	for _, kv := range kvs {
+		m[kv.Key] = kv.Value
+	}
+	return m
+}
+
 func TestPipelineRun(t *testing.T) {
 	modelA := datalayer.NewModel("model-a")
 	modelB := datalayer.NewModel("model-b")