P2: real engine-state feed replaces stale shadow counters for migration targeting
vLLM scheduler publishes real state (running/waiting, KV free, and the max-in-progress-prefill signal /metrics lacks) to a tmpfs/redis store ~20Hz; router reads it and avoids GIL-stall (mid-large-prefill) + KV-capacity-wall targets, using real load over 30s-stale shadow counters. Components: engine_state.py (canonical+reader), instrument_engine_state.py (scheduler patch, file/redis writer), migration_target.py (scorer), proxy wiring (--engine-state-uri, off=unchanged). All unit-tested without GPU; not yet run live. See P2_ENGINE_STATE.md. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
@@ -111,6 +111,13 @@ class Settings:
|
||||
v3_recent_mig_weight: float = 1.0 # how many "virtual requests" each
|
||||
# recent migration counts as
|
||||
|
||||
# P2: real engine-state feed (replaces 30s-stale shadow counters for
|
||||
# migration target selection). Empty = disabled (use shadow only).
|
||||
engine_state_uri: str = "" # file:///dev/shm/... or redis://...
|
||||
engine_state_period_ms: int = 50 # router poll period
|
||||
es_big_prefill_threshold: int = 16000 # target mid-prefill >= this => avoid (GIL stall)
|
||||
es_kv_wall_frac: float = 0.90 # target KV usage >= this => avoid (capacity wall)
|
||||
|
||||
# Direction B knob: LMetric fallback adds decode-token penalty to score.
|
||||
# score = (pending_prefill + new + lmetric_decode_weight * ongoing_decode_tok) * num_req
|
||||
# Empirical iter-time slope on H100 + Qwen3-30B-A3B: each decode token in
|
||||
@@ -206,6 +213,9 @@ class InstanceState:
|
||||
# recent-migration count over a sliding window, preventing back-to-back
|
||||
# decisions from clustering on the same dst.
|
||||
self.recent_mig_targeted_at: deque[float] = deque(maxlen=64)
|
||||
# P2: latest real engine state (from the engine-state feed), or None
|
||||
# when the feed is disabled/stale. Set by _engine_state_poll_loop.
|
||||
self.real_state: dict | None = None
|
||||
|
||||
def estimate_cache_hit(self, token_ids: list[int] | None) -> int:
|
||||
if not token_ids or len(token_ids) < BLOCK_SIZE:
|
||||
@@ -672,20 +682,29 @@ def pick_instance_unified_v3(
|
||||
now_mono = _time.monotonic()
|
||||
cutoff = now_mono - SETTINGS.v3_recent_mig_window_s
|
||||
|
||||
def _real_load(inst):
|
||||
# P2: prefer REAL engine state (running+waiting) over the proxy's
|
||||
# 30s-stale shadow num_requests, when the engine-state feed is fresh.
|
||||
rs = getattr(inst, "real_state", None)
|
||||
if rs is not None:
|
||||
return rs.get("num_running", 0) + rs.get("num_waiting", 0)
|
||||
return inst.num_requests
|
||||
|
||||
def effective_load(inst):
|
||||
# Drop expired entries lazily.
|
||||
while inst.recent_mig_targeted_at and inst.recent_mig_targeted_at[0] < cutoff:
|
||||
inst.recent_mig_targeted_at.popleft()
|
||||
recent = len(inst.recent_mig_targeted_at)
|
||||
return inst.num_requests + recent * SETTINGS.v3_recent_mig_weight
|
||||
return _real_load(inst) + recent * SETTINGS.v3_recent_mig_weight
|
||||
|
||||
ph_load = _real_load(prefill_host)
|
||||
threshold_loaded = max(1,
|
||||
int(prefill_host.num_requests * SETTINGS.v3_target_load_ratio))
|
||||
int(ph_load * SETTINGS.v3_target_load_ratio))
|
||||
candidates = [
|
||||
(i, inst) for i, inst in enumerate(instances)
|
||||
if i != prefill_idx
|
||||
and effective_load(inst) < threshold_loaded
|
||||
and effective_load(inst) <= prefill_host.num_requests - SETTINGS.v3_min_load_gap
|
||||
and effective_load(inst) <= ph_load - SETTINGS.v3_min_load_gap
|
||||
]
|
||||
if not candidates:
|
||||
decision["v3_reason"] = (
|
||||
@@ -700,11 +719,23 @@ def pick_instance_unified_v3(
|
||||
# cache_hit DESC (more cache = less KV to transfer), then by effective_load
|
||||
# (which includes recent-migration penalty), then by ongoing_tokens.
|
||||
if SETTINGS.v3_prefer_cache_target:
|
||||
decode_target_idx, decode_target = min(
|
||||
candidates,
|
||||
key=lambda x: (-x[1].estimate_cache_hit(token_ids),
|
||||
effective_load(x[1]),
|
||||
x[1].ongoing_tokens))
|
||||
def _tgt_key(x):
|
||||
# P2: avoid a target that is mid-large-prefill (holds the GIL,
|
||||
# stalls the mooncake receiver_loop = the ~45% control-plane
|
||||
# residual layer-wise can't fix) or near the KV capacity wall,
|
||||
# before ranking by cache-richness and real load.
|
||||
inst = x[1]
|
||||
ch = inst.estimate_cache_hit(token_ids)
|
||||
rs = getattr(inst, "real_state", None)
|
||||
stalls = near_wall = 0
|
||||
if rs is not None:
|
||||
if int(rs.get("max_prefill_remaining", 0)) >= SETTINGS.es_big_prefill_threshold:
|
||||
stalls = 1
|
||||
f = rs.get("gpu_kv_used_frac", 0.0) or 0.0
|
||||
if float(f) >= SETTINGS.es_kv_wall_frac:
|
||||
near_wall = 1
|
||||
return (stalls, near_wall, -ch, effective_load(inst), inst.ongoing_tokens)
|
||||
decode_target_idx, decode_target = min(candidates, key=_tgt_key)
|
||||
else:
|
||||
decode_target_idx, decode_target = min(
|
||||
candidates, key=lambda x: (effective_load(x[1]), x[1].ongoing_tokens))
|
||||
@@ -857,6 +888,57 @@ async def _fetch_vllm_inflight(inst: "InstanceState") -> tuple[int, int] | None:
|
||||
return running, waiting
|
||||
|
||||
|
||||
def _engine_state_read_all(uri: str, max_age_s: float = 2.0) -> dict:
|
||||
"""P2 reader (inlined; mirrors engine_state.StateReader). Returns
|
||||
{engine_id: state}, dropping records older than max_age_s."""
|
||||
now = _time.time()
|
||||
out: dict = {}
|
||||
try:
|
||||
if uri.startswith("file://"):
|
||||
import glob
|
||||
d = uri[len("file://"):]
|
||||
for p in glob.glob(os.path.join(d, "*.json")):
|
||||
try:
|
||||
s = json.load(open(p))
|
||||
except Exception:
|
||||
continue
|
||||
if now - s.get("ts", 0) <= max_age_s:
|
||||
out[s.get("engine_id", os.path.basename(p)[:-5])] = s
|
||||
elif uri.startswith("redis://"):
|
||||
import redis
|
||||
r = redis.Redis.from_url(uri)
|
||||
for k in r.scan_iter("engine_state:*"):
|
||||
v = r.get(k)
|
||||
if not v:
|
||||
continue
|
||||
s = json.loads(v)
|
||||
if now - s.get("ts", 0) <= max_age_s:
|
||||
out[s.get("engine_id")] = s
|
||||
except Exception:
|
||||
pass
|
||||
return out
|
||||
|
||||
|
||||
async def _engine_state_poll_loop():
|
||||
"""P2: poll the engine-state feed and attach real_state to each instance.
|
||||
Instance i is keyed engine_{i} (matches AGENTIC_WORKER_ID in the launcher).
|
||||
"""
|
||||
uri = SETTINGS.engine_state_uri
|
||||
if not uri:
|
||||
return
|
||||
period = max(0.01, SETTINGS.engine_state_period_ms / 1000.0)
|
||||
insts = combined_instances or (prefill_instances + decode_instances)
|
||||
print(f"[engine-state] polling {uri} every {period*1000:.0f}ms for {len(insts)} instances")
|
||||
while True:
|
||||
try:
|
||||
await asyncio.sleep(period)
|
||||
except asyncio.CancelledError:
|
||||
return
|
||||
states = await asyncio.to_thread(_engine_state_read_all, uri)
|
||||
for i, inst in enumerate(insts):
|
||||
inst.real_state = states.get(f"engine_{i}")
|
||||
|
||||
|
||||
async def _reconcile_loop():
|
||||
"""Periodic shadow-state reconciliation against vLLM /metrics truth.
|
||||
|
||||
@@ -960,6 +1042,7 @@ async def lifespan(app: FastAPI):
|
||||
_verify_vllm_patch()
|
||||
|
||||
reconcile_task = asyncio.create_task(_reconcile_loop())
|
||||
engine_state_task = asyncio.create_task(_engine_state_poll_loop())
|
||||
|
||||
if global_args.combined:
|
||||
is_pd_sep = False
|
||||
@@ -1720,6 +1803,10 @@ def parse_args():
|
||||
" penalised. 0 = original behavior; 0.01 is a reasonable start.")
|
||||
p.add_argument("--overload-factor", type=float, default=2.0,
|
||||
help="Break session affinity when instance load > factor * avg")
|
||||
p.add_argument("--engine-state-uri", type=str, default="",
|
||||
help="P2: real engine-state feed for migration target "
|
||||
"selection (file:///dev/shm/... or redis://...). "
|
||||
"Empty=disabled (shadow counters only).")
|
||||
# The four flags below are accepted for bench.sh backward compatibility but
|
||||
# have no effect after the PD-sep offload path was retired (REPORT §3.9,
|
||||
# commits 4c583f2 / cc6e562). Removing them would break scripts/bench.sh and
|
||||
@@ -1752,6 +1839,7 @@ if __name__ == "__main__":
|
||||
global_args = parse_args()
|
||||
SETTINGS.heavy_threshold = global_args.heavy_threshold
|
||||
SETTINGS.overload_factor = global_args.overload_factor
|
||||
SETTINGS.engine_state_uri = getattr(global_args, 'engine_state_uri', '') or ''
|
||||
SETTINGS.max_offload_inflight = global_args.max_offload_inflight
|
||||
SETTINGS.cache_gate_ratio = global_args.cache_gate_ratio
|
||||
SETTINGS.decode_iteration_s = getattr(global_args, 'decode_iteration_s', 0.05)
|
||||
|
||||
Reference in New Issue
Block a user