fix(kvbm): handle typed KvCacheConnectorConfig in consolidator check (#8117)

Signed-off-by: Alec Flowers <aflowers@nvidia.com>

lib/bindings/kvbm/python/kvbm/trtllm_integration/consolidator_config.py

@@ -48,19 +48,23 @@ def should_enable_consolidator(arg_map) -> bool:
         )
         return False
 
-    # Check if KVBM connector is enabled
-    if not isinstance(arg_map, dict):
-        logger.warning("KV Event Consolidator is not enabled: arg_map is not a dict")
-        return False
-
-    kv_connector_config = arg_map.get("kv_connector_config", {})
-    if not isinstance(kv_connector_config, dict):
+    # Check if KVBM connector is enabled by extracting connector_module
+    # from kv_connector_config (works whether arg_map holds raw dicts or typed objects)
+    kv_connector_config = (
+        arg_map.get("kv_connector_config") if isinstance(arg_map, dict) else None
+    )
+    if kv_connector_config is None:
         logger.warning(
-            "KV Event Consolidator is not enabled: kv_connector_config is not a dict"
+            "KV Event Consolidator is not enabled: no kv_connector_config found"
         )
         return False
 
+    if isinstance(kv_connector_config, dict):
         connector_module = kv_connector_config.get("connector_module", "")
+    else:
+        # Access directly so AttributeError surfaces if the contract changes
+        connector_module = kv_connector_config.connector_module or ""
+
     has_kvbm_connector = "kvbm.trtllm_integration.connector" in connector_module
 
     if not has_kvbm_connector:

tests/kvbm_integration/test_consolidator_config_unit.py

+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Unit tests for consolidator_config.should_enable_consolidator.
+
+Covers both raw dict and typed object paths for kv_connector_config,
+and asserts that the production code path (build_kv_connector_config)
+produces a type that the consolidator check handles correctly.
+"""
+
+import os
+from unittest.mock import patch
+
+import pytest
+
+kvbm = pytest.importorskip("kvbm", reason="kvbm package not installed")
+from kvbm.trtllm_integration.consolidator_config import (  # noqa: E402
+    should_enable_consolidator,
+)
+
+KVBM_MODULE = "kvbm.trtllm_integration.connector"
+
+
+@pytest.mark.unit
+@pytest.mark.pre_merge
+@pytest.mark.kvbm
+@pytest.mark.gpu_0
+class TestShouldEnableConsolidatorDict:
+    """Tests that only need kvbm (no GPU, no trtllm)."""
+
+    def test_dict_config_with_kvbm(self):
+        """Raw dict path (--extra-engine-args YAML, used in tests)."""
+        arg_map = {"kv_connector_config": {"connector_module": KVBM_MODULE}}
+        assert should_enable_consolidator(arg_map) is True
+
+    def test_non_kvbm_connector(self):
+        arg_map = {"kv_connector_config": {"connector_module": "other.connector"}}
+        assert should_enable_consolidator(arg_map) is False
+
+    def test_no_connector_config(self):
+        assert should_enable_consolidator({"backend": "pytorch"}) is False
+
+    def test_env_var_disables(self):
+        arg_map = {"kv_connector_config": {"connector_module": KVBM_MODULE}}
+        with patch.dict(
+            os.environ, {"DYN_KVBM_KV_EVENTS_ENABLE_CONSOLIDATOR": "false"}
+        ):
+            assert should_enable_consolidator(arg_map) is False
+
+
+@pytest.mark.unit
+@pytest.mark.pre_merge
+@pytest.mark.kvbm
+@pytest.mark.trtllm
+@pytest.mark.gpu_1
+class TestShouldEnableConsolidatorTyped:
+    """Tests that need trtllm (requires GPU for import)."""
+
+    def test_typed_config_with_kvbm(self):
+        """Typed object path (DYN_CONNECTOR=kvbm, used in production)."""
+        from tensorrt_llm.llmapi.llm_args import KvCacheConnectorConfig
+
+        config = KvCacheConnectorConfig(
+            connector_module=KVBM_MODULE,
+            connector_scheduler_class="DynamoKVBMConnectorLeader",
+            connector_worker_class="DynamoKVBMConnectorWorker",
+        )
+        arg_map = {"kv_connector_config": config}
+        assert should_enable_consolidator(arg_map) is True
+
+    def test_production_build_kv_connector_config(self):
+        """Assert build_kv_connector_config output works with consolidator check.
+
+        This is the key regression test: if llm_worker.py changes the type
+        returned by build_kv_connector_config, this test will catch it.
+        """
+        from dynamo.trtllm.workers.llm_worker import build_kv_connector_config
+
+        class FakeConfig:
+            connector = ("kvbm",)
+
+        config = build_kv_connector_config(FakeConfig())
+        assert config is not None
+        assert hasattr(config, "connector_module")
+        assert KVBM_MODULE in config.connector_module
+
+        # The consolidator check must work with this exact object
+        arg_map = {"kv_connector_config": config}
+        assert should_enable_consolidator(arg_map) is True