From e6bbc3dae85e9ea3057e10a4fc59209d9bbb8ecb Mon Sep 17 00:00:00 2001
From: Teng Ma <teng-ma@linux.alibaba.com>
Date: Wed, 18 Jun 2025 19:04:09 +0800
Subject: [PATCH] [Build] add TE bench into wheel package (#514)

---
 .gitignore                           |  1 +
 mooncake-wheel/mooncake/cli_bench.py | 28 ++++++++++++++++++++++++++++
 mooncake-wheel/pyproject.toml        |  3 ++-
 mooncake-wheel/tests/test_cli.py     | 11 +++++++++++
 scripts/build_wheel.sh               |  1 +
 scripts/test_installation.sh         |  5 +++++
 6 files changed, 48 insertions(+), 1 deletion(-)
 create mode 100644 mooncake-wheel/mooncake/cli_bench.py

diff --git a/.gitignore b/.gitignore
index 9aba795e..4c206afa 100644
--- a/.gitignore
+++ b/.gitignore
@@ -193,3 +193,4 @@ cmake-build
 libetcd_wrapper.h
 
 mooncake-wheel/mooncake/mooncake_master
+mooncake-wheel/mooncake/transfer_engine_bench
diff --git a/mooncake-wheel/mooncake/cli_bench.py b/mooncake-wheel/mooncake/cli_bench.py
new file mode 100644
index 00000000..2eb5ea98
--- /dev/null
+++ b/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())
diff --git a/mooncake-wheel/pyproject.toml b/mooncake-wheel/pyproject.toml
index df1419ec..7fabe224 100644
--- a/mooncake-wheel/pyproject.toml
+++ b/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"]
diff --git a/mooncake-wheel/tests/test_cli.py b/mooncake-wheel/tests/test_cli.py
index 268cfda7..27e22ddb 100644
--- a/mooncake-wheel/tests/test_cli.py
+++ b/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}")
diff --git a/scripts/build_wheel.sh b/scripts/build_wheel.sh
index 2f875d68..0711d1cc 100755
--- a/scripts/build_wheel.sh
+++ b/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..."
diff --git a/scripts/test_installation.sh b/scripts/test_installation.sh
index 21a95304..6bb0eaf7 100755
--- a/scripts/test_installation.sh
+++ b/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!"