feat(session-server): multi-backend vanilla + per-request observability

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):

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")

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,