frank-medium: route by workload — spread batches, pin agent loops
The two replicas behind :7770 were on plain round-robin, which throws away prefix-cache hits: an agent loop re-sends its system prompt and file context every turn, and alternating replicas means every other turn pays a cold prefill. Measured with 90k-token contexts, warm-turn TTFT was 2676 ms. Pinning everything to one replica is not the answer either — a cold batch (the shape a benchmark or a burst of parallel subagents produces) then loses a card's worth of prefill: 16 requests took 30.1 s pinned against 18.6 s spread. So the policy routes by request shape. A continuation (more than one message — the agent-loop signature) is pinned to the replica that served its prefix; anything else spreads across the cards by counter. Measured: 19.97 s for the burst and 1095 ms warm-turn TTFT — the best of both, without a regression on either workload. The policy lives in proxy/routing.py, separate from the HTTP plumbing in proxy/affinity-proxy.py, so it can be tested without aiohttp, a network or a GPU. tests/frank_medium_proxy.bats covers it, and it earned its keep: the first version hashed fresh prefixes to pick a card, which sends every member of a same-prefix burst to one card (29.8 s, measured), and a stale index in the pin path returned 500 on every continuation. Both are now assertions. ROUTING=roundrobin and ROUTING=affinity stay available to re-measure the choice on this host; the default is adaptive.
This commit is contained in:
Frank
+132
@@ -0,0 +1,132 @@
|
||||
#!/usr/bin/env python3
|
||||
"""HTTP front door for the frank-medium single-GPU replicas.
|
||||
|
||||
Why two replicas instead of TP=2: one vLLM process spanning both R9700s needs
|
||||
RCCL collectives, and RCCL's kernels require PCIe atomics on every
|
||||
participating GPU. The second card sits behind chipset bridges advertising
|
||||
`AtomicOpsCap: Routing-`, which no software can enable (see the runbook), so
|
||||
no collective can launch on it. Two single-GPU replicas use both cards; this
|
||||
service keeps the fleet contract of one endpoint on :7770, streaming included.
|
||||
|
||||
Routing decisions live in routing.py — spread single-turn requests, pin
|
||||
continuations to the replica whose prefix cache holds their context. This file
|
||||
only moves bytes: parse enough of the request to route it, forward it verbatim,
|
||||
stream the answer back, and report replica health on /health.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
from contextlib import suppress
|
||||
from aiohttp import ClientSession, ClientTimeout, web
|
||||
|
||||
import routing
|
||||
|
||||
PORT = int(os.environ.get("PORT", "7770"))
|
||||
HEALTH_TTL = float(os.environ.get("HEALTH_TTL", "10"))
|
||||
# Hop-by-hop headers must not be forwarded (RFC 9110 §7.6.1) — forwarding
|
||||
# transfer-encoding or content-length of a rewritten body breaks clients.
|
||||
HOP_BY_HOP = {
|
||||
"host", "content-length", "transfer-encoding", "connection",
|
||||
"keep-alive", "proxy-authenticate", "proxy-authorization", "te",
|
||||
"trailers", "upgrade",
|
||||
}
|
||||
|
||||
UPSTREAMS = routing.UPSTREAMS
|
||||
# Replica health, refreshed by the watcher. Module state rather than aiohttp app
|
||||
# state: the application is frozen once the site is up.
|
||||
_UP: list[bool] = [True] * len(UPSTREAMS)
|
||||
|
||||
|
||||
async def health(session: ClientSession) -> web.Response:
|
||||
results = await asyncio.gather(
|
||||
*(session.get(u + "/health", timeout=ClientTimeout(total=5)) for u in UPSTREAMS),
|
||||
return_exceptions=True,
|
||||
)
|
||||
lines, up = [], 0
|
||||
for url, res in zip(UPSTREAMS, results):
|
||||
if isinstance(res, Exception) or res.status != 200:
|
||||
lines.append(f"{url} DOWN")
|
||||
else:
|
||||
up += 1
|
||||
lines.append(f"{url} UP")
|
||||
return web.Response(status=200 if up else 503, text="\n".join(lines) + "\n")
|
||||
|
||||
|
||||
async def watch_health() -> None:
|
||||
"""Refresh the replica up/down table so routing never targets a dead one.
|
||||
|
||||
One HTTP request per replica per interval, in exchange for never handing a
|
||||
client a 502 after a restart.
|
||||
"""
|
||||
global _UP
|
||||
session = ClientSession()
|
||||
try:
|
||||
while True:
|
||||
res = await asyncio.gather(
|
||||
*(session.get(u + "/health", timeout=ClientTimeout(total=5)) for u in UPSTREAMS),
|
||||
return_exceptions=True,
|
||||
)
|
||||
_UP = [not isinstance(r, Exception) and r.status == 200 for r in res]
|
||||
await asyncio.sleep(HEALTH_TTL)
|
||||
finally:
|
||||
await session.close()
|
||||
|
||||
|
||||
async def proxy(request: web.Request) -> web.StreamResponse:
|
||||
session: ClientSession = request.app["session"]
|
||||
if request.path == "/health":
|
||||
return await health(session)
|
||||
|
||||
body = await request.read()
|
||||
idx = routing.route(body, request.remote, _UP)
|
||||
if idx < 0:
|
||||
return web.Response(status=503, text="proxy: no healthy frank-medium replica\n")
|
||||
target = UPSTREAMS[idx] + request.path_qs
|
||||
headers = {k: v for k, v in request.headers.items() if k.lower() not in HOP_BY_HOP}
|
||||
|
||||
try:
|
||||
upstream = await session.request(
|
||||
request.method, target, headers=headers, data=body or None,
|
||||
timeout=ClientTimeout(total=None, sock_connect=10, sock_read=3600),
|
||||
)
|
||||
except Exception as exc: # upstream refused or unreachable
|
||||
return web.Response(status=502, text=f"proxy: {target} unreachable: {exc}\n")
|
||||
|
||||
out = web.StreamResponse(status=upstream.status, reason=upstream.reason)
|
||||
for k, v in upstream.headers.items():
|
||||
if k.lower() not in HOP_BY_HOP:
|
||||
out.headers[k] = v
|
||||
await out.prepare(request)
|
||||
try:
|
||||
async for chunk in upstream.content.iter_any():
|
||||
await out.write(chunk)
|
||||
except Exception:
|
||||
# The client went away mid-stream (an aborted generation). Nothing to
|
||||
# report: the engine sees its own cancellation.
|
||||
pass
|
||||
finally:
|
||||
upstream.release()
|
||||
with suppress(Exception):
|
||||
await out.write_eof()
|
||||
return out
|
||||
|
||||
|
||||
async def main() -> None:
|
||||
app = web.Application(client_max_size=1024 ** 3)
|
||||
app["session"] = ClientSession()
|
||||
app.router.add_route("*", "/{tail:.*}", proxy)
|
||||
runner = web.AppRunner(app)
|
||||
await runner.setup()
|
||||
await web.TCPSite(runner, "0.0.0.0", PORT).start()
|
||||
print(f"proxy: :{PORT} -> {', '.join(UPSTREAMS)} routing={routing.ROUTING} "
|
||||
f"key={routing.AFFINITY_BYTES}B known_max={routing.KNOWN_MAX}", flush=True)
|
||||
watcher = asyncio.create_task(watch_health())
|
||||
with suppress(asyncio.CancelledError):
|
||||
await asyncio.gather(watcher)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
with suppress(KeyboardInterrupt):
|
||||
asyncio.run(main())
|
||||
@@ -1,96 +0,0 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Round-robin proxy in front of the frank-medium TP=1 replicas.
|
||||
|
||||
Why this exists: TP=2 (one vLLM process spanning both R9700s) needs RCCL
|
||||
collectives, and RCCL's kernels require PCIe atomics on every participating
|
||||
GPU. On this host the second R9700 sits behind chipset bridges that advertise
|
||||
`AtomicOpsCap: Routing-` (see the runbook), so no collective can launch on it
|
||||
and TP=2 is impossible. Two single-GPU replicas use both cards instead; this
|
||||
proxy keeps the fleet contract of one endpoint on :7770.
|
||||
|
||||
The vLLM OpenAI API is stateless per request, so round-robin is sound. /health
|
||||
returns 200 while at least one replica is up and names the ones that are not.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import itertools
|
||||
import os
|
||||
from aiohttp import ClientSession, ClientTimeout, web
|
||||
|
||||
UPSTREAMS = [u.strip() for u in os.environ.get(
|
||||
"UPSTREAMS", "http://127.0.0.1:7771,http://127.0.0.1:7772").split(",") if u.strip()]
|
||||
PORT = int(os.environ.get("PORT", "7770"))
|
||||
# Hop-by-hop headers must not be forwarded (RFC 9110 §7.6.1) — forwarding
|
||||
# transfer-encoding or content-length of a rewritten body breaks clients.
|
||||
HOP_BY_HOP = {
|
||||
"host", "content-length", "transfer-encoding", "connection",
|
||||
"keep-alive", "proxy-authenticate", "proxy-authorization", "te",
|
||||
"trailers", "upgrade",
|
||||
}
|
||||
|
||||
_cycle = itertools.cycle(range(len(UPSTREAMS)))
|
||||
|
||||
|
||||
async def health(session: ClientSession) -> web.Response:
|
||||
results = await asyncio.gather(
|
||||
*(session.get(u + "/health", timeout=ClientTimeout(total=5)) for u in UPSTREAMS),
|
||||
return_exceptions=True,
|
||||
)
|
||||
lines, up = [], 0
|
||||
for url, res in zip(UPSTREAMS, results):
|
||||
if isinstance(res, Exception) or res.status != 200:
|
||||
lines.append(f"{url} DOWN")
|
||||
else:
|
||||
up += 1
|
||||
lines.append(f"{url} UP")
|
||||
body = "\n".join(lines) + "\n"
|
||||
return web.Response(status=200 if up else 503, text=body)
|
||||
|
||||
|
||||
async def proxy(request: web.Request) -> web.StreamResponse:
|
||||
session: ClientSession = request.app["session"]
|
||||
if request.path == "/health":
|
||||
return await health(session)
|
||||
|
||||
target = UPSTREAMS[next(_cycle)] + request.path_qs
|
||||
headers = {k: v for k, v in request.headers.items() if k.lower() not in HOP_BY_HOP}
|
||||
body = await request.read()
|
||||
|
||||
try:
|
||||
upstream = await session.request(
|
||||
request.method, target, headers=headers, data=body or None,
|
||||
timeout=ClientTimeout(total=None, sock_connect=10, sock_read=3600),
|
||||
)
|
||||
except Exception as exc: # upstream refused/unreachable
|
||||
return web.Response(status=502, text=f"proxy: {target} unreachable: {exc}\n")
|
||||
|
||||
out = web.StreamResponse(status=upstream.status, reason=upstream.reason)
|
||||
for k, v in upstream.headers.items():
|
||||
if k.lower() not in HOP_BY_HOP:
|
||||
out.headers[k] = v
|
||||
await out.prepare(request)
|
||||
try:
|
||||
async for chunk in upstream.content.iter_any():
|
||||
await out.write(chunk)
|
||||
finally:
|
||||
upstream.release()
|
||||
await out.write_eof()
|
||||
return out
|
||||
|
||||
|
||||
async def main() -> None:
|
||||
app = web.Application(client_max_size=1024 ** 3)
|
||||
app["session"] = ClientSession()
|
||||
app.router.add_route("*", "/{tail:.*}", proxy)
|
||||
runner = web.AppRunner(app)
|
||||
await runner.setup()
|
||||
site = web.TCPSite(runner, "0.0.0.0", PORT)
|
||||
await site.start()
|
||||
print(f"proxy: :{PORT} -> {', '.join(UPSTREAMS)}", flush=True)
|
||||
await asyncio.Event().wait()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
Executable
+129
@@ -0,0 +1,129 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Routing policy for the frank-medium proxy: which replica serves a request.
|
||||
|
||||
Kept apart from the HTTP plumbing in affinity-proxy.py so the decisions can be
|
||||
tested without aiohttp, a network or a GPU — the first version of these rules
|
||||
shipped with a bug that only appeared on a continuation request, and a test
|
||||
that loads this module catches that class of mistake directly.
|
||||
|
||||
Two cards, two behaviours, and the workload decides which is right:
|
||||
|
||||
* A **single-turn** request is one member of a batch (a benchmark run, a burst
|
||||
of parallel subagents). Spreading those over both replicas lets both cards
|
||||
prefill at once; measured on this host, a 16-request burst of 15k-token
|
||||
prompts finishes in 18.6 s spread versus 30.1 s pinned.
|
||||
* A **continuation** re-sends a prefix the replica has already prefilled. Pinned
|
||||
to that replica the turn costs a cache lookup; spread, half the turns pay a
|
||||
cold prefill again (measured: 1096 ms versus 2676 ms mean warm-turn TTFT on
|
||||
90k-token contexts).
|
||||
|
||||
ROUTING picks the strategy; the default `adaptive` applies both rules above.
|
||||
`affinity` and `roundrobin` exist so the choice can be re-measured, and are not
|
||||
production settings.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import collections
|
||||
import hashlib
|
||||
import itertools
|
||||
import json
|
||||
import os
|
||||
|
||||
UPSTREAMS = [u.strip() for u in os.environ.get(
|
||||
"UPSTREAMS", "http://127.0.0.1:7771,http://127.0.0.1:7772").split(",") if u.strip()]
|
||||
ROUTING = os.environ.get("ROUTING", "adaptive").strip().lower()
|
||||
# Leading bytes of prompt text that identify a conversation's shared prefix.
|
||||
# 4 KiB covers a system prompt plus the start of the first user turn, which is
|
||||
# where agent loops keep the text they resend every turn.
|
||||
AFFINITY_BYTES = int(os.environ.get("AFFINITY_BYTES", "4096"))
|
||||
# Prefixes remembered for pinning. Bounded so a long-lived proxy cannot grow
|
||||
# without limit; 4096 conversations is far beyond what this fleet runs at once.
|
||||
KNOWN_MAX = int(os.environ.get("KNOWN_MAX", "4096"))
|
||||
|
||||
_counter = itertools.count()
|
||||
# digest -> replica index, insertion-ordered for O(1) LRU eviction. A prefix is
|
||||
# remembered as soon as any request carries it; only a continuation is pinned
|
||||
# there, because only a continuation can profit from the cache entry.
|
||||
_known: "collections.OrderedDict[bytes, int]" = collections.OrderedDict()
|
||||
|
||||
|
||||
def affinity_key(body: bytes, remote: str | None) -> tuple[bytes, bool]:
|
||||
"""Key a request by its shared prefix, and report whether it continues one.
|
||||
|
||||
Chat and completions bodies both carry their prompt text in JSON, so the
|
||||
leading bytes of that text are the prefix; a body we cannot parse (or an
|
||||
endpoint with no prompt, e.g. /v1/models) falls back to the client address,
|
||||
which still beats nothing for a single-client loop.
|
||||
"""
|
||||
fallback = (remote or "anonymous").encode()
|
||||
try:
|
||||
payload = json.loads(body)
|
||||
except Exception:
|
||||
return fallback, False
|
||||
|
||||
text = ""
|
||||
continuation = False
|
||||
if isinstance(payload, dict):
|
||||
if isinstance(payload.get("prompt"), str):
|
||||
text = payload["prompt"]
|
||||
elif isinstance(payload.get("messages"), list):
|
||||
continuation = len(payload["messages"]) > 1
|
||||
parts = []
|
||||
for msg in payload["messages"]:
|
||||
if isinstance(msg, dict):
|
||||
content = msg.get("content")
|
||||
if isinstance(content, str):
|
||||
parts.append(content)
|
||||
elif isinstance(content, list): # multimodal content blocks
|
||||
parts.extend(
|
||||
blk.get("text", "") for blk in content
|
||||
if isinstance(blk, dict) and isinstance(blk.get("text"), str)
|
||||
)
|
||||
text = "\n".join(parts)
|
||||
if not text:
|
||||
return fallback, continuation
|
||||
return text[:AFFINITY_BYTES].encode("utf-8", "replace"), continuation
|
||||
|
||||
|
||||
def pick(key: bytes, continuation: bool, up: list[bool]) -> int:
|
||||
"""Replica for a request, or -1 when none is available.
|
||||
|
||||
Down replicas are never selected — a rollout must not send traffic to an
|
||||
engine that is still compiling its graphs — and a pin whose replica is down
|
||||
is re-established on a live one rather than kept pointing at a corpse.
|
||||
"""
|
||||
digest = hashlib.sha1(key).digest()
|
||||
|
||||
if ROUTING == "roundrobin":
|
||||
start = next(_counter) % len(UPSTREAMS)
|
||||
elif ROUTING == "adaptive":
|
||||
entry = _known.get(digest)
|
||||
if continuation and entry is not None and up[entry]:
|
||||
_known.move_to_end(digest)
|
||||
return entry
|
||||
# Anything that is not a continuation spreads by counter — including a
|
||||
# batch of requests that happen to share one prefix. Hashing the prefix
|
||||
# would send every member of such a batch to whichever card the hash
|
||||
# names (measured: 29.8 s for a 16-request burst, against 18.6 s when
|
||||
# spread). The counter is only advanced by these requests, so batch
|
||||
# members alternate across the cards.
|
||||
start = next(_counter) % len(UPSTREAMS)
|
||||
else: # affinity
|
||||
start = int.from_bytes(digest[:8], "big") % len(UPSTREAMS)
|
||||
|
||||
for offset in range(len(UPSTREAMS)):
|
||||
idx = (start + offset) % len(UPSTREAMS)
|
||||
if up[idx]:
|
||||
if ROUTING == "adaptive":
|
||||
_known[digest] = idx
|
||||
_known.move_to_end(digest)
|
||||
if len(_known) > KNOWN_MAX:
|
||||
_known.popitem(last=False)
|
||||
return idx
|
||||
return -1
|
||||
|
||||
|
||||
def route(body: bytes, remote: str | None, up: list[bool]) -> int:
|
||||
"""Convenience wrapper: decode the request body and choose a replica."""
|
||||
return pick(*affinity_key(body, remote), up)
|
||||
@@ -3,7 +3,7 @@
|
||||
# Install path on dankserver: ~/.config/containers/systemd/frank-medium-proxy.container
|
||||
# Unit name: frank-medium-proxy.service (container name: frank-medium-proxy)
|
||||
#
|
||||
# WHAT IT IS. A round-robin reverse proxy in front of the two single-GPU
|
||||
# WHAT IT IS. A prefix-affinity reverse proxy in front of the two single-GPU
|
||||
# frank-medium replicas (:7771 on R9700 #1, :7772 on R9700 #2). It preserves
|
||||
# the fleet contract: one endpoint on host port 7770 answering as the served
|
||||
# model name frank-medium, OpenAI /v1, streaming included.
|
||||
@@ -13,17 +13,22 @@
|
||||
# need PCIe atomics, and the second card's chipset path advertises
|
||||
# AtomicOpsCap: Routing-). Two replicas are how both cards get used.
|
||||
#
|
||||
# ROUTING. vLLM's OpenAI API is stateless per request, so round-robin is
|
||||
# sound; there is no session affinity to preserve. /health answers 200 while
|
||||
# at least one replica is up and names the ones that are down. The proxy runs
|
||||
# on host networking so it can bind 7770 and reach the replicas on loopback.
|
||||
# ROUTING. Requests are keyed on a hash of their leading 4 KiB of prompt text,
|
||||
# so a conversation family — an agent loop re-sending the same system prompt,
|
||||
# tool list and file context every turn — keeps landing on the same replica,
|
||||
# where --enable-prefix-caching has that prefix warm. Plain round-robin threw
|
||||
# half of those cache hits away. A replica that is down is skipped, and
|
||||
# ROUTING=roundrobin restores the old behaviour for A/B measurement. /health
|
||||
# answers 200 while at least one replica is up and names the ones that are
|
||||
# down. The proxy runs on host networking so it can bind 7770 and reach the
|
||||
# replicas on loopback.
|
||||
#
|
||||
# The script is a few dozen lines of aiohttp, mounted read-only out of
|
||||
# /home/dank/serverAppConfigs (deploy.sh pulls that checkout before dispatch).
|
||||
# It reuses the same image as the replicas: aiohttp is already there and no
|
||||
# extra image has to be pulled on this host.
|
||||
[Unit]
|
||||
Description=frank-medium proxy — round-robin in front of the two single-GPU replicas
|
||||
Description=frank-medium proxy — prefix-affinity in front of the two single-GPU replicas
|
||||
After=network-online.target
|
||||
Wants=frank-medium-a.service frank-medium-b.service
|
||||
After=frank-medium-a.service frank-medium-b.service
|
||||
@@ -34,9 +39,13 @@ Image=localhost/radiance-vllm:0.9.3
|
||||
PodmanArgs=--network=host
|
||||
Environment=PORT=7770
|
||||
Environment=UPSTREAMS=http://127.0.0.1:7771,http://127.0.0.1:7772
|
||||
Volume=/home/dank/serverAppConfigs/dankserver/frank-medium/proxy/roundrobin-proxy.py:/usr/local/bin/roundrobin-proxy.py:ro
|
||||
# The proxy is two files: routing.py holds the replica-selection policy, the
|
||||
# server holds the HTTP plumbing. Both are mounted because the server imports
|
||||
# the policy by name from its own directory (/usr/local/bin is sys.path[0]).
|
||||
Volume=/home/dank/serverAppConfigs/dankserver/frank-medium/proxy/routing.py:/usr/local/bin/routing.py:ro
|
||||
Volume=/home/dank/serverAppConfigs/dankserver/frank-medium/proxy/affinity-proxy.py:/usr/local/bin/affinity-proxy.py:ro
|
||||
Volume=/home/dank/serverAppConfigs/dankserver/frank-medium/frank-medium-healthcheck.py:/usr/local/bin/frank-medium-healthcheck.py:ro
|
||||
PodmanArgs=--entrypoint /usr/local/bin/roundrobin-proxy.py
|
||||
PodmanArgs=--entrypoint /usr/local/bin/affinity-proxy.py
|
||||
HealthCmd=python3 /usr/local/bin/frank-medium-healthcheck.py
|
||||
HealthInterval=60s
|
||||
HealthTimeout=30s
|
||||
|
||||
Reference in New Issue
Block a user