- fixed kvevents and made api/example more intuitive (#88)

- dockerfile fix
- deployment fix

Signed-off-by: Maroon Ayoub <[email protected]>

Author: vMaroon

Dockerfile

@@ -26,10 +26,6 @@ RUN dnf install -y 'https://dl.fedoraproject.org/pub/epel/epel-release-latest-8.
     dnf install -y gcc-c++ libstdc++ libstdc++-devel clang zeromq-devel pkgconfig && \
     dnf clean all
 
-# Install Python 3.9+ for chat template functionality
-RUN dnf install -y python39 python39-pip python39-devel && \
-    dnf clean all
-
 # Copy the Go Modules manifests
 COPY go.mod go.mod
 COPY go.sum go.sum
@@ -41,9 +37,6 @@ RUN go mod download
 COPY examples/kv_events examples/kv_events
 COPY . .
 
-# Install Python dependencies for chat template functionality
-RUN python3.9 -m pip install -r pkg/preprocessing/chat_completions_template/requirements.txt
-
 # HuggingFace tokenizer bindings
 RUN mkdir -p lib
 ARG RELEASE_VERSION=v1.22.1
@@ -66,14 +59,9 @@ WORKDIR /
 # The final image is UBI9, so we need epel-release-9.
 USER root
 RUN dnf install -y 'https://dl.fedoraproject.org/pub/epel/epel-release-latest-9.noarch.rpm' && \
-    dnf install -y zeromq python39 python39-pip
-
-# Install Python dependencies in the final image
-COPY pkg/preprocessing/chat_completions_template/requirements.txt /tmp/requirements.txt
-RUN python3.9 -m pip install -r /tmp/requirements.txt
+    dnf install -y zeromq
 
 COPY --from=builder /workspace/bin/kv-cache-manager /app/kv-cache-manager
-COPY --from=builder /workspace/pkg/preprocessing/chat_completions_template/chat_template_wrapper.py /app/chat_template_wrapper.py
 USER 65532:65532
 
 # Set the entrypoint to the kv-cache-manager binary

examples/kv_events/offline/main.go

@@ -17,7 +17,6 @@ limitations under the License.
 package main
 
 import (
-	"bytes"
 	"context"
 	_ "embed"
 	"fmt"
@@ -166,22 +165,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)
+	blockStoredEvent := kvevents.BlockStored{
+		BlockHashes:     testdata.PromptHashes,
+		ParentBlockHash: nil,
+		TokenIds:        []uint32{1, 2, 3},
+		BlockSize:       256,
+		LoraID:          nil,
+	}
 
 	//nolint // won't fail
-	enc.Encode(&kvevents.BlockStoredEvent{
-		TypeField:   "BlockStored",
-		BlockStored: &kvevents.BlockStored{BlockHashes: testdata.PromptHashes},
-	})
-
-	dpRank := 0
+	blockStoredPayload, _ := msgpack.Marshal(blockStoredEvent.ToTaggedUnion())
 
 	eventBatch := kvevents.EventBatch{
 		TS:               float64(time.Now().UnixNano()) / 1e9,
-		Events:           []msgpack.RawMessage{blockStoredPayload.Bytes()},
-		DataParallelRank: &dpRank,
+		Events:           []msgpack.RawMessage{blockStoredPayload},
+		DataParallelRank: nil,
 	}
 
 	topic := fmt.Sprintf("kv@vllm-pod1@%s", testdata.ModelName)
@@ -204,20 +202,17 @@ 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)
+	blockRemovedEvent := kvevents.BlockRemoved{
+		BlockHashes: testdata.PromptHashes[2:], // Remove last blocks
+	}
 
 	//nolint // won't fail
