mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-12 03:48:56 +00:00
feat(benchmark): add RPS suite (10 runs 1k–10k RPS, one at a time, separate files)
- scripts/run_benchmark_rps_suite.py: run benchmark at 1k, 2k, ..., 10k RPS sequentially - Each run writes to benchmark_results_<rps>rps.txt (e.g. benchmark_results_1000rps.txt) - First run waits for proxy; runs 2–10 use --no-wait. --output-dir, --requests supported - proxy_cli: LITELLM_AUTO_BENCHMARK_RPS_SUITE=1 runs suite when proxy is up - Optional: LITELLM_AUTO_BENCHMARK_OUTPUT_DIR, LITELLM_AUTO_BENCHMARK_REQUESTS
This commit is contained in:
Alexsander Hamir
parent
9da08ced8f
commit
69a65b4430
2 changed files with 144 additions and 12 deletions
|
|
@ -173,11 +173,16 @@ class ProxyInitializationHelpers:
|
|||
_proxy_cli_dir = os.path.dirname(os.path.abspath(__file__))
|
||||
_litellm_root = os.path.dirname(_proxy_cli_dir)
|
||||
_repo_root = os.path.dirname(_litellm_root)
|
||||
script_path = os.path.join(_repo_root, "scripts", "run_benchmark_after_proxy.py")
|
||||
rps_suite = os.getenv("LITELLM_AUTO_BENCHMARK_RPS_SUITE", "").strip().lower() in ("1", "true", "yes")
|
||||
if rps_suite:
|
||||
script_path = os.path.join(_repo_root, "scripts", "run_benchmark_rps_suite.py")
|
||||
else:
|
||||
script_path = os.path.join(_repo_root, "scripts", "run_benchmark_after_proxy.py")
|
||||
if not os.path.isfile(script_path):
|
||||
return
|
||||
|
||||
output = os.getenv("LITELLM_AUTO_BENCHMARK_OUTPUT", "").strip()
|
||||
output_dir = os.getenv("LITELLM_AUTO_BENCHMARK_OUTPUT_DIR", "").strip()
|
||||
extra_args_str = os.getenv("LITELLM_AUTO_BENCHMARK_ARGS", "--rps-control 10000 --requests 500000").strip()
|
||||
extra_args = extra_args_str.split() if extra_args_str else ["--rps-control", "10000", "--requests", "500000"]
|
||||
|
||||
|
|
@ -207,24 +212,38 @@ class ProxyInitializationHelpers:
|
|||
else:
|
||||
print("LiteLLM: Auto-benchmark timed out waiting for proxy health", file=sys.stderr)
|
||||
return
|
||||
# Proxy is up: run benchmark and save to file
|
||||
# Proxy is up: run benchmark (single run or RPS suite)
|
||||
from datetime import datetime
|
||||
|
||||
if not output:
|
||||
out_path = os.path.join(_repo_root, f"benchmark_results_{datetime.now().strftime('%Y%m%d_%H%M%S')}.txt")
|
||||
else:
|
||||
out_path = output if os.path.isabs(output) else os.path.join(_repo_root, output)
|
||||
mock_base, _ = ProxyInitializationHelpers._get_mock_chat_endpoint()
|
||||
provider_url_default = f"{mock_base.rstrip('/')}/chat/completions"
|
||||
env = os.environ.copy()
|
||||
env["LITELLM_PROXY_URL"] = f"{proxy_base_url}/chat/completions"
|
||||
env.setdefault("PROVIDER_URL", provider_url_default)
|
||||
cmd = [sys.executable, script_path, "--no-wait", "--output", out_path, "--proxy-url", env["LITELLM_PROXY_URL"], "--provider-url", env["PROVIDER_URL"]] + extra_args
|
||||
try:
|
||||
print(f"LiteLLM: Auto-benchmark starting, output -> {out_path}", flush=True)
|
||||
subprocess.run(cmd, env=env, cwd=_repo_root)
|
||||
except Exception as e:
|
||||
print(f"LiteLLM: Auto-benchmark failed: {e}", file=sys.stderr)
|
||||
|
||||
if rps_suite:
|
||||
cmd = [sys.executable, script_path, "--no-wait", "--proxy-url", env["LITELLM_PROXY_URL"], "--provider-url", env["PROVIDER_URL"]]
|
||||
if output_dir:
|
||||
cmd.extend(["--output-dir", output_dir if os.path.isabs(output_dir) else os.path.join(_repo_root, output_dir)])
|
||||
requests_env = os.getenv("LITELLM_AUTO_BENCHMARK_REQUESTS", "").strip()
|
||||
if requests_env.isdigit():
|
||||
cmd.extend(["--requests", requests_env])
|
||||
try:
|
||||
print("LiteLLM: Auto-benchmark RPS suite starting (1k–10k RPS, 10 runs)", flush=True)
|
||||
subprocess.run(cmd, env=env, cwd=_repo_root)
|
||||
except Exception as e:
|
||||
print(f"LiteLLM: Auto-benchmark RPS suite failed: {e}", file=sys.stderr)
|
||||
else:
|
||||
if not output:
|
||||
out_path = os.path.join(_repo_root, f"benchmark_results_{datetime.now().strftime('%Y%m%d_%H%M%S')}.txt")
|
||||
else:
|
||||
out_path = output if os.path.isabs(output) else os.path.join(_repo_root, output)
|
||||
cmd = [sys.executable, script_path, "--no-wait", "--output", out_path, "--proxy-url", env["LITELLM_PROXY_URL"], "--provider-url", env["PROVIDER_URL"]] + extra_args
|
||||
try:
|
||||
print(f"LiteLLM: Auto-benchmark starting, output -> {out_path}", flush=True)
|
||||
subprocess.run(cmd, env=env, cwd=_repo_root)
|
||||
except Exception as e:
|
||||
print(f"LiteLLM: Auto-benchmark failed: {e}", file=sys.stderr)
|
||||
|
||||
t = threading.Thread(target=_run, daemon=True)
|
||||
t.start()
|
||||
|
|
|
|||
113
scripts/run_benchmark_rps_suite.py
Normal file
113
scripts/run_benchmark_rps_suite.py
Normal file
|
|
@ -0,0 +1,113 @@
|
|||
#!/usr/bin/env python3
|
||||
"""
|
||||
Run the proxy-vs-provider benchmark 10 times: 1k, 2k, ..., 10k RPS, one at a time.
|
||||
Each run writes to a separate results file (e.g. benchmark_results_1000rps.txt, ...).
|
||||
|
||||
Usage:
|
||||
# Wait for proxy, then run 10 tests (1k–10k RPS), 10k requests each:
|
||||
python scripts/run_benchmark_rps_suite.py
|
||||
|
||||
# Skip proxy wait (proxy already up):
|
||||
python scripts/run_benchmark_rps_suite.py --no-wait
|
||||
|
||||
# Custom output directory and request count:
|
||||
python scripts/run_benchmark_rps_suite.py --output-dir ./results --requests 5000
|
||||
|
||||
# Custom proxy/provider URLs:
|
||||
python scripts/run_benchmark_rps_suite.py --proxy-url http://localhost:4000/chat/completions --provider-url http://localhost:8090/chat/completions
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
|
||||
# RPS values for the 10 runs: 1k, 2k, ..., 10k
|
||||
RPS_LEVELS = [1000, 2000, 3000, 4000, 5000, 6000, 7000, 8000, 9000, 10000]
|
||||
|
||||
|
||||
def main() -> int:
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Run benchmark 10 times at 1k–10k RPS, one at a time, each to a separate file.",
|
||||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||||
epilog=__doc__,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output-dir",
|
||||
"-o",
|
||||
default=None,
|
||||
help="Directory for result files (default: current directory). Files: benchmark_results_1000rps.txt, ...",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--proxy-url",
|
||||
default=os.getenv("LITELLM_PROXY_URL", "http://localhost:4000/chat/completions"),
|
||||
help="Proxy chat completions URL",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--provider-url",
|
||||
default=os.getenv("PROVIDER_URL", "http://localhost:8090/chat/completions"),
|
||||
help="Direct provider chat completions URL",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--requests",
|
||||
type=int,
|
||||
default=10000,
|
||||
help="Number of requests per run (default: 10000)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--no-wait",
|
||||
action="store_true",
|
||||
help="Do not wait for proxy; assume it is already up.",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
script_dir = os.path.dirname(os.path.abspath(__file__))
|
||||
runner_script = os.path.join(script_dir, "run_benchmark_after_proxy.py")
|
||||
if not os.path.isfile(runner_script):
|
||||
print(f"Error: {runner_script} not found", file=sys.stderr)
|
||||
return 1
|
||||
|
||||
output_dir = args.output_dir or os.getcwd()
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
|
||||
env = os.environ.copy()
|
||||
env["LITELLM_PROXY_URL"] = args.proxy_url
|
||||
env["PROVIDER_URL"] = args.provider_url
|
||||
|
||||
print(f"RPS suite: 10 runs at {RPS_LEVELS[0]}–{RPS_LEVELS[-1]} RPS, {args.requests} requests each")
|
||||
print(f"Output dir: {os.path.abspath(output_dir)}")
|
||||
print()
|
||||
|
||||
for i, rps in enumerate(RPS_LEVELS, 1):
|
||||
out_file = os.path.join(output_dir, f"benchmark_results_{rps}rps.txt")
|
||||
cmd = [
|
||||
sys.executable,
|
||||
runner_script,
|
||||
"--output",
|
||||
out_file,
|
||||
"--proxy-url",
|
||||
args.proxy_url,
|
||||
"--provider-url",
|
||||
args.provider_url,
|
||||
"--rps-control",
|
||||
str(rps),
|
||||
"--requests",
|
||||
str(args.requests),
|
||||
]
|
||||
# Only wait for proxy on first run (unless --no-wait)
|
||||
if i > 1 or args.no_wait:
|
||||
cmd.append("--no-wait")
|
||||
print(f"[{i}/10] RPS={rps} -> {os.path.basename(out_file)}")
|
||||
cwd = os.path.dirname(script_dir) # repo root so benchmark script finds benchmark_proxy_vs_provider.py
|
||||
ret = subprocess.run(cmd, env=env, cwd=cwd)
|
||||
if ret.returncode != 0:
|
||||
print(f"Run failed with exit code {ret.returncode}", file=sys.stderr)
|
||||
return ret.returncode
|
||||
print()
|
||||
|
||||
print("All 10 runs completed.")
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
Loading…
Add table
Reference in a new issue