mooncake-transfer-engine 0.1.0__tar.gz

mooncake_transfer_engine-0.1.0/PKG-INFO

@@ -0,0 +1,20 @@
+Metadata-Version: 2.4
+Name: mooncake-transfer-engine
+Version: 0.1.0
+Summary: Python binding of a Mooncake library using pybind11
+Home-page: https://github.com/kvcache-ai/Mooncake
+Author: Mooncake Authors
+Project-URL: Documentation, https://github.com/kvcache-ai/Mooncake/tree/main/doc
+Project-URL: Source, https://github.com/kvcache-ai/Mooncake
+Project-URL: Issues, https://github.com/kvcache-ai/Mooncake/issues
+Keywords: mooncake,data transfer,kv cache,llm inference
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: C++
+Classifier: Operating System :: POSIX :: Linux
+Classifier: License :: OSI Approved :: Apache Software License
+Dynamic: author
+Dynamic: classifier
+Dynamic: home-page
+Dynamic: keywords
+Dynamic: project-url
+Dynamic: summary

mooncake_transfer_engine-0.1.0/mooncake/__init__.py

@@ -0,0 +1,3 @@
+from .mooncake_vllm_adaptor import MooncakeDistributedStore
+from .mooncake_vllm_adaptor import mooncake_vllm_adaptor
+

mooncake_transfer_engine-0.1.0/mooncake_transfer_engine.egg-info/PKG-INFO

@@ -0,0 +1,20 @@
+Metadata-Version: 2.4
+Name: mooncake-transfer-engine
+Version: 0.1.0
+Summary: Python binding of a Mooncake library using pybind11
+Home-page: https://github.com/kvcache-ai/Mooncake
+Author: Mooncake Authors
+Project-URL: Documentation, https://github.com/kvcache-ai/Mooncake/tree/main/doc
+Project-URL: Source, https://github.com/kvcache-ai/Mooncake
+Project-URL: Issues, https://github.com/kvcache-ai/Mooncake/issues
+Keywords: mooncake,data transfer,kv cache,llm inference
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: C++
+Classifier: Operating System :: POSIX :: Linux
+Classifier: License :: OSI Approved :: Apache Software License
+Dynamic: author
+Dynamic: classifier
+Dynamic: home-page
+Dynamic: keywords
+Dynamic: project-url
+Dynamic: summary

mooncake_transfer_engine-0.1.0/mooncake_transfer_engine.egg-info/SOURCES.txt

@@ -0,0 +1,13 @@
+pyproject.toml
+setup.py
+mooncake/__init__.py
+mooncake/mooncake_master
+mooncake/mooncake_sglang_adaptor.cpython-310-x86_64-linux-gnu.so
+mooncake/mooncake_vllm_adaptor.cpython-310-x86_64-linux-gnu.so
+mooncake/lib_so/libetcd-cpp-api.so
+mooncake_transfer_engine.egg-info/PKG-INFO
+mooncake_transfer_engine.egg-info/SOURCES.txt
+mooncake_transfer_engine.egg-info/dependency_links.txt
+mooncake_transfer_engine.egg-info/not-zip-safe
+mooncake_transfer_engine.egg-info/top_level.txt
+tests/test_distributed_object_store.py

mooncake_transfer_engine-0.1.0/mooncake_transfer_engine.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

mooncake_transfer_engine-0.1.0/mooncake_transfer_engine.egg-info/not-zip-safe

@@ -0,0 +1 @@
+

mooncake_transfer_engine-0.1.0/mooncake_transfer_engine.egg-info/top_level.txt

@@ -0,0 +1 @@
+mooncake

mooncake_transfer_engine-0.1.0/pyproject.toml

@@ -0,0 +1,3 @@
+[build-system]
+requires = ["setuptools", "wheel"]
+build-backend = "setuptools.build_meta"

mooncake_transfer_engine-0.1.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

mooncake_transfer_engine-0.1.0/setup.py

