purewhiter--mobilegym
wehub-resource-sync
128 行
4.2 KiB
Python
128 行
4.2 KiB
Python
"""Hierarchical stopwatch for profiling episode lifecycle."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import threading
|
|
import time
|
|
from contextlib import contextmanager
|
|
from dataclasses import dataclass, field
|
|
from typing import Any, Optional
|
|
|
|
|
|
@dataclass
|
|
class _Entry:
|
|
name: str
|
|
elapsed: float
|
|
children: list[_Entry] = field(default_factory=list)
|
|
|
|
|
|
class StopWatch:
|
|
"""Nestable stopwatch for async episode profiling.
|
|
|
|
Usage::
|
|
|
|
sw = StopWatch()
|
|
with sw.phase("setup"):
|
|
with sw.phase("reset"):
|
|
...
|
|
with sw.phase("warm"):
|
|
...
|
|
print(sw.summary())
|
|
# setup=3.21s { reset=2.10s warm=1.05s } | run=25.3s
|
|
"""
|
|
|
|
def __init__(self) -> None:
|
|
self._stack: list[_Entry] = []
|
|
self._roots: list[_Entry] = []
|
|
|
|
@contextmanager
|
|
def phase(self, name: str):
|
|
t0 = time.monotonic()
|
|
entry = _Entry(name=name, elapsed=0.0)
|
|
if self._stack:
|
|
self._stack[-1].children.append(entry)
|
|
else:
|
|
self._roots.append(entry)
|
|
self._stack.append(entry)
|
|
try:
|
|
yield
|
|
finally:
|
|
entry.elapsed = time.monotonic() - t0
|
|
self._stack.pop()
|
|
|
|
def record(self, name: str, elapsed: float) -> None:
|
|
"""Insert a pre-measured phase entry under the current stack top.
|
|
|
|
Use when the duration was measured outside ``phase()`` — e.g. across an
|
|
``asyncio.to_thread`` boundary, where the sync call is dispatched to a
|
|
worker thread and the queue/exec split is captured by the wrapper.
|
|
Appends as a child of the currently-open phase (or as a root entry if
|
|
none is open). Append-only: thread-safe under CPython GIL.
|
|
"""
|
|
entry = _Entry(name=name, elapsed=elapsed)
|
|
if self._stack:
|
|
self._stack[-1].children.append(entry)
|
|
else:
|
|
self._roots.append(entry)
|
|
|
|
def summary(self, roots: list[_Entry] | None = None, depth: int = 0) -> str:
|
|
parts: list[str] = []
|
|
for e in (roots or self._roots):
|
|
s = f"{e.name}={e.elapsed:.2f}s"
|
|
if e.children:
|
|
child_s = " ".join(
|
|
self.summary([c], depth + 1) for c in e.children
|
|
)
|
|
s += f" {{ {child_s} }}"
|
|
parts.append(s)
|
|
return " | ".join(parts)
|
|
|
|
def to_tree(self, roots: list[_Entry] | None = None) -> list[dict[str, Any]]:
|
|
return [self._entry_to_tree(entry) for entry in (roots or self._roots)]
|
|
|
|
def to_flat(self, roots: list[_Entry] | None = None) -> dict[str, float]:
|
|
flat: dict[str, float] = {}
|
|
for entry in (roots or self._roots):
|
|
self._entry_to_flat(entry, flat, prefix="")
|
|
return flat
|
|
|
|
def _entry_to_tree(self, entry: _Entry) -> dict[str, Any]:
|
|
return {
|
|
"name": entry.name,
|
|
"elapsed_s": entry.elapsed,
|
|
"children": [self._entry_to_tree(child) for child in entry.children],
|
|
}
|
|
|
|
def _entry_to_flat(self, entry: _Entry, flat: dict[str, float], prefix: str) -> None:
|
|
key = f"{prefix}.{entry.name}" if prefix else entry.name
|
|
flat[key] = entry.elapsed
|
|
for child in entry.children:
|
|
self._entry_to_flat(child, flat, prefix=key)
|
|
|
|
@property
|
|
def total(self) -> float:
|
|
return sum(e.elapsed for e in self._roots)
|
|
|
|
def reset(self) -> None:
|
|
self._stack.clear()
|
|
self._roots.clear()
|
|
|
|
|
|
# Thread-local "current stopwatch" so sync code dispatched into a worker thread
|
|
# (typically via ``asyncio.to_thread(agent.act, ...)``) can locate the env's
|
|
# stopwatch without threading it through every API. The runner sets this for
|
|
# the duration of one ``to_thread`` call; LLMClient and other thread-side code
|
|
# call ``current_stopwatch()`` and ``StopWatch.record()`` to attribute their
|
|
# internal timings to the same tree.
|
|
_thread_local = threading.local()
|
|
|
|
|
|
def set_current_stopwatch(sw: Optional["StopWatch"]) -> None:
|
|
"""Bind ``sw`` (or unbind with ``None``) to the calling thread."""
|
|
_thread_local.sw = sw
|
|
|
|
|
|
def current_stopwatch() -> Optional["StopWatch"]:
|
|
"""Return the calling thread's bound StopWatch, or ``None`` if unset."""
|
|
return getattr(_thread_local, "sw", None)
|