[Build] add TE bench into wheel package (#514)

.gitignore

@@ -193,3 +193,4 @@ cmake-build
 libetcd_wrapper.h
 
 mooncake-wheel/mooncake/mooncake_master
+mooncake-wheel/mooncake/transfer_engine_bench

mooncake-wheel/mooncake/cli_bench.py

@@ -0,0 +1,28 @@
+#!/usr/bin/env python3
+"""
+Minimal CLI module for transfer_engine_bench.
+"""
+
+import os
+import sys
+import subprocess
+
+
+def main():
+    """
+    Main entry point for the transfer_engine_bench command.
+    Simply runs the transfer_engine_bench binary with all arguments passed through.
+    """
+    # Get the path to the transfer_engine_bench binary
+    package_dir = os.path.dirname(os.path.abspath(__file__))
+    bin_path = os.path.join(package_dir, "transfer_engine_bench")
+    
+    # Make sure the binary is executable
+    os.chmod(bin_path, 0o755)
+    
+    # Run the binary with all arguments passed through
+    return subprocess.call([bin_path] + sys.argv[1:])
+
+
+if __name__ == "__main__":
+    sys.exit(main())

mooncake-wheel/pyproject.toml

@@ -27,6 +27,7 @@ Issues = "https://github.com/kvcache-ai/Mooncake/issues"
 
 [project.scripts]
 mooncake_master = "mooncake.cli:main"
+transfer_engine_bench = "mooncake.cli_bench:main"
 mooncake_http_metadata_server = "mooncake.http_metadata_server:main"
 mc_store_rest_server = "mooncake.mooncake_store_service:main"
 
@@ -39,4 +40,4 @@ where = ["."]
 include = ["mooncake*"]
 
 [tool.setuptools.package-data]
-mooncake = ["*.so", "mooncake_master"]
+mooncake = ["*.so", "mooncake_master", "transfer_engine_bench"]

mooncake-wheel/tests/test_cli.py

@@ -24,6 +24,17 @@ def test_entry_point_installed():
             return False
             
         print(f"✅ mooncake_master entry point found at: {result.stdout.strip()}")
+        result = subprocess.run(
+            ["which", "transfer_engine_bench"], 
+            capture_output=True, 
+            text=True
+        )
+        
+        if result.returncode != 0:
+            print("❌ transfer_engine_bench entry point not found in PATH")
+            return False
+            
+        print(f"✅ transfer_engine_bench entry point found at: {result.stdout.strip()}")
         return True
     except Exception as e:
         print(f"❌ Error checking for entry point: {e}")

scripts/build_wheel.sh

@@ -38,6 +38,7 @@ fi
 echo "Copying master binary and shared libraries..."
 # Copy master binary and shared libraries
 cp build/mooncake-store/src/mooncake_master mooncake-wheel/mooncake/
+cp build/mooncake-transfer-engine/example/transfer_engine_bench mooncake-wheel/mooncake/
 
 
 echo "Building wheel package..."

scripts/test_installation.sh

@@ -47,6 +47,11 @@ echo "Verifying mooncake_master entry point..."
 which mooncake_master || { echo "ERROR: mooncake_master entry point not found!"; exit 1; }
 echo "Success: mooncake_master entry point found"
 
+echo "Verifying transfer_engine_bench entry point..."
+# Check if the transfer_engine_bench entry point is installed and executable
+which transfer_engine_bench || { echo "ERROR: transfer_engine_bench entry point not found!"; exit 1; }
+echo "Success: transfer_engine_bench entry point found"
+
 cd ..
 
 echo "Installation test completed successfully!"