@@ -0,0 +1,52 @@
+import sys
+import os
+from setuptools import setup, find_packages
+from setuptools.dist import Distribution
+from wheel.bdist_wheel import bdist_wheel
+
+class BinaryDistribution(Distribution):
+    def has_ext_modules(self):
+        return True
+
+class CustomBdistWheel(bdist_wheel):
+    def finalize_options(self):
+        bdist_wheel.finalize_options(self)
+        self.root_is_pure = False
+        self.plat_name_supplied = True
+        self.plat_name = "manylinux2014_x86_64"
+
+python_version = f">={sys.version_info.major}.{sys.version_info.minor}"
+
+VERSION = os.environ.get("VERSION", "0.1.0")
+
+setup(
+    name="mooncake-transfer-engine",
+    version=VERSION,
+    packages=find_packages(),
+    package_data={"mooncake": [
+        "*.so",
+        "mooncake_master",
+        "lib_so/libetcd-cpp-api.so",
+    ]},
+    include_package_data=True,
+    zip_safe=False,
+    distclass=BinaryDistribution,
+    cmdclass={
+        'bdist_wheel': CustomBdistWheel,
+    },
+    author="Mooncake Authors",
+    description="Python binding of a Mooncake library using pybind11",
+    url="https://github.com/kvcache-ai/Mooncake",
+    project_urls={
+        "Documentation": "https://github.com/kvcache-ai/Mooncake/tree/main/doc",
+        "Source": "https://github.com/kvcache-ai/Mooncake",
+        "Issues": "https://github.com/kvcache-ai/Mooncake/issues",
+    },
+    keywords=["mooncake", "data transfer", "kv cache", "llm inference"],
+    classifiers=[
+        "Programming Language :: Python :: 3",
+        "Programming Language :: C++",
+        "Operating System :: POSIX :: Linux",
+        "License :: OSI Approved :: Apache Software License",
+    ]
+)

mooncake_transfer_engine-0.1.0/tests/test_distributed_object_store.py

