@@ -20,11 +20,15 @@ import (
2020 "context"
2121 "fmt"
2222
23+ "go.opentelemetry.io/otel/attribute"
24+ "go.opentelemetry.io/otel/codes"
25+ "go.opentelemetry.io/otel/trace"
2326 "k8s.io/apimachinery/pkg/util/sets"
2427 "sigs.k8s.io/controller-runtime/pkg/log"
2528
2629 "github.com/llm-d/llm-d-kv-cache-manager/pkg/kvcache/kvblock"
2730 preprocessing "github.com/llm-d/llm-d-kv-cache-manager/pkg/preprocessing/chat_completions"
31+ "github.com/llm-d/llm-d-kv-cache-manager/pkg/telemetry"
2832 "github.com/llm-d/llm-d-kv-cache-manager/pkg/tokenization"
2933 "github.com/llm-d/llm-d-kv-cache-manager/pkg/tokenization/prefixstore"
3034 "github.com/llm-d/llm-d-kv-cache-manager/pkg/utils/logging"
@@ -132,6 +136,19 @@ func (k *Indexer) KVBlockIndex() kvblock.Index {
132136func (k * Indexer ) GetPodScores (ctx context.Context , renderReq * preprocessing.RenderJinjaTemplateRequest , prompt , modelName string ,
133137 podIdentifiers []string ,
134138) (map [string ]float64 , error ) {
139+ // Start tracing span for main operation
140+ tracer := telemetry .Tracer ()
141+ ctx , span := tracer .Start (ctx , "kvcache.manager.get_scores" ,
142+ trace .WithSpanKind (trace .SpanKindServer ),
143+ )
144+ defer span .End ()
145+
146+ // Set initial attributes
147+ span .SetAttributes (
148+ attribute .String ("gen_ai.request.model" , modelName ),
149+ attribute .Int ("kvcache.pod_count" , len (podIdentifiers )),
150+ )
151+
135152 traceLogger := log .FromContext (ctx ).V (logging .TRACE ).WithName ("kvcache.GetPodScores" )
136153
137154 // 1. tokenize prompt
@@ -141,30 +158,142 @@ func (k *Indexer) GetPodScores(ctx context.Context, renderReq *preprocessing.Ren
141158 blockKeys := k .tokensProcessor .TokensToKVBlockKeys (tokens , modelName )
142159 if len (blockKeys ) == 0 {
143160 traceLogger .Info ("no block keys found, returning empty scores" )
161+ span .SetAttributes (attribute .Int ("kvcache.block_keys.count" , 0 ))
162+ span .SetStatus (codes .Ok , "" )
144163 //nolint:nilnil // no need to return an error
145164 return nil , nil
146165 }
147166
167+ span .SetAttributes (attribute .Int ("kvcache.block_keys.count" , len (blockKeys )))
148168 traceLogger .Info ("found tokens" , "tokens" , tokens , "block-keys" , blockKeys )
149169
150- // 3. query kvblock indexer for pods
151- keyToPods , err := k .kvBlockIndex . Lookup (ctx , blockKeys , sets .New (podIdentifiers ... ))
170+ // 3. query kvblock indexer for pods (with child span)
171+ keyToPods , err := k .lookupWithSpan (ctx , blockKeys , sets .New (podIdentifiers ... ))
152172 if err != nil {
173+ span .RecordError (err )
174+ span .SetStatus (codes .Error , err .Error ())
153175 return nil , fmt .Errorf ("failed to query kvblock indexer: %w" , err )
154176 }
155177 traceLogger .Info ("found block keys" , "block-keys" , blockKeys ,
156178 "pods" , podsPerKeyPrintHelper (keyToPods ))
157179
158- // 4. score pods
159- podScores , err := k .kvBlockScorer .Score (blockKeys , keyToPods )
180+ // Calculate total blocks available
181+ totalBlocksAvailable := 0
182+ for _ , pods := range keyToPods {
183+ totalBlocksAvailable += len (pods )
184+ }
185+ span .SetAttributes (attribute .Int ("kvcache.total_blocks_available" , totalBlocksAvailable ))
186+
187+ // 4. score pods (with child span)
188+ podScores , err := k .scoreWithSpan (ctx , blockKeys , keyToPods )
160189 if err != nil {
190+ span .RecordError (err )
191+ span .SetStatus (codes .Error , err .Error ())
161192 return nil , fmt .Errorf ("failed to query kvblock scorer: %w" , err )
162193 }
163194 traceLogger .Info ("found pod scores" , "pod-scores" , podScores )
164195
196+ // Calculate hit ratio (pods with non-zero scores / total pods)
197+ podsWithHits := 0
198+ for _ , score := range podScores {
199+ if score > 0 {
200+ podsWithHits ++
201+ }
202+ }
203+ hitRatio := 0.0
204+ if len (podIdentifiers ) > 0 {
205+ hitRatio = float64 (podsWithHits ) / float64 (len (podIdentifiers ))
206+ }
207+ span .SetAttributes (
208+ attribute .Float64 ("kvcache.hit_ratio" , hitRatio ),
209+ attribute .Int ("kvcache.pods_with_hits" , podsWithHits ),
210+ )
211+
212+ span .SetStatus (codes .Ok , "" )
165213 return podScores , nil
166214}
167215
216+ // lookupWithSpan wraps kvBlockIndex.Lookup with a tracing span
217+ func (k * Indexer ) lookupWithSpan (ctx context.Context , blockKeys []kvblock.Key , podSet sets.Set [string ]) (map [kvblock.Key ][]kvblock.PodEntry , error ) {
218+ tracer := telemetry .Tracer ()
219+ ctx , span := tracer .Start (ctx , "kvcache.storage.lookup" ,
220+ trace .WithSpanKind (trace .SpanKindInternal ),
221+ )
222+ defer span .End ()
223+
224+ span .SetAttributes (
225+ attribute .Int ("kvcache.lookup.block_count" , len (blockKeys )),
226+ attribute .Int ("kvcache.lookup.pod_filter_count" , podSet .Len ()),
227+ )
228+
229+ result , err := k .kvBlockIndex .Lookup (ctx , blockKeys , podSet )
230+ if err != nil {
231+ span .RecordError (err )
232+ span .SetStatus (codes .Error , err .Error ())
233+ return nil , err
234+ }
235+
236+ // Calculate cache hit metrics
237+ blocksFound := 0
238+ for _ , pods := range result {
239+ if len (pods ) > 0 {
240+ blocksFound ++
241+ }
242+ }
243+ cacheHit := blocksFound > 0
244+
245+ span .SetAttributes (
246+ attribute .Bool ("kvcache.lookup.cache_hit" , cacheHit ),
247+ attribute .Int ("kvcache.lookup.blocks_found" , blocksFound ),
248+ )
249+
250+ span .SetStatus (codes .Ok , "" )
251+ return result , nil
252+ }
253+
254+ // scoreWithSpan wraps kvBlockScorer.Score with a tracing span
255+ func (k * Indexer ) scoreWithSpan (ctx context.Context , keys []kvblock.Key , keyToPods map [kvblock.Key ][]kvblock.PodEntry ) (map [string ]float64 , error ) {
256+ tracer := telemetry .Tracer ()
257+ ctx , span := tracer .Start (ctx , "kvcache.scorer.compute" ,
258+ trace .WithSpanKind (trace .SpanKindInternal ),
259+ )
260+ defer span .End ()
261+
262+ span .SetAttributes (
263+ attribute .String ("kvcache.scorer.algorithm" , string (k .kvBlockScorer .Strategy ())),
264+ attribute .Int ("kvcache.scorer.key_count" , len (keys )),
265+ )
266+
267+ scores , err := k .kvBlockScorer .Score (keys , keyToPods )
268+ if err != nil {
269+ span .RecordError (err )
270+ span .SetStatus (codes .Error , err .Error ())
271+ return nil , err
272+ }
273+
274+ // Calculate score distribution
275+ if len (scores ) > 0 {
276+ maxScore := 0.0
277+ totalScore := 0.0
278+ for _ , score := range scores {
279+ if score > maxScore {
280+ maxScore = score
281+ }
282+ totalScore += score
283+ }
284+ avgScore := totalScore / float64 (len (scores ))
285+
286+ span .SetAttributes (
287+ attribute .Float64 ("kvcache.score.max" , maxScore ),
288+ attribute .Float64 ("kvcache.score.avg" , avgScore ),
289+ attribute .Int ("kvcache.scorer.pods_scored" , len (scores )),
290+ )
291+ }
292+
293+ span .SetStatus (codes .Ok , "" )
294+ return scores , nil
295+ }
296+
168297// podsPerKeyPrintHelper formats a map of keys to pod entries for printing.
169298func podsPerKeyPrintHelper (ks map [kvblock.Key ][]kvblock.PodEntry ) string {
170299 flattened := ""
0 commit comments