Fix kv_events offline example (#82)

* Added new structs for events to contain TagField, changed marshaling to encoding as array

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

* Fixed the sent hashes, changed publisher endpoint, fixed readme

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

---------

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

Author: irar2

examples/kv_events/README.md

@@ -18,7 +18,7 @@ Download tokenizer bindings:
 
 ### Running the Example
 ```
-    go run -ldflags="-extldflags '-L$(pwd)/lib'" examples/kv-events/offline/main.go examples/kv-events/offline/publisher.go
+    go run -ldflags="-extldflags '-L$(pwd)/lib'" examples/kv_events/offline/main.go examples/kv_events/offline/publisher.go
 ```
 
 The example will start the KV-Cache Manager (indexer) and a dummy publisher that simulates KV-Events. 

examples/kv_events/offline/main.go

@@ -17,6 +17,7 @@ limitations under the License.
 package main
 
 import (
+	"bytes"
 	"context"
 	_ "embed"
 	"fmt"
@@ -135,11 +136,10 @@ func setupEventsPool(ctx context.Context, kvBlockIndex kvblock.Index) *kvevents.
 func setupPublisher(ctx context.Context) (*Publisher, error) {
 	logger := klog.FromContext(ctx)
 
-	cfg := kvevents.DefaultConfig()
-
-	logger.Info("Creating ZMQ publisher (simulating vLLM engines)", "endpoint", cfg.ZMQEndpoint)
+	endpoint := "tcp://localhost:5557"
+	logger.Info("Creating ZMQ publisher (simulating vLLM engines)", "endpoint", endpoint)
 
-	publisher, err := NewPublisher(cfg.ZMQEndpoint)
+	publisher, err := NewPublisher(endpoint)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create ZMQ publisher: %w", err)
 	}
@@ -166,13 +166,21 @@ func runEventsDemo(ctx context.Context, kvCacheIndexer *kvcache.Indexer, publish
 	// Simulate vLLM engine publishing BlockStored events
 	logger.Info("@@@ Simulating vLLM engine publishing BlockStored events...")
 
+	var blockStoredPayload bytes.Buffer
+	enc := msgpack.NewEncoder(&blockStoredPayload)
+	enc.UseArrayEncodedStructs(true)
+
 	//nolint // won't fail
-	blockStoredPayloadBytes, _ := msgpack.Marshal(kvevents.BlockStored{BlockHashes: testdata.PromptHashes})
+	enc.Encode(&kvevents.BlockStoredEvent{
+		TypeField:   "BlockStored",
+		BlockStored: &kvevents.BlockStored{BlockHashes: testdata.PromptHashes},
+	})
+
 	dpRank := 0
 
 	eventBatch := kvevents.EventBatch{
 		TS:               float64(time.Now().UnixNano()) / 1e9,
-		Events:           []msgpack.RawMessage{blockStoredPayloadBytes},
+		Events:           []msgpack.RawMessage{blockStoredPayload.Bytes()},
 		DataParallelRank: &dpRank,
 	}
 
@@ -196,14 +204,19 @@ func runEventsDemo(ctx context.Context, kvCacheIndexer *kvcache.Indexer, publish
 	// Simulate removing some blocks
 	logger.Info("@@@ Simulating vLLM engine removing some blocks...")
 
+	var blockRemovedPayload bytes.Buffer
+	enc = msgpack.NewEncoder(&blockRemovedPayload)
+	enc.UseArrayEncodedStructs(true)
+
 	//nolint // won't fail
-	blockRemovedPayloadBytes, _ := msgpack.Marshal(kvevents.BlockRemoved{
-		BlockHashes: testdata.PromptHashes[:2],
+	enc.Encode(&kvevents.BlockRemovedEvent{
+		TypeField:    "BlockRemoved",
+		BlockRemoved: &kvevents.BlockRemoved{BlockHashes: testdata.PromptHashes[2:]},
 	})
 
 	removeEventBatch := kvevents.EventBatch{
 		TS:               float64(time.Now().UnixNano()) / 1e9,
-		Events:           []msgpack.RawMessage{blockRemovedPayloadBytes},
+		Events:           []msgpack.RawMessage{blockRemovedPayload.Bytes()},
 		DataParallelRank: &dpRank,
 	}
 

examples/kv_events/offline/publisher.go

@@ -17,6 +17,7 @@ limitations under the License.
 package main
 
 import (
+	"bytes"
 	"context"
 	"encoding/binary"
 	"fmt"
@@ -58,7 +59,11 @@ func NewPublisher(endpoint string) (*Publisher, error) {
 func (p *Publisher) PublishEvent(ctx context.Context, topic string, batch interface{}) error {
 	logger := klog.FromContext(ctx).V(0)
 
-	payload, err := msgpack.Marshal(batch)
+	// Use an encoder configured for struct as array
+	var payload bytes.Buffer
+	enc := msgpack.NewEncoder(&payload)
+	enc.UseArrayEncodedStructs(true)
+	err := enc.Encode(batch)
 	if err != nil {
 		return fmt.Errorf("failed to marshal event batch: %w", err)
 	}
@@ -69,7 +74,7 @@ func (p *Publisher) PublishEvent(ctx context.Context, topic string, batch interf
 	binary.BigEndian.PutUint64(seqBytes, seq)
 
 	// send topic, sequence, payload
-	if _, err := p.socket.SendMessage(topic, seqBytes, payload); err != nil {
+	if _, err := p.socket.SendMessage(topic, seqBytes, payload.Bytes()); err != nil {
 		return fmt.Errorf("failed to send message to topic %s: %w", topic, err)
 	}
 

pkg/kvcache/kvevents/events.go

@@ -63,6 +63,12 @@ type BlockStored struct {
 
 func (BlockStored) isEvent() {}
 
+type BlockStoredEvent struct {
+	_         struct{} `msgpack:",array"`
+	TypeField string
+	*BlockStored
+}
+
 // BlockRemoved event.
 type BlockRemoved struct {
 	_           struct{} `msgpack:",array"`
@@ -71,6 +77,12 @@ type BlockRemoved struct {
 
 func (BlockRemoved) isEvent() {}
 
+type BlockRemovedEvent struct {
+	_         struct{} `msgpack:",array"`
+	TypeField string
+	*BlockRemoved
+}
+
 // AllBlocksCleared event.
 type AllBlocksCleared struct {
 	_ struct{} `msgpack:",array"`

pkg/kvcache/kvevents/pool.go

@@ -193,13 +193,13 @@ func (p *Pool) processEvent(ctx context.Context, msg *Message) {
 		var unmarshalErr error
 		switch tag {
 		case "BlockStored":
-			var bs BlockStored
-			unmarshalErr = msgpack.Unmarshal(payloadBytes, &bs)
-			event = bs
+			var bs BlockStoredEvent
+			unmarshalErr = msgpack.Unmarshal(rawEvent, &bs)
+			event = bs.BlockStored
 		case "BlockRemoved":
-			var br BlockRemoved
-			unmarshalErr = msgpack.Unmarshal(payloadBytes, &br)
-			event = br
+			var br BlockRemovedEvent
+			unmarshalErr = msgpack.Unmarshal(rawEvent, &br)
+			event = br.BlockRemoved
 		case "AllBlocksCleared":
 			var ac AllBlocksCleared
 			unmarshalErr = msgpack.Unmarshal(payloadBytes, &ac)