fix(v1,serve): env-worker memory — step slimming, registry eviction, intercept discard, worker trim - #1608
Conversation
Each v1 trajectory step stored serializable(response) wholesale: its message content restates the step's completion and its message.tokens restate the step's tokens (full-prefix prompt_ids/mask + per-token prompt_attribution) — doubling every step in env-worker memory, on the wire, and in the orchestrator's group buffers. On long uncapped-tool browser rollouts (worldsims perplexity, 75 turns, ~57k ctx) this put env pods at ~200-400MB per active rollout and OOM-killed workers by step 5. Keep identity/usage only: no consumer reads the heavy fields (all step-level Response readers isinstance-check the live object, which a serialized dict never passes; usage is recorded on state separately at this call site via record_response_usage). Port of 93b7e5a (worldsims/ephemeral-mm-pixels-vf-model), refitted to this branch's runtime structure. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
![macroscopeapp[bot]](https://avatars.githubusercontent.com/in/900172?s=60&v=4){kind=small}
ApprovabilityVerdict: Approved Memory optimization fixes that remove redundant data from trajectory steps and clean up completed rollouts from a registry. Both changes are well-documented, limited in scope, and don't alter functional behavior. You can customize Macroscope's approvability policy. Learn more. |
register_trajectory stores every rollout's live trajectory list in Runtime.trajectories so handle-borrowing sub-runtime states can resolve it — but nothing ever removed the entry, and the Runtime lives for the env worker process lifetime: every COMPLETED rollout's full trajectory stayed referenced forever (~50-400MB leaked per rollout on browser envs; OOM-killed worldsims perplexity env pods within ~5 training steps even after per-rollout slimming). Pop it in cleanup_rollout (runs in the harness finally, so error paths evict too). Borrowers' lifetimes are bounded by the owning rollout, so no consumer can resolve the entry after cleanup. Port of 038cb29 (worldsims/ephemeral-mm-pixels-vf-model). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Dismissing prior approval to re-evaluate c8e2b90
hubert-marek
changed the title
![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: c8e2b90087
ℹ️ About Codex in GitHub
Your team has set up Codex to 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 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| # lifetime is within the owning rollout — without this pop the | ||
| # long-lived Runtime retains every completed rollout's full | ||
| # trajectory and the env worker leaks ~50-400MB per rollout. | ||
| self.trajectories.pop(str(state["trajectory_id"]), None) |
There was a problem hiding this comment.
For grouped rollouts this cleanup runs before the group phase: Env._run_group_states awaits all harness.run(...) calls and only then calls harness.score_group(...), while Harness.run skips cleanup_group when group_key is present. Popping the live trajectory here therefore makes any group update/reward/cleanup that runs a stored child state with transcript="append" fail in resolve_trajectory with No live trajectory registered..., even though those runtime handles are intentionally retained until group cleanup. This eviction should be delayed for grouped states (or repeated in cleanup_group) so group-stage borrowers can still append to the parent trajectory.
Useful? React with 👍 / 👎.
| # this call site via ``record_response_usage``. | ||
| response_meta = serializable(response) | ||
| if isinstance(response_meta, dict): | ||
| response_meta.pop("message", None) |
There was a problem hiding this comment.
When renderer rollouts hit training-budget truncation, renderer_client._step_token_ids intentionally falls back from empty step["tokens"].completion_ids to step["response"]["message"]["tokens"] (see verifiers/clients/renderer_client.py and the fallback tests in tests/test_renderer_client.py). Removing message here makes that fallback return None, so subsequent turns are fully re-rendered instead of bridged and can lose the cache/multimodal placeholder continuity the renderer path relies on for long multimodal rollouts. If the goal is to drop duplicated content, keep a minimal raw-token/is-truncated sidecar rather than deleting the whole message.
Useful? React with 👍 / 👎.
hubert-marek
changed the title
1 participant
Four env-worker memory fixes for the v1 trajectory/endpoint path, found and production-validated incrementally on worldsims browser runs (each predicted/confirmed by a measurement harness; details below):
1. Drop the response message dump from trajectory steps
Each step stored
serializable(response)wholesale — itsmessagerestates the step'scompletionand itstokens(full-prefix ids/mask plus per-tokenprompt_attribution), doubling every step in worker memory, on the wire, and in prime-rl's group buffers. Steps now keep identity/usage only. No v1-path consumer reads the dropped fields (all step-level readersisinstance(…, Response)-check, which serialized dicts fail).2. Evict completed rollouts from
Runtime.trajectories(true leak)prepare_state → register_trajectorystores every rollout's live trajectory for handle-borrowing sub-runtime states, and nothing ever removed it — the process-lifetime Runtime retained every completed rollout's full trajectory (~50–400MB each; OOM-killed env pods by training step ~5). Popped incleanup_rollout(runs in the harnessfinally).3. Discard delivered intercepts (dominant term: −74% worker peak)
InterceptionServer.interceptsretained every request's raw body — the full message history including in-sandbox base64 screenshots — until rollout unregister, so a long rollout held every turn's request simultaneously. (The renderer's image offload rewrites a normalized copy; the intercept's base64 never left.)forward_requestnow discards after delivery; the HTTP handler keeps its local reference, discard is idempotent, unregister still sweeps undelivered entries.4.
malloc_trimat the worker stats cadence (resting RSS ÷3)Workers had no trim anywhere (the per-batch trim is orchestrator-only); freed arena pages ratcheted RSS to ~3× the live set. Trim every 10s via ctypes (GIL released). Deliberately not
gc.collect— full collections on fat heaps are what caused the worker heartbeat-timeout kills.Measurement (churn simulation, 4 concurrent 45-turn rollouts/worker, real subprocess RSS)
Production lineage on the worldsims env pin (
worldsims/ephemeral-mm-pixels-vf-model): no-fixes OOMed at step 5 (~115GB env pod); fixes 1+2 still climbed to ~90GB and failed; all four shipped ase522d36e— runoh7jsedjxix7rbk23sdvsujdtraining on it now.Commits 1–2 are ports refitted to this branch's runtime structure; 3–4 cherry-picked clean.
py_compile+ruffclean. Structural successor for 1–2: #1606's delta-native message graph.🤖 Generated with Claude Code