Skip to content
69 changes: 29 additions & 40 deletions __perf_gate__.sh
Original file line number Diff line number Diff line change
@@ -1,57 +1,46 @@
#!/usr/bin/env bash
set -uo pipefail

script_dir=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)
bench_py="$script_dir/examples/qwen3/bench_allreduce.py"
if [[ ! -f "$bench_py" ]]; then
echo "ERROR: missing benchmark: $bench_py" >&2
exit 1
fi

: "${ARK_ROOT:=$PWD}"
export ARK_ROOT
export PYTHONPATH="${PYTHONPATH:-$ARK_ROOT/python}"

target_ms=$(python3 - <<'PY'
import importlib.util
import pathlib

path = pathlib.Path("../examples/qwen3/bench_allreduce.py")
spec = importlib.util.spec_from_file_location("bench_allreduce", path)
module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(module)
print(f"{module._DECODE_TARGET_MS:.4f}")
PY
)

tmpdir=$(mktemp -d)
trap 'rm -rf "$tmpdir"' EXIT
status=0
python3 ../examples/qwen3/bench_allreduce.py --world-size 2 --shape decode >"$tmpdir/tp2.log" 2>"$tmpdir/tp2.err" || status=1
python3 ../examples/qwen3/bench_allreduce.py --world-size 8 --shape decode >"$tmpdir/tp8.log" 2>"$tmpdir/tp8.err" || status=1
python3 "$bench_py" --world-size 2 --shape all --input-mode all \
>"$tmpdir/tp2.log" 2>"$tmpdir/tp2.err" || status=1
python3 "$bench_py" --world-size 8 --shape all --input-mode all \
>"$tmpdir/tp8.log" 2>"$tmpdir/tp8.err" || status=1

ark_ms=$(python3 - "$tmpdir/tp2.log" "$tmpdir/tp8.log" "$status" <<'PY'
import re
python3 - "$status" "$bench_py" "$tmpdir/tp2.log" "$tmpdir/tp8.log" <<'PY'
import importlib.util
import pathlib
import sys

values = []
for name in sys.argv[1:3]:
text = open(name, encoding="utf-8").read()
match = re.search(r"PERF_GATE name=allreduce\s+ark_ms=([0-9.]+)", text)
if match:
values.append(float(match.group(1)))
if int(sys.argv[3]) or len(values) != 2:
print("999999.0000")
else:
print(f"{max(values):.4f}")
PY
status = int(sys.argv[1])
path = pathlib.Path(sys.argv[2])
logs = sys.argv[3:]
spec = importlib.util.spec_from_file_location("bench_allreduce", path)
module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(module)
target_ms = module._DECODE_TARGET_MS
texts = [pathlib.Path(log).read_text(encoding="utf-8") for log in logs]
ark_ms = (
999999.0 if status else module._decode_gate_ark_ms_from_logs(texts)
)
ratio=$(python3 - "$ark_ms" "$target_ms" <<'PY'
import sys

print(f"{float(sys.argv[1]) / float(sys.argv[2]):.4f}")
PY
ratio = ark_ms / target_ms
print(
f"PERF_GATE name={module._DECODE_GATE_NAME} ark_ms={ark_ms:.4f} "
f"sglang_ms={target_ms:.4f} ratio={ratio:.4f}"
)
printf 'PERF_GATE name=allreduce ark_ms=%s sglang_ms=%s ratio=%s\n' "$ark_ms" "$target_ms" "$ratio"
python3 - "$ark_ms" "$target_ms" "$status" <<'PY'
import sys

ark_ms = float(sys.argv[1])
target_ms = float(sys.argv[2])
status = int(sys.argv[3])
if status or ark_ms >= target_ms:
if status or ark_ms >= 999999.0 or ark_ms >= target_ms:
raise SystemExit(1)
PY
Loading