init v0.23.0

Signed-off-by: Sun Ruoxi <sunruoxi@4paradigm.com>

tests/ut/distributed/mooncake/test_mooncake_config_data.py

@@ -0,0 +1,74 @@
+import sys
+import types
+import unittest
+from unittest.mock import MagicMock
+
+fake_engine = types.ModuleType("mooncake.engine")
+fake_engine.TransferEngine = MagicMock()  # type: ignore[attr-defined]
+sys.modules["mooncake.engine"] = fake_engine
+fake_store = types.ModuleType("mooncake.store")
+fake_store.ReplicateConfig = MagicMock()  # type: ignore[attr-defined]
+sys.modules["mooncake.store"] = fake_store
+
+from vllm_ascend.distributed.kv_transfer.kv_pool.ascend_store.backend.mooncake_backend import (  # noqa: E402
+    _convert_to_bytes,
+    _parse_global_segment_size,
+)
+
+
+class TestParseGlobalSegmentSize(unittest.TestCase):
+    def test_int_input(self):
+        self.assertEqual(_parse_global_segment_size(1024), 1024)
+        self.assertEqual(_parse_global_segment_size(0), 0)
+
+    def test_gb_unit(self):
+        self.assertEqual(_parse_global_segment_size("2GB"), 2 * 1024**3)
+        self.assertEqual(_parse_global_segment_size("1.5GB"), int(1.5 * 1024**3))
+        self.assertEqual(_parse_global_segment_size(" 2 GB "), 2 * 1024**3)
+
+    def test_gb_unit_edge_cases(self):
+        with self.assertRaises(ValueError):
+            _parse_global_segment_size("GB")
+        with self.assertRaises(ValueError):
+            _parse_global_segment_size("abcGB")
+
+    def test_mb_unit(self):
+        self.assertEqual(_parse_global_segment_size("512MB"), 512 * 1024**2)
+        self.assertEqual(_parse_global_segment_size("0.5MB"), int(0.5 * 1024**2))
+        self.assertEqual(_parse_global_segment_size("1024MB"), 1024 * 1024**2)
+
+    def test_kb_unit(self):
+        self.assertEqual(_parse_global_segment_size("256KB"), 256 * 1024)
+        self.assertEqual(_parse_global_segment_size("1.25KB"), int(1.25 * 1024))
+
+    def test_b_unit(self):
+        self.assertEqual(_parse_global_segment_size("4096B"), 4096)
+        self.assertEqual(_parse_global_segment_size("1024b"), 1024)
+
+    def test_no_unit(self):
+        self.assertEqual(_parse_global_segment_size("2048"), 2048)
+        self.assertEqual(_parse_global_segment_size("0"), 0)
+
+    def test_non_string_non_int_input(self):
+        self.assertEqual(_parse_global_segment_size(2048.0), 2048)
+        self.assertEqual(_parse_global_segment_size(True), 1)
+
+        with self.assertRaises(TypeError):
+            _parse_global_segment_size(None)
+
+        with self.assertRaises(TypeError):
+            _parse_global_segment_size({"size": 1024})
+
+
+class TestConvertToBytes(unittest.TestCase):
+    def test_valid_conversion(self):
+        self.assertEqual(_convert_to_bytes("10", 1, "10"), 10)
+        self.assertEqual(_convert_to_bytes("1.5", 1024, "1.5KB"), int(1.5 * 1024))
+        self.assertEqual(_convert_to_bytes("0", 1024**3, "0GB"), 0)
+
+    def test_invalid_numbers(self):
+        with self.assertRaises(ValueError):
+            _convert_to_bytes("abc", 1, "abc")
+
+        with self.assertRaises(ValueError):
+            _convert_to_bytes("1.2.3", 1024, "1.2.3KB")

tests/ut/distributed/mooncake/test_mooncake_kv_transfer.py

@@ -0,0 +1,73 @@
+import threading
+import unittest
+from types import SimpleNamespace
+
+import torch
+
+if not hasattr(torch, "npu"):
+    torch.npu = SimpleNamespace(Event=object)  # type: ignore[attr-defined]
+
+from vllm_ascend.distributed.kv_transfer.kv_pool.ascend_store.config_data import (
+    ChunkedTokenDatabase,
+    KeyMetadata,
+    ReqMeta,
+)
+from vllm_ascend.distributed.kv_transfer.kv_pool.ascend_store.kv_transfer import (
+    KVCacheStoreSendingThread,
+)
+
+
+class _FakeStore:
+    def __init__(self, exists_result: list[int]):
+        self.exists_result = exists_result
+        self.put_calls: list[tuple[list[str], list[list[int]], list[list[int]]]] = []
+
+    def set_device(self):
+        return None
+
+    def exists(self, keys: list[str]) -> list[int]:
+        # Return exact number of states for requested keys.
+        return self.exists_result[: len(keys)]
+
+    def put(self, keys, addrs, sizes):
+        self.put_calls.append((list(keys), list(addrs), list(sizes)))
+
+
+class TestKVTransferMissingKeyPut(unittest.TestCase):
+    def test_sending_thread_only_puts_missing_keys(self):
+        store = _FakeStore(exists_result=[1, 0, 1, 0])
+        token_db = ChunkedTokenDatabase([KeyMetadata("m", 0, 0, 0, 0)], [16], None)
+        token_db.set_group_buffers({0: [1000]}, {0: [16]}, {0: [1]})
+        thread = KVCacheStoreSendingThread(
+            m_store=store,
+            token_database=token_db,
+            block_size=16,
+            tp_rank=0,
+            dcp_size=1,
+            put_step=1,
+            kv_role="kv_producer",
+            ready_event=threading.Event(),
+            group_uses_align_state=[False],
+            enable_kv_event=False,
+        )
+
+        req_meta = ReqMeta(
+            req_id="req-1",
+            token_len_chunk=64,
+            block_ids=[0, 1, 2, 3],
+            block_hashes=[b"h0", b"h1", b"h2", b"h3"],  # type: ignore[arg-type]
+            current_event=None,
+        )
+        thread.add_stored_request("req-1")
+        thread.request_queue.put(req_meta)
+        thread._handle_request(req_meta)
+
+        self.assertEqual(len(store.put_calls), 1)
+        put_keys, put_addrs, put_sizes = store.put_calls[0]
+        self.assertEqual(len(put_keys), 2)
+        self.assertEqual(put_addrs, [[1001], [1003]])
+        self.assertEqual(put_sizes, [[16], [16]])
+
+
+if __name__ == "__main__":
+    unittest.main()