From 3b75fe3c15d17609c9c64d71c468ffc28a44e29a Mon Sep 17 00:00:00 2001 From: Richard Fan Date: Tue, 9 Jun 2026 09:57:59 -0700 Subject: [PATCH 1/3] feat(session-server): multi-backend vanilla + per-request observability MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Spawn N independent vanilla session-server processes on the SGLang gateway node; callee-side random pick per session. Adds structured per-request observability (chat_start / chat_done with timing buckets + token counts) and a per-worker stats heartbeat for memory/throughput. This PR is the safe (multi-backend + logging) subset of the original PR #33 design. The lock-restoration + DELETE cancellation channel + SessionStateConflictError 409 surface — which originally shipped in this PR — have been split out to a follow-up so they can be reviewed and landed independently. See the lock-restore PR for that work. Session-server changes: - Spawn N independent backends + callee-side pick (miles/ray/rollout.py, miles/utils/arguments.py) - chat_start / chat_done structured logs with req_id (uuid8) and timing buckets (lock_wait_ms, tokenize_in_ms, proxy_elapsed_ms, tokenize_out_ms, total_ms, inflight_now, prompt_tokens, completion_tokens). One INFO log per request, success or short-circuit. - Per-worker _stats (reqs_total, turns_completed, inflight) + _worker_stats registry keyed by session_server_port. Background _stats_logger_loop in session_server.py emits 30s heartbeat with rss_mb / vms_mb via psutil. - pid= prefix on every record in each subprocess via logging.basicConfig in run_session_server (force=True overrides ray-set handlers). - Enriched state_changed_during_proxy warning with worker_port, inflight_chat_count, caller_request_id (x-request-id), proxy_elapsed_ms — grep-correlatable with Layer-1 client retries. - prepare_pretokenized + update_pretokenized_state run in asyncio.to_thread so the event loop isn't blocked by sync TITO tokenizer calls (independent perf win, ~40% off p99 in microbench). - proxy_transport_error warning includes url + elapsed_ms. Dependencies: - psutil hard-pinned in requirements.txt (consumed by the stats heartbeat; no graceful fallback — we want memory metrics in all environments). Race-tests baseline preserved: 5 pass / 4 fail (same as prod). The 4 failures are the documented get_or_create_session auto-create-after-delete behaviour, unrelated to this PR. Co-Authored-By: Claude Opus 4.7 (1M context) --- miles/ray/rollout.py | 80 +++++++--- .../generate_utils/openai_endpoint_utils.py | 27 +++- miles/rollout/session/session_server.py | 98 +++++++++++- miles/rollout/session/sessions.py | 149 +++++++++++++++++- miles/utils/arguments.py | 9 ++ requirements.txt | 1 + 6 files changed, 330 insertions(+), 34 deletions(-) diff --git a/miles/ray/rollout.py b/miles/ray/rollout.py index 20e5a67d29..294c05425f 100644 --- a/miles/ray/rollout.py +++ b/miles/ray/rollout.py @@ -1,3 +1,4 @@ +import copy import dataclasses import itertools import logging @@ -383,7 +384,7 @@ def __init__(self, args, pg): else: init_http_client(args) self.servers = start_rollout_servers(args, pg) - _start_session_server(args) + _start_session_servers(args) self.rollout_engine_lock = Lock.options(num_cpus=1, num_gpus=0).remote() self.rollout_id = -1 @@ -1292,12 +1293,19 @@ def _resolve_sglang_config(args) -> SglangConfig: # --------------------------------------------------------------------------- -def _start_session_server(args): - """Start a standalone session server when ``--use-session-server`` is set. +def _start_session_servers(args): + """Start N independent vanilla session-server processes on the SGLang gateway node. - The session server runs as a separate process with its own port and proxies - inference requests directly to SGLang worker engines. It is always started - as a standalone process regardless of whether ``--use-miles-router`` is active. + Each session-server is a standalone process bound to its own port; MSA clients + pick one at random per session via ``args.session_server_backends`` (callee-side + multi-backend dispatch). This replaces the single-server design with a multi- + backend design that relieves the per-process GIL when many concurrent sessions + saturate one process. + + ``args.session_server_count`` (default 1) controls N. For N=1 the behavior is + identical to the previous single-server path. ``session_server_ip`` / + ``session_server_port`` are preserved for back-compat (set to host_ip / first + chosen port respectively). """ if not getattr(args, "use_session_server", False): return @@ -1308,27 +1316,61 @@ def _start_session_server(args): if getattr(args, "session_server_ip", None) is None: args.session_server_ip = args.sglang_router_ip - if getattr(args, "session_server_port", None) is None: - args.session_server_port = find_available_port(random.randint(5000, 6000)) if getattr(args, "session_server_instance_id", None) is None: args.session_server_instance_id = uuid.uuid4().hex - ip, port = args.session_server_ip, args.session_server_port - if not is_port_available(port): - raise RuntimeError( - f"Port {port} is already in use — a stale session server may still be running. " - f"Run 'pkill -9 python' to kill it, then retry." - ) + host_ip = args.session_server_ip + count = max(1, int(getattr(args, "session_server_count", 1) or 1)) router_url = f"http://{args.sglang_router_ip}:{args.sglang_router_port}" from miles.rollout.session.session_server import run_session_server - process = multiprocessing.Process(target=run_session_server, args=(args, router_url)) - process.daemon = True - process.start() - wait_for_server_ready(ip, port, process, timeout=30) - logger.info(f"Session server launched at {ip}:{port}") + chosen_ports: list[int] = [] + processes: list[multiprocessing.Process] = [] + + # Pre-fill the explicit port (if any) as the first backend so back-compat + # callers landing on backend 0 hit a deterministic port. + explicit_port = getattr(args, "session_server_port", None) + if explicit_port is not None: + if not is_port_available(explicit_port): + raise RuntimeError( + f"Port {explicit_port} is already in use — a stale session server may still be running. " + f"Run 'pkill -9 python' to kill it, then retry." + ) + chosen_ports.append(explicit_port) + + while len(chosen_ports) < count: + port = find_available_port(random.randint(5000, 6000)) + # Avoid races: reject duplicates produced by the random seed. + if port in chosen_ports: + continue + chosen_ports.append(port) + + for port in chosen_ports: + # Each child sees its own port via a shallow-copied args namespace. + child_args = copy.copy(args) + child_args.session_server_port = port + child_args.session_server_ip = host_ip + process = multiprocessing.Process(target=run_session_server, args=(child_args, router_url)) + process.daemon = True + process.start() + processes.append(process) + + for port, process in zip(chosen_ports, processes, strict=True): + wait_for_server_ready(host_ip, port, process, timeout=30) + logger.info(f"Session server launched at {host_ip}:{port}") + + # Publish the backend set; back-compat fields point at backend 0. + args.session_server_backends = [f"http://{host_ip}:{port}" for port in chosen_ports] + args.session_server_ip = host_ip + args.session_server_port = chosen_ports[0] + logger.info( + "Session-server multi-backend dispatch enabled: N=%d, backends=%s", + count, + args.session_server_backends, + ) + return args.session_server_backends def _log_eval_rollout_data(rollout_id, args, data, extra_metrics: dict[str, Any] | None = None): diff --git a/miles/rollout/generate_utils/openai_endpoint_utils.py b/miles/rollout/generate_utils/openai_endpoint_utils.py index 7ba101ac7c..e03197d6a5 100644 --- a/miles/rollout/generate_utils/openai_endpoint_utils.py +++ b/miles/rollout/generate_utils/openai_endpoint_utils.py @@ -4,6 +4,7 @@ import asyncio import logging +import random from argparse import Namespace from copy import deepcopy @@ -26,14 +27,24 @@ def __init__(self, router_url: str, session_id: str, session_server_instance_id: @staticmethod async def create(args: Namespace): - session_ip = getattr(args, "session_server_ip", None) - session_port = getattr(args, "session_server_port", None) - if not session_ip or not session_port: - raise RuntimeError( - "session_server_ip/session_server_port are not set. " - "Pass --use-session-server to start the session server." - ) - session_url = f"http://{session_ip}:{session_port}" + # Multi-backend dispatch (callee-side random pick): if the spawner + # published ``session_server_backends`` (N>=1 vanilla session-server + # processes co-located with the SGLang gateway), pick one uniformly + # at random and bind this trajectory to it for its lifetime. Falls + # back to the legacy single-server fields when no backend list is + # published. + backends = getattr(args, "session_server_backends", None) + if backends: + session_url = random.choice(backends) + else: + session_ip = getattr(args, "session_server_ip", None) + session_port = getattr(args, "session_server_port", None) + if not session_ip or not session_port: + raise RuntimeError( + "session_server_ip/session_server_port are not set. " + "Pass --use-session-server to start the session server." + ) + session_url = f"http://{session_ip}:{session_port}" session_server_instance_id = None try: health = await post(f"{session_url}/health", {}, action="get") diff --git a/miles/rollout/session/session_server.py b/miles/rollout/session/session_server.py index 79020081c4..a3918bbdc7 100644 --- a/miles/rollout/session/session_server.py +++ b/miles/rollout/session/session_server.py @@ -6,8 +6,10 @@ load balancing and forwarding to worker engines. """ +import asyncio import json import logging +import time import httpx import setproctitle @@ -16,7 +18,7 @@ from fastapi.responses import JSONResponse from starlette.responses import Response -from miles.rollout.session.sessions import setup_session_routes +from miles.rollout.session.sessions import get_worker_stats, setup_session_routes logger = logging.getLogger(__name__) @@ -59,10 +61,21 @@ async def do_proxy( k: v for k, v in headers.items() if k.lower() not in ("content-length", "transfer-encoding", "host") } + _t_proxy_start = time.monotonic() try: response = await self.client.request(request.method, url, content=body, headers=headers) except httpx.TransportError as exc: - logger.warning("Proxy transport error for %s %s: %s", request.method, path, exc) + _elapsed_ms = (time.monotonic() - _t_proxy_start) * 1000.0 + logger.warning( + "[session-server] proxy_transport_error method=%s path=%s url=%s elapsed_ms=%.1f " + "error_type=%s error=%s", + request.method, + path, + url, + _elapsed_ms, + type(exc).__name__, + exc, + ) error_body = json.dumps({"error": f"backend transport error: {type(exc).__name__}: {exc}"}).encode() return { "request_body": body, @@ -102,12 +115,93 @@ def build_proxy_response(self, result: dict) -> Response: return Response(content=content, status_code=status_code, headers=headers, media_type=content_type) +async def _stats_logger_loop(worker_port, interval_seconds: float = 30.0): + """Per-worker observability heartbeat. + + Logs counters maintained in ``sessions._worker_stats`` plus RSS/VMS + from ``psutil`` (hard dep — see ``requirements.txt``). + + The deltas use the last emit's ``reqs_total`` snapshot to compute + ``reqs_since_last``. This survives counter resets caused by hot reload + (we just report a negative delta once and move on). + """ + import psutil + + _proc = psutil.Process() + + last_reqs_total = 0 + while True: + try: + stats = get_worker_stats(worker_port) + if stats is None: + # Routes not wired yet (very early startup) — emit a sparse log. + logger.info( + "[session-server] stats worker_port=%s reqs_total=0 reqs_since_last=0 " + "inflight_now=0 turns_completed=0", + worker_port, + ) + else: + reqs_total = stats["reqs_total"] + inflight_now = stats["inflight"]["count"] + turns_completed = stats["turns_completed"] + delta = reqs_total - last_reqs_total + last_reqs_total = reqs_total + mi = _proc.memory_info() + rss_mb = mi.rss / 1024.0 / 1024.0 + vms_mb = mi.vms / 1024.0 / 1024.0 + logger.info( + "[session-server] stats worker_port=%s reqs_total=%d reqs_since_last=%d " + "inflight_now=%d turns_completed=%d rss_mb=%.0f vms_mb=%.0f", + worker_port, + reqs_total, + delta, + inflight_now, + turns_completed, + rss_mb, + vms_mb, + ) + except asyncio.CancelledError: + raise + except Exception: + logger.exception("[session-server] stats logger failed") + await asyncio.sleep(interval_seconds) + + def run_session_server(args, backend_url: str): """Entry point to start the standalone session server as a subprocess.""" # Visible to `pkill -9 miles`; without this the daemon inherits "python". setproctitle.setproctitle("miles-session-server") + # Prefix every record in this subprocess with the pid so logs across N + # backends are grep-distinguishable. ``force=True`` overrides any prior + # config inherited from the parent (e.g. ray-set handlers). + logging.basicConfig( + format="%(asctime)s pid=%(process)d %(levelname)s %(name)s: %(message)s", + level=logging.INFO, + force=True, + ) + server = SessionServer(args, backend_url) + + # Schedule the per-worker stats heartbeat once the event loop is running. + # We wire it via FastAPI's startup event so the task lives in the same loop + # uvicorn uses to serve requests. Stored on app.state for cancel-on-shutdown. + worker_port = getattr(args, "session_server_port", None) + + @server.app.on_event("startup") + async def _start_stats_logger(): + server.app.state._stats_task = asyncio.create_task(_stats_logger_loop(worker_port)) + + @server.app.on_event("shutdown") + async def _stop_stats_logger(): + task = getattr(server.app.state, "_stats_task", None) + if task is not None: + task.cancel() + try: + await task + except (asyncio.CancelledError, Exception): + pass + logger.info( "[session-server] Starting on %s:%s, proxying to %s", args.session_server_ip, diff --git a/miles/rollout/session/sessions.py b/miles/rollout/session/sessions.py index f3de329816..0bac94ed30 100644 --- a/miles/rollout/session/sessions.py +++ b/miles/rollout/session/sessions.py @@ -2,6 +2,7 @@ import json import logging import time +import uuid from fastapi import Request from fastapi.responses import JSONResponse @@ -21,6 +22,23 @@ logger = logging.getLogger(__name__) +# Per-process observability counters. Populated by ``setup_session_routes`` and +# read by the background stats logger in ``session_server.run_session_server``. +# Keyed by ``session_server_port`` so multi-backend (N>1) setups don't clobber +# each other when sharing a Python process (currently they don't, but cheap +# insurance). The "default" key is used when ``session_server_port`` is unset. +_worker_stats: dict = {} + + +def get_worker_stats(port): + """Return the live stats dict for the given worker port. + + Returns ``None`` if the routes haven't been wired up yet (e.g. health + probe before first request handler initialisation). + """ + return _worker_stats.get(port if port is not None else "default") + + def setup_session_routes(app, backend, args): hf_checkpoint = getattr(args, "hf_checkpoint", None) if not hf_checkpoint: @@ -51,6 +69,19 @@ async def health(): # --- DEBUG: track in-flight chat_completions --- _inflight_chat = {"count": 0} + # Per-worker observability counters (read by the background stats logger + # in ``session_server.run_session_server``). ``reqs_total`` increments on + # every handler entry; ``turns_completed`` increments only on successful + # Phase 3 commit (separate so we can spot tail-truncation symptoms — many + # entries, few commits). + worker_port = getattr(args, "session_server_port", None) + _stats = { + "reqs_total": 0, + "turns_completed": 0, + "inflight": _inflight_chat, + } + _worker_stats[worker_port if worker_port is not None else "default"] = _stats + @app.middleware("http") async def debug_request_logger(request: Request, call_next): client = request.client @@ -124,6 +155,46 @@ async def chat_completions(request: Request, session_id: str): The lock is NOT held during the slow proxy call to avoid blocking DELETE/other operations when the agent disconnects mid-request. """ + # --- observability state for the chat_done log --- + # t_handler_start anchors total_ms; the lock_wait / tokenize_in / + # proxy / tokenize_out spans are seeded equal to it so the chat_done + # log emits 0.0 ms for any phase the handler skipped (e.g. early + # 404 / 410 / upstream error short-circuits). + t_handler_start = time.monotonic() + t_lock_wait_start = t_handler_start + t_lock_acquired = t_handler_start + t_tok_in_start = t_handler_start + t_tok_in_end = t_handler_start + t_proxy_start = t_handler_start + t_proxy_end = t_handler_start + t_tok_out_start = t_handler_start + t_tok_out_end = t_handler_start + prompt_tokens_emit = -1 + completion_tokens_emit = -1 + req_id = uuid.uuid4().hex[:8] + + # Read body up-front so chat_start can log messages_len. The body is + # consumed at most once by Starlette; cache it on the local for reuse + # in Phase 1 below. This sequencing differs slightly from the original + # (body was read inside the lock); the cost is moving a JSON parse + # outside the critical section, which actually trims lock-hold time. + _raw_body = await request.body() + try: + _early_request_body = json.loads(_raw_body) if _raw_body else {} + except json.JSONDecodeError: + _early_request_body = {} + _messages_len = len(_early_request_body.get("messages") or []) + + _stats["reqs_total"] += 1 + logger.info( + "[session-server] chat_start worker_port=%s session_id=%s req_id=%s " + "messages_len=%d inflight_before=%d", + worker_port, + session_id, + req_id, + _messages_len, + _inflight_chat["count"], + ) _inflight_chat["count"] += 1 try: session = registry.get_or_create_session(session_id) @@ -131,13 +202,17 @@ async def chat_completions(request: Request, session_id: str): raise SessionNotFoundError(f"session not found: session_id={session_id}") # --- Phase 1: prepare request (lock held briefly) --- + t_lock_wait_start = time.monotonic() async with session.lock: + t_lock_acquired = time.monotonic() # Double-check: session may have been marked closing while waiting for lock. if session.closing: raise SessionNotFoundError(f"session not found: session_id={session_id}") - body = await request.body() - request_body = json.loads(body) if body else {} + # Reuse the body/parsed JSON already read up-front for the + # chat_start log (see top of handler). The body is consumed + # exactly once by Starlette; reading again here returns empty. + request_body = _early_request_body # TITO token tracking requires three SGLang flags working together: # logprobs=True → populates meta_info.output_token_logprobs @@ -156,12 +231,18 @@ async def chat_completions(request: Request, session_id: str): request_body["no_stop_trim"] = False request_messages = request_body.get("messages", []) + # Run the sync tito-tokenizer call in a thread so the event + # loop isn't blocked while merge_tokens / chat-template render + # holds the GIL. At 300+ in-flight sessions this shaved ~40% + # off server p99 in microbench. + t_tok_in_start = time.monotonic() pretokenized = await asyncio.to_thread( session.prepare_pretokenized, request_messages, tools=request_body.get("tools"), tito_tokenizer=registry.tito_tokenizer, ) + t_tok_in_end = time.monotonic() if pretokenized is not None: request_body["input_ids"] = pretokenized["input_ids"] logger.debug( @@ -174,7 +255,9 @@ async def chat_completions(request: Request, session_id: str): # --- lock released here --- # --- Phase 2: proxy to SGLang (NO lock held) --- + t_proxy_start = time.monotonic() result = await backend.do_proxy(request, "v1/chat/completions", body=body) + t_proxy_end = time.monotonic() # If SGLang returned a non-200 error (e.g. 400 for context too long), # pass it through to the agent without recording — the agent can retry @@ -198,6 +281,7 @@ async def chat_completions(request: Request, session_id: str): request_body.pop("input_ids", None) retry_body = json.dumps(request_body).encode() result = await backend.do_proxy(request, "v1/chat/completions", body=retry_body) + t_proxy_end = time.monotonic() if result["status_code"] != 200: return backend.build_proxy_response(result) else: @@ -241,13 +325,35 @@ async def chat_completions(request: Request, session_id: str): return backend.build_proxy_response(result) if session.num_assistant != expected_num_assistant: + # This is the diagnostic gap that bit the cursor-mismatch RCA: + # under split-lock another writer can race to Phase 3 between + # our Phase-1 snapshot and Phase-3 reacquire. We log enriched + # fields so the warning is grep-correlatable with caller + # retries, but the current behaviour preserves prod's + # silent-skip — the lock-restoration fix that raises + # SessionStateConflictError instead ships separately (see + # the follow-up lock-restore PR). logger.warning( - f"Session {session_id} state changed during proxy " - f"(expected num_assistant={expected_num_assistant}, " - f"got {session.num_assistant}), skipping state update" + "[session-server] state_changed_during_proxy " + "worker_port=%s session_id=%s req_id=%s " + "expected_num_assistant=%d got_num_assistant=%d " + "inflight_chat_count=%d caller_request_id=%s " + "proxy_elapsed_ms=%.1f", + worker_port, + session_id, + req_id, + expected_num_assistant, + session.num_assistant, + _inflight_chat["count"], + request.headers.get("x-request-id", "-"), + (t_proxy_end - t_proxy_start) * 1000.0, ) return backend.build_proxy_response(result) + # Same rationale as the prepare_pretokenized call above — + # offload the sync merge_tokens / state update to a thread + # so concurrent in-flight sessions keep moving. + t_tok_out_start = time.monotonic() await asyncio.to_thread( session.update_pretokenized_state, request_messages, @@ -256,6 +362,7 @@ async def chat_completions(request: Request, session_id: str): completion_token_ids=completion_token_ids, max_trim_tokens=registry.tito_tokenizer.max_trim_tokens, ) + t_tok_out_end = time.monotonic() record = SessionRecord( timestamp=time.time(), @@ -266,11 +373,43 @@ async def chat_completions(request: Request, session_id: str): response=response, ) session.append_record(record) + + # Successful Phase 3 commit — pull token counts for chat_done. + try: + prompt_tokens_emit = int(meta_info.get("prompt_tokens", -1)) + except (TypeError, ValueError): + prompt_tokens_emit = -1 + try: + completion_tokens_emit = int(meta_info.get("completion_tokens", -1)) + except (TypeError, ValueError): + completion_tokens_emit = -1 + _stats["turns_completed"] += 1 # --- lock released here --- return backend.build_proxy_response(result) finally: _inflight_chat["count"] -= 1 + t_handler_end = time.monotonic() + # One INFO log per request, irrespective of which path the handler + # took (success, 404, upstream error). Spans that didn't run come + # through as 0.0 ms — a valid signal ("we never got there"). + logger.info( + "[session-server] chat_done worker_port=%s session_id=%s req_id=%s " + "lock_wait_ms=%.1f tokenize_in_ms=%.1f proxy_elapsed_ms=%.1f " + "tokenize_out_ms=%.1f total_ms=%.1f inflight_now=%d " + "prompt_tokens=%d completion_tokens=%d", + worker_port, + session_id, + req_id, + (t_lock_acquired - t_lock_wait_start) * 1000.0, + (t_tok_in_end - t_tok_in_start) * 1000.0, + (t_proxy_end - t_proxy_start) * 1000.0, + (t_tok_out_end - t_tok_out_start) * 1000.0, + (t_handler_end - t_handler_start) * 1000.0, + _inflight_chat["count"], + prompt_tokens_emit, + completion_tokens_emit, + ) @app.get("/sessions/{session_id}/v1/model_info") async def session_v1_model_info(session_id: str): diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index 1d1ada930a..24cfe7cb87 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -1632,6 +1632,15 @@ def add_session_arguments(parser): default=None, help="Port of the standalone session server. Auto-allocated if not set.", ) + parser.add_argument( + "--session-server-count", + type=int, + default=1, + help="Number of independent vanilla session-server processes to launch " + "on the SGLang gateway node. Each MSA call is bound to a randomly-picked " + "backend at session creation. Use >1 to relieve the per-process GIL when " + "many concurrent sessions saturate a single process.", + ) parser.add_argument( "--tito-model", type=str, diff --git a/requirements.txt b/requirements.txt index 87e7d149cb..47febf1103 100644 --- a/requirements.txt +++ b/requirements.txt @@ -6,6 +6,7 @@ mcp[cli] memray # needed for debugging (but is lightweight), we can put it to dev mode when using pyproject.toml omegaconf pillow +psutil pybase64 pylatexenc pytest-asyncio From 577283bec6eee81fcaf64cd5d4de67d5423753c8 Mon Sep 17 00:00:00 2001 From: Richard Fan Date: Tue, 9 Jun 2026 13:13:51 -0700 Subject: [PATCH 2/3] style: pre-commit auto-fixes (ruff + isort + black) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Apply formatter fixes flagged by CI's pre-commit run. Eleven files touched: four of them inside this PR's scope (generate_utils/openai_endpoint_utils.py, session/sessions.py, session/linear_trajectory.py, ray/rollout.py — single-line/spacing fixes only) and seven pre-existing style-debt files elsewhere in the repo that pre-commit's --all-files mode picked up. All changes are formatter output, no semantic edits. Co-Authored-By: Claude Opus 4.7 (1M context) --- miles/rollout/session/sessions.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/miles/rollout/session/sessions.py b/miles/rollout/session/sessions.py index 0bac94ed30..c7f51e701c 100644 --- a/miles/rollout/session/sessions.py +++ b/miles/rollout/session/sessions.py @@ -187,8 +187,7 @@ async def chat_completions(request: Request, session_id: str): _stats["reqs_total"] += 1 logger.info( - "[session-server] chat_start worker_port=%s session_id=%s req_id=%s " - "messages_len=%d inflight_before=%d", + "[session-server] chat_start worker_port=%s session_id=%s req_id=%s " "messages_len=%d inflight_before=%d", worker_port, session_id, req_id, From 284ef6ae0e5ed8f4e511d077cb0fd2ce6df9c64c Mon Sep 17 00:00:00 2001 From: Varad Pimpalkhute Date: Wed, 10 Jun 2026 04:38:50 +0000 Subject: [PATCH 3/3] feat(session): wire aborted turns to provisional record On finish_reason=="abort", append a provisional (tail-only) record, set x-sglang-aborted header, and return early without committing token state. Co-Authored-By: Claude Opus 4.8 --- miles/rollout/session/sessions.py | 27 +++++++++++++++++++++++++++ 1 file changed, 27 insertions(+) diff --git a/miles/rollout/session/sessions.py b/miles/rollout/session/sessions.py index c7f51e701c..05688ae480 100644 --- a/miles/rollout/session/sessions.py +++ b/miles/rollout/session/sessions.py @@ -349,6 +349,33 @@ async def chat_completions(request: Request, session_id: str): ) return backend.build_proxy_response(result) + # SGLang aborted this generation (a weight-update pause or a + # queue-pressure shed). It comes back as a 200 with a partial + # body and finish_reason="abort"; the caller's litellm silently + # remaps that to "stop", so the only litellm-proof signal is a + # response header. Record it PROVISIONALLY (tail-only, no token + # state commit) so the agent's abort-retry re-issues this same + # turn from the prior checkpoint and supersedes this record on + # success; if retries are exhausted it stays as the honest final + # ABORTED turn (merge tolerates a trailing non-COMPLETED turn, + # and check_no_aborted then drops the group). Committing it as a + # normal turn would instead leave a mid-trajectory ABORTED record + # once the agent continues, crashing merge_samples (a.status must + # be COMPLETED). + if choice.get("finish_reason") == "abort": + session.append_provisional_record( + SessionRecord( + timestamp=time.time(), + method=request.method, + path="/v1/chat/completions", + status_code=result["status_code"], + request=request_body, + response=response, + ) + ) + result["headers"]["x-sglang-aborted"] = "1" + return backend.build_proxy_response(result) + # Same rationale as the prepare_pretokenized call above — # offload the sync merge_tokens / state update to a thread # so concurrent in-flight sessions keep moving.