@@ -0,0 +1,271 @@
+import unittest
+import os
+import time
+import threading
+import random
+from mooncake import MooncakeDistributedStore
+
+
+def get_client(store):
+    """Initialize and setup the distributed store client."""
+    protocol = os.getenv("PROTOCOL", "tcp")
+    device_name = os.getenv("DEVICE_NAME", "ibp6s0")
+    local_hostname = os.getenv("LOCAL_HOSTNAME", "localhost")
+    metadata_server = os.getenv("MC_METADATA_SERVER", "127.0.0.1:2379")
+    global_segment_size = 3200 * 1024 * 1024  # 3200 MB
+    local_buffer_size = 512 * 1024 * 1024     # 512 MB
+    master_server_address = os.getenv("MASTER_SERVER", "127.0.0.1:50051")
+    
+    retcode = store.setup(
+        local_hostname, 
+        metadata_server, 
+        global_segment_size,
+        local_buffer_size, 
+        protocol, 
+        device_name,
+        master_server_address
+    )
+    
+    if retcode:
+        raise RuntimeError(f"Failed to setup store client. Return code: {retcode}")
+
+
+class TestDistributedObjectStore(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        """Initialize the store once for all tests."""
+        cls.store = MooncakeDistributedStore()
+        get_client(cls.store)
+    
+    def test_client_tear_down(self):
+        """Test client tear down and re-initialization."""
+        test_data = b"Hello, World!"
+        key = "test_teardown_key"
+        
+        # Put data and verify teardown clears it
+        self.assertEqual(self.store.put(key, test_data), 0)
+        self.assertEqual(self.store.close(), 0)
+        time.sleep(1)  # Allow time for teardown to complete
+        
+        # Re-initialize the store
+        get_client(self.store)
+        
+        # Verify data is gone after teardown
+        retrieved_data = self.store.get(key)
+        self.assertEqual(retrieved_data, b"")
+        
+        # Verify store is functional after re-initialization
+        self.assertEqual(self.store.put(key, test_data), 0)
+        retrieved_data = self.store.get(key)
+        self.assertEqual(retrieved_data, test_data)
+
+    def test_basic_put_get_exist_operations(self):
+        """Test basic Put/Get/Exist operations through the Python interface."""
+        test_data = b"Hello, World!"
+        key = "test_basic_key"
+
+        # Test Put operation
+        self.assertEqual(self.store.put(key, test_data), 0)
+
+        # Verify data through Get operation
+        self.assertEqual(self.store.getSize(key), len(test_data))
+        retrieved_data = self.store.get(key)
+        self.assertEqual(retrieved_data, test_data)
+
+        # Put again with the same key, should succeed
+        self.assertEqual(self.store.put(key, test_data), 0)
+
+        # Remove the key
+        self.assertEqual(self.store.remove(key), 0)
+
+        # Get after remove should return empty bytes
+        self.assertLess(self.store.getSize(key), 0)
+        empty_data = self.store.get(key)
+        self.assertEqual(empty_data, b"")
+
+        # Test isExist functionality
+        test_data_2 = b"Testing exists!"
+        key_2 = "test_exist_key"
+        
+        # Should not exist initially
+        self.assertLess(self.store.getSize(key_2), 0)
+        self.assertEqual(self.store.isExist(key_2), 0)
+        
+        # Should exist after put
+        self.assertEqual(self.store.put(key_2, test_data_2), 0)
+        self.assertEqual(self.store.isExist(key_2), 1)
+        self.assertEqual(self.store.getSize(key_2), len(test_data_2))
+        
+        # Should not exist after remove
+        self.assertEqual(self.store.remove(key_2), 0)
+        self.assertLess(self.store.getSize(key_2), 0)
+        self.assertEqual(self.store.isExist(key_2), 0)
+
+    def test_concurrent_stress_with_barrier(self):
+        """Test concurrent Put/Get operations with multiple threads using barrier."""
+        NUM_THREADS = 8
+        VALUE_SIZE = 1024 * 1024  # 1MB
+        OPERATIONS_PER_THREAD = 100
+        
+        # Create barriers for synchronization
+        start_barrier = threading.Barrier(NUM_THREADS + 1)  # +1 for main thread
+        put_barrier = threading.Barrier(NUM_THREADS + 1)    # Barrier after put operations
+        get_barrier = threading.Barrier(NUM_THREADS + 1)    # Barrier after get operations
+        
+        # Statistics for system-wide timing
+        system_stats = {
+            'put_start': 0,
+            'put_end': 0,
+            'get_start': 0,
+            'get_end': 0
+        }
+        thread_exceptions = []
+        
+        def worker(thread_id):
+            try:
+                # Generate test data (1MB)
+                test_data = os.urandom(VALUE_SIZE)
+                thread_keys = [f"key_{thread_id}_{i}" for i in range(OPERATIONS_PER_THREAD)]
+                
+                # Wait for all threads to be ready
+                start_barrier.wait()
+                
+                # Put operations
+                for key in thread_keys:
+                    result = self.store.put(key, test_data)
+                    self.assertEqual(result, 0, f"Put operation failed for key {key}")
+                
+                # Wait for all threads to complete put operations
+                put_barrier.wait()
+                
+                # Get operations
+                for key in thread_keys:
+                    retrieved_data = self.store.get(key)
+                    self.assertEqual(len(retrieved_data), VALUE_SIZE, 
+                                    f"Retrieved data size mismatch for key {key}")
+                    self.assertEqual(retrieved_data, test_data, 
+                                    f"Retrieved data content mismatch for key {key}")
+                
+                # Wait for all threads to complete get operations
+                get_barrier.wait()
+                
+                # Remove all keys
+                for key in thread_keys:
+                    self.assertEqual(self.store.remove(key), 0)
+                
+                
+            except Exception as e:
+                thread_exceptions.append(f"Thread {thread_id} failed: {str(e)}")
+        
+        # Create and start threads
+        threads = []
+        for i in range(NUM_THREADS):
+            t = threading.Thread(target=worker, args=(i,), name=f"Worker-{i}")
+            threads.append(t)
+            t.start()
+        
+        # Wait for all threads to be ready and start the test
+        start_barrier.wait()
+        
+        # Record put start time
+        system_stats['put_start'] = time.time()
+        
+        # Wait for all put operations to complete
+        put_barrier.wait()
+        system_stats['put_end'] = time.time()
+        
+        # Record get start time
+        system_stats['get_start'] = time.time()
+        
+        # Wait for all get operations to complete
+        get_barrier.wait()
+        system_stats['get_end'] = time.time()
+        
+        
+        # Join all threads
+        for t in threads:
+            t.join()
+        
+        # Check for any exceptions
+        self.assertEqual(len(thread_exceptions), 0, "\n".join(thread_exceptions))
+        
+        # Calculate system-wide statistics
+        total_operations = NUM_THREADS * OPERATIONS_PER_THREAD
+        put_duration = system_stats['put_end'] - system_stats['put_start']
+        get_duration = system_stats['get_end'] - system_stats['get_start']
+        total_data_size_gb = (VALUE_SIZE * total_operations) / (1024**3)
+        
+        print(f"\nConcurrent Stress Test Results:")
+        print(f"Total threads: {NUM_THREADS}")
+        print(f"Operations per thread: {OPERATIONS_PER_THREAD}")
+        print(f"Total operations: {total_operations}")
+        print(f"Data block size: {VALUE_SIZE/1024/1024:.2f}MB")
+        print(f"Total data processed: {total_data_size_gb:.2f}GB")
+        print(f"Put duration: {put_duration:.2f} seconds")
+        print(f"Get duration: {get_duration:.2f} seconds")
+        print(f"System Put throughput: {total_operations/put_duration:.2f} ops/sec")
+        print(f"System Get throughput: {total_operations/get_duration:.2f} ops/sec")
+        print(f"System Put bandwidth: {total_data_size_gb/put_duration:.2f} GB/sec")
+        print(f"System Get bandwidth: {total_data_size_gb/get_duration:.2f} GB/sec")
+
+    def test_dict_fuzz_e2e(self):
+         """End-to-end fuzz test comparing distributed store behavior with dict.
+         Performs ~1000 random operations (put, get, remove) with random value sizes between 1KB and 64MB.
+         After testing, all keys are removed.
+         """
+         import random
+         # Local reference dict to simulate expected dict behavior
+         reference = {}
+         operations = 1000
+         # Use a pool of keys to limit memory consumption
+         keys_pool = [f"key_{i}" for i in range(100)]
+         # Track which keys have values assigned to ensure consistency
+         key_values = {}
+         # Fuzz record for debugging in case of errors
+         fuzz_record = []
+         try:
+             for i in range(operations):
+                 op = random.choice(["put", "get", "remove"])
+                 key = random.choice(keys_pool)
+                 if op == "put":
+                     # If key already exists, use the same value to ensure consistency
+                     if key in key_values:
+                         value = key_values[key]
+                         size = len(value)
+                     else:
+                         size = random.randint(1, 64 * 1024 * 1024)
+                         value = os.urandom(size)
+                         key_values[key] = value
+                     
+                     fuzz_record.append(f"{i}: put {key} [size: {size}]")
+                     error_code = self.store.put(key, value)
+                     if error_code == -200: 
+                         # The space is not enough, continue to next operation
+                         continue
+                     elif error_code == 0:
+                         reference[key] = value
+                     else:
+                         raise RuntimeError(f"Put operation failed for key {key}. Error code: {error_code}")
+                 elif op == "get":
+                     fuzz_record.append(f"{i}: get {key}")
+                     retrieved = self.store.get(key)
+                     expected = reference.get(key, b"")
+                     self.assertEqual(retrieved, expected)
+                 elif op == "remove":
+                     fuzz_record.append(f"{i}: remove {key}")
+                     self.store.remove(key)
+                     reference.pop(key, None)
+                     # Also remove from key_values to allow new value if key is reused
+                     key_values.pop(key, None)
+         except Exception as e:
+             print(f"Error: {e}")
+             print('\nFuzz record (operations so far):')
+             for record in fuzz_record:
+                 print(record)
+             raise e
+         # Cleanup: ensure all remaining keys are removed
+         for key in list(reference.keys()):
+             self.store.remove(key)
+             
+if __name__ == '__main__':
+    unittest.main()