Reconnect and send the observation again when the inference socket drops - #695
Reconnect and send the observation again when the inference socket drops#695v-positronic wants to merge 1 commit into
Conversation
A backend that scales to zero drops the websocket when its container recycles. The rig read that as fatal: one dropped send ended the World mid-episode and the run recorded nothing. `InferenceSession.infer` now reconnects once through the client's own cold-start retries and sends the same observation on the new socket. The handshake moves out of `InferenceSession.__init__` into `_handshake`, so `InferenceClient._open` can connect, handshake and retry as one step and serve both a session open and a reconnect. A stall keeps raising: the server may still be computing the observation, so a second send doubles the work on a backend already too slow. So does a second drop, which is a backend that cannot serve the observation at all. The reconnect gets 45s of wall clock against `connect_deadline`'s 900s, because it is spent with the arm holding its last setpoint and an attended trial carries no episode deadline behind it. The connect loop now clips its backoff sleep to what is left, so that deadline bounds the wall clock spent rather than the instant the last attempt starts at. Ticket: none — unblocks a customer rollout that died on the rig
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 9b008323c3
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| self._websocket = self._reopen() | ||
| # The server built this session a moment ago, so it holds none of the dropped one's state. A second | ||
| # drop is a backend that cannot serve this observation at all, and reaches the caller. | ||
| return self._round_trip(serialised) |
There was a problem hiding this comment.
Rule hidden-dependency violated:
InferenceSession.infer replaces the socket with a newly opened server session and immediately resends only the current observation. If any earlier inference completed, this silently discards the per-episode state that Policy.Session explicitly owns—such as trajectory buffers or model history—so stateful policies can return incorrect robot actions after a drop. Restrict this recovery to the first inference, or explicitly restore/replay session state before continuing.
AGENTS.md reference: AGENTS.md:L7-L8
Useful? React with 👍 / 👎.
| ping_interval=20.0, | ||
| ) | ||
| return InferenceSession(ws, infer_timeout=self.infer_timeout) | ||
| return ws, _handshake(ws) |
There was a problem hiding this comment.
Rule misleading-name violated:
reconnect_deadline is not actually a deadline because _open invokes connect with the full open_timeout and _handshake with an independent recurring 30-second timeout. An attempt started near the cutoff can overrun it, and a backend that keeps sending loading updates can remain in _handshake indefinitely while the live robot holds its last setpoint. Pass the remaining budget into both operations and enforce a total handshake deadline.
AGENTS.md reference: AGENTS.md:L7-L8
Useful? React with 👍 / 👎.
1 participant
A customer rollout died on the rig this morning having recorded nothing. The endpoint accepted the
websocket and answered at the application level, then dropped it while still loading; the first
inferof episode 1 found it dead and the World came down under a live arm.
A backend that scales to zero drops the socket when its container recycles, and a blind round samples
two endpoints, so whichever is dialled second cold-starts mid-run and dies the same way.
The change
InferenceSession.infertreats a dropped socket as recoverable: it reconnects through the client's owncold-start retries and sends the same observation on the new socket. The status handshake moves out of
InferenceSession.__init__into a module-level_handshake, soInferenceClient._open(deadline)canconnect, handshake and retry as one step and serve both a session open and a reconnect.
InferenceSessionkeeps its constructor shape;metadataandreopenare new keyword arguments.The connect loop now clips its backoff sleep to the budget left. Without that, a 45s deadline with
backoff at 32s returns after ~63s — the deadline bounded the instant the last attempt starts at, not
the wall clock spent.
The calls
In the client, not
remote.round_trip. The session owns the socket and is the only layer that canrebuild it;
round_tripholds anInferenceSessionand no way to get a new one, so a retry therere-sends on the same dead socket. It also sits inside the
policy.inferspan, so a reconnect reads asa slow inference rather than disappearing.
45s of wall clock, against
connect_deadline's 900s, spanning one reconnect. A session opens withthe arm parked and can afford the long deadline; a drop mid-episode is spent with the arm live, and the
harness cannot bound it — an attended droid trial is
timeout_sec=None(cfg/eval/real/droid.py), sono episode deadline stands behind it. It also bounds teardown: the runtime waits the round trip out
before the world comes down, so a 900s reconnect would make Abort take 900s. 45s allows attempts at
t=0, 1, 3, 7, 15, 31 — enough for a recycle onto a warm replacement, short of an operator watching a
still arm and reaching for Abort.
The arm holds its last setpoint meanwhile.
round_tripruns on an Executor worker, so the harnessloop keeps running and
_playkeeps emitting from the schedule in hand; once that drains no commandgoes out.
inference.waitkeeps pollingshould_stop, so Abort works throughout.One retry. The reconnected session is one the server has just built, holding none of the dropped
one's state, so re-sending is safe. A second drop on a freshly handshaked socket is a backend that
cannot serve this observation at all, and it reaches the caller.
A stall keeps raising. A recv timeout leaves the observation with a server that may still be
computing it, so a second send doubles the work on a backend already too slow.
[loading]needs no new handling._handshakealready loops on it untilready, andnew_sessionalready retries a handshake that drops. This drop happened afterready.Tests
TestInferDropsTheConnection, over a fake connection: a dropped send reconnects and serves theobservation (fails without the fix); a stall does not reconnect (fails if the fix over-reaches);
a session with no
reopenraises; a second drop reaches the caller. Both were verified by breaking thefix each way.
Five connect-path tests in
test_remote_policy.pypatchedInferenceSessionto skip the handshake;they now patch
_handshake, which is what they were stubbing.