diff --git a/test/manual/kv_transfer/test_mooncake_transfer_engine.py b/test/manual/kv_transfer/test_mooncake_transfer_engine.py new file mode 100755 index 000000000..f0ed315e8 --- /dev/null +++ b/test/manual/kv_transfer/test_mooncake_transfer_engine.py @@ -0,0 +1,64 @@ +#!/usr/bin/env python3 +import argparse + +try: + import mooncake + + BENCH_TOOL_PATH = f"{mooncake.__path__[0]}/transfer_engine_bench" + print(f"Mooncake is installed. Bench tool path:\n{BENCH_TOOL_PATH}") +except ImportError: + BENCH_TOOL_PATH = None + print("Mooncake is not installed.") + exit(0) + + +def run_cmd(args): + cmd = [BENCH_TOOL_PATH] + if args.initiator: + cmd += ["--mode=initiator"] + elif args.target: + cmd += ["--mode=target"] + + if args.metadata_server: + cmd += [f"--metadata_server={args.metadata_server}"] + if args.mc_segment_id: + cmd += [f"--segment_id={args.mc_segment_id}"] + if args.device: + cmd += [f"--device_name={args.device}"] + + if args.bench_h2h: + cmd += ["--use_vram=false"] + + cmd += ["--auto_discovery"] + print(f"Executing command: {' '.join(cmd)}") + import subprocess + + try: + subprocess.run(cmd, check=True) + except subprocess.CalledProcessError as e: + print(f"Command failed with error: {e}") + exit(1) + + +def main(): + parser = argparse.ArgumentParser() + group = parser.add_mutually_exclusive_group() + group.add_argument("--initiator", action="store_true", help="Run as initiator") + group.add_argument("--target", action="store_true", help="Run as target") + parser.add_argument("--metadata-server", type=str, default="P2PHANDSHAKE") + parser.add_argument("--mc-segment-id", type=str, default=None) + parser.add_argument("--bench-h2h", action="store_true") + parser.add_argument("--device", type=str, default="mlx5_0") + args = parser.parse_args() + + print("Running Mooncake transfer engine benchmark...") + if not args.initiator and not args.target: + parser.error("Please specify --initiator or --target") + if args.initiator and args.mc_segment_id is None: + parser.error("Please specify --mc-segment-id for initiator") + + run_cmd(args) + + +if __name__ == "__main__": + main()