-	enc.Encode(&kvevents.BlockRemovedEvent{
-		TypeField:    "BlockRemoved",
-		BlockRemoved: &kvevents.BlockRemoved{BlockHashes: testdata.PromptHashes[2:]},
-	})
+	blockRemovedPayload, _ := msgpack.Marshal(blockRemovedEvent.ToTaggedUnion())
 
 	removeEventBatch := kvevents.EventBatch{
 		TS:               float64(time.Now().UnixNano()) / 1e9,
-		Events:           []msgpack.RawMessage{blockRemovedPayload.Bytes()},
-		DataParallelRank: &dpRank,
+		Events:           []msgpack.RawMessage{blockRemovedPayload},
+		DataParallelRank: nil,
 	}
 
 	if err := publisher.PublishEvent(ctx, topic, removeEventBatch); err != nil {

pkg/kvcache/kvevents/events.go

@@ -1,14 +1,22 @@
 package kvevents
 
 import (
-	"fmt"
-
 	"github.com/vmihailenco/msgpack/v5"
 )
 
+const (
+	// BlockStoredEventTag is the tag for BlockStored events.
+	BlockStoredEventTag = "BlockStored"
+	// BlockRemovedEventTag is the tag for BlockRemoved events.
+	BlockRemovedEventTag = "BlockRemoved"
+	// AllBlocksClearedEventTag is the tag for AllBlocksCleared events.
+	AllBlocksClearedEventTag = "AllBlocksCleared"
+)
+
 // event is a marker interface for KV-cache events.
 type event interface {
 	isEvent()
+	ToTaggedUnion() []any
 }
 
 // EventBatch represents a batch of events.
@@ -20,37 +28,6 @@ type EventBatch struct {
 	DataParallelRank *int `msgpack:",omitempty"`
 }
 
-// DecodeMsgpack allows 2- or 3-element array-encoded batches.
-func (e *EventBatch) DecodeMsgpack(decoder *msgpack.Decoder) error {
-	length, err := decoder.DecodeArrayLen()
-	if err != nil {
-		return err
-	}
-	if length < 2 {
-		return fmt.Errorf("EventBatch: expected at least 2 fields, got %d", length)
-	}
-	if e.TS, err = decoder.DecodeFloat64(); err != nil {
-		return err
-	}
-	if err := decoder.Decode(&e.Events); err != nil {
-		return err
-	}
-	if length > 2 {
-		var rank int
-		if err := decoder.Decode(&rank); err != nil {
-			return err
-		}
-		e.DataParallelRank = &rank
-	}
-	// skip any extra
-	for i := 3; i < length; i++ {
-		if err := decoder.Skip(); err != nil {
-			return err
-		}
-	}
-	return nil
-}
-
 // BlockStored event.
 type BlockStored struct {
 	_               struct{} `msgpack:",array"`
@@ -61,31 +38,43 @@ type BlockStored struct {
 	LoraID          *int
 }
 
-func (BlockStored) isEvent() {}
-
-type BlockStoredEvent struct {
-	_         struct{} `msgpack:",array"`
-	TypeField string
-	*BlockStored
+func (bs BlockStored) ToTaggedUnion() []any {
+	return []any{
+		BlockStoredEventTag,
+		bs.BlockHashes,
+		bs.ParentBlockHash,
+		bs.TokenIds,
+		bs.BlockSize,
+		bs.LoraID,
+	}
 }
 
+func (BlockStored) isEvent() {}
+
 // BlockRemoved event.
 type BlockRemoved struct {
 	_           struct{} `msgpack:",array"`
 	BlockHashes []uint64
 }
 
-func (BlockRemoved) isEvent() {}
-
-type BlockRemovedEvent struct {
-	_         struct{} `msgpack:",array"`
-	TypeField string
-	*BlockRemoved
+func (br BlockRemoved) ToTaggedUnion() []any {
+	return []any{
+		BlockRemovedEventTag,
+		br.BlockHashes,
+	}
 }
 
+func (BlockRemoved) isEvent() {}
+
 // AllBlocksCleared event.
 type AllBlocksCleared struct {
 	_ struct{} `msgpack:",array"`
 }
 
+func (ac AllBlocksCleared) ToTaggedUnion() []any {
+	return []any{
+		AllBlocksClearedEventTag,
+	}
+}
+
 func (AllBlocksCleared) isEvent() {}

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 BlockStoredEvent
-			unmarshalErr = msgpack.Unmarshal(rawEvent, &bs)
-			event = bs.BlockStored
+			var bs BlockStored
+			unmarshalErr = msgpack.Unmarshal(payloadBytes, &bs)
+			event = bs
 		case "BlockRemoved":
-			var br BlockRemovedEvent
-			unmarshalErr = msgpack.Unmarshal(rawEvent, &br)
-			event = br.BlockRemoved
+			var br BlockRemoved
+			unmarshalErr = msgpack.Unmarshal(payloadBytes, &br)
+			event = br
 		case "AllBlocksCleared":
 			var ac AllBlocksCleared
 			unmarshalErr = msgpack.Unmarshal(payloadBytes, &ac)

vllm-setup-helm/templates/deployment.yaml

@@ -91,6 +91,8 @@ spec:
           startupProbe:
             {{- toYaml .Values.vllm.startupProbe | nindent 12 }}
           env:
+            - name: HOME
+              value: {{ .Values.persistence.mountPath }}
             - name: HF_HOME
               value: {{ .Values.persistence.mountPath }}
             - name: POD_IP