File size: 7,548 Bytes
932bc69 | 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 | """Benchmark-only launcher: one/two same-GPU engines behind a local HTTP proxy.
Use as LAUNCH_VLLM in the single-node recipe. Both engines retain the normal
hidden-state connector and their own launch provenance. This is not a default.
"""
import asyncio
import contextlib
import itertools
import os
import signal
import sys
from pathlib import Path
import aiohttp
from aiohttp import web
def option(arguments, name, default=None):
return arguments[arguments.index(name) + 1] if name in arguments else default
def replace_option(arguments, name, value):
result = list(arguments)
if name in result:
result[result.index(name) + 1] = str(value)
else:
result.extend((name, str(value)))
return result
async def main():
arguments = sys.argv[1:]
boundary = arguments.index("--")
launch_args, engine_args = arguments[:boundary], arguments[boundary + 1 :]
provenance = Path(option(launch_args, "--provenance-dir"))
port = int(option(engine_args, "--port", "8000"))
host = option(engine_args, "--host", "127.0.0.1")
if host != "127.0.0.1":
raise ValueError("The benchmark proxy must remain loopback-only")
if option(engine_args, "--data-parallel-size", "1") != "1":
raise ValueError("This benchmark requires DP=1 per engine")
engine_count = int(os.environ.get("VLLM_BENCHMARK_ENGINES", "2"))
if engine_count not in (1, 2):
raise ValueError("Only one or two engines fit this benchmark topology")
endpoints = [
f"http://127.0.0.1:{port + offset}" for offset in range(1, engine_count + 1)
]
children, logs = [], []
stop = asyncio.Event()
for sig in (signal.SIGINT, signal.SIGTERM):
asyncio.get_running_loop().add_signal_handler(sig, stop.set)
runner = None
try:
timeout = aiohttp.ClientTimeout(total=1800)
async with aiohttp.ClientSession(
timeout=timeout, auto_decompress=False, trust_env=False
) as client:
# Sequential startup avoids cross-process interference in vLLM's
# activation-memory profiler. Runtime inference remains concurrent.
for index, endpoint in enumerate(endpoints):
engine_provenance = provenance / f"engine{index}"
engine_provenance.mkdir(parents=True, exist_ok=True)
child_launch = replace_option(
launch_args, "--provenance-dir", engine_provenance
)
child_engine = replace_option(engine_args, "--port", port + index + 1)
child_engine = replace_option(
child_engine,
"--gpu-memory-utilization",
"0.44" if engine_count == 2 else "0.92",
)
child_engine = replace_option(child_engine, "--api-server-count", "1")
log = (engine_provenance / "server.log").open("w")
logs.append(log)
child = await asyncio.create_subprocess_exec(
sys.executable,
str(Path(__file__).resolve().parents[1] / "launch_vllm.py"),
*child_launch,
"--",
*child_engine,
env={**os.environ, "VLLM_PORT": str(port + index + 101)},
stdout=log,
stderr=asyncio.subprocess.STDOUT,
)
children.append(child)
print(
f"Starting engine {index}: pid={child.pid}, {endpoint}, log={log.name}",
flush=True,
)
deadline = asyncio.get_running_loop().time() + 1500
while not stop.is_set():
if child.returncode is not None:
raise RuntimeError(
f"Engine {index} exited {child.returncode}; see {log.name}"
)
try:
async with client.get(
endpoint + "/health", timeout=aiohttp.ClientTimeout(total=3)
) as response:
if response.status == 200:
break
except (aiohttp.ClientError, TimeoutError):
pass
if asyncio.get_running_loop().time() >= deadline:
raise TimeoutError(f"Engine {index} startup timed out")
with contextlib.suppress(TimeoutError):
await asyncio.wait_for(stop.wait(), timeout=2)
if stop.is_set():
return
print(f"Engine {index} healthy", flush=True)
rotation = itertools.cycle(endpoints)
hop_headers = {
"host",
"connection",
"transfer-encoding",
"keep-alive",
"content-length",
}
async def forward(request):
endpoint = next(rotation) if request.method == "POST" else endpoints[0]
headers = {
key: value
for key, value in request.headers.items()
if key.lower() not in hop_headers
}
async with client.request(
request.method,
endpoint + request.path_qs,
headers=headers,
data=await request.read(),
allow_redirects=False,
) as upstream:
response = web.StreamResponse(
status=upstream.status,
headers={
key: value
for key, value in upstream.headers.items()
if key.lower() not in hop_headers
},
)
await response.prepare(request)
async for chunk in upstream.content.iter_chunked(65536):
await response.write(chunk)
await response.write_eof()
return response
app = web.Application(client_max_size=64 * 1024 * 1024)
app.router.add_route("*", "/{path:.*}", forward)
runner = web.AppRunner(app, access_log=None)
await runner.setup()
await web.TCPSite(runner, host, port).start()
print(f"{engine_count}-engine proxy healthy on {host}:{port}", flush=True)
while not stop.is_set():
for index, child in enumerate(children):
if child.returncode is not None:
raise RuntimeError(f"Engine {index} exited {child.returncode}")
with contextlib.suppress(TimeoutError):
await asyncio.wait_for(stop.wait(), timeout=2)
finally:
if runner is not None:
await runner.cleanup()
for child in children:
if child.returncode is None:
with contextlib.suppress(ProcessLookupError):
child.terminate()
for child in children:
try:
await asyncio.wait_for(child.wait(), timeout=20)
except TimeoutError:
with contextlib.suppress(ProcessLookupError):
child.kill()
await child.wait()
for log in logs:
log.close()
if __name__ == "__main__":
asyncio.run(main())
|