feat: add OpenTelemetry spans for tokenization and KV-cache event processing

Signed-off-by: satyamg1620 <Satyam.Gupta.3@ibm.com>

Author: satyamg1620

pkg/epp/framework/plugins/requestcontrol/dataproducer/tokenizer/tokenizer.go

@@ -31,12 +31,18 @@ import (
 	"github.com/llm-d/llm-d-router/pkg/kvcache/kvblock"
 	"github.com/llm-d/llm-d-router/pkg/kvcache/tokenization"
 	tokenizerTypes "github.com/llm-d/llm-d-router/pkg/kvcache/tokenization/types"
+	"go.opentelemetry.io/otel/attribute"
+	"go.opentelemetry.io/otel/codes"
+	"go.opentelemetry.io/otel/trace"
 	"sigs.k8s.io/controller-runtime/pkg/log"
 
+	"github.com/llm-d/llm-d-router/pkg/common/observability/tracing"
 	"github.com/llm-d/llm-d-router/pkg/epp/framework/interface/plugin"
 	"github.com/llm-d/llm-d-router/pkg/epp/framework/interface/requestcontrol"
 	fwkrh "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/requesthandling"
 	"github.com/llm-d/llm-d-router/pkg/epp/framework/interface/scheduling"
+	mmobs "github.com/llm-d/llm-d-router/pkg/epp/framework/observability/multimodal"
+	rcplugins "github.com/llm-d/llm-d-router/pkg/epp/framework/plugins/requestcontrol"
 )
 
 type tokenizer interface {
@@ -57,6 +63,13 @@ const (
 	tokenizedPromptKeyID = "TokenizedPrompt"
 )
 
+// Backend identifiers reported on the tokenize span.
+const (
+	backendUDS      = "uds"
+	backendVLLM     = "vllm"
+	backendEstimate = "estimate"
+)
+
 var TokenizedPromptDataKey = plugin.NewDataKey(tokenizedPromptKeyID, PluginType)
 
 // tokenizerPluginConfig holds the configuration for the tokenizer plugin.
@@ -264,6 +277,7 @@ func LegacyPluginFactory(name string, rawParameters *json.Decoder, handle plugin
 // (the default when no backend is set).
 func NewPlugin(ctx context.Context, name string, config *tokenizerPluginConfig) (*Plugin, error) {
 	var backend tokenInputProducer
+	var backendName string
 	switch {
 	case config.TokenizerConfig.IsEnabled():
 		log.FromContext(ctx).Info(
@@ -275,6 +289,7 @@ func NewPlugin(ctx context.Context, name string, config *tokenizerPluginConfig)
 			return nil, fmt.Errorf("failed to initialize UDS tokenizer for '%s' plugin - %w", PluginType, err)
 		}
 		backend = renderBackend{tk: uds}
+		backendName = backendUDS
 	case config.VLLM != nil || config.ModelName != "":
 		cfg := config.VLLM
 		if cfg == nil {
@@ -285,14 +300,17 @@ func NewPlugin(ctx context.Context, name string, config *tokenizerPluginConfig)
 			return nil, fmt.Errorf("failed to initialize vLLM HTTP renderer for '%s' plugin - %w", PluginType, err)
 		}
 		backend = renderBackend{tk: renderer}
+		backendName = backendVLLM
 	default:
 		backend = estimateBackend{img: newImageEstimator(config.Estimate), vid: newVideoEstimator(config.Estimate)}
+		backendName = backendEstimate
 	}
 
 	p := &Plugin{
-		typedName: plugin.TypedName{Type: PluginType, Name: name},
-		backend:   backend,
-		dk:        TokenizedPromptDataKey.WithNonEmptyProducerName(name),
+		typedName:   plugin.TypedName{Type: PluginType, Name: name},
+		backend:     backend,
+		backendName: backendName,
+		dk:          TokenizedPromptDataKey.WithNonEmptyProducerName(name),
 	}
 	if w, ok := backend.(warmer); ok {
 		go w.warmup(ctx)
@@ -305,7 +323,9 @@ func NewPlugin(ctx context.Context, name string, config *tokenizerPluginConfig)
 type Plugin struct {
 	typedName plugin.TypedName
 	backend   tokenInputProducer
-	dk        plugin.DataKey
+	// backendName identifies the configured backend on the tokenize span.
+	backendName string
+	dk          plugin.DataKey
 }
 
 // compile-time assertions.
@@ -337,6 +357,9 @@ func (p *Plugin) ProduceTimeout() time.Duration {
 // Produce derives the request's TokenizedPrompt via the configured backend and
 // stores it on the body. Skips when one is already present; errors propagate to
 // the Director, which logs and continues.
+//
+// The tokenize span covers the backend call only, so the skip path stays
+// untraced and the span always represents tokenization work.
 func (p *Plugin) Produce(ctx context.Context, request *scheduling.InferenceRequest, _ []scheduling.Endpoint) error {
 	if request.Body == nil {
 		return errors.New("request body is nil")
@@ -350,16 +373,33 @@ func (p *Plugin) Produce(ctx context.Context, request *scheduling.InferenceReque
 		return nil
 	}
 
+	ctx, span := tracing.Tracer(rcplugins.TracerScope).Start(ctx, "tokenize",
+		trace.WithSpanKind(trace.SpanKindInternal),
+	)
+	defer span.End()
+
+	span.SetAttributes(attribute.String("llm_d.epp.token_producer.backend", p.backendName))
+	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))
+	}
+
 	ctx = withMMMetadata(ctx, parseMMMetadataHeaders(request.Headers))
 	tp, err := p.backend.produce(ctx, request.Body)
 	if err != nil {
+		span.SetStatus(codes.Error, err.Error())
 		return err
 	}
 	if tp == nil || tp.TokenCount() == 0 {
 		return nil
 	}
 	tp.CacheSalt = CacheSaltFromBody(request.Body)
 	request.Body.TokenizedPrompt = tp
+
+	span.SetAttributes(attribute.Int("llm_d.epp.token_producer.token_count", tp.TokenCount()))
+	span.SetAttributes(mmobs.SpanAttributes(request)...)
 	return nil
 }
 

pkg/epp/framework/plugins/requestcontrol/dataproducer/tokenizer/tokenizer_test.go

@@ -50,8 +50,9 @@ func (m *mockTokenizer) RenderChat(_ context.Context, payload fwkrh.RequestPaylo
 
 func newTestPlugin(tok tokenizer) *Plugin {
 	return &Plugin{
-		typedName: plugin.TypedName{Type: PluginType, Name: "test"},
-		backend:   renderBackend{tk: tok},
+		typedName:   plugin.TypedName{Type: PluginType, Name: "test"},
+		backend:     renderBackend{tk: tok},
+		backendName: backendVLLM,
 	}
 }
 

pkg/epp/framework/plugins/requestcontrol/dataproducer/tokenizer/tokenizer_tracing_test.go

@@ -0,0 +1,163 @@
+/*
+Copyright 2026 The llm-d 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 tokenizer
+
+import (
+	"context"
+	"errors"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"go.opentelemetry.io/otel"
+	"go.opentelemetry.io/otel/attribute"
+	"go.opentelemetry.io/otel/codes"
+	sdktrace "go.opentelemetry.io/otel/sdk/trace"
+	"go.opentelemetry.io/otel/sdk/trace/tracetest"
+
+	fwkrh "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/requesthandling"
+	"github.com/llm-d/llm-d-router/pkg/epp/framework/interface/scheduling"
+	"github.com/llm-d/llm-d-router/pkg/kvcache/kvblock"
+	"github.com/llm-d/llm-d-router/pkg/kvcache/tokenization"
+)
+
+// setupSpanRecorder installs an in-memory span recorder as the global tracer
+// provider and returns it, restoring the previous provider on cleanup.
+func setupSpanRecorder(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 tokenizeSpan(t *testing.T, recorder *tracetest.SpanRecorder) sdktrace.ReadOnlySpan {
+	t.Helper()
+	for _, s := range recorder.Ended() {
+		if s.Name() == "tokenize" {
+			return s
+		}
+	}
+	t.Fatal("no tokenize span recorded")
+	return nil
+}
+
+func spanAttrs(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 chatRequest() *scheduling.InferenceRequest {
+	return &scheduling.InferenceRequest{
+		RequestID:   "req-1",
+		TargetModel: "model-a",
+		Body: &fwkrh.InferenceRequestBody{
+			ChatCompletions: &fwkrh.ChatCompletionsRequest{
+				Messages: []fwkrh.Message{{Role: "user", Content: fwkrh.Content{Raw: "hi"}}},
+			},
+			Payload: fwkrh.PayloadMap{},
+		},
+	}
+}
+
+func TestProduce_EmitsTokenizeSpan(t *testing.T) {
+	recorder := setupSpanRecorder(t)
+	p := newTestPlugin(&mockTokenizer{
+		renderChatFunc: func(_ fwkrh.RequestPayload) ([]uint32, *tokenization.MultiModalFeatures, error) {
+			return []uint32{1, 2, 3}, nil, nil
+		},
+	})
+
+	require.NoError(t, p.Produce(context.Background(), chatRequest(), nil))
+
+	attrs := spanAttrs(tokenizeSpan(t, recorder))
+	assert.Equal(t, backendVLLM, attrs["llm_d.epp.token_producer.backend"].AsString())
+	assert.Equal(t, int64(3), attrs["llm_d.epp.token_producer.token_count"].AsInt64())
+	assert.Equal(t, "model-a", attrs["gen_ai.request.model"].AsString())
+	assert.Equal(t, "req-1", attrs["gen_ai.request.id"].AsString())
+	assert.Equal(t, "none", attrs["mm.modality"].AsString())
+	assert.Equal(t, int64(0), attrs["mm.hash_count"].AsInt64())
+}
+
+// The span must parent to the caller's span so tokenization appears under the
+// request trace rather than as a detached root.
+func TestProduce_TokenizeSpanParentedToCaller(t *testing.T) {
+	recorder := setupSpanRecorder(t)
+	p := newTestPlugin(&mockTokenizer{
+		renderChatFunc: func(_ fwkrh.RequestPayload) ([]uint32, *tokenization.MultiModalFeatures, error) {
+			return []uint32{1}, nil, nil
+		},
+	})
+
+	ctx, parent := otel.Tracer("test").Start(context.Background(), "parent")
+	require.NoError(t, p.Produce(ctx, chatRequest(), nil))
+	parent.End()
+
+	span := tokenizeSpan(t, recorder)
+	assert.Equal(t, parent.SpanContext().TraceID(), span.SpanContext().TraceID())
+	assert.Equal(t, parent.SpanContext().SpanID(), span.Parent().SpanID())
+}
+
+func TestProduce_TokenizeSpanRecordsMultiModal(t *testing.T) {
+	recorder := setupSpanRecorder(t)
+	p := newTestPlugin(&mockTokenizer{
+		renderChatFunc: func(_ fwkrh.RequestPayload) ([]uint32, *tokenization.MultiModalFeatures, error) {
+			return []uint32{1, 2}, &tokenization.MultiModalFeatures{
+				MMHashes:       map[string][]string{"image": {"h1"}},
+				MMPlaceholders: map[string][]kvblock.PlaceholderRange{"image": {{Offset: 0, Length: 2}}},
+			}, nil
+		},
+	})
+
+	require.NoError(t, p.Produce(context.Background(), chatRequest(), nil))
+
+	attrs := spanAttrs(tokenizeSpan(t, recorder))
+	assert.Equal(t, "image", attrs["mm.modality"].AsString())
+	assert.Equal(t, int64(1), attrs["mm.hash_count"].AsInt64())
+}
+
+func TestProduce_TokenizeSpanRecordsError(t *testing.T) {
+	recorder := setupSpanRecorder(t)
+	p := newTestPlugin(&mockTokenizer{
+		renderChatFunc: func(_ fwkrh.RequestPayload) ([]uint32, *tokenization.MultiModalFeatures, error) {
+			return nil, nil, errors.New("render failed")
+		},
+	})
+
+	require.Error(t, p.Produce(context.Background(), chatRequest(), nil))
+
+	span := tokenizeSpan(t, recorder)
+	assert.Equal(t, codes.Error, span.Status().Code)
+	assert.Contains(t, span.Status().Description, "render failed")
+}
+
+// The skip path does no tokenization work, so it must not emit a span.
+func TestProduce_NoSpanWhenAlreadyTokenized(t *testing.T) {
+	recorder := setupSpanRecorder(t)
+	p := newTestPlugin(&mockTokenizer{})
+
+	req := chatRequest()
+	req.Body.TokenizedPrompt = &fwkrh.TokenizedPrompt{PerPromptTokens: [][]uint32{{7}}}
+	require.NoError(t, p.Produce(context.Background(), req, nil))
+
+	assert.Empty(t, recorder.Ended())
+}

pkg/kvevents/engineadapter/sglang_adapter.go

@@ -17,6 +17,7 @@ limitations under the License.
 package engineadapter
 
 import (
+	"context"
 	"fmt"
 
 	"github.com/vmihailenco/msgpack/v5"
@@ -68,7 +69,7 @@ func (s *SGLangAdapter) ShardingKey(msg *kvevents.RawMessage) string {
 // and decodes the msgpack payload into an EventBatch.
 //
 //nolint:gocritic // unnamedResult: named returns conflict with nonamedreturns linter
-func (s *SGLangAdapter) ParseMessage(msg *kvevents.RawMessage) (string, string, kvevents.EventBatch, error) {
+func (s *SGLangAdapter) ParseMessage(_ context.Context, msg *kvevents.RawMessage) (string, string, kvevents.EventBatch, error) {
 	podID, modelName := parseTopic(msg.Topic)
 
 	var batch msgpackSGLangEventBatch

pkg/kvevents/engineadapter/sglang_adapter_test.go

@@ -17,6 +17,7 @@ limitations under the License.
 package engineadapter //nolint:testpackage // Tests access unexported functions
 
 import (
+	"context"
 	"testing"
 
 	"github.com/llm-d/llm-d-router/pkg/kvevents"
@@ -61,7 +62,7 @@ func TestSGLangParseMessage_Valid(t *testing.T) {
 		Payload:  payload,
 	}
 
-	podID, modelName, eventBatch, err := adapter.ParseMessage(msg)
+	podID, modelName, eventBatch, err := adapter.ParseMessage(context.Background(), msg)
 	require.NoError(t, err)
 	assert.Equal(t, "pod-1", podID)
 	assert.Equal(t, "llama-2-7b", modelName)
@@ -82,7 +83,7 @@ func TestSGLangParseMessage_InvalidPayload(t *testing.T) {
 		Payload: []byte{0xFF, 0xFF, 0xFF},
 	}
 
-	_, _, _, err := adapter.ParseMessage(msg)
+	_, _, _, err := adapter.ParseMessage(context.Background(), msg)
 	assert.Error(t, err)
 }
 

pkg/kvevents/engineadapter/vllm_adapter.go

@@ -17,6 +17,7 @@ limitations under the License.
 package engineadapter
 
 import (
+	"context"
 	"fmt"
 
 	"github.com/vmihailenco/msgpack/v5"
@@ -61,7 +62,7 @@ func (v *VLLMAdapter) ShardingKey(msg *kvevents.RawMessage) string {
 // and decodes the msgpack payload into an EventBatch.
 //
 //nolint:gocritic // unnamedResult: named returns conflict with nonamedreturns linter
-func (v *VLLMAdapter) ParseMessage(msg *kvevents.RawMessage) (string, string, kvevents.EventBatch, error) {
+func (v *VLLMAdapter) ParseMessage(_ context.Context, msg *kvevents.RawMessage) (string, string, kvevents.EventBatch, error) {
 	podID, modelName := parseTopic(msg.Topic)
 
 	var vllmBatch msgpackVLLMEventBatch

pkg/kvevents/engineadapter/vllm_adapter_bench_test.go

@@ -17,6 +17,7 @@ limitations under the License.
 package engineadapter //nolint:testpackage // Benchmarks access unexported functions
 
 import (
+	"context"
 	"fmt"
 	"testing"
 
@@ -154,7 +155,7 @@ func BenchmarkParseMessage_Batch(b *testing.B) {
 			b.SetBytes(int64(len(batchPayload)))
 			b.ResetTimer()
 			for i := 0; i < b.N; i++ {
-				_, _, _, err := adapter.ParseMessage(msg)
+				_, _, _, err := adapter.ParseMessage(context.Background(), msg)
 				if err != nil {
 					b.Fatal(err)
 				}