From 9c098ca00c756998f42707867ec7b4e18ca1863d Mon Sep 17 00:00:00 2001 From: jingshenghang Date: Wed, 27 May 2026 03:59:14 +0000 Subject: [PATCH] [coding_agent_rl] middleware: shutdown_session closes sid and drains in-flight handlers before sglang release_memory_occupation When a SWE rollout sample returns, slime's train loop will later call release_memory_occupation, which hard-asserts sglang's `is_fully_idle()` and SIGQUITs the scheduler on miss. Straggler claude-cli requests from the just-torn-down sandbox can still reach sglang via middleware's `setdefault(sid, Session())` at that moment, tripping the assert. Every multi-step SWE RL run was crashing on the first offload after a rollout step. Two-layer close in the middleware (only examples/ touched, slime backend untouched): - module-level `_inflight: dict[sid, set[Task]]` tracks every running `_handle_request`; lives outside the Session dataclass so it survives `pop_session_split` - module-level `_closed: set[sid]` tombstones drained sids; `_handle_request` early-returns 503 for them so stragglers cannot silently `setdefault` a fresh Session and reach sglang - `shutdown_session(sid)`: `_closed.add(sid)`, then `asyncio.wait` the in-flight bucket, cancel + gather pending (cancel fires `/abort_request` to sglang via `_generate`'s cleanup) - `pop_session_split` deliberately does NOT discard from `_closed` (permanent tombstone). sids are unique per rollout step so the set only grows linearly per process - `generate.py` finally-block calls `shutdown_session` before `pop_session_split`; idempotent on the happy path Does NOT wait for sglang idle here -- slime's `sglang_engine.release_memory_occupation` already calls `flush_cache()` with 60x1s polling (returns 200 only when scheduler is idle), strictly stronger than anything we could probe from the middleware. Once `_closed` is set and local handlers are drained, no new request can reach sglang, so the backend's existing flush_cache covers the actual idle wait. Validated 2026-05-26 on 8-node Qwen3.6-35B-A3B SWE RL: step 0 (rollout + train + weight update + offload) and step 1 (rollout + train, loss -1.03, grad_norm 3.36) both completed cleanly. --- examples/coding_agent_rl/generate.py | 5 +++ examples/coding_agent_rl/middleware.py | 50 ++++++++++++++++++++------ 2 files changed, 44 insertions(+), 11 deletions(-) diff --git a/examples/coding_agent_rl/generate.py b/examples/coding_agent_rl/generate.py index 2ff2c2e2b..b059fe87f 100644 --- a/examples/coding_agent_rl/generate.py +++ b/examples/coding_agent_rl/generate.py @@ -437,6 +437,11 @@ async def generate(args, sample: Sample, sampling_params: dict[str, Any]): traceback.format_exc(), ) return _abort_result(sample, f"exception:{type(e).__name__}") + finally: + # Close the sid before next train step's release_memory_occupation; + # stragglers from this trajectory would otherwise race its idle assert. + await middleware.shutdown_session(session_id) + middleware.pop_session_split(state.store, session_id) # idempotent def _log_timeout_diagnostic(t0: float) -> None: diff --git a/examples/coding_agent_rl/middleware.py b/examples/coding_agent_rl/middleware.py index 4698bb194..06f64f17c 100644 --- a/examples/coding_agent_rl/middleware.py +++ b/examples/coding_agent_rl/middleware.py @@ -18,6 +18,7 @@ start(...) build the aiohttp app + store open_session(store, sid, ...) register a session with sampling defaults pop_session_split(store, sid) drain a session's trajectory for training + shutdown_session(sid, ...) tombstone sid + drain in-flight handlers Chain, Session dataclasses (exposed for type hints) Read `_handle_request` (§3) top-to-bottom -- that's the whole turn: @@ -143,6 +144,12 @@ class Session: _Store = dict[str, Session] +# Drain state for shutdown_session. Module-level so it survives +# pop_session_split; _closed is a permanent tombstone (late requests 503). +_inflight: dict[str, set[asyncio.Task]] = {} +_closed: set[str] = set() + + def _make_segment(chain: Chain, kind: str) -> tuple: """Snapshot the chain's current token state into a segment tuple (kind, prompt_ids, response_ids, loss_mask, meta) for the train loop.""" @@ -720,19 +727,24 @@ def _start_sub_chain(s: Session, dispatch_id: str) -> None: async def _handle_request(request: web.Request) -> web.StreamResponse: body = await request.json() sid = request.headers["Authorization"].removeprefix("Bearer ").strip() + if sid in _closed: # session drained; refuse stragglers + return web.Response(status=503, text="session closed") app = request.app s = app["store"].setdefault(sid, Session()) - - async with s.lock: # same sid -> serialized - target, is_sub, kind = _select_chain(s, body) - ideal_ids = _build_prompt(target, body, kind, app["tokenizer"]) - output_ids, finish = await _generate(target, ideal_ids, s, body, app) - blocks, stop, did = _build_reply(target, output_ids, finish, app) - if did and not is_sub: # sub doesn't nest - _start_sub_chain(s, did) - in_tok, out_tok = len(ideal_ids), len(output_ids) - - return await _stream_response(request, blocks, stop, in_tok, out_tok) + task = asyncio.current_task() + _inflight.setdefault(sid, set()).add(task) + try: + async with s.lock: # same sid -> serialized + target, is_sub, kind = _select_chain(s, body) + ideal_ids = _build_prompt(target, body, kind, app["tokenizer"]) + output_ids, finish = await _generate(target, ideal_ids, s, body, app) + blocks, stop, did = _build_reply(target, output_ids, finish, app) + if did and not is_sub: # sub doesn't nest + _start_sub_chain(s, did) + in_tok, out_tok = len(ideal_ids), len(output_ids) + return await _stream_response(request, blocks, stop, in_tok, out_tok) + finally: + _inflight.get(sid, set()).discard(task) async def _stream_response(request, blocks, stop_reason, in_tok, out_tok) -> web.StreamResponse: @@ -836,6 +848,22 @@ def pop_session_split(store: _Store, sid: str) -> list[tuple]: ] +async def shutdown_session(sid: str, *, wait_timeout: float = 5.0) -> None: + """Tombstone sid (late requests 503) and drain in-flight local handlers + (cancel fires /abort_request to sglang). Does NOT wait for sglang idle -- + sglang_engine.release_memory_occupation already calls flush_cache() with + 60×1s polling. Idempotent.""" + _closed.add(sid) + tasks = [t for t in _inflight.pop(sid, ()) if not t.done()] + if not tasks: + return + _, pending = await asyncio.wait(tasks, timeout=wait_timeout) + for t in pending: + t.cancel() + if pending: + await asyncio.gather(*pending, return_exceptions=True) + + # Trivial endpoints claude-code probes during a session: count_tokens runs # every turn (return 0 -- client uses it as a hint, not a hard budget), # healthz/v1/models are startup readiness checks.