File size: 8,856 Bytes
4a28d4d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 | import os
import subprocess
import time
import fire
import yaml
def get_launching_server_cmd(model_path, backend, server_config):
if backend in ['turbomind', 'pytorch']:
cmd = ['lmdeploy', 'serve', 'api_server', model_path, '--backend', backend]
elif backend == 'sglang':
cmd = ['python3', '-m', 'sglang.launch_server', '--model-path', model_path]
elif backend == 'vllm':
cmd = ['vllm', 'serve', model_path]
else:
raise ValueError(f'unknown backend: {backend}')
for key, value in server_config.items():
# Convert snake_case to kebab-case for command line args
key = key.replace('_', '-')
cmd.append(f'--{key}')
if str(value):
cmd.append(str(value))
# Special handling for proxy server case
if server_config.get('proxy_url') and server_config.get('dp'):
cmd.append('--allow-terminate-by-client')
return cmd
def get_output_file(model_path, backend, server_config):
"""Generate the benchmark output filename."""
model_name = server_config.get('model_name', None) or os.path.basename(model_path)
if backend not in ['turbomind', 'pytorch', 'sglang', 'vllm']:
raise ValueError(f'Unknown backend: {backend}')
if backend in ['sglang', 'vllm']:
return f'benchmark_{model_name}_{backend}.csv'
# For turbomind/pytorch backends
params = [
('bs', server_config['max_batch_size']),
('tp', server_config.get('tp', 1)),
('dp', server_config.get('dp', '')),
('ep', server_config.get('ep', '')),
('cache', server_config.get('cache_max_entry_count', 0.8)),
('mptk', server_config.get('max_prefill_token_num', '')),
]
params_str = '_'.join(f'{k}{v}' for k, v in params if v != '')
# Turbomind-specific additions
if backend == 'turbomind' and (comm := server_config.get('communicator')):
params_str += f'_{comm}'
return f'benchmark_{model_name}_{backend}_{params_str}.csv'
def get_server_ip_port(backend: str, server_config: dict) -> tuple[str, int]:
if backend in ['turbomind', 'pytorch']:
if server_config.get('proxy_url'):
# If proxy_url is set, we use the proxy server's IP and port
parts = server_config['proxy_url'].split(':')
server_ip = parts[1].lstrip('//')
server_port = int(parts[2])
else:
# Default to the server IP and port specified in the config
server_ip = server_config.get('server_ip', '0.0.0.0')
server_port = server_config.get('server_port', 23333)
elif backend == 'sglang':
return (server_config.get('server_ip', '0.0.0.0'), server_config.get('port', 30000))
elif backend == 'vllm':
return (server_config.get('server_ip', '0.0.0.0'), server_config.get('port', 8000))
else:
raise ValueError(f'unknown backend: {backend}')
return server_ip, server_port
def wait_server_ready(server_ip: str, server_port: int) -> bool:
"""Wait for the API server to become ready."""
from openai import OpenAI
while True:
try:
client = OpenAI(api_key='DUMMPY', base_url=f'http://{server_ip}:{server_port}/v1')
model_name = client.models.list().data[0].id
if model_name:
print('Server is ready.')
return True
except Exception as e:
print(f'connect to server http://{server_ip}:{server_port} failed {e}')
time.sleep(5)
def get_client_cmd(backend: str, server_ip: str, server_port: int, client_config: dict) -> list[str]:
"""Generate the client benchmark command."""
current_dir = os.path.dirname(os.path.abspath(__file__))
if backend in ['turbomind', 'pytorch']:
backend = 'lmdeploy'
cmd = [
'python3', f'{current_dir}/profile_restful_api.py', '--backend', backend, '--host', server_ip, '--port',
str(server_port)
]
for key, value in client_config.items():
# change the key like 'dataset_path' to 'dataset-path' to suit the optional when performing
# "python3 benchmark/profile_restful_api.py"
key = key.replace('_', '-')
if key == 'disable-warmup':
if str(value).lower() == 'true':
cmd.append(f'--{key}')
continue
cmd.append(f'--{key}')
cmd.append(str(value))
return cmd
def benchmark(model_path: str, backend: str, server_config: dict, data_config: dict | list[dict]):
"""Benchmark the server with the given configuration.
Args:
model_path: Path to the model.
backend: Backend to use.
server_config: Configuration for the server and the inference engine.
data_config: Configuration for the data.
"""
if isinstance(data_config, dict):
data_config = [data_config]
if not (isinstance(data_config, list) and all(isinstance(d, dict) for d in data_config)):
raise ValueError('data_config must be a dict or list of dicts')
server_cmd = get_launching_server_cmd(model_path, backend, server_config)
server_ip, server_port = get_server_ip_port(backend, server_config)
proc = None
try:
print(f"Starting api_server: {' '.join(server_cmd)}", flush=True)
proc = subprocess.Popen(server_cmd)
# Wait for the server to be ready
wait_server_ready(server_ip, server_port)
# Run benchmarks
output_file = get_output_file(model_path, backend, server_config)
for data in data_config:
data = data.copy()
data['output_file'] = output_file
client_cmd = get_client_cmd(backend, server_ip, server_port, data)
print(f"Running benchmark: {' '.join(client_cmd)}")
subprocess.run(client_cmd, check=True)
except Exception as e:
print(f'Unexpected error: {e}')
raise
finally:
# Clean up server process
if proc and proc.poll() is None:
if server_config.get('proxy_url') and server_config.get('dp'):
# Sending termination request to proxy_server. The request will be broadcasted to
# api_server on each dp_rank by proxy server
# Note that api_server is supposed to be launched with --allow-terminate-by-client
print('Sending termination request to proxy server')
subprocess.run(['curl', '-X', 'POST', f'{server_config["proxy_url"]}/nodes/terminate_all'],
check=True,
timeout=10)
proc.terminate()
try:
proc.wait(timeout=30)
except subprocess.TimeoutExpired:
print('Server did not terminate gracefully - killing')
proc.kill()
def validate_config(config: dict) -> None:
"""Validate the configuration structure.
Args:
config: Loaded configuration dictionary
Raises:
BenchmarkConfigError: If configuration is invalid
"""
required_sections = ['api_server', 'engine', 'data']
for section in required_sections:
if section not in config:
raise ValueError(f'Missing required config section: {section}')
if not isinstance(config['engine'], (dict, list)):
raise ValueError('engine config must be a dict or list of dicts')
if not isinstance(config['data'], (dict, list)):
raise ValueError('data config must be a dict or list of dicts')
def main(backend: str, config_path: str, model_path: str | None = None):
"""Main entry point for the benchmark script.
Args:
backend: Backend to use
config_path: Path to config file
model_path: Optional override for model path
Raises:
BenchmarkConfigError: If required parameters are missing or config is invalid
"""
with open(config_path) as f:
config = yaml.safe_load(f)
server_config = config['server']
engine_configs = config['engine']
data_config = config['data']
if isinstance(engine_configs, dict):
engine_configs = [engine_configs]
assert isinstance(engine_configs, list) and all(isinstance(s, dict) for s in engine_configs)
for engine_config in engine_configs:
server_config = server_config.copy()
server_config.update(engine_config) # Merge engine config with server config
# The model_path provided by the user will override the model_path in the config file.
model_path = model_path or server_config.pop('model_path')
# Remove model_path from server_config to avoid passing it to the server command
server_config.pop('model_path', None)
benchmark(model_path, backend, server_config, data_config)
if __name__ == '__main__':
fire.Fire(main)
|