ZMQ publisher (llm-d#119)

* ZMQ publisher Signed-off-by: Ira <IRAR@il.ibm.com> * Fix lint errors Signed-off-by: Ira <IRAR@il.ibm.com> * Fix lint errors Signed-off-by: Ira <IRAR@il.ibm.com> * Use sleep instead of ticker Signed-off-by: Ira <IRAR@il.ibm.com> --------- Signed-off-by: Ira <IRAR@il.ibm.com>

Author: irar2

.github/workflows/ci-pr-checks.yaml

@@ -21,6 +21,14 @@ jobs:
  go-version: '1.24.0'
  cache-dependency-path: ./go.sum
 
+ - name: Install libzmq dependencies (kvcache/kvevents)
+ run: |
+ sudo apt-get update
+ sudo apt-get install -y libzmq3-dev pkg-config
+
+ - name: Set PKG_CONFIG_PATH
+ run: echo "PKG_CONFIG_PATH=/usr/lib/pkgconfig" >> $GITHUB_ENV
+
  - name: Run lint checks
  uses: golangci/golangci-lint-action@v8
  with:

Dockerfile

@@ -41,6 +41,14 @@ RUN go build -a -o bin/llm-d-inference-sim -ldflags="-extldflags '-L$(pwd)/lib'"
 FROM registry.access.redhat.com/ubi9/ubi-minimal:latest
 
 WORKDIR /
+
+# Install zeromq runtime library needed by the manager.
+# The final image is UBI9, so we need epel-release-9.
+USER root
+RUN microdnf install -y dnf && \
+ dnf install -y 'https://dl.fedoraproject.org/pub/epel/epel-release-latest-9.noarch.rpm' && \
+ dnf install -y zeromq
+
 COPY --from=builder /workspace/bin/llm-d-inference-sim /app/llm-d-inference-sim
 
 # USER 65532:65532

Makefile

@@ -63,7 +63,7 @@ format: ## Format Go source files
 @gofmt -l -w $(SRC)
 
 .PHONY: test
-test: check-ginkgo download-tokenizer ## Run tests
+test: check-ginkgo download-tokenizer download-zmq ## Run tests
 @printf "\033[33;1m==== Running tests ====\033[0m\n"
 CGO_ENABLED=1 ginkgo -ldflags="$(LDFLAGS)" -v -r
 
@@ -80,7 +80,7 @@ lint: check-golangci-lint ## Run lint
 ##@ Build
 
 .PHONY: build
-build: check-go download-tokenizer ##
+build: check-go download-tokenizer download-zmq 
 @printf "\033[33;1m==== Building ====\033[0m\n"
 go build -ldflags="$(LDFLAGS)" -o bin/$(PROJECT_NAME) cmd/$(PROJECT_NAME)/main.go
 
@@ -193,3 +193,35 @@ print-project-name: ## Print the current project name
 .PHONY: install-hooks
 install-hooks: ## Install git hooks
 git config core.hooksPath hooks
+
+##@ ZMQ Setup
+
+.PHONY: download-zmq
+download-zmq: ## Install ZMQ dependencies based on OS/ARCH
+@echo "Checking if ZMQ is already installed..."
+@if pkg-config --exists libzmq; then \
+ echo "✅ ZMQ is already installed."; \
+else \
+ echo "Installing ZMQ dependencies..."; \
+ if [ "$(TARGETOS)" = "linux" ]; then \
+ if [ -x "$(command -v apt)" ]; then \
+ apt update && apt install -y libzmq3-dev; \
+ elif [ -x "$(command -v dnf)" ]; then \
+ dnf install -y zeromq-devel; \
+ else \
+ echo "Unsupported Linux package manager. Install libzmq manually."; \
+ exit 1; \
+ fi; \
+ elif [ "$(TARGETOS)" = "darwin" ]; then \
+ if [ -x "$(command -v brew)" ]; then \
+ brew install zeromq; \
+ else \
+ echo "Homebrew is not installed and is required to install zeromq. Install it from https://brew.sh/"; \
+ exit 1; \
+ fi; \
+ else \
+ echo "Unsupported OS: $(TARGETOS). Install libzmq manually - check https://zeromq.org/download/ for guidance."; \
+ exit 1; \
+ fi; \
+ echo "✅ ZMQ dependencies installed."; \
+fi

go.mod

@@ -12,10 +12,12 @@ require (
 github.com/onsi/ginkgo/v2 v2.23.4
 github.com/onsi/gomega v1.37.0
 github.com/openai/openai-go v0.1.0-beta.10
+github.com/pebbe/zmq4 v1.4.0
 github.com/prometheus/client_golang v1.22.0
 github.com/santhosh-tekuri/jsonschema/v5 v5.3.1
 github.com/spf13/pflag v1.0.6
 github.com/valyala/fasthttp v1.59.0
+github.com/vmihailenco/msgpack v4.0.4+incompatible
 gopkg.in/yaml.v3 v3.0.1
 k8s.io/klog/v2 v2.130.1
 )
@@ -34,6 +36,7 @@ require (
 github.com/go-openapi/swag v0.23.0 // indirect
 github.com/go-task/slim-sprig/v3 v3.0.0 // indirect
 github.com/gogo/protobuf v1.3.2 // indirect
+github.com/golang/protobuf v1.5.2 // indirect
 github.com/google/gnostic-models v0.6.9 // indirect
 github.com/google/go-cmp v0.7.0 // indirect
 github.com/google/pprof v0.0.0-20250403155104-27863c87afa6 // indirect
@@ -64,6 +67,7 @@ require (
 golang.org/x/text v0.23.0 // indirect
 golang.org/x/time v0.9.0 // indirect
 golang.org/x/tools v0.31.0 // indirect
+google.golang.org/appengine v1.6.8 // indirect
 google.golang.org/protobuf v1.36.5 // indirect
 gopkg.in/evanphx/json-patch.v4 v4.12.0 // indirect
 gopkg.in/inf.v0 v0.9.1 // indirect

go.sum

@@ -40,6 +40,10 @@ github.com/gogo/protobuf v1.3.2/go.mod h1:P1XiOD3dCwIKUDQYPy72D8LYyHL2YPYrpS2s69
 github.com/google/gnostic-models v0.6.9 h1:MU/8wDLif2qCXZmzncUQ/BOfxWfthHi63KqpoNbWqVw=
 github.com/google/gnostic-models v0.6.9/go.mod h1:CiWsm0s6BSQd1hRn8/QmxqB6BesYcbSZxsz9b0KuDBw=
 github.com/google/go-cmp v0.5.9/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY=
+github.com/golang/protobuf v1.5.0/go.mod h1:FsONVRAS9T7sI+LIUmWTfcYkHO4aIWwzhcaSAoJOfIk=
+github.com/golang/protobuf v1.5.2 h1:ROPKBNFfQgOUMifHyP+KYbvpjbdoFNs+aK7DXlji0Tw=
+github.com/golang/protobuf v1.5.2/go.mod h1:XVQd3VNwM+JqD3oG2Ue2ip4fOMUkwXdXDdiuN0vRsmY=
+github.com/google/go-cmp v0.5.5/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE=
 github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8=
 github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU=
 github.com/google/gofuzz v1.0.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/M65Eg=
@@ -86,6 +90,8 @@ github.com/openai/openai-go v0.1.0-beta.10/go.mod h1:g461MYGXEXBVdV5SaR/5tNzNbSf
 github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4=
 github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
+github.com/pebbe/zmq4 v1.4.0 h1:gO5P92Ayl8GXpPZdYcD62Cwbq0slSBVVQRIXwGSJ6eQ=
+github.com/pebbe/zmq4 v1.4.0/go.mod h1:nqnPueOapVhE2wItZ0uOErngczsJdLOGkebMxaO8r48=
 github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 h1:Jamvg5psRIccs7FGNTlIRMkT8wgtp5eCXdBlqhYGL6U=
 github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
 github.com/prashantv/gostub v1.1.0 h1:BTyx3RfQjRHnUWaGF9oQos79AlQ5k8WNktv7VGvVH4g=
@@ -180,6 +186,49 @@ golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8T
 golang.org/x/xerrors v0.0.0-20191011141410-1b5146add898/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
 golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
 golang.org/x/xerrors v0.0.0-20200804184101-5ec99f83aff1/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
+github.com/vmihailenco/msgpack v4.0.4+incompatible h1:dSLoQfGFAo3F6OoNhwUmLwVgaUXK79GlxNBwueZn0xI=
+github.com/vmihailenco/msgpack v4.0.4+incompatible/go.mod h1:fy3FlTQTDXWkZ7Bh6AcGMlsjHatGryHQYUTf1ShIgkk=
+github.com/xyproto/randomstring v1.0.5 h1:YtlWPoRdgMu3NZtP45drfy1GKoojuR7hmRcnhZqKjWU=
+github.com/xyproto/randomstring v1.0.5/go.mod h1:rgmS5DeNXLivK7YprL0pY+lTuhNQW3iGxZ18UQApw/E=
+github.com/yuin/goldmark v1.4.13/go.mod h1:6yULJ656Px+3vBD8DxQVa3kxgyrAnzto9xy5taEt/CY=
+go.uber.org/automaxprocs v1.6.0 h1:O3y2/QNTOdbF+e/dpXNNW7Rx2hZ4sTIPyybbxyNqTUs=
+go.uber.org/automaxprocs v1.6.0/go.mod h1:ifeIMSnPZuznNm6jmdzmU3/bfk01Fe2fotchwEFJ8r8=
+golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w=
+golang.org/x/crypto v0.0.0-20210921155107-089bfa567519/go.mod h1:GvvjBRRGRdwPK5ydBHafDWAxML/pGHZbMvKqRZ5+Abc=
+golang.org/x/mod v0.6.0-dev.0.20220419223038-86c51ed26bb4/go.mod h1:jJ57K6gSWd91VN4djpZkiMVwK6gcyfeH4XE8wZrZaV4=
+golang.org/x/net v0.0.0-20190620200207-3b0461eec859/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s=
+golang.org/x/net v0.0.0-20210226172049-e18ecbb05110/go.mod h1:m0MpNAwzfU5UDzcl9v0D8zg8gWTRqZa9RBIspLL5mdg=
+golang.org/x/net v0.0.0-20220722155237-a158d28d115b/go.mod h1:XRhObCWvk6IyKnWLug+ECip1KBveYUHfp+8e9klMJ9c=
+golang.org/x/net v0.38.0 h1:vRMAPTMaeGqVhG5QyLJHqNDwecKTomGeqbnfZyKlBI8=
+golang.org/x/net v0.38.0/go.mod h1:ivrbrMbzFq5J41QOQh0siUuly180yBYtLp+CKbEaFx8=
+golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
+golang.org/x/sync v0.0.0-20220722155255-886fb9371eb4/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
+golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
+golang.org/x/sys v0.0.0-20201119102817-f84b799fce68/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
+golang.org/x/sys v0.0.0-20210615035016-665e8c7367d1/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
+golang.org/x/sys v0.0.0-20220520151302-bc2c85ada10a/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
+golang.org/x/sys v0.0.0-20220722155257-8c9f86f7a55f/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
+golang.org/x/sys v0.32.0 h1:s77OFDvIQeibCmezSnk/q6iAfkdiQaJi4VzroCFrN20=
+golang.org/x/sys v0.32.0/go.mod h1:BJP2sWEmIv4KK5OTEluFJCKSidICx8ciO85XgH3Ak8k=
+golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo=
+golang.org/x/term v0.0.0-20210927222741-03fcf44c2211/go.mod h1:jbD1KX2456YbFQfuXm/mYQcufACuNUgVhRMnK/tPxf8=
+golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
+golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
+golang.org/x/text v0.3.7/go.mod h1:u+2+/6zg+i71rQMx5EYifcz6MCKuco9NR6JIITiCfzQ=
+golang.org/x/text v0.3.8/go.mod h1:E6s5w1FMmriuDzIBO73fBruAKo1PCIq6d2Q6DHfQ8WQ=
+golang.org/x/text v0.23.0 h1:D71I7dUrlY+VX0gQShAThNGHFxZ13dGLBHQLVl1mJlY=
+golang.org/x/text v0.23.0/go.mod h1:/BLNzu4aZCJ1+kcD0DNRotWKage4q2rGVAg4o22unh4=
+golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
+golang.org/x/tools v0.0.0-20191119224855-298f0cb1881e/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo=
+golang.org/x/tools v0.1.12/go.mod h1:hNGJHUnrk76NpqgfD5Aqm5Crs+Hm0VOH/i9J2+nxYbc=
+golang.org/x/tools v0.31.0 h1:0EedkvKDbh+qistFTd0Bcwe/YLh4vHwWEkiI0toFIBU=
+golang.org/x/tools v0.31.0/go.mod h1:naFTU+Cev749tSJRXJlna0T3WxKvb1kWEx15xA4SdmQ=
+golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
+golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
+google.golang.org/appengine v1.6.8 h1:IhEN5q69dyKagZPYMSdIjS2HqprW324FRQZJcGqPAsM=
+google.golang.org/appengine v1.6.8/go.mod h1:1jJ3jBArFh5pcgW8gCtRJnepW8FzD1V44FJffLiz/Ds=
+google.golang.org/protobuf v1.26.0-rc.1/go.mod h1:jlhhOSvTdKEhbULTjvd4ARK9grFBp09yW+WbY/TyQbw=
+google.golang.org/protobuf v1.26.0/go.mod h1:9q0QmTI4eRPtz6boOQmLYwt+qCgq0jsYwAQnmE0givc=
 google.golang.org/protobuf v1.36.5 h1:tPhr+woSbjfYvY6/GPufUoYizxw1cF/yFoxJ2fmpwlM=
 google.golang.org/protobuf v1.36.5/go.mod h1:9fA7Ob0pmnwhb644+1+CVWFRbNajQ6iRojtC/QF5bRE=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=

pkg/common/config.go

@@ -313,6 +313,7 @@ func ParseCommandParamsAndLoadConfig() (*Configuration, error) {
 f.IntVar(&config.MinToolCallArrayParamLength, "min-tool-call-array-param-length", config.MinToolCallArrayParamLength, "Minimum possible length of array parameters in a tool call")
 f.IntVar(&config.ToolCallNotRequiredParamProbability, "tool-call-not-required-param-probability", config.ToolCallNotRequiredParamProbability, "Probability to add a parameter, that is not required, in a tool call")
 f.IntVar(&config.ObjectToolCallNotRequiredParamProbability, "object-tool-call-not-required-field-probability", config.ObjectToolCallNotRequiredParamProbability, "Probability to add a field, that is not required, in an object in a tool call")
+f.BoolVar(&config.EnableKVCache, "enable-kvcache", config.EnableKVCache, "Defines if KV cache feature is enabled")
 
 // These values were manually parsed above in getParamValueFromArgs, we leave this in order to get these flags in --help
 var dummyString string

pkg/common/publisher.go

@@ -0,0 +1,90 @@
+/*
+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 (
+"context"
+"encoding/binary"
+"errors"
+"fmt"
+"sync/atomic"
+
+zmq "github.com/pebbe/zmq4"
+"github.com/vmihailenco/msgpack"
+"k8s.io/klog/v2"
+)
+
+// Publisher sends events to a ZMQ endpoint.
+type Publisher struct {
+socket *zmq.Socket
+endpoint string
+seqNum uint64
+}
+
+// NewPublisher creates a new ZMQ publisher.
+// endpoint is the ZMQ address to bind to (e.g., "tcp://*:5557").
+func NewPublisher(endpoint string) (*Publisher, error) {
+socket, err := zmq.NewSocket(zmq.PUB)
+if err != nil {
+return nil, fmt.Errorf("failed to create ZMQ PUB socket: %w", err)
+}
+
+if err := socket.Connect(endpoint); err != nil {
+errClose := socket.Close()
+return nil, errors.Join(
+fmt.Errorf("failed to connect to %s: %w", endpoint, err),
+errClose,
+)
+}
+
+return &Publisher{
+socket: socket,
+endpoint: endpoint,
+}, nil
+}
+
+// PublishEvent publishes a KV cache event batch to the ZMQ topic.
+// topic should include the pod identifier (e.g., "kv.pod1").
+func (p *Publisher) PublishEvent(ctx context.Context, topic string, batch interface{}) error {
+logger := klog.FromContext(ctx).V(0)
+
+payload, err := msgpack.Marshal(batch)
+if err != nil {
+return fmt.Errorf("failed to marshal event batch: %w", err)
+}
+
+// sequence number for ordering
+seq := atomic.AddUint64(&p.seqNum, 1)
+seqBytes := make([]byte, 8)
+binary.BigEndian.PutUint64(seqBytes, seq)
+
+// send topic, sequence, payload
+if _, err := p.socket.SendMessage(topic, seqBytes, payload); err != nil {
+return fmt.Errorf("failed to send message to topic %s: %w", topic, err)
+}
+
+logger.Info("Published event batch", "topic", topic, "seq", seq)
+return nil
+}
+
+// Close closes the publisher and cleans up resources.
+func (p *Publisher) Close() error {
+if p.socket != nil {
+return p.socket.Close()
+}
+return nil
+}

pkg/common/publisher_test.go

@@ -0,0 +1,78 @@
+/*
+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 (
+"context"
+"encoding/binary"
+
+"time"
+
+. "github.com/onsi/ginkgo/v2"
+. "github.com/onsi/gomega"
+zmq "github.com/pebbe/zmq4"
+"github.com/vmihailenco/msgpack"
+)
+
+const (
+topic = "test-topic"
+endpoint = "tcp://localhost:5557"
+data = "Hello"
+)
+
+var _ = Describe("Publisher", func() {
+It("should publish and receive correct message", func() {
+zctx, err := zmq.NewContext()
+Expect(err).NotTo(HaveOccurred())
+sub, err := zctx.NewSocket(zmq.SUB)
+Expect(err).NotTo(HaveOccurred())
+err = sub.Bind(endpoint)
+Expect(err).NotTo(HaveOccurred())
+err = sub.SetSubscribe(topic)
+Expect(err).NotTo(HaveOccurred())
+
+time.Sleep(100 * time.Millisecond)
+
+pub, err := NewPublisher(endpoint)
+Expect(err).NotTo(HaveOccurred())
+
+ctx, cancel := context.WithCancel(context.Background())
+defer cancel()
+
+go func() {
+// Make sure that sub.RecvMessageBytes is called before pub.PublishEvent
+time.Sleep(time.Second)
+err := pub.PublishEvent(ctx, topic, data)
+Expect(err).NotTo(HaveOccurred())
+}()
+
+// The message should be [topic, seq, payload]
+parts, err := sub.RecvMessageBytes(0)
+Expect(err).NotTo(HaveOccurred())
+Expect(parts).To(HaveLen(3))
+
+Expect(string(parts[0])).To(Equal(topic))
+
+seq := binary.BigEndian.Uint64(parts[1])
+Expect(seq).To(Equal(uint64(1)))
+
+var payload string
+err = msgpack.Unmarshal(parts[2], &payload)
+Expect(err).NotTo(HaveOccurred())
+Expect(payload).To(Equal(data))
+})
+})