Sleep mode (#252)

* Sleep mode

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

* lint

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

* Disable the new sleep test to check if it causes the problem

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

* Wait before sending request

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

* Review comments and test fix

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

* Log levels and lint

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

---------

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

Author: irar2

pkg/common/config.go

@@ -220,6 +220,9 @@ type Configuration struct {
 	DatasetURL string `yaml:"dataset-url" json:"dataset-url"`
 	// DatasetInMemory defines whether to load the entire dataset into memory for faster access.
 	DatasetInMemory bool `yaml:"dataset-in-memory" json:"dataset-in-memory"`
+
+	// EnableSleepMode enables sleep mode
+	EnableSleepMode bool `yaml:"enable-sleep-mode" json:"enable-sleep-mode"`
 }
 
 type Metrics struct {
@@ -741,6 +744,8 @@ func ParseCommandParamsAndLoadConfig() (*Configuration, error) {
 	f.StringVar(&config.DatasetURL, "dataset-url", config.DatasetURL, "URL to download the sqlite db file for response generation from a dataset")
 	f.BoolVar(&config.DatasetInMemory, "dataset-in-memory", config.DatasetInMemory, "Load the entire dataset into memory for faster access")
 
+	f.BoolVar(&config.EnableSleepMode, "enable-sleep-mode", config.EnableSleepMode, "Enable sleep mode")
+
 	f.IntVar(&config.FailureInjectionRate, "failure-injection-rate", config.FailureInjectionRate, "Probability (0-100) of injecting failures")
 	failureTypes := getParamValueFromArgs("failure-types")
 	var dummyFailureTypes multiString

pkg/common/test_utils.go

@@ -0,0 +1,40 @@
+/*
+Copyright 2025 The llm-d-inference-sim 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 common
+
+import (
+	"github.com/onsi/gomega"
+	zmq "github.com/pebbe/zmq4"
+)
+
+// CreateSub creates a ZMQ sub, subscribes to the provided topic, and returns the
+// sub and the endpoint to publish events on
+func CreateSub(topic string) (*zmq.Socket, string) {
+	wildcardEndpoint := "tcp://*:*"
+	zctx, err := zmq.NewContext()
+	gomega.Expect(err).NotTo(gomega.HaveOccurred())
+	sub, err := zctx.NewSocket(zmq.SUB)
+	gomega.Expect(err).NotTo(gomega.HaveOccurred())
+	err = sub.Bind(wildcardEndpoint)
+	gomega.Expect(err).NotTo(gomega.HaveOccurred())
+	// get the actual port
+	endpoint, err := sub.GetLastEndpoint()
+	gomega.Expect(err).NotTo(gomega.HaveOccurred())
+	err = sub.SetSubscribe(topic)
+	gomega.Expect(err).NotTo(gomega.HaveOccurred())
+	return sub, endpoint
+}

pkg/kv-cache/block_cache.go

@@ -24,6 +24,7 @@ import (
 
 	"github.com/go-logr/logr"
 	"github.com/llm-d/llm-d-inference-sim/pkg/common"
+	"github.com/llm-d/llm-d-inference-sim/pkg/common/logging"
 )
 
 const (
@@ -42,6 +43,7 @@ type blockCache struct {
 	eventChan       chan EventData       // channel for asynchronous event processing
 	usageChan       chan float64         // channel for usage reporting
 	logger          logr.Logger
+	disabled        bool // indicated whether the cache is disabled
 }
 
 // newBlockCache creates a new blockCache with the specified maximum number of blocks
@@ -58,31 +60,66 @@ func newBlockCache(config *common.Configuration, logger logr.Logger, usageChan c
 		}
 	}
 
+	eventSender := NewKVEventSender(publisher, CreateKVEventsTopic(config.Port, config.Model),
+		eChan, config.EventBatchSize, delay, logger)
+
 	return &blockCache{
 		requestToBlocks: make(map[string][]uint64),
 		usedBlocks:      make(map[uint64]int),
 		unusedBlocks:    make(map[uint64]time.Time),
 		maxBlocks:       config.KVCacheSize,
 		eventChan:       eChan,
 		usageChan:       usageChan,
-		eventSender:     NewKVEventSender(publisher, createTopic(config), eChan, config.EventBatchSize, delay, logger),
+		eventSender:     eventSender,
 		logger:          logger,
 	}, nil
 }
 
 func (bc *blockCache) start(ctx context.Context) {
+	bc.logger.V(logging.INFO).Info("Starting KV cache")
 	err := bc.eventSender.Run(ctx)
 	if err != nil {
 		bc.logger.Error(err, "Sender stopped with error")
 	}
 }
 
+func (bc *blockCache) discard() {
+	bc.logger.V(logging.INFO).Info("Discarding KV cache")
+
+	bc.mu.Lock()
+	defer bc.mu.Unlock()
+
+	bc.disabled = true
+
+	bc.requestToBlocks = make(map[string][]uint64)
+	bc.usedBlocks = make(map[uint64]int)
+	bc.unusedBlocks = make(map[uint64]time.Time)
+
+	common.WriteToChannel(bc.eventChan,
+		EventData{action: eventActionAllBlocksCleared},
+		bc.logger, "block cache eventChan")
+}
+
+func (bc *blockCache) activate() {
+	bc.logger.V(logging.INFO).Info("Activating KV cache")
+
+	bc.mu.Lock()
+	defer bc.mu.Unlock()
+
+	bc.disabled = false
+}
+
 // startRequest adds a request with its associated block hashes to the cache
 // and returns the number of blocks that were already in the cache
 func (bc *blockCache) startRequest(requestID string, blocks []uint64) (int, error) {
 	bc.mu.Lock()
 	defer bc.mu.Unlock()
 
+	if bc.disabled {
+		bc.logger.V(logging.TRACE).Info("KV cache is disabled, request is not added to the kv cache")
+		return 0, nil
+	}
+
 	if _, exists := bc.requestToBlocks[requestID]; exists {
 		// request with the same id already exists
 		return 0, fmt.Errorf("request already exists for id %s", requestID)
@@ -167,6 +204,11 @@ func (bc *blockCache) finishRequest(requestID string) error {
 	bc.mu.Lock()
 	defer bc.mu.Unlock()
 
+	if bc.disabled {
+		bc.logger.V(logging.TRACE).Info("KV cache is disabled, request completion is not processed by the kv cache")
+		return nil
+	}
+
 	// Get blocks associated with this request
 	blockHashes, exists := bc.requestToBlocks[requestID]
 	if !exists {
@@ -239,6 +281,6 @@ func (bc *blockCache) getBlockInfo(blockHash uint64) (int, bool) {
 	return 0, false
 }
 
-func createTopic(config *common.Configuration) string {
-	return fmt.Sprintf("kv@$localhost:%d@%s", config.Port, config.Model)
+func CreateKVEventsTopic(port int, model string) string {
+	return fmt.Sprintf("kv@$localhost:%d@%s", port, model)
 }

pkg/kv-cache/kv_cache.go

@@ -63,6 +63,14 @@ func (h *KVCacheHelper) Run(ctx context.Context) {
 	h.blockCache.start(ctx)
 }
 
+func (h *KVCacheHelper) Discard() {
+	h.blockCache.discard()
+}
+
+func (h *KVCacheHelper) Activate() {
+	h.blockCache.activate()
+}
+
 func (h *KVCacheHelper) OnRequestStart(vllmReq openaiserverapi.CompletionRequest) error {
 	h.logger.V(logging.TRACE).Info("KV cache - process request")
 

pkg/kv-cache/kv_cache_sender.go

@@ -32,6 +32,7 @@ type EventAction int
 const (
 	eventActionStore EventAction = iota
 	eventActionRemove
+	eventActionAllBlocksCleared
 )
 
 type EventData struct {
@@ -97,6 +98,8 @@ func (s *KVEventSender) Run(ctx context.Context) error {
 				payload, err = msgpack.Marshal(kvevents.BlockStored{BlockHashes: eventData.hashValues}.ToTaggedUnion())
 			case eventActionRemove:
 				payload, err = msgpack.Marshal(kvevents.BlockRemoved{BlockHashes: eventData.hashValues}.ToTaggedUnion())
+			case eventActionAllBlocksCleared:
+				payload, err = msgpack.Marshal(kvevents.AllBlocksCleared{}.ToTaggedUnion())
 			default:
 				return fmt.Errorf("invalid event action %d", eventData.action)
 			}

pkg/kv-cache/kv_cache_test.go

@@ -18,16 +18,11 @@ package kvcache
 
 import (
 	"context"
-	"encoding/binary"
 	"fmt"
 	"sync"
 	"time"
 
-	zmq "github.com/pebbe/zmq4"
-	"github.com/vmihailenco/msgpack/v5"
-
 	"github.com/llm-d/llm-d-inference-sim/pkg/common"
-	"github.com/llm-d/llm-d-kv-cache-manager/pkg/kvcache/kvevents"
 	. "github.com/onsi/ginkgo/v2"
 	. "github.com/onsi/gomega"
 )
@@ -207,7 +202,9 @@ var _ = Describe("KV cache", Ordered, func() {
 					EventBatchSize:        1,
 				}
 
-				sub, topic := createSub(config)
+				topic := CreateKVEventsTopic(config.Port, config.Model)
+				sub, endpoint := common.CreateSub(topic)
+				config.ZMQEndpoint = endpoint
 				//nolint
 				defer sub.Close()
 
@@ -289,7 +286,7 @@ var _ = Describe("KV cache", Ordered, func() {
 				for i := range test.expectedRemovedBlocks + test.expectedStoredBlocks {
 					parts, err := sub.RecvMessageBytes(0)
 					Expect(err).NotTo(HaveOccurred())
-					stored, removed := parseEvent(parts, topic, uint64(i+1))
+					stored, removed, _ := ParseKVEvent(parts, topic, uint64(i+1))
 					storedCount += len(stored)
 					removedCount += len(removed)
 				}
@@ -309,7 +306,9 @@ var _ = Describe("KV cache", Ordered, func() {
 				ZMQMaxConnectAttempts: 3,
 			}
 
-			sub, topic := createSub(config)
+			topic := CreateKVEventsTopic(config.Port, config.Model)
+			sub, endpoint := common.CreateSub(topic)
+			config.ZMQEndpoint = endpoint
 			//nolint
 			defer sub.Close()
 
@@ -378,7 +377,7 @@ var _ = Describe("KV cache", Ordered, func() {
 			for {
 				parts, err := sub.RecvMessageBytes(0)
 				Expect(err).NotTo(HaveOccurred())
-				stored, removed := parseEvent(parts, topic, count)
+				stored, removed, _ := ParseKVEvent(parts, topic, count)
 				storedBlocks = append(storedBlocks, stored...)
 				removedBlocks = append(removedBlocks, removed...)
 				count++
@@ -484,67 +483,3 @@ func createRandomArray(minArrLen, maxArrLen int, maxValue uint64, random *common
 
 	return arr
 }
-
-func parseEvent(parts [][]byte, expectedTopic string, expectedSeq uint64) ([]uint64, []uint64) {
-	// The message should be [topic, seq, payload]
-	Expect(parts).To(HaveLen(3))
-
-	Expect(string(parts[0])).To(Equal(expectedTopic))
-
-	seq := binary.BigEndian.Uint64(parts[1])
-	Expect(seq).To(Equal(expectedSeq))
-
-	removed := make([]uint64, 0)
-	stored := make([]uint64, 0)
-
-	var eventBatch kvevents.EventBatch
-	err := msgpack.Unmarshal(parts[2], &eventBatch)
-	Expect(err).NotTo(HaveOccurred())
-	for _, rawEvent := range eventBatch.Events {
-		var taggedUnion []msgpack.RawMessage
-		err := msgpack.Unmarshal(rawEvent, &taggedUnion)
-		Expect(err).NotTo(HaveOccurred())
-		Expect(len(taggedUnion)).To(BeNumerically(">", 1))
-
-		payloadBytes, err := msgpack.Marshal(taggedUnion[1:])
-		Expect(err).NotTo(HaveOccurred())
-
-		var tag string
-		err = msgpack.Unmarshal(taggedUnion[0], &tag)
-		Expect(err).NotTo(HaveOccurred())
-
-		switch tag {
-		case kvevents.BlockStoredEventTag:
-			var bs kvevents.BlockStored
-			err = msgpack.Unmarshal(payloadBytes, &bs)
-			stored = append(stored, bs.BlockHashes...)
-		case kvevents.BlockRemovedEventTag:
-			var br kvevents.BlockRemoved
-			err = msgpack.Unmarshal(payloadBytes, &br)
-			removed = append(removed, br.BlockHashes...)
-
-		default:
-			Fail("unexpected tag " + tag)
-			continue
-		}
-		Expect(err).NotTo(HaveOccurred())
-	}
-	return stored, removed
-}
-
-func createSub(config *common.Configuration) (*zmq.Socket, string) {
-	zctx, err := zmq.NewContext()
-	Expect(err).NotTo(HaveOccurred())
-	sub, err := zctx.NewSocket(zmq.SUB)
-	Expect(err).NotTo(HaveOccurred())
-	err = sub.Bind(wildcardEndpoint)
-	Expect(err).NotTo(HaveOccurred())
-	// get the actual port
-	endpoint, err := sub.GetLastEndpoint()
-	Expect(err).NotTo(HaveOccurred())
-	config.ZMQEndpoint = endpoint
-	topic := createTopic(config)
-	err = sub.SetSubscribe(topic)
-	Expect(err).NotTo(HaveOccurred())
-	return sub, topic
-}