Fix bench_test

AlibabaPAI · Oct 11, 2024 · 83911f1 · 83911f1
1 parent 73bea56
commit 83911f1
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 3 deletions.
diff --git a/Makefile b/Makefile
@@ -45,6 +45,10 @@ offline_test:
 e2e_test:
 	@pytest -v tests/e2e_test/test_e2e.py
 
+.PHONY: bench_test
+bench_test:
+	@pytest -v ./tests/e2e_test/test_bench.py
+
 #################### pygloo install for gloo migration backend begin ####################
 
 BAZEL_CMD = bazel

diff --git a/tests/e2e_test/test_bench.py b/tests/e2e_test/test_bench.py
@@ -99,8 +99,8 @@ def get_markdown_data(key: str, head_name: str):
 
         return data
 
-    decode_data = get_markdown_data('decode_latencies', 'decode')
-    prefill_data = get_markdown_data('prefill_latencies', 'prefill')
+    decode_data = get_markdown_data('decode_token_latencies', 'decode')
+    prefill_data = get_markdown_data('prefill_token_latencies', 'prefill')
 
     return to_markdown_table(prefill_data) + "\n\n" + to_markdown_table(decode_data)
 

diff --git a/tools/bench_test.sh b/tools/bench_test.sh
@@ -3,4 +3,4 @@ set -ex
 
 nvidia-docker run --rm -t --net host --ipc host -v ${PWD}:/workspace -v /mnt:/mnt -w /workspace \
   registry.cn-beijing.aliyuncs.com/llumnix/llumnix-dev:20240909_action_678a439 \
-  bash -c "pip install -e . > /dev/null && pytest -v ./tests/e2e_test/test_bench.py"
+  bash -c "pip install -e . > /dev/null && make bench_test"