issue/340 - enhance validation with some checks

Author: pengcheng888

python/infinilm/config/kv_transfer.py

@@ -41,14 +41,16 @@ def __post_init__(self) -> None:
         if self.engine_id is None:
             self.engine_id = f"{self.kv_role}_" + str(uuid.uuid4())
 
-        if not self.kv_connector_extra_config:
-            self.kv_connector_extra_config = dict(self.kv_connector_extra_config or {})
-            self.kv_connector_extra_config.setdefault("mooncake_protocol", "rdma")
-
-        assert all(
-            key in ["mooncake_protocol", "num_workers"]
-            for key in self.kv_connector_extra_config.keys()
-        )
+        self.kv_connector_extra_config = dict(self.kv_connector_extra_config or {})
+        self.kv_connector_extra_config.setdefault("mooncake_protocol", "rdma")
+
+        allowed_extra_config_keys = frozenset({"mooncake_protocol", "num_workers"})
+        unknown_keys = set(self.kv_connector_extra_config.keys()) - allowed_extra_config_keys
+        if unknown_keys:
+            raise ValueError(
+                f"Unsupported kv_connector_extra_config keys: {sorted(unknown_keys)}. "
+                f"Supported keys are {sorted(allowed_extra_config_keys)}"
+            )
 
         mooncake_protocol = self.kv_connector_extra_config["mooncake_protocol"]
         if mooncake_protocol not in ["tcp", "rdma"]:

python/infinilm/kv_connector/mooncake/mooncake_connector_worker.py

@@ -1,7 +1,7 @@
 try:
     from mooncake.engine import TransferEngine
 except ImportError as e:
-    raise ImportError("Please install mooncake") from e
+    raise ImportError("Please pip install mooncake-transfer-engine") from e
 
 import asyncio
 import logging
@@ -933,6 +933,7 @@ async def receive_kv_from_single_worker(
 
         except zmq.ContextTerminated:
             logger.debug("ZMQ context terminated, exiting Mooncake receiver thread.")
+            # TODO: handle this error
         except Exception as e:
             logger.error("MooncakeXferMetadata transfer failed for %s: %s", req_ids, e)
             return

python/infinilm/llm/cache_manager.py

@@ -62,9 +62,9 @@ def compute_hash(
         return h.intdigest()
 
     def __init__(self, num_blocks: int, block_size: int):
-        assert (
-            num_blocks > 0 and block_size > 0
-        ), "num_blocks and block_size must be positive"
+        assert num_blocks > 0 and block_size > 0, (
+            "num_blocks and block_size must be positive"
+        )
         self.num_blocks = num_blocks
         self.block_size = block_size
 
@@ -105,9 +105,9 @@ def _allocate_full_block(self) -> Block:
     def _deallocate_block(self, block_id: int):
         """Deallocate a block and return it to free list."""
         block = self.blocks[block_id]
-        assert (
-            block.ref_count == 0
-        ), f"Block {block_id} ref_count not zero, cannot deallocate"
+        assert block.ref_count == 0, (
+            f"Block {block_id} ref_count not zero, cannot deallocate"
+        )
 
         if block.hash != -1 and self.hash_to_block_id.get(block.hash) == block_id:
             del self.hash_to_block_id[block.hash]
@@ -396,6 +396,7 @@ def free_blocks(self, block_table: List[int]):
         immediately freed to allow reuse."""
         for block_id in reversed(block_table):
             block = self.blocks[block_id]
+            assert block.ref_count > 0, "block ref_count must be greater than 0"
             block.ref_count -= 1
 
     def try_free_blocks(self, num_required: int) -> bool:
@@ -425,9 +426,9 @@ def update_blocks_hash(self, block_table: List[int], num_local_cached_tokens: in
             num_local_cached_tokens: Number of locally cached tokens (must be a multiple of
                 block_size).
         """
-        assert (
-            num_local_cached_tokens % self.block_size == 0
-        ), "num_local_cached_tokens must be multiple of block_size"
+        assert num_local_cached_tokens % self.block_size == 0, (
+            "num_local_cached_tokens must be multiple of block_size"
+        )
         for idx in range(num_local_cached_tokens // self.block_size, len(block_table)):
             block_id = block_table[idx]
             block = self.blocks[block_id]

python/infinilm/llm/model_runner/model_runner.py

@@ -45,7 +45,7 @@ class ModelRunner:
     def __init__(self, config: EngineConfig):
         self.config = config
         self.kv_transfer_config = config.kv_transfer_config
-        print(f"kv_transfer_config: {self.kv_transfer_config}")
+        logger.info(f"kv_transfer_config: {self.kv_transfer_config}")
 
         self._init_device()
 

python/infinilm/llm/sampling_params.py

@@ -13,7 +13,7 @@ class SamplingParams:
     temperature: float = 1.0
     top_p: float = 0.8
     top_k: int = 1
-    max_tokens: Optional[int] = None
+    max_tokens: int = 512
     stop: Optional[List[str]] = None
     stop_token_ids: Optional[List[int]] = (
         None  # Placeholder for future usage, not currently handled