refactor: Update BlockStored KVEvent

examples/helper/events.go

@@ -32,12 +32,15 @@ const topic = "kv@vllm-pod1@" + testdata.ModelName
 func SimulateProduceEvent(ctx context.Context, publisher *Publisher) error {
 	logger := log.FromContext(ctx)
 	logger.Info("@@@ Simulating vLLM engine publishing BlockStored events...")
+	medium := "GPU"
 	blockStoredEvent := kvevents.BlockStored{
 		BlockHashes:     utils.SliceMap(testdata.PromptHashes, func(h uint64) any { return h }),
 		ParentBlockHash: nil,
 		TokenIds:        []uint32{1, 2, 3},
 		BlockSize:       256,
 		LoraID:          nil,
+		Medium:          &medium,
+		LoraName:        nil,
 	}
 
 	//nolint // won't fail

pkg/kvcache/kvevents/events.go

@@ -15,6 +15,8 @@
 package kvevents
 
 import (
+	"fmt"
+
 	"github.com/vmihailenco/msgpack/v5"
 )
 
@@ -53,6 +55,7 @@
 	BlockSize       int
 	LoraID          *int    `msgpack:",omitempty"`
 	Medium          *string `msgpack:",omitempty"`
+	LoraName        *string `msgpack:",omitempty"`
 }
 
 // ToTaggedUnion converts the BlockStored event to a tagged union format.
@@ -67,6 +70,7 @@
 		bs.BlockSize,
 		bs.LoraID,
 		bs.Medium,
+		bs.LoraName,
 	}
 }
 
@@ -102,3 +106,66 @@
 }
 
 func (AllBlocksCleared) isEvent() {}
+
+// unmarshalKVEvent unmarshals a raw msgpack event into the event interface
+func unmarshalKVEvent(rawEvent msgpack.RawMessage) (event, error) {
+	var taggedUnion []msgpack.RawMessage
+	if err := msgpack.Unmarshal(rawEvent, &taggedUnion); err != nil {
+		return nil, fmt.Errorf("failed to unmarshal tagged union: %w", err)
+	}
+
+	if len(taggedUnion) < 1 {
+		return nil, fmt.Errorf("malformed tagged union: no tag")
+	}
+
+	var tag string
+	if err := msgpack.Unmarshal(taggedUnion[0], &tag); err != nil {
+		return nil, fmt.Errorf("failed to unmarshal tag: %w", err)
+	}
+
+	// The 'payload' starts from index 1 of the tagged union
+	payloadParts := taggedUnion[1:]
+
+	switch tag {
+	case BlockStoredEventTag:
+		// Mandatory fields: BlockHashes, Parent, TokenIds, BlockSize (indices 0-3 of payload)
+		if len(payloadParts) < 5 {
+			return nil, fmt.Errorf("BlockStored missing mandatory fields: got %d", len(payloadParts))
+		}
+
+		var bs BlockStored
+		// Manual mapping to bypass the strict "array-encoded struct" length check
+		_ = msgpack.Unmarshal(payloadParts[0], &bs.BlockHashes)
+		_ = msgpack.Unmarshal(payloadParts[1], &bs.ParentBlockHash)
+		_ = msgpack.Unmarshal(payloadParts[2], &bs.TokenIds)
+		_ = msgpack.Unmarshal(payloadParts[3], &bs.BlockSize)
+		_ = msgpack.Unmarshal(payloadParts[4], &bs.LoraID)
+
+		// Optional fields (Indices 5 and 6 of payload)
+		if len(payloadParts) > 5 {
+			_ = msgpack.Unmarshal(payloadParts[5], &bs.Medium)
+		}
+		if len(payloadParts) > 6 {
+			_ = msgpack.Unmarshal(payloadParts[6], &bs.LoraName)
+		}
+
+		return bs, nil
+
+	case BlockRemovedEventTag:
+		if len(payloadParts) < 1 {
+			return nil, fmt.Errorf("BlockRemoved missing mandatory BlockHashes")
+		}
+		var br BlockRemoved
+		_ = msgpack.Unmarshal(payloadParts[0], &br.BlockHashes)
+		if len(payloadParts) > 1 {
+			_ = msgpack.Unmarshal(payloadParts[1], &br.Medium)
+		}
+		return br, nil
+
+	case AllBlocksClearedEventTag:
+		return AllBlocksCleared{}, nil
+
+	default:
+		return nil, fmt.Errorf("unknown event tag: %s", tag)
+	}
+}

