From 326551854c5f75ecc33af77a0b3c100e4689e3df Mon Sep 17 00:00:00 2001 From: caoxiaoyu <1844118046@qq.com> Date: Fri, 30 Oct 2026 06:58:37 +0000 Subject: [PATCH] capture: share image history strings across turns Each chat request re-sends the whole conversation and is JSON-decoded on its own, so turn k holds a fresh copy of every image and long text from turns 1..k. RolloutGraph kept all of them, which makes memory grow quadratically with the number of turns in multimodal rollouts. add_turn now replaces long strings (>= 1024 chars) in request_messages, request_tools and response_message with one graph-local copy. Only immutable str objects are shared; dicts and lists are rebuilt per turn, so values and message structure are unchanged and exports are byte-identical. Co-Authored-By: Claude Opus 5.5 --- src/openenv/core/harness/capture/graph.py | 16 +++ tests/envs/test_capture_graph_shared_text.py | 117 +++++++++++++++++++ 2 files changed, 133 insertions(+) create mode 100644 tests/envs/test_capture_graph_shared_text.py diff --git a/src/openenv/core/harness/capture/graph.py b/src/openenv/core/harness/capture/graph.py index 61cc88bf2..820689b54 100644 --- a/src/openenv/core/harness/capture/graph.py +++ b/src/openenv/core/harness/capture/graph.py @@ -248,9 +248,25 @@ def __init__(self) -> None: self._nodes: dict[str, TurnNode] = {} self._order: list[str] = [] self._children: dict[str, list[str]] = {} + # Every request re-sends the whole history, so each turn decodes a fresh copy of every + # earlier image and long text. Keep one copy per graph. Only immutable strings are shared; + # message dicts and lists stay per turn. + self._shared_text: dict[str, str] = {} + + def _share_text(self, value: Any) -> Any: + if isinstance(value, str) and len(value) >= 1024: + return self._shared_text.setdefault(value, value) + if isinstance(value, list): + return [self._share_text(item) for item in value] + if isinstance(value, dict): + return {key: self._share_text(item) for key, item in value.items()} + return value # --- construction -------------------------------------------------- def add_turn(self, node: TurnNode) -> TurnNode: + node.request_messages = self._share_text(node.request_messages) + node.request_tools = self._share_text(node.request_tools) + node.response_message = self._share_text(node.response_message) node.index = len(self._order) node.parent_id = self._find_parent(node) self._nodes[node.node_id] = node diff --git a/tests/envs/test_capture_graph_shared_text.py b/tests/envs/test_capture_graph_shared_text.py new file mode 100644 index 000000000..42ebf0031 --- /dev/null +++ b/tests/envs/test_capture_graph_shared_text.py @@ -0,0 +1,117 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Repeated history in a multi-turn rollout is held once, and exports are unchanged. + +Every chat request re-sends the whole conversation, and each request body is decoded separately, so +turn k carries its own fresh copy of every image attached in turns 1..k. Retaining those copies makes +graph memory grow quadratically with the number of turns. +""" + +from __future__ import annotations + +import json + +import pytest + +graph_mod = pytest.importorskip("openenv.core.harness.capture.graph") +export_mod = pytest.importorskip("openenv.core.harness.capture.export") + +RolloutGraph = graph_mod.RolloutGraph +TurnNode = graph_mod.TurnNode + +N_TURNS = 12 + + +class FakeSession: + def __init__(self, graph): + self.session_id = "s1" + self.metadata: dict = {} + self.findings: list[str] = [] + self.graph = graph + + +def image(i: int) -> str: + return "data:image/png;base64," + chr(ord("A") + i) * (64 * 1024) + + +def request_body(turn: int) -> list[dict]: + """The messages of request `turn`, decoded from the wire like the capture server does.""" + messages: list[dict] = [{"role": "system", "content": "You are a helpful agent."}] + for i in range(turn + 1): + messages.append( + { + "role": "user", + "content": [ + {"type": "text", "text": f"step {i}"}, + {"type": "image_url", "image_url": {"url": image(i)}}, + ], + } + ) + if i < turn: + messages.append({"role": "assistant", "content": f"reply {i}"}) + return json.loads(json.dumps(messages)) + + +def build_graph() -> RolloutGraph: + graph = RolloutGraph() + for turn in range(N_TURNS): + graph.add_turn( + TurnNode( + node_id=f"t{turn}", + prompt_ids=[], + sampled_ids=[], + request_messages=request_body(turn), + response_message={"role": "assistant", "content": f"reply {turn}"}, + ) + ) + return graph + + +def image_urls(node: TurnNode) -> list[str]: + return [ + part["image_url"]["url"] + for message in node.request_messages + if isinstance(message["content"], list) + for part in message["content"] + if part["type"] == "image_url" + ] + + +def test_each_image_is_retained_once_across_turns(): + graph = build_graph() + retained = {id(url): url for node in graph.nodes() for url in image_urls(node)} + # Without sharing this is N_TURNS * (N_TURNS + 1) / 2 copies. + assert len(retained) == N_TURNS + first_seen = image_urls(graph.get("t0"))[0] + for node in graph.nodes(): + assert image_urls(node)[0] is first_seen + + +def test_sharing_strings_does_not_alias_message_containers(): + graph = build_graph() + graph.get("t1").request_messages[1]["content"][1]["image_url"]["url"] = "edited" + assert graph.get("t0").request_messages[1]["content"][1]["image_url"]["url"] == ( + image(0) + ) + assert graph.get("t2").request_messages[1]["content"][1]["image_url"]["url"] == ( + image(0) + ) + + +def test_export_is_byte_identical_to_unshared_graph(monkeypatch): + def export(graph: RolloutGraph) -> bytes: + document = export_mod.export_session( + FakeSession(graph), include_messages=True, capture_level="text" + ) + return json.dumps(document, sort_keys=True).encode() + + shared = export(build_graph()) + monkeypatch.setattr(RolloutGraph, "_share_text", lambda self, value: value) + unshared = export(build_graph()) + + assert shared == unshared + assert len(json.loads(shared)["turns"]) == N_TURNS