How RL now keeps track of a finished Gym rollout until Gym is allowed to forget it
NVIDIA-NeMo/RL #4265
(stack part 2 of 3, on top of #4264).
RL links pin head ed14402 (base: ad3c9cb); Gym links pin f4fcf8c.
Before
Each Gym try got a random ID. RL kept no record of which finished Gym results it owned, and had no way to tell Gym "I have this one, you can delete your stored copy".
This PR
Gives each rollout a stable ID plus a try number, fetches a receipt from Gym for every finished row, and stores "I still owe Gym an ACK for this receipt" in the saved recovery state.
Status
New in this PR — under review. The bug in prompt-group mode is on the second page.
Nine pages on this stack — this one: #4265: how a finished Gym rollout is tracked.
What /run is, what Gym keeps, and what an ACK deletes
/run is the HTTP call RL makes to a Gym agent to play out one rollout (one sibling)
from start to finish: every model call, every tool call, and the final scoring. It is a single
blocking POST: it returns only when that rollout is done.
RL sends all siblings' /run calls at once and
handles each reply as it arrives. There is no polling.
response: the output messages and tool calls, as text
reward: the score the environment gave this rollout
mask_sample: "leave this sample out", set when the environment itself failed
a few IDs, and the original request fields
This reply is the only way the reward reaches RL. RL reads it straight from the reply
(nemo_gym.py:2053). So RL must hear back from every
/run. Without the reply, the rollout has no reward and cannot be trained on.
The reply has no tokens. The token IDs and logprobs were already written to
TQ
during the rollout, one row per model call.
Why: Gym's HTTP client re-sends the same /run if the
connection drops. If the first reply was lost after the agent finished, the re-sent call
gets the stored reply back instead of running the rollout
again. This copy lives only in memory: a Gym checkpoint
does not save it, so it does not survive a restore.
An ACK is RL saying "I have saved this result". Gym
checks the receipt matches exactly, then deletes the stored
reply and marks that ID and try as done for good. It does not touch TQ or anything RL still needs.
Dotted words have a short explanation: hover or tap them.
Gym runs each rollout as an agent. With checkpointing on, Gym keeps every finished result
until RL sends an ACK for it (a short "I have saved this result, you can delete your copy" message). Its code says so:
results "remain replayable until their receipt is acknowledged", and Gym
will not commit a checkpoint
while any finished result is still un-ACKed. So RL must ACK each result, but only after the
result is safely in RL's own saved state. Otherwise a crash between the two loses the rollout.
This PR builds that hand-off.
One rollout, before and after
The same three siblings, step by step
Where the two paths are the same: the row still goes through
run_rollouts to Gym's /run, and RL still seals it with
mark_sibling_sealed.
What is new sits on either side of that seal: the receipt before it, and the ACK debt after it.
before (#4264) after (#4265)
ID sent for g2, try 1 "p7_g2_a9c41…" (random) "p7_g2" + attempt 1 → Gym key "p7_g2-a1"
g2 retried again new random ID attempt 2 → "p7_g2-a2"
g0 finishes seal g0 GET receipt, then in one locked step:
seal g0 + debt {("p7_g0", 0)}
saved recovery state sealed siblings only sealed siblings + ACK debts (schema v3)
ACK to Gym none — no ACK pathPOST /acknowledge, then remove the debt
Gym prepare 409 while any result un-ACKed same rule; now RL can clear it
The ACK debt is written in the same locked step as the seal
(rollout_manager.py:2445-2465),
and the network send happens after the lock is released. That order is required: if RL crashed after
sending the ACK but before saving, Gym would have deleted its copy of a result RL never saved. Debts are removed only
after Gym accepts the ACK
(mark_completed_executions_acknowledged).
Two things called "receipt", and the order RL does things
Completion receipt
Rollout receipt
Made by
Gym, for one finished run
RL, from Gym's token-capture list for that run
Holds
ID, try number, a hash of the /run reply, which token list to read, the last model call
which TQ rows make up this rollout, which call was the last one, the reward
Used for
the body of the ACK (and picking the right token list)
By step 5, RL already has everything it needs. The ACK only lets Gym free its copy.
So what. After this PR, a finished Gym rollout has one owner at every moment: Gym until
RL has saved it, and RL after. A retry never reuses an old try number. The one path that breaks this
is prompt-group mode, where RL waits for the whole group before recording any debt. See
page 2.