pkg/kvcache/kvevents/pool.go

@@ -188,51 +188,9 @@ func (p *Pool) processEvent(ctx context.Context, msg *Message) {
 
 	events := make([]event, 0, len(eventBatch.Events))
 	for _, rawEvent := range eventBatch.Events {
-		var taggedUnion []msgpack.RawMessage
-		if err := msgpack.Unmarshal(rawEvent, &taggedUnion); err != nil {
-			debugLogger.Error(err, "Failed to unmarshal tagged union, skipping event")
-			continue
-		}
-
-		// Handle array_like tagged union: re-marshall tail parts into a payload array
-		if len(taggedUnion) < 1 {
-			debugLogger.Error(nil, "Malformed tagged union, no tag element", "parts", len(taggedUnion))
-			continue
-		}
-		payloadBytes, err := msgpack.Marshal(taggedUnion[1:])
+		event, err := unmarshalKVEvent(rawEvent)
 		if err != nil {
-			debugLogger.Error(err, "Failed to re-marshal payload parts, skipping event")
-			continue
-		}
-
-		var tag string
-		if err := msgpack.Unmarshal(taggedUnion[0], &tag); err != nil {
-			debugLogger.Error(err, "Failed to unmarshal tag from tagged union, skipping event")
-			continue
-		}
-
-		var event event
-		var unmarshalErr error
-		switch tag {
-		case "BlockStored":
-			var bs BlockStored
-			unmarshalErr = msgpack.Unmarshal(payloadBytes, &bs)
-			event = bs
-		case "BlockRemoved":
-			var br BlockRemoved
-			unmarshalErr = msgpack.Unmarshal(payloadBytes, &br)
-			event = br
-		case "AllBlocksCleared":
-			var ac AllBlocksCleared
-			unmarshalErr = msgpack.Unmarshal(payloadBytes, &ac)
-			event = ac
-		default:
-			debugLogger.Info("Unknown event tag", "tag", tag)
-			continue
-		}
-
-		if unmarshalErr != nil {
-			debugLogger.Error(unmarshalErr, "Failed to unmarshal event value", "tag", tag)
+			debugLogger.Error(err, "Failed to unmarshal event, skipping")
 			continue
 		}
 		events = append(events, event)

pkg/kvcache/kvevents/process_event_test.go

@@ -0,0 +1,158 @@
+/*
+Copyright 2025 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 kvevents
+
+import (
+	"testing"
+
+	"github.com/vmihailenco/msgpack/v5"
+)
+
+// Helper function to create BlockStored raw msgpack message
+func createBlockStoredRaw(t *testing.T, fields []any) msgpack.RawMessage {
+	data, err := msgpack.Marshal(fields)
+	if err != nil {
+		t.Fatalf("Failed to marshal fields: %v", err)
+	}
+	return msgpack.RawMessage(data)
+}
+
+func TestBlockStoredMissingMediumAndLoraName(t *testing.T) {
+	rawMsg := createBlockStoredRaw(t, []any{
+		BlockStoredEventTag,               // Event tag
+		[]any{uint64(1001), uint64(1002)}, // BlockHashes
+		nil,                               // ParentBlockHash
+		[]uint32{1, 2, 3},                 // TokenIds
+		256,                               // BlockSize
+		42,                                // LoraID
+		// Medium and LoraName are missing
+	})
+
+	event, err := unmarshalKVEvent(rawMsg)
+	if err != nil {
+		t.Fatalf("Failed to process BlockStored event: %v", err)
+	}
+
+	if event == nil {
+		t.Error("Expected event to be non-nil")
+	}
+
+	blockStored, ok := event.(BlockStored)
+	if !ok {
+		t.Fatalf("Expected BlockStored event, got %T", event)
+	}
+
+	if blockStored.Medium != nil {
+		t.Errorf("Expected Medium to be nil, got %v", *blockStored.Medium)
+	}
+	if blockStored.LoraName != nil {
+		t.Errorf("Expected LoraName to be nil, got %v", *blockStored.LoraName)
+	}
+}
+
+func TestBlockStoredMissingLoraName(t *testing.T) {
+	rawMsg := createBlockStoredRaw(t, []any{
+		BlockStoredEventTag,               // Event tag
+		[]any{uint64(1001), uint64(1002)}, // BlockHashes
+		nil,                               // ParentBlockHash
+		[]uint32{1, 2, 3},                 // TokenIds
+		256,                               // BlockSize
+		42,                                // LoraID
+		"GPU",                             // Medium
+		// LoraName is missing
+	})
+
+	event, err := unmarshalKVEvent(rawMsg)
+	if err != nil {
+		t.Fatalf("Failed to process BlockStored event: %v", err)
+	}
+
+	if event == nil {
+		t.Error("Expected event to be non-nil")
+	}
+
+	blockStored, ok := event.(BlockStored)
+	if !ok {
+		t.Fatalf("Expected BlockStored event, got %T", event)
+	}
+
+	if blockStored.Medium == nil || *blockStored.Medium != "cpu" {
+		t.Errorf("Expected Medium to be 'cpu', got %v", blockStored.Medium)
+	}
+	if blockStored.LoraName != nil {
+		t.Errorf("Expected LoraName to be nil, got %v", *blockStored.LoraName)
+	}
+}
+
+func TestBlockStoredAllFieldsPresent(t *testing.T) {
+	rawMsg := createBlockStoredRaw(t, []any{
+		BlockStoredEventTag,               // Event tag
+		[]any{uint64(1001), uint64(1002)}, // BlockHashes
+		nil,                               // ParentBlockHash
+		[]uint32{1, 2, 3},                 // TokenIds
+		256,                               // BlockSize
+		42,                                // LoraID
+		"gpu",                             // Medium
+		"test-lora",                       // LoraName
+	})
+
+	event, err := unmarshalKVEvent(rawMsg)
+	if err != nil {
+		t.Fatalf("Failed to process BlockStored event: %v", err)
+	}
+
+	if event == nil {
+		t.Error("Expected event to be non-nil")
+	}
+
+	blockStored, ok := event.(BlockStored)
+	if !ok {
+		t.Fatalf("Expected BlockStored event, got %T", event)
+	}
+
+	if blockStored.Medium == nil || *blockStored.Medium != "gpu" {
+		t.Errorf("Expected Medium to be 'gpu', got %v", blockStored.Medium)
+	}
+	if blockStored.LoraName == nil || *blockStored.LoraName != "test-lora" {
+		t.Errorf("Expected LoraName to be 'test-lora', got %v", blockStored.LoraName)
+	}
+}
+
+func TestUnmarshalKVEventErrors(t *testing.T) {
+	// Test invalid msgpack
+	_, err := unmarshalKVEvent(msgpack.RawMessage([]byte{0x01, 0x02, 0x03}))
+	if err == nil {
+		t.Error("Expected error for invalid msgpack")
+	}
+
+	// Test unknown event tag
+	rawMsg := createBlockStoredRaw(t, []any{
+		"UnknownEvent",
+		[]any{uint64(1001)},
+	})
+	_, err = unmarshalKVEvent(rawMsg)
+	if err == nil {
+		t.Error("Expected error for unknown event tag")
+	}
+
+	// Test malformed union (empty array)
+	emptyRawMsg := createBlockStoredRaw(t, []any{})
+	_, err = unmarshalKVEvent(emptyRawMsg)
+	if err == nil {
+		t.Error("Expected error for malformed union")
+	}
+}