Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
85 changes: 84 additions & 1 deletion rlm/environments/local_repl.py
Original file line number Diff line number Diff line change
Expand Up @@ -177,6 +177,15 @@ def __init__(
self.original_cwd = os.getcwd()
self.temp_dir = tempfile.mkdtemp(prefix=f"repl_env_{uuid.uuid4()}_")
self._lock = threading.Lock()
# Inter-turn idle instrumentation (RLM Lab, BR002-WO002): mirrors
# PyWasmEnv's timing fields so native and wasm idle structure compare
# directly. Mid-execute park is deliberately NOT instrumented here
# (BR002-CL005) — active_seconds therefore includes nested llm_query
# waits and is an upper bound on true active time.
self._timing_lock = threading.Lock()
self._timing = self._new_idle_timing()
self._last_execute_ended_at: float | None = None
self._inter_turn_idle_started_at: float | None = None
self._context_count: int = 0
self._history_count: int = 0
self.compaction = compaction
Expand Down Expand Up @@ -597,15 +606,89 @@ def execute_code(self, code: str) -> REPLResult:
final_answer = self._last_final_answer
self._last_final_answer = None

end_time = time.perf_counter()
with self._timing_lock:
wall = max(0.0, end_time - start_time)
self._timing["execute_wall_seconds"] += wall
self._timing["active_seconds"] += wall
self._timing["execute_count"] += 1
self._last_execute_ended_at = end_time

return REPLResult(
stdout=stdout,
stderr=stderr,
locals=self.locals.copy(),
execution_time=time.perf_counter() - start_time,
execution_time=end_time - start_time,
rlm_calls=self._pending_llm_calls.copy(),
final_answer=final_answer,
)

# -- Inter-turn idle instrumentation (RLM Lab, BR002-WO002) ------------- #

@staticmethod
def _new_idle_timing() -> dict[str, float | int]:
return {
"execute_wall_seconds": 0.0,
"active_seconds": 0.0,
"mid_execute_park_seconds": 0.0, # not instrumented (BR002-CL005)
"inter_turn_idle_seconds": 0.0,
"execute_count": 0,
"mid_execute_park_count": 0,
"inter_turn_idle_count": 0,
}

def reset_idle_timing(self) -> None:
with self._timing_lock:
self._timing = self._new_idle_timing()
self._last_execute_ended_at = None
self._inter_turn_idle_started_at = None

def begin_inter_turn_idle(self) -> None:
"""Mark that the session is awaiting the next root-model response.

Idle is backdated to the last execute's end so the interval matches
PyWasmEnv's boundary (execute-return -> next code execution). The
initial root wait (no prior REPL state) is deliberately not counted.
"""
with self._timing_lock:
if (
self._last_execute_ended_at is not None
and self._inter_turn_idle_started_at is None
):
self._inter_turn_idle_started_at = self._last_execute_ended_at

def end_inter_turn_idle(self) -> None:
with self._timing_lock:
if self._inter_turn_idle_started_at is None:
return
now = time.perf_counter()
self._timing["inter_turn_idle_seconds"] += max(
0.0, now - self._inter_turn_idle_started_at
)
self._timing["inter_turn_idle_count"] += 1
self._inter_turn_idle_started_at = None
self._last_execute_ended_at = None

def runtime_timing(self) -> dict[str, Any]:
with self._timing_lock:
timing = dict(self._timing)
if self._inter_turn_idle_started_at is not None:
timing["inter_turn_idle_seconds"] += max(
0.0, time.perf_counter() - self._inter_turn_idle_started_at
)
active = float(timing["active_seconds"])
inter_turn = float(timing["inter_turn_idle_seconds"])
mid_execute = float(timing["mid_execute_park_seconds"])
observed = active + inter_turn + mid_execute
timing["observed_seconds"] = observed
timing["active_fraction"] = active / observed if observed else 0.0
timing["inter_turn_idle_fraction"] = inter_turn / observed if observed else 0.0
timing["mid_execute_park_fraction"] = (
mid_execute / observed if observed else 0.0
)
timing["mid_execute_park_instrumented"] = False
return timing

def __enter__(self):
return self

Expand Down
6 changes: 6 additions & 0 deletions rlm/environments/warm_reset.py
Original file line number Diff line number Diff line change
Expand Up @@ -53,3 +53,9 @@ def reset_for_request(env: "LocalREPL") -> None:
if env.compaction:
env._compaction_history = []
env.locals["history"] = env._compaction_history

# Reset per-request idle timing (RLM Lab, BR002-WO002) so a reused warm env
# reports each task's idle structure like a fresh environment.
reset_timing = getattr(env, "reset_idle_timing", None)
if callable(reset_timing):
reset_timing()
67 changes: 67 additions & 0 deletions tests/test_local_repl.py
Original file line number Diff line number Diff line change
Expand Up @@ -291,3 +291,70 @@ def test_simulated_rlm_completions_functions_not_preserved(self):
assert "NameError" in result.stderr
assert "my_helper" in result.stderr
completion_2_env.cleanup()


class TestLocalREPLIdleTiming:
"""Inter-turn idle instrumentation (RLM Lab, BR002-WO002) — mirrors the
PyWasmEnv timing contract on the native env."""

def test_execute_accumulates_active_time(self):
repl = LocalREPL()
try:
repl.execute_code("x = 1")
timing = repl.runtime_timing()
assert timing["execute_count"] == 1
assert timing["active_seconds"] > 0.0
assert timing["execute_wall_seconds"] == timing["active_seconds"]
assert timing["mid_execute_park_instrumented"] is False
finally:
repl.cleanup()

def test_inter_turn_idle_backdates_to_last_execute_end(self):
import time as _time

repl = LocalREPL()
try:
repl.execute_code("x = 1")
_time.sleep(0.05)
repl.begin_inter_turn_idle() # backdates to the execute end
repl.end_inter_turn_idle()
timing = repl.runtime_timing()
assert timing["inter_turn_idle_count"] == 1
assert timing["inter_turn_idle_seconds"] >= 0.04
finally:
repl.cleanup()

def test_initial_root_wait_is_uncounted(self):
repl = LocalREPL()
try:
repl.begin_inter_turn_idle() # no execute yet -> not an idle window
repl.end_inter_turn_idle()
assert repl.runtime_timing()["inter_turn_idle_count"] == 0
finally:
repl.cleanup()

def test_open_idle_window_included_in_runtime_timing(self):
repl = LocalREPL()
try:
repl.execute_code("x = 1")
repl.begin_inter_turn_idle()
timing = repl.runtime_timing() # window still open
assert timing["inter_turn_idle_seconds"] >= 0.0
assert timing["active_fraction"] <= 1.0
finally:
repl.cleanup()

def test_reset_idle_timing_clears_all_counters(self):
repl = LocalREPL()
try:
repl.execute_code("x = 1")
repl.begin_inter_turn_idle()
repl.end_inter_turn_idle()
repl.reset_idle_timing()
timing = repl.runtime_timing()
assert timing["execute_count"] == 0
assert timing["active_seconds"] == 0.0
assert timing["inter_turn_idle_count"] == 0
assert timing["inter_turn_idle_seconds"] == 0.0
finally:
repl.cleanup()