Install python3-dev and pkg-config

Extract template wrapper from kv-cache-manager

Signed-off-by: Pierangelo Di Pilato <[email protected]>

Author: pierDipi

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

@@ -21,22 +21,51 @@ jobs:
           go-version-file: ./go.mod
           cache-dependency-path: ./go.sum
 
-      - name: Install libzmq dependencies (kvcache/kvevents)
+      - name: Install libzmq and Python dependencies (kvcache/kvevents)
         run: |
           sudo apt-get update
+          sudo apt-get install -y pkg-config python3-dev python3-pip
           make download-zmq
+          pip3 install transformers --break-system-packages
 
-      - name: Set PKG_CONFIG_PATH
-        run: echo "PKG_CONFIG_PATH=/usr/lib/pkgconfig" >> $GITHUB_ENV
+      - name: Configure CGO for Python
+        run: |
+          PYTHON_INCLUDE=$(python3 -c "import sysconfig; print(sysconfig.get_path('include'))")
+          echo "CPATH=${PYTHON_INCLUDE}:${CPATH}" >> $GITHUB_ENV
+          echo "CGO_ENABLED=1" >> $GITHUB_ENV
+          echo "CGO_CFLAGS=$(python3-config --cflags --embed)" >> $GITHUB_ENV
+          echo "CGO_LDFLAGS=$(python3-config --ldflags --embed)" >> $GITHUB_ENV
+
+      - name: Set PKG_CONFIG_PATH and PYTHONPATH
+        run: |
+          echo "PKG_CONFIG_PATH=/usr/lib/pkgconfig" >> $GITHUB_ENV
+          GOMODCACHE=$(go env GOMODCACHE)
+          # Extract kv-cache-manager version from go.mod
+          KV_CACHE_MGR_VERSION=$(go list -m -f '{{.Version}}' github.com/llm-d/llm-d-kv-cache-manager)
+          KV_CACHE_MGR_PATH="${GOMODCACHE}/github.com/llm-d/llm-d-kv-cache-manager@${KV_CACHE_MGR_VERSION}/pkg/preprocessing/chat_completions"
+          echo "PYTHONPATH=${KV_CACHE_MGR_PATH}:${PYTHONPATH}" >> $GITHUB_ENV
 
       - name: Run lint checks
         uses: golangci/golangci-lint-action@v8
         with:
           version: 'v2.4.0'
           args: "--config=./.golangci.yml"
+        env:
+          CGO_ENABLED: ${{ env.CGO_ENABLED }}
+          CGO_CFLAGS: ${{ env.CGO_CFLAGS }}
+          CGO_LDFLAGS: ${{ env.CGO_LDFLAGS }}
+          CPATH: ${{ env.CPATH }}
+          PKG_CONFIG_PATH: ${{ env.PKG_CONFIG_PATH }}
 
       - name: Run go test
-        shell: bash 
+        shell: bash
         run: |
           echo "Running tests with Ginkgo..."
           make test
+        env:
+          CGO_ENABLED: ${{ env.CGO_ENABLED }}
+          CGO_CFLAGS: ${{ env.CGO_CFLAGS }}
+          CGO_LDFLAGS: ${{ env.CGO_LDFLAGS }}
+          CPATH: ${{ env.CPATH }}
+          PKG_CONFIG_PATH: ${{ env.PKG_CONFIG_PATH }}
+          PYTHONPATH: ${{ env.PYTHONPATH }}

Dockerfile

@@ -6,7 +6,7 @@ ARG TARGETARCH
 # Install build tools
 # The builder is based on UBI8, so we need epel-release-8.
 RUN dnf install -y 'https://dl.fedoraproject.org/pub/epel/epel-release-latest-8.noarch.rpm' && \
-    dnf install -y gcc-c++ libstdc++ libstdc++-devel clang zeromq-devel pkgconfig && \
+    dnf install -y gcc-c++ libstdc++ libstdc++-devel clang zeromq-devel pkgconfig python3-devel && \
     dnf clean all
 
 WORKDIR /workspace
