add test & fix bugs

Author: stmatengss

tests/kv_transfer/test_mooncake_kv.py

@@ -0,0 +1,33 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import os
+import random
+
+import torch
+from tqdm import tqdm
+
+from vllm.config import KVTransferConfig
+from vllm.config import VllmConfig
+from vllm.distributed.kv_transfer.kv_lookup_buffer.base import MooncakeStore
+
+
+if __name__ == "__main__":
+
+    # my_rank = int(os.environ['RANK'])
+
+    # print(f"initialized! My rank is {my_rank}")
+
+    # config = KVTransferConfig(
+    #     kv_connector='PyNcclConnector',
+    #     kv_buffer_device='cuda',
+    #     kv_buffer_size=1e9,
+    #     kv_rank=my_rank,
+    #     kv_role="kv_both",  # this arg doesn't matter in this test
+    #     kv_parallel_size=2,
+    #     kv_ip="127.0.0.1",
+    #     kv_port=12345,
+    # )
+    vllm_config = VllmConfig()
+    mooncake = MooncakeStore("localhost", 0, vllm_config)
+
+    print('Done')

vllm/distributed/kv_transfer/kv_lookup_buffer/base.py

@@ -14,6 +14,9 @@
 from vllm.distributed.kv_transfer.kv_pipe.mooncake_pipe import MooncakeTransferEngineConfig
 
 import torch
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
 
 
 class KVLookupBufferBase(ABC):
@@ -157,6 +160,9 @@ def drop_select(
             roi: Optional[torch.Tensor]) -> List[Optional[torch.Tensor]]:  
         # V1 (get and unpack all tensors): consume tensors from MooncakeStore's buffer
         raise NotImplementedError("Insert method is not implemented")
+    
+    def close(self):
+        pass
 
     def put(
         self,