Merge branch 'llm-d:main' into unit-tests

Signed-off-by: sagiahrac <[email protected]>
Signed-off-by: Sage Ahrac <[email protected]>

Author: sagiahrac

Makefile

@@ -75,19 +75,70 @@ verify-boilerplate: $(TOOLS_DIR)/verify_boilerplate.py
 	$(TOOLS_DIR)/verify_boilerplate.py --boilerplate-dir=hack/boilerplate --skip docs
 
 .PHONY: unit-test
-unit-test: download-tokenizer download-zmq
+unit-test: download-tokenizer detect-python install-python-deps download-zmq
 	@printf "\033[33;1m==== Running unit tests ====\033[0m\n"
 	go test -ldflags="$(LDFLAGS)" ./pkg/...
+	@printf "\033[33;1m==== Running chat template tests ====\033[0m\n"
+	go test -tags=exclude -v -ldflags="$(LDFLAGS)" ./pkg/preprocessing/chat_completions_template/
+	@printf "\033[33;1m==== Running chat template benchmarks ====\033[0m\n"
+	go test -tags=exclude -bench=. -benchmem -ldflags="$(LDFLAGS)" ./pkg/preprocessing/chat_completions_template/
 
 .PHONY: e2e-test
-e2e-test: download-tokenizer download-zmq
-	@printf "\033[33;1m==== Running unit tests ====\033[0m\n"
+e2e-test: download-tokenizer detect-python install-python-deps download-zmq
+	@printf "\033[33;1m==== Running e2e tests ====\033[0m\n"
 	go test -v -ldflags="$(LDFLAGS)" ./tests/...
 
 ##@ Build
 
+# Python detection and build configuration
+PYTHON_VERSION := 3.11.7
+PYTHON_DIR = build/python-$(PYTHON_VERSION)
+
+.PHONY: detect-python
+detect-python:
+	@printf "\033[33;1m==== Detecting Python installation ====\033[0m\n"
+	@if python3 -c "import sys; print(sys.version)" >/dev/null 2>&1; then \
+		echo "Using system Python"; \
+		PYTHON_PATH=$$(python3 -c "import sys; print(sys.prefix)"); \
+		PYTHON_VERSION=$$(python3 -c "import sys; print(f'{sys.version_info.major}.{sys.version_info.minor}')"); \
+		sed -i.bak "s|{{PYTHON_PATH}}|$$PYTHON_PATH|g; s|{{PYTHON_VERSION}}|$$PYTHON_VERSION|g" \
+			pkg/preprocessing/chat_completions_template/cgo_functions.go; \
+		rm -f pkg/preprocessing/chat_completions_template/cgo_functions.go.bak; \
+	else \
+		echo "System Python not found, downloading..."; \
+		$(MAKE) download-python; \
+	fi
+
+.PHONY: download-python
+download-python:
+	@printf "\033[33;1m==== Downloading Python $(PYTHON_VERSION) ====\033[0m\n"
+	@mkdir -p build
+	@if [ ! -d "$(PYTHON_DIR)" ]; then \
+		if [ "$(TARGETOS)" = "darwin" ]; then \
+			curl -L https://www.python.org/ftp/python/$(PYTHON_VERSION)/python-$(PYTHON_VERSION)-macos11.pkg -o build/python.pkg; \
+			sudo installer -pkg build/python.pkg -target /; \
+		elif [ "$(TARGETOS)" = "linux" ]; then \
+			curl -L https://www.python.org/ftp/python/$(PYTHON_VERSION)/Python-$(PYTHON_VERSION).tgz -o build/python.tgz; \
+			tar -xzf build/python.tgz -C build/; \
+			cd build/Python-$(PYTHON_VERSION) && ./configure --prefix=$(PWD)/$(PYTHON_DIR) && make && make install; \
+		fi; \
+	fi
+	@# Update CGo flags with downloaded Python path
+	@sed -i.bak "s|{{PYTHON_PATH}}|$(PWD)/$(PYTHON_DIR)|g; s|{{PYTHON_VERSION}}|$(PYTHON_VERSION)|g" \
+		pkg/preprocessing/chat_completions_template/cgo_functions.go
+	@rm -f pkg/preprocessing/chat_completions_template/cgo_functions.go.bak
+
+.PHONY: install-python-deps
+install-python-deps: detect-python
+	@printf "\033[33;1m==== Installing Python dependencies ====\033[0m\n"
+	@if [ -d "$(PYTHON_DIR)" ]; then \
+		$(PYTHON_DIR)/bin/pip install -r pkg/preprocessing/chat_completions_template/requirements.txt; \
+	else \
+		python3 -m pip install -r pkg/preprocessing/chat_completions_template/requirements.txt; \
+	fi
+
 .PHONY: build
-build: check-go download-tokenizer download-zmq
+build: check-go download-tokenizer detect-python install-python-deps download-zmq
 	@printf "\033[33;1m==== Building ====\033[0m\n"
 	go build -ldflags="$(LDFLAGS)" -o bin/$(PROJECT_NAME) examples/kv_cache_index/main.go
 
@@ -351,6 +402,14 @@ print-namespace: ## Print the current namespace
 print-project-name: ## Print the current project name
 	@echo "$(PROJECT_NAME)"
 
+.PHONY: clean
+clean: ## Clean build artifacts and restore placeholders
+	@printf "\033[33;1m==== Cleaning build artifacts ====\033[0m\n"
+	@rm -rf build/
+	@# Restore original placeholders
+	@git checkout pkg/preprocessing/chat_completions_template/cgo_functions.go
+	@echo "✅ Build artifacts cleaned and placeholders restored"
+
 .PHONY: install-hooks
 install-hooks: ## Install git hooks
 	git config core.hooksPath hooks

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)

pkg/preprocessing/chat_completions_template/README.md

@@ -62,16 +62,16 @@ The templating process (steps 1.1-1.4) handles the conversion from structured re
     └── cgo_functions.go:NewChatTemplateCGoWrapper()
         └── Creates ChatTemplateCGoWrapper struct with initialized=false
 
-1.2. **Template Fetching**: wrapper.GetModelChatTemplate(getReq)
-    ├── cgo_functions.go:GetModelChatTemplate(req)
+1.2. **Template Fetching**: wrapper.GetModelChatTemplate(ctx, getReq)
+    ├── cgo_functions.go:GetModelChatTemplate(ctx, req)
     │   ├── Initialize() Python interpreter via CGO
     │   ├── executePythonCode() - **CGO Binding** to Python
     │   └── **Python Wrapper**: chat_template_wrapper.py:get_model_chat_template()
     │       └── Uses Hugging Face AutoTokenizer to fetch model template
     └── Returns: (template, template_vars)
 
-1.3. **Template Rendering**: wrapper.RenderChatTemplate(req)
-    ├── cgo_functions.go:RenderChatTemplate(req)
+1.3. **Template Rendering**: wrapper.RenderChatTemplate(ctx, req)
+    ├── cgo_functions.go:RenderChatTemplate(ctx, req)
     │   ├── Initialize() Python interpreter via CGO (if not already done)
     │   ├── executePythonCode() - **CGO Binding** to Python
     │   └── **Python Wrapper**: chat_template_wrapper.py:render_jinja_template()