@@ -28,6 +28,13 @@ ARG TOKENIZER_VERSION=v1.22.1
 RUN curl -L https://github.com/daulet/tokenizers/releases/download/${TOKENIZER_VERSION}/libtokenizers.${TARGETOS}-${TARGETARCH}.tar.gz | tar -xz -C lib
 RUN ranlib lib/*.a
 
+# Copy Python wrapper from kv-cache-manager dependency
+# Extract version dynamically and copy to a known location
+RUN KV_CACHE_MGR_VERSION=$(go list -m -f '{{.Version}}' github.com/llm-d/llm-d-kv-cache-manager) && \
+    mkdir -p /workspace/kv-cache-manager-wrapper && \
+    cp /go/pkg/mod/github.com/llm-d/llm-d-kv-cache-manager@${KV_CACHE_MGR_VERSION}/pkg/preprocessing/chat_completions/render_jinja_template_wrapper.py \
+       /workspace/kv-cache-manager-wrapper/
+
 # Build
 # the GOARCH has not a default value to allow the binary be built according to the host where the command
 # was called. For example, if we call make image-build in a local env which has the Apple Silicon M1 SO
@@ -38,20 +45,38 @@ ENV GOOS=${TARGETOS:-linux}
 ENV GOARCH=${TARGETARCH}
 RUN go build -a -o bin/llm-d-inference-sim -ldflags="-extldflags '-L$(pwd)/lib'" cmd/cmd.go
 
+# Runtime stage
 # Use ubi9 as a minimal base image to package the manager binary
 # Refer to https://catalog.redhat.com/software/containers/ubi9/ubi-minimal/615bd9b4075b022acc111bf5 for more details
 FROM registry.access.redhat.com/ubi9/ubi-minimal:latest
 
 WORKDIR /
 
-# Install zeromq runtime library needed by the manager.
+# Install zeromq runtime library and Python runtime needed by the manager.
 # The final image is UBI9, so we need epel-release-9.
+# Using microdnf for minimal image size
 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 && \
-    dnf clean all && \
-    rm -rf /var/cache/dnf /var/lib/dnf
+RUN curl -L -o /tmp/epel-release.rpm https://dl.fedoraproject.org/pub/epel/epel-release-latest-9.noarch.rpm && \
+    rpm -i /tmp/epel-release.rpm && \
+    rm /tmp/epel-release.rpm && \
+    microdnf install -y --setopt=install_weak_deps=0 zeromq python3 python3-libs python3-pip && \
+    microdnf clean all && \
+    rm -rf /var/cache/yum /var/lib/yum
+
+# Install wrapper as a module in site-packages
+# Extract the kv-cache-manager version dynamically from go.mod in the builder stage
+RUN mkdir -p /usr/local/lib/python3.9/site-packages/
+COPY --from=builder /workspace/kv-cache-manager-wrapper/render_jinja_template_wrapper.py /usr/local/lib/python3.9/site-packages/
+
+# Python deps (no cache, single target) – install transformers
+ENV PIP_NO_CACHE_DIR=1 PIP_DISABLE_PIP_VERSION_CHECK=1
+RUN python3 -m pip install --no-cache-dir --upgrade pip setuptools wheel && \
+    python3 -m pip install --no-cache-dir --target /usr/local/lib/python3.9/site-packages transformers && \
+    rm -rf /root/.cache/pip
+
+# Python env
+ENV PYTHONPATH="/usr/local/lib/python3.9/site-packages:/usr/lib/python3.9/site-packages"
+ENV PYTHON=python3
 
 COPY --from=builder /workspace/bin/llm-d-inference-sim /app/llm-d-inference-sim
 

pkg/llm-d-inference-sim/simulator.go

@@ -347,7 +347,7 @@ func (s *VllmSimulator) initializeSim(ctx context.Context) error {
 
 	s.tokenizer, err = tokenization.NewCachedHFTokenizer(tokenizationConfig.HFTokenizerConfig)
 	if err != nil {
-		return fmt.Errorf("failed to create tokenizer: %w", err)
+		return fmt.Errorf("failed to create hf tokenizer: %w", err)
 	}
 
 	if s.config.EnableKVCache {