diff --git a/src/agent_env/task/step_journal.py b/src/agent_env/task/step_journal.py index 1eda5806..db68fc60 100644 --- a/src/agent_env/task/step_journal.py +++ b/src/agent_env/task/step_journal.py @@ -29,11 +29,15 @@ def _union(existing: list, items: list, *, move_to_end: bool = False, path: str """Union by value (deployed envs by instance), order-stable. ``move_to_end`` is the re-record case: that step's completion moves to the end of the replay order, so its items must move with it.""" key = _item_key(path) - fresh = [x for i, x in enumerate(items) if key(x) not in [key(y) for y in items[:i]]] - fresh_keys = [key(x) for x in fresh] + fresh, fresh_keys = [], [] + for item in items: + item_key = key(item) + if item_key not in fresh_keys: + fresh.append(item) + fresh_keys.append(item_key) + existing_keys = [key(item) for item in existing] if move_to_end: - return [x for x in existing if key(x) not in fresh_keys] + fresh - existing_keys = [key(x) for x in existing] + return [item for item, item_key in zip(existing, existing_keys) if item_key not in fresh_keys] + fresh return list(existing) + [x for x, k in zip(fresh, fresh_keys) if k not in existing_keys] diff --git a/tst/benchmarks/journal_union.py b/tst/benchmarks/journal_union.py new file mode 100644 index 00000000..fb0e2ad4 --- /dev/null +++ b/tst/benchmarks/journal_union.py @@ -0,0 +1,80 @@ +"""Compare journal list union with a git revision using the production function bodies. + +Run from the repository root, for example: + python tst/benchmarks/journal_union.py --base origin/main --repeats 5 +""" + +from __future__ import annotations + +import argparse +import ast +import statistics +import subprocess +import time +from pathlib import Path + +ROOT = Path(__file__).resolve().parents[2] +TARGET = "src/agent_env/task/step_journal.py" + + +def positive_int(value: str) -> int: + try: + result = int(value) + except ValueError as exc: + raise argparse.ArgumentTypeError("must be a positive integer") from exc + if result <= 0: + raise argparse.ArgumentTypeError("must be a positive integer") + return result + + +def load_union(source: str): + tree = ast.parse(source) + wanted = {"_union", "_item_key"} + nodes = [node for node in tree.body if isinstance(node, ast.FunctionDef) and node.name in wanted] + if {node.name for node in nodes} != wanted: + raise ValueError("source does not define both _union and _item_key") + namespace = {} + module = ast.Module(body=nodes, type_ignores=[]) + exec(compile(ast.fix_missing_locations(module), "", "exec"), namespace) + return namespace["_union"] + + +def median_ms(fn, repeats: int) -> float: + samples = [] + for _ in range(repeats): + start = time.perf_counter() + fn() + samples.append((time.perf_counter() - start) * 1000) + return statistics.median(samples) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--base", default="origin/main", help="git revision to compare against (default: origin/main)") + parser.add_argument("--repeats", type=positive_int, default=5, help="timing samples per implementation and size (default: 5)") + args = parser.parse_args() + + baseline_source = subprocess.check_output( + ["git", "show", f"{args.base}:{TARGET}"], cwd=ROOT, text=True + ) + candidate_source = (ROOT / TARGET).read_text() + baseline, candidate = load_union(baseline_source), load_union(candidate_source) + + print(f"Median runtime; production _union/_item_key function bodies extracted from {args.base} and the worktree") + print("items baseline_ms worktree_ms speedup") + for size in (100, 1000, 3000): + items = [ + {"instance_id": str(i), "sandbox_id": str(i), "card": {"tools": ["a", "b"]}} + for i in range(size) + ] + expected = baseline([], items, path="context.deployed_envs") + actual = candidate([], items, path="context.deployed_envs") + if actual != expected: + raise AssertionError(f"union output differs at {size} items") + before = median_ms(lambda: baseline([], items, path="context.deployed_envs"), args.repeats) + after = median_ms(lambda: candidate([], items, path="context.deployed_envs"), args.repeats) + print(f"{size:>5} {before:>11.3f} {after:>11.3f} {before / after:>6.2f}x") + + +if __name__ == "__main__": + main() diff --git a/tst/unit/task/step_journal_test.py b/tst/unit/task/step_journal_test.py index 04143f7c..d1be8f39 100644 --- a/tst/unit/task/step_journal_test.py +++ b/tst/unit/task/step_journal_test.py @@ -9,6 +9,7 @@ import pytest +import agent_env.task.step_journal as journal_mod import agent_env.task.store as store_mod from agent_env.store import Filter, LocalSqliteDocumentStore, Sort from agent_env.task.step_journal import _union, commit_ordered, replay_context @@ -574,6 +575,30 @@ def test_a_deployed_env_written_by_two_kernels_stays_one_entry(): assert _union([old], [new], path="context.prompt_responses") == [old, new] # every other list stays by value +def test_union_projects_each_item_key_once(monkeypatch): + calls = 0 + + def item_key(path): + def project(item): + nonlocal calls + calls += 1 + return item + + return project + + monkeypatch.setattr(journal_mod, "_item_key", item_key) + existing = [["old"], ["shared"]] + items = [["shared"], ["new"], ["new"]] + + assert _union(existing, items) == [["old"], ["shared"], ["new"]] + assert calls == len(existing) + len(items) + + calls = 0 + moved_items = [["old"]] + assert _union(existing, moved_items, move_to_end=True) == [["shared"], ["old"]] + assert calls == len(existing) + len(moved_items) + + def test_a_completion_that_re_adds_a_deployed_env_keeps_one_entry(store): iid = _new(store) old = {"env_id": "e", "instance_id": "i1", "mcp_server_name": None}