fix: restoring KVBM wheel generation for `framework=none` and adding CI validation (#4772)

Signed-off-by: Olga Andreeva <oandreeva@nvidia.com>

.github/workflows/container-validation-dynamo.yml

@@ -44,7 +44,7 @@ jobs:
         env:
           GITHUB_TOKEN: ${{ secrets.CI_TOKEN }}
         run: |
-          ./container/build.sh --tag ${{ steps.define_image_tag.outputs.image_tag }} --target dev --framework none
+          ./container/build.sh --tag ${{ steps.define_image_tag.outputs.image_tag }} --target dev --framework none --enable-kvbm
       - name: Start services with docker-compose
         working-directory: ./deploy
         run: |

container/build.sh

@@ -125,6 +125,10 @@ NIXL_UCX_EFA_REF=9d2b88a1f67faf9876f267658bd077b379b8bb76
 
 NO_CACHE=""
 
+# KVBM (KV Cache Block Manager) - default disabled, enabled automatically for VLLM/TRTLLM
+# or can be explicitly enabled via --enable-kvbm flag
+ENABLE_KVBM=false
+
 # sccache configuration for S3
 USE_SCCACHE=""
 SCCACHE_BUCKET=""
@@ -802,15 +806,15 @@ fi
 
 # ENABLE_KVBM: Used in base Dockerfile for block-manager feature.
 #              Declared but not currently used in Dockerfile.{vllm,trtllm}.
+# Force KVBM to be enabled for VLLM and TRTLLM frameworks
 if [[ $FRAMEWORK == "VLLM" ]] || [[ $FRAMEWORK == "TRTLLM" ]]; then
     echo "Forcing enable_kvbm to true in ${FRAMEWORK} image build"
     ENABLE_KVBM=true
-else
-    ENABLE_KVBM=false
 fi
+# For other frameworks, ENABLE_KVBM defaults to false unless --enable-kvbm flag was provided
 
-if [  ! -z ${ENABLE_KVBM} ]; then
-    echo "Enabling the KVBM in the dynamo image"
+if [[ ${ENABLE_KVBM} == "true" ]]; then
+    echo "Enabling KVBM in the dynamo image"
     BUILD_ARGS+=" --build-arg ENABLE_KVBM=${ENABLE_KVBM} "
 fi
 

tests/dependencies/test_kvbm_imports.py

+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Unit tests to verify KVBM package and wheels are properly installed."""
+
+import subprocess
+
+import pytest
+
+
+# Helper functions for KVBM verification
+def _check_kvbm_wheel_exists():
+    """Helper to verify KVBM wheel file exists in expected location."""
+    result = subprocess.run(
+        ["bash", "-c", "ls /opt/dynamo/wheelhouse/kvbm*.whl"],
+        capture_output=True,
+        text=True,
+    )
+
+    assert result.returncode == 0, (
+        f"KVBM wheel not found in /opt/dynamo/wheelhouse/\n"
+        f"stdout: {result.stdout}\n"
+        f"stderr: {result.stderr}"
+    )
+    assert (
+        "kvbm" in result.stdout
+    ), f"Expected kvbm wheel in output, got: {result.stdout}"
+
+
+def _check_kvbm_imports():
+    """Helper to verify KVBM package and core classes can be imported."""
+    try:
+        import kvbm
+        from kvbm import BlockManager, KvbmLeader, KvbmWorker
+
+        assert kvbm is not None, "kvbm module is None"
+        assert BlockManager is not None, "BlockManager class not available"
+        assert KvbmLeader is not None, "KvbmLeader class not available"
+        assert KvbmWorker is not None, "KvbmWorker class not available"
+    except ImportError as e:
+        pytest.fail(f"Failed to import KVBM package or core classes: {e}")
+
+
+# Base tests (no framework markers) - run in main job with --framework none --enable-kvbm
+@pytest.mark.pre_merge
+def test_kvbm_wheel_exists():
+    """Verify KVBM wheel file exists in expected location."""
+    _check_kvbm_wheel_exists()
+
+
+@pytest.mark.pre_merge
+def test_kvbm_imports():
+    """Verify KVBM package and core classes can be imported."""
+    _check_kvbm_imports()
+
+
+# vLLM-specific tests - run in vLLM job (vLLM auto-enables KVBM)
+@pytest.mark.pre_merge
+@pytest.mark.vllm
+def test_kvbm_wheel_exists_vllm():
+    """Verify KVBM wheel exists in vLLM image."""
+    _check_kvbm_wheel_exists()
+
+
+@pytest.mark.pre_merge
+@pytest.mark.vllm
+def test_kvbm_imports_vllm():
+    """Verify KVBM package and core classes can be imported in vLLM image."""
+    _check_kvbm_imports()
+
+
+# TRT-LLM-specific tests - run in TRT-LLM job (TRT-LLM auto-enables KVBM)
+@pytest.mark.pre_merge
+@pytest.mark.trtllm
+def test_kvbm_wheel_exists_trtllm():
+    """Verify KVBM wheel exists in TRT-LLM image."""
+    _check_kvbm_wheel_exists()
+
+
+@pytest.mark.pre_merge
+@pytest.mark.trtllm
+def test_kvbm_imports_trtllm():
+    """Verify KVBM package and core classes can be imported in TRT-LLM image."""
+    _check_kvbm_imports()