nemo_gym.server_runtime.request_handling.session_routing
nemo_gym.server_runtime.request_handling.session_routing
Route each session’s requests to the uvicorn worker that holds its state.
uvicorn’s workers share one listening socket (unless uvicorn_reuse_port gives each its own),
so a request reaches whichever worker accepts the connection,
while a session’s state lives in the memory of its current owner.
With more than one worker, each worker also serves the same app on a private Unix socket,
in a directory the main process creates for the server, and requests are routed by one rule:
- Find the request’s session ID.
- Look up the session’s owner, the worker that currently holds the session.
- Forward the request to that worker’s private socket, unchanged, and stream its reply back unchanged. If this worker owns the session, or the session has no owner, the request is handled here.
A request that arrives over a private socket is always handled where it arrives, so it is forwarded at most once. If the owner’s socket is gone, because its worker exited, the request gets a 410 instead of silently running against empty state.
Where the session ID comes from:
The body is read only on the seed and close routes, where it names the session before any cookie exists.
If the cookie and the body name different sessions, the reply is 409.
An episode’s resources session uses its resources_session_id as its cookie’s session_id too,
so the seed, its close and every request in between name the session by the same ID.
A resources request with no session ID gets a new ID, which this router mints.
A session start, such as a seed, takes ownership of the session before its handler runs,
since its reply’s cookie names the session.
Any other request takes ownership only if the handler used the session, before the reply leaves the worker.
A request that keeps nothing, such as a health check or /aggregate_metrics, takes no ownership and gets no cookie,
so no cookie names a session nobody owns.
A request whose cookie names a session whose ownership ended makes the worker that serves it the owner the same way.
Ownership is recorded in an owner file in the socket directory, named by a hash of the session ID and holding the owner’s worker ID. The worker that creates the session creates the owner file, atomically, so two workers can never own one ID. Seeds and closes take ownership, so a retried seed, a cleanup close, or a delayed seed after a close all reach the worker that saw the ID first. Ownership ends once its session’s last record on the owner is gone:
- an agent session’s ownership ends when the agent drops the session’s last record,
which is its fence or failed seed,
session_close_retry_window_secondsafter it was written; - a resources session’s ownership lasts while the session is open,
from a successful seed until its close or
/verify, and expiressession_close_retry_window_secondsafter the last request for it otherwise. A/verifythat says a close follows, as an environment server’s does, leaves the session open for that close.
An agent worker also owns /rollouts/<rollout id> from the start of a session’s seed until the session closes,
so a rollout stop, which names only the rollout, reaches the worker that runs it.
That request is forwarded to the key’s owner but never takes ownership of the key,
so a stop for a rollout no worker runs is handled where it arrives and leaves no owner file.
The owner files of a worker that exited stay until the server exits, so its sessions keep getting a 410. The whole directory goes with the server’s runtime directory when the server exits. With one worker there is no private socket and no owner.
A request without a session ID that belongs to one worker, such as an agent’s call to itself within one /run,
names that worker in the x-ng-session-owner header, and is forwarded there, or gets a 410 if that worker is gone.
A partial-rollout checkpoint restore installs each restored session on one live worker,
which rewrites the session’s owner file to name itself (see SessionOwners.rewrite).
A cookie or MCP token from before the restore names only the session ID, so its requests reach the new holder.
Module Contents
Classes
Functions
Data
_PRIVATE_SERVER_KEEPALIVE_SECONDS
API
An agent server, whose sessions are owned by agent_session_id across workers.
Drop the fences and failed seeds that expired, releasing ownership of their sessions.
Whether this worker still keeps any record of the session, open, failed or closed.
Keep the session owners, to release this worker’s ownership when it drops a session’s last record.
How long a resources server’s worker keeps ownership of a session, judged from the requests it served.
A resources server keeps its per-session state in its own subclass, so the router judges by the requests:
- A successful seed, or another declared session start such as a Gymnasium
/reset, opens the session: its ownership lasts until the session ends. - A successful
/close_sessionor/verify, orSimpleResourcesServer.checkpoint_session_ended, ends it, unless the request says a close follows, as an environment server’s/verifydoes. - A session that is not open keeps its owner
close_retry_window_secondsafter the last request for it. The caller’s retried or cleanup close reaches the owner within that window. A delayed seed after a close reaches it too, where the server refuses it while the session is fenced. - A failed seed or close of an open session leaves it open.
So ownership outlives the session’s last request by at most the close retry window, except for a session never closed, which keeps its owner as long as the server keeps the session.
End a session this worker owns, from a route that ends a session under another name.
Drop a session whose install a failed restore undid, so its ownership can be released at once.
Open a session a checkpoint restore installed on this worker, as a successful seed would.
Keep ownership of a session a checkpoint retire released until ended, when the retire’s refusal ends.
Until then a late request that carries only the session’s cookie reaches this worker, which refuses it.
Where one kind of server’s requests name their session.
How long a worker keeps ownership of the sessions it created.
Whether this worker still keeps any record of the session.
Record a request for a session this worker owns, once it is served; status is None if it failed.
Whether a request on path starts a session, so it takes ownership whatever its handler does.
Drop the records whose time has passed, and return the sessions whose ownership can be released.
Which worker currently owns each session, shared by the workers of one server.
Each session’s owner is recorded in an owner file in the owners directory, named by a hash of the session ID and holding the owner’s worker ID. It is created with a hard link from a complete temporary file, which fails if the owner file exists, so two workers cannot both own one ID and a reader never sees a partly written owner file.
A worker owns a session as long as holds says it keeps anything of the session.
It does not release ownership while a request for the session is being served here,
since that request may be about to create a record.
Only the owner removes its owner file, so a removal cannot delete an owner file another worker wrote since.
A checkpoint restore can rewrite an owner file for another worker at any time, though.
The removal’s check of the owner and its unlink, and every rewrite, hold one lock file shared by the workers,
so a removal never deletes an owner file a rewrite has just written.
Hold the workers’ owners lock; it is held only for a few file operations, never across an await.
Replace the owner file of each ID with one naming this worker, and append each ID to replaced once done.
It touches only files, so it can run in a thread.
Make this worker the ID’s owner unless another worker already owns it; return the owner.
Returns None for an owner file that names no valid worker, which nothing here writes.
Stop owning the ID here, and leave its owner file to the worker a checkpoint restore moved it to.
The worker that owns the ID, or None if nobody does.
Release this worker’s ownership of the ID, unless a request for it is in progress or the session remains.
Make this worker the ID’s owner whoever owns it now, for a session a checkpoint restore installs here.
The owner file is replaced by one rename, so a reader finds either the old owner or this one.
A live worker that owned the ID before must disown it, since a worker trusts its own ownership first.
rewrite each ID, in a thread, so a restore of many sessions does not hold up the event loop.
The thread only replaces files, and holds the workers’ owners lock once per _REWRITE_BATCH renames.
The IDs this worker owns change only on the event loop, once the thread has ended,
even if the caller was cancelled meanwhile.
They then name exactly the owner files the thread rewrote, so a release after a failure misses none of them.
Until then this worker does not count them as its own, so a release here leaves them alone.
The caller holds the sessions (see holds) from before this call, so no release afterwards removes them.
Keep ownership while a request for the ID is served here, and release it after if nothing is left.
Find a request’s session ID, and forward the request to the worker that owns the session.
It sits outside SessionMiddleware and reads the cookie by decoding it with the same secret. A forwarded request therefore never reaches this worker’s SessionMiddleware, so the reply carries only the owner’s Set-Cookie, built from the session as the owner left it. Sitting inside SessionMiddleware instead would add this worker’s stale copy of the session as a second Set-Cookie.
With one worker there is no ownership, and it only reads the session ID of seeds and closes: it refuses a request whose cookie and body name different sessions, and gives an episode’s resources session its ID in the cookie.
Serve a request for a session this worker owns, or may take, and record it for the ownership’s lifetime.
An agent’s own records decide its ownership: it releases a session when it drops the session’s last record.
Bases: HttpToolsProtocol
uvicorn’s httptools protocol without access logs, which it would otherwise write whenever the logger has handlers.
Match any of paths, bare or under a rollout prefix.
Read the request body, or None if the client disconnected.
A receive that gives the app the body already read, then passes later messages, such as a disconnect, on.
The session ID in field of a seed or close body, or None for a body that has no non-empty string one.
Release ownership of the sessions that have ended, so ownership ends even when no more requests arrive.
Return the session stored in a session cookie, or None if Starlette’s SessionMiddleware would reject it.
Decoded exactly as SessionMiddleware decodes it. A cookie it rejects (badly signed or expired) starts a fresh session.
The directory for one server’s private worker sockets and owner files: this process’s runtime directory.
Route this worker’s requests by session ID, and serve forwarded requests on a private socket.
Call once the app’s middleware is otherwise in place: the router must sit outside SessionMiddleware.
Without a socket_dir (one worker), it only reads the session ID of seeds and closes, and owns nothing.
With one, it records session owners in it and returns this worker’s ID,
also stored in app.state.nemo_gym_routing_id.
Pass the agent for an agent server, whose own records decide how long it owns its sessions,
or the close retry window and the session start and end routes for a resources server
(see ResourcesSessionLifetimes).
Return the session ID in an MCP session token, or None for a token with a bad signature.
A rejected token is handled locally, where the MCP endpoint’s own token check refuses it.
Serve app on a Unix socket for requests other workers forward here, on the running event loop.