add more spans

Signed-off-by: sallyom <[email protected]>

Author: sallyom

pkg/kvcache/indexer.go

@@ -67,13 +67,24 @@ type Indexer struct {
 
 // NewKVCacheIndexer creates a KVCacheIndex given a Config.
 func NewKVCacheIndexer(ctx context.Context, config *Config) (*Indexer, error) {
+	tracer := otel.GetTracerProvider().Tracer("llm-d-kv-cache-manager")
+	ctx, span := tracer.Start(ctx, "llm_d.kv_cache_manager.initialization")
+	defer span.End()
+
+	span.SetAttributes(
+		attribute.String("component", "llm-d-kv-cache-manager"),
+		attribute.String("operation", "initialization"),
+	)
+
 	logger := klog.FromContext(ctx)
 	if config != nil && config.TokenProcessorConfig != nil {
 		logger.Info("NewKVCacheIndexer config", "blockSize", config.TokenProcessorConfig.BlockSize)
+		span.SetAttributes(attribute.Int("llm_d.kv_cache_manager.block_size", config.TokenProcessorConfig.BlockSize))
 	}
 
 	tokensIndexer, err := prefixstore.NewLRUTokenStore(config.PrefixStoreConfig)
 	if err != nil {
+		span.SetAttributes(attribute.String("operation.outcome", "error"))
 		return nil, fmt.Errorf("failed to create prefixstore.Indexer: %w", err)
 	}
 
@@ -91,9 +102,11 @@ func NewKVCacheIndexer(ctx context.Context, config *Config) (*Indexer, error) {
 
 	tokenizersPool, err := tokenization.NewTokenizationPool(config.TokenizersPoolConfig, tokensIndexer)
 	if err != nil {
+		span.SetAttributes(attribute.String("operation.outcome", "error"))
 		return nil, fmt.Errorf("failed to create tokenizers pool: %w", err)
 	}
 
+	span.SetAttributes(attribute.String("operation.outcome", "success"))
 	return &Indexer{
 		config:          config,
 		tokensIndexer:   tokensIndexer,
@@ -124,47 +137,107 @@ func (k *Indexer) KVBlockIndex() kvblock.Index {
 func (k *Indexer) GetPodScores(ctx context.Context, prompt, modelName string,
 	podIdentifiers []string,
 ) (map[string]int, error) {
-	tracer := otel.GetTracerProvider().Tracer("llm-d-epp")
-	ctx, span := tracer.Start(ctx, "kv-cache-manager.GetPodScores")
+	tracer := otel.GetTracerProvider().Tracer("llm-d-kv-cache-manager")
+	ctx, span := tracer.Start(ctx, "llm_d.kv_cache_manager.GetPodScores")
 	defer span.End()
 
 	span.SetAttributes(
-		attribute.String("component", "llm-d-kv-cache-manager"),
-		attribute.String("operation", "get_pod_scores"),
 		attribute.String("gen_ai.request.model", modelName),
-		attribute.Int("llm_d.kv_cache.pod_count", len(podIdentifiers)),
+		attribute.Int("llm_d.kv_cache_manager.pod_count", len(podIdentifiers)),
 	)
 
 	traceLogger := klog.FromContext(ctx).V(logging.TRACE).WithName("kvcache.GetPodScores")
 
 	// 1. tokenize prompt
+	// 1. get available tokens of longest prefix
+	_, tokenSpan := tracer.Start(ctx, "llm_d.kv_cache_manager.find_tokens")
+	tokenSpan.SetAttributes(
+		attribute.String("gen_ai.request.model", modelName),
+	)
 	tokens := k.tokenizersPool.Tokenize(prompt, modelName)
+	if len(tokens) == 0 {
+		tokenSpan.SetAttributes(
+			attribute.Int("llm_d.kv_cache_manager.tokens_found", 0),
+			attribute.String("operation.outcome", "success"),
+		)
+		tokenSpan.End()
+		//nolint:nilnil // no need to return an error
+		return nil, nil
+	}
+	tokenSpan.SetAttributes(
+		attribute.Int("llm_d.kv_cache_manager.tokens_found", len(tokens)),
+		attribute.String("operation.outcome", "success"),
+	)
+	tokenSpan.End()
 
 	// 2. get block keys
+	_, blockSpan := tracer.Start(ctx, "llm_d.kv_cache_manager.tokens_to_block_keys")
+	blockSpan.SetAttributes(
+		attribute.String("gen_ai.request.model", modelName),
+		attribute.Int("llm_d.kv_cache_manager.input_tokens", len(tokens)),
+	)
 	blockKeys := k.tokensProcessor.TokensToKVBlockKeys(tokens, modelName)
 	if len(blockKeys) == 0 {
+		blockSpan.SetAttributes(
+			attribute.Int("llm_d.kv_cache_manager.block_keys_generated", 0),
+			attribute.String("operation.outcome", "success"),
+		)
+		blockSpan.End()
 		traceLogger.Info("no block keys found, returning empty scores")
 		//nolint:nilnil // no need to return an error
 		return nil, nil
 	}
+	blockSpan.SetAttributes(
+		attribute.Int("llm_d.kv_cache_manager.block_keys_generated", len(blockKeys)),
+		attribute.String("operation.outcome", "success"),
+	)
+	blockSpan.End()
 
 	traceLogger.Info("found tokens", "tokens", tokens, "block-keys", blockKeys)
 
 	// 3. query kvblock indexer for pods
+	_, lookupSpan := tracer.Start(ctx, "llm_d.kv_cache_manager.lookup_pods")
+	lookupSpan.SetAttributes(
+		attribute.String("gen_ai.request.model", modelName),
+		attribute.Int("llm_d.kv_cache_manager.block_keys_count", len(blockKeys)),
+	)
 	keyToPods, err := k.kvBlockIndex.Lookup(ctx, blockKeys, sets.New(podIdentifiers...))
 	if err != nil {
+		lookupSpan.RecordError(err)
+		lookupSpan.SetAttributes(attribute.String("operation.outcome", "error"))
+		lookupSpan.End()
 		span.RecordError(err)
+		span.SetAttributes(attribute.String("operation.outcome", "error"))
 		return nil, fmt.Errorf("failed to query kvblock indexer: %w", err)
 	}
+	lookupSpan.SetAttributes(
+		attribute.Int("llm_d.kv_cache_manager.lookup_results", len(keyToPods)),
+		attribute.String("operation.outcome", "success"),
+	)
+	lookupSpan.End()
 	traceLogger.Info("found block keys", "block-keys", blockKeys,
 		"pods", podsPerKeyPrintHelper(keyToPods))
 
 	// 4. score pods
+	_, scoreSpan := tracer.Start(ctx, "llm_d.kv_cache_manager.score_pods")
+	scoreSpan.SetAttributes(
+		attribute.String("gen_ai.request.model", modelName),
+		attribute.Int("llm_d.kv_cache_manager.block_keys_count", len(blockKeys)),
+	)
 	podScores, err := k.kvBlockScorer.Score(blockKeys, keyToPods)
 	if err != nil {
+		scoreSpan.RecordError(err)
+		scoreSpan.SetAttributes(attribute.String("operation.outcome", "error"))
+		scoreSpan.End()
 		span.RecordError(err)
+		span.SetAttributes(attribute.String("operation.outcome", "error"))
 		return nil, fmt.Errorf("failed to query kvblock scorer: %w", err)
 	}
+	scoreSpan.SetAttributes(
+		attribute.Int("llm_d.kv_cache_manager.scored_pods", len(podScores)),
+		attribute.String("operation.outcome", "success"),
+	)
+	scoreSpan.End()
 	traceLogger.Info("found pod scores", "pod-scores", podScores)
 
 	// Calculate hit ratio for observability
@@ -180,7 +253,8 @@ func (k *Indexer) GetPodScores(ctx context.Context, prompt, modelName string,
 	}
 
 	span.SetAttributes(
-		attribute.Float64("llm_d.kv_cache.hit_ratio", hitRatio),
+		attribute.Float64("llm_d.kv_cache_manager.hit_ratio", hitRatio),
+		attribute.String("operation.outcome", "success"),
 	)
 
 	return podScores, nil