From b3a0a2b0305ccaac16ff2586edb3fd8baf1da0a3 Mon Sep 17 00:00:00 2001 From: "jinli.yl" Date: Fri, 31 Jul 2026 22:09:06 +0800 Subject: [PATCH] fix(evaluation): clarify usage semantics and atomic counters --- benchmark/beam/run.py | 2 +- benchmark/longmemeval/run.py | 2 +- .../agent_wrapper/base_agent_wrapper.py | 17 ++-- reme/components/job/background_job.py | 2 +- reme/components/job/base_job.py | 18 +--- reme/components/job/cron_job.py | 2 +- reme/components/job/stream_job.py | 2 +- reme/schema/token_usage.py | 10 +- reme/utils/__init__.py | 9 +- reme/utils/counter.py | 94 +++++++++++++------ reme/utils/evaluation_interface.py | 12 +-- tests/unit/test_evaluation_interface.py | 4 +- tests/unit/test_job.py | 10 +- tests/unit/test_token_usage.py | 18 ++++ tests/unit/test_utils.py | 75 +++++++++++++++ 15 files changed, 198 insertions(+), 79 deletions(-) diff --git a/benchmark/beam/run.py b/benchmark/beam/run.py index 351658f2..97c620e5 100644 --- a/benchmark/beam/run.py +++ b/benchmark/beam/run.py @@ -737,7 +737,7 @@ def main( # pylint: disable=too-many-statements print(f" {'OVERALL':<40s}: {overall:.3f} binary={binary_overall:.3f} ({len(all_scores)} Qs)") tool_call_mean, tool_call_std = _mean_and_std(all_tool_call_totals) print(f" Tool calls/query: mean={tool_call_mean:.2f} std={tool_call_std:.2f}") - print(" Bench tokens/query:") + print(" Bench reported tokens/query:") for metric in _TOKEN_USAGE_METRICS: values = [usage[metric] for usage in all_token_usages if usage.get(metric) is not None] if values: diff --git a/benchmark/longmemeval/run.py b/benchmark/longmemeval/run.py index 0e682bf1..baf386ae 100644 --- a/benchmark/longmemeval/run.py +++ b/benchmark/longmemeval/run.py @@ -735,7 +735,7 @@ def _print_summary(results: list[dict], start_time: float) -> None: tool_call_mean, tool_call_std = _mean_and_std(tool_call_totals) print(f" Tool calls/query: mean={tool_call_mean:.2f} std={tool_call_std:.2f}") token_usages = [r.get("agentic_token_usage", {}) for r in results] - print(" Bench tokens/query:") + print(" Bench reported tokens/query:") for metric in _TOKEN_USAGE_METRICS: values = [usage[metric] for usage in token_usages if usage.get(metric) is not None] if values: diff --git a/reme/components/agent_wrapper/base_agent_wrapper.py b/reme/components/agent_wrapper/base_agent_wrapper.py index 078d81bb..69f1ace6 100644 --- a/reme/components/agent_wrapper/base_agent_wrapper.py +++ b/reme/components/agent_wrapper/base_agent_wrapper.py @@ -12,7 +12,7 @@ from ..base_component import BaseComponent from ..outbound_proxy import BaseOutboundProxy from ...enumeration import ChunkEnum, ComponentEnum from ...schema import StreamChunk, TokenUsage -from ...utils import global_counter_add +from ...utils import global_counter_add_many if TYPE_CHECKING: from ..job.base_job import BaseJob @@ -236,12 +236,15 @@ class BaseAgentWrapper(BaseComponent): counters = getattr(self.app_context, "metadata", None) if not isinstance(counters, dict): return - for field in ("input_tokens", "output_tokens", "total_tokens"): - global_counter_add( - counters, - [self.TOKEN_COUNTER_PREFIX, self.name, field], - getattr(usage, field), - ) + prefix = (self.TOKEN_COUNTER_PREFIX, self.name) + global_counter_add_many( + counters, + { + (*prefix, "input_tokens"): usage.input_tokens, + (*prefix, "output_tokens"): usage.output_tokens, + (*prefix, "total_tokens"): usage.total_tokens, + }, + ) @abstractmethod async def reply(self, inputs: Any, **kwargs) -> dict: diff --git a/reme/components/job/background_job.py b/reme/components/job/background_job.py index 2f124bdc..662840e9 100644 --- a/reme/components/job/background_job.py +++ b/reme/components/job/background_job.py @@ -123,7 +123,7 @@ class BackgroundJob(BaseJob): async def __call__(self, **kwargs) -> Response: """Default body: run steps in order; errors propagate to supervisor.""" - self._record_call(BackgroundJob) + self._record_call() merged = {**self.kwargs, **kwargs} context = RuntimeContext(stop_event=self._stop_event, **merged) for step in self._build_steps(): diff --git a/reme/components/job/base_job.py b/reme/components/job/base_job.py index c19134a0..5133111d 100644 --- a/reme/components/job/base_job.py +++ b/reme/components/job/base_job.py @@ -58,27 +58,15 @@ class BaseJob(BaseComponent): # dict(params) copies kwargs so steps cannot mutate the shared spec. return [step_cls(**dict(params)) for step_cls, params in self.step_specs] - def _counter_key(self, entry_class: type["BaseJob"]) -> list[str]: - """Build this job's application-lifetime call counter key. - - The counter path groups calls by the built-in job implementation that - accepted the call, then preserves any project-specific subclasses - between that implementation and the concrete job class. - """ - mro = type(self).mro() - entry_index = mro.index(entry_class) - subclass_path = [cls.__name__ for cls in mro[:entry_index]] - return ["__job_counter", entry_class.__name__, *subclass_path, self.name] - - def _record_call(self, entry_class: type["BaseJob"]) -> None: + def _record_call(self) -> None: """Increment this job's application-lifetime call counter.""" metadata = getattr(self.app_context, "metadata", None) if isinstance(metadata, dict): - global_counter_inc(metadata, self._counter_key(entry_class)) + global_counter_inc(metadata, ["__job_counter", self.name]) async def __call__(self, **kwargs) -> Response: """Run all steps in order, capturing any failure into the response.""" - self._record_call(BaseJob) + self._record_call() merged = {**self.kwargs, **kwargs} context = RuntimeContext(**merged) try: diff --git a/reme/components/job/cron_job.py b/reme/components/job/cron_job.py index 49a117d6..a72a7878 100644 --- a/reme/components/job/cron_job.py +++ b/reme/components/job/cron_job.py @@ -47,7 +47,7 @@ class CronJob(BackgroundJob): if self._stop_event.is_set(): break try: - self._record_call(CronJob) + self._record_call() await self._execute_steps() except Exception as exc: self.logger.exception(f"Cron job '{self.name}' failed: {exc}") diff --git a/reme/components/job/stream_job.py b/reme/components/job/stream_job.py index 6de561a0..c48bda75 100644 --- a/reme/components/job/stream_job.py +++ b/reme/components/job/stream_job.py @@ -12,7 +12,7 @@ class StreamJob(BaseJob): async def __call__(self, **kwargs) -> None: """Run steps; emit failures as ERROR chunks, then a terminal DONE marker.""" - self._record_call(StreamJob) + self._record_call() merged = {**self.kwargs, **kwargs} context = RuntimeContext(**merged) try: diff --git a/reme/schema/token_usage.py b/reme/schema/token_usage.py index f0ba3587..df178152 100644 --- a/reme/schema/token_usage.py +++ b/reme/schema/token_usage.py @@ -6,10 +6,12 @@ from pydantic import BaseModel, Field, model_validator class TokenUsage(BaseModel): - """Token usage for one completed agent invocation. + """Portable token usage reported for one completed agent invocation. - ``input_tokens`` and ``output_tokens`` are the complete usage reported by - one backend invocation. ``total_tokens`` is their derived sum. + Only the provider's top-level input and output counters are retained. + Provider-specific cache and reasoning breakdowns are intentionally + excluded, so values may not share identical billing semantics across + providers. ``total_tokens`` is always their derived sum. """ input_tokens: int = Field(default=0, ge=0) @@ -26,7 +28,7 @@ class TokenUsage(BaseModel): cls, usage: Any, ) -> "TokenUsage": - """Normalize a provider usage object or mapping to input/output totals.""" + """Keep only a provider's portable top-level input/output counters.""" def get(*names: str) -> int | None: for name in names: diff --git a/reme/utils/__init__.py b/reme/utils/__init__.py index 6fbbe767..a41e66e9 100644 --- a/reme/utils/__init__.py +++ b/reme/utils/__init__.py @@ -15,7 +15,13 @@ from .service_utils import find_reme, locate_reme, precheck_start, cli_find_reme from .similarity_utils import cosine_similarity, batch_cosine_similarity from .token_utils import estimate_token_count from .agent_state_io import AsStateHandler -from .counter import global_counter_add, global_counter_get, global_counter_get_all, global_counter_inc +from .counter import ( + global_counter_add, + global_counter_add_many, + global_counter_get, + global_counter_get_all, + global_counter_inc, +) __all__ = [ "hash_text", @@ -39,6 +45,7 @@ __all__ = [ "estimate_token_count", "AsStateHandler", "global_counter_add", + "global_counter_add_many", "global_counter_get", "global_counter_get_all", "global_counter_inc", diff --git a/reme/utils/counter.py b/reme/utils/counter.py index cd4c53d0..a2f578af 100644 --- a/reme/utils/counter.py +++ b/reme/utils/counter.py @@ -2,10 +2,70 @@ import copy import threading +from collections.abc import Mapping from typing import Any COUNTER_TREE_KEY = "_counter_tree" COUNTER_LOCK_KEY = "_counter_tree_lock" +_COUNTER_INIT_LOCK = threading.Lock() + + +def _get_counter_lock(metadata: dict[str, Any]) -> Any: + """Return the metadata-scoped lock, creating it once when needed.""" + lock = metadata.get(COUNTER_LOCK_KEY) + if lock is not None: + return lock + + # Two threads may reach the first counter operation concurrently. Guard + # initialization so they cannot install and then use different locks. + with _COUNTER_INIT_LOCK: + lock = metadata.get(COUNTER_LOCK_KEY) + if lock is None: + lock = threading.Lock() + metadata[COUNTER_LOCK_KEY] = lock + return lock + + +def global_counter_add_many( + metadata: dict[str, Any], + updates: Mapping[tuple[str, ...], int], +) -> dict[tuple[str, ...], int]: + """Atomically fetch-and-add multiple counter paths. + + All paths are validated before the counter tree is mutated. The returned + mapping contains each path's value immediately before its increment. + """ + normalized = dict(updates) + for path, value in normalized.items(): + if not isinstance(path, tuple) or not all(isinstance(part, str) for part in path): + raise TypeError("counter paths must be tuples of strings") + if not isinstance(value, int): + raise TypeError("counter increments must be integers") + if not normalized: + return {} + + lock = _get_counter_lock(metadata) + with lock: + tree = metadata.get(COUNTER_TREE_KEY) + if tree is None: + tree = {"value": 0, "children": {}} + metadata[COUNTER_TREE_KEY] = tree + + nodes: dict[tuple[str, ...], dict[str, Any]] = {} + for path in normalized: + node = tree + for part in path: + child = node["children"].get(part) + if child is None: + child = {"value": 0, "children": {}} + node["children"][part] = child + node = child + nodes[path] = node + + previous = {path: node["value"] for path, node in nodes.items()} + for path, value in normalized.items(): + nodes[path]["value"] += value + return previous def global_counter_add(metadata: dict[str, Any], key: list[str], val: int) -> int: @@ -23,28 +83,8 @@ def global_counter_add(metadata: dict[str, Any], key: list[str], val: int) -> in If they are missing they are created lazily so the function is safe to call with a plain ``dict``. """ - lock = metadata.get(COUNTER_LOCK_KEY) - if lock is None: - lock = threading.Lock() - metadata[COUNTER_LOCK_KEY] = lock - - with lock: - tree = metadata.get(COUNTER_TREE_KEY) - if tree is None: - tree = {"value": 0, "children": {}} - metadata[COUNTER_TREE_KEY] = tree - - node: dict[str, Any] = tree - for part in key: - assert isinstance(part, str) - tmp = node["children"].get(part, None) - if tmp is None: - tmp = {"value": 0, "children": {}} - node["children"][part] = tmp - node = tmp - res = node["value"] - node["value"] = res + val - return res + path = tuple(key) + return global_counter_add_many(metadata, {path: val})[path] def global_counter_inc(metadata: dict[str, Any], key: list[str]) -> int: @@ -63,10 +103,7 @@ def global_counter_get(metadata: dict[str, Any], key: list[str]) -> int: path that does not exist yet is reported as 0, matching the value the node would hold right before its first increment. """ - lock = metadata.get(COUNTER_LOCK_KEY) - if lock is None: - lock = threading.Lock() - metadata[COUNTER_LOCK_KEY] = lock + lock = _get_counter_lock(metadata) with lock: tree = metadata.get(COUNTER_TREE_KEY) @@ -91,10 +128,7 @@ def global_counter_get_all(metadata: dict[str, Any], key: list[str]) -> dict[str Returns ``None`` when the tree or any part of ``key`` does not exist. An empty ``key`` returns a copy of the whole tree. """ - lock = metadata.get(COUNTER_LOCK_KEY) - if lock is None: - lock = threading.Lock() - metadata[COUNTER_LOCK_KEY] = lock + lock = _get_counter_lock(metadata) with lock: tree = metadata.get(COUNTER_TREE_KEY) diff --git a/reme/utils/evaluation_interface.py b/reme/utils/evaluation_interface.py index 49e226e9..1f6819d4 100644 --- a/reme/utils/evaluation_interface.py +++ b/reme/utils/evaluation_interface.py @@ -9,14 +9,12 @@ work sharing its Application instance. from typing import TYPE_CHECKING -from ..components.job import BackgroundJob, BaseJob, CronJob, StreamJob from .counter import global_counter_get, global_counter_get_all if TYPE_CHECKING: from ..components.application_context import ApplicationContext -_JOB_ENTRY_CLASSES = {CronJob, StreamJob, BackgroundJob, BaseJob} _TOKEN_METRICS = ( "input_tokens", "output_tokens", @@ -31,15 +29,9 @@ def check_job_count(job_name: str, app_context: "ApplicationContext") -> int: current Application instance. Unknown job names use the same ``KeyError`` contract as :meth:`Application.run_job`. """ - job = app_context.jobs.get(job_name) - if job is None: + if job_name not in app_context.jobs: raise KeyError(f"Job '{job_name}' not found") - - entry_class = next((cls for cls in type(job).mro() if cls in _JOB_ENTRY_CLASSES), None) - if entry_class is None: - raise TypeError(f"Job '{job_name}' does not inherit from a supported job implementation") - # pylint: disable-next=protected-access - return global_counter_get(app_context.metadata, job._counter_key(entry_class)) + return global_counter_get(app_context.metadata, ["__job_counter", job_name]) class JobCountTracker: diff --git a/tests/unit/test_evaluation_interface.py b/tests/unit/test_evaluation_interface.py index 086528a7..fd58b2a9 100644 --- a/tests/unit/test_evaluation_interface.py +++ b/tests/unit/test_evaluation_interface.py @@ -32,8 +32,8 @@ def test_check_job_count_reads_registered_base_job_count(): asyncio.run(run()) -def test_check_job_count_resolves_custom_job_inheritance_path(): - """check_job_count resolves counters recorded under subclassed job paths.""" +def test_check_job_count_reads_custom_job_by_name(): + """check_job_count reads subclassed job counters without depending on inheritance.""" async def run(): class ProjectStreamJob(StreamJob): diff --git a/tests/unit/test_job.py b/tests/unit/test_job.py index 3531fcd8..b1749975 100644 --- a/tests/unit/test_job.py +++ b/tests/unit/test_job.py @@ -144,12 +144,12 @@ def test_base_job_records_calls_by_name(): await job() await job() - assert global_counter_get(app_context.metadata, ["__job_counter", "BaseJob", "search"]) == 2 + assert global_counter_get(app_context.metadata, ["__job_counter", "search"]) == 2 asyncio.run(run()) -def test_stream_job_records_concrete_subclass_path(): +def test_stream_job_subclass_records_calls_by_job_name(): async def run(): class ProjectStreamJob(StreamJob): pass @@ -159,7 +159,7 @@ def test_stream_job_records_concrete_subclass_path(): await job(stream_queue=asyncio.Queue()) - assert global_counter_get(app_context.metadata, ["__job_counter", "StreamJob", "ProjectStreamJob", "chat"]) == 1 + assert global_counter_get(app_context.metadata, ["__job_counter", "chat"]) == 1 asyncio.run(run()) @@ -171,7 +171,7 @@ def test_background_job_records_calls_by_name(): await job() - assert global_counter_get(app_context.metadata, ["__job_counter", "BackgroundJob", "watch"]) == 1 + assert global_counter_get(app_context.metadata, ["__job_counter", "watch"]) == 1 asyncio.run(run()) @@ -195,7 +195,7 @@ def test_cron_job_records_each_triggered_execution(): await job() - assert global_counter_get(app_context.metadata, ["__job_counter", "CronJob", "nightly"]) == 1 + assert global_counter_get(app_context.metadata, ["__job_counter", "nightly"]) == 1 asyncio.run(run()) diff --git a/tests/unit/test_token_usage.py b/tests/unit/test_token_usage.py index 5f7f5c04..cd505815 100644 --- a/tests/unit/test_token_usage.py +++ b/tests/unit/test_token_usage.py @@ -50,6 +50,24 @@ def test_provider_usage_uses_reported_input_without_cache_adjustment(): assert usage.total_tokens == 64 +def test_provider_usage_accepts_prompt_and_completion_aliases(): + """Portable OpenAI-style aliases normalize to the common counters.""" + usage = TokenUsage.from_provider({"prompt_tokens": 3, "completion_tokens": 2}) + + assert usage.model_dump() == { + "input_tokens": 3, + "output_tokens": 2, + "total_tokens": 5, + } + + +def test_total_tokens_is_always_derived_from_input_and_output(): + """Caller-supplied totals cannot violate the portable usage invariant.""" + usage = TokenUsage(input_tokens=3, output_tokens=2, total_tokens=999) + + assert usage.total_tokens == 5 + + def test_agentscope_usage_keeps_portable_input_and_output(tmp_path): """AgentScope usage has the same input/output-only contract.""" context = ApplicationContext(workspace_dir=str(tmp_path)) diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index 34eff188..6b21eef1 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -9,8 +9,10 @@ import pytest from reme.utils import common_utils from reme.utils.counter import ( + COUNTER_LOCK_KEY, COUNTER_TREE_KEY, global_counter_add, + global_counter_add_many, global_counter_get, global_counter_get_all, global_counter_inc, @@ -94,6 +96,79 @@ def test_add_returns_old_value_and_adds_val(): assert global_counter_get(metadata, ["a"]) == 16 +def test_add_many_returns_old_values_and_updates_all_paths(): + """``add_many`` updates sibling metrics under one counter-tree lock.""" + metadata: dict = {} + global_counter_add(metadata, ["usage", "input"], 2) + + previous = global_counter_add_many( + metadata, + { + ("usage", "input"): 3, + ("usage", "output"): 4, + ("usage", "total"): 7, + }, + ) + + assert previous == { + ("usage", "input"): 2, + ("usage", "output"): 0, + ("usage", "total"): 0, + } + assert global_counter_get(metadata, ["usage", "input"]) == 5 + assert global_counter_get(metadata, ["usage", "output"]) == 4 + assert global_counter_get(metadata, ["usage", "total"]) == 7 + + +def test_add_many_holds_the_counter_lock_once_for_the_batch(): + """A multi-metric update is one critical section, not several writes.""" + + class CountingLock: + """Lock test double that counts entered critical sections.""" + + def __init__(self): + self.lock = threading.Lock() + self.entries = 0 + + def __enter__(self): + self.lock.acquire() + self.entries += 1 + return self + + def __exit__(self, *_args): + self.lock.release() + + lock = CountingLock() + metadata = {COUNTER_LOCK_KEY: lock} + + global_counter_add_many( + metadata, + { + ("usage", "input"): 3, + ("usage", "output"): 4, + ("usage", "total"): 7, + }, + ) + + assert lock.entries == 1 + + +def test_add_many_validates_every_update_before_mutating(): + """One invalid update leaves every valid counter unchanged.""" + metadata: dict = {} + + with pytest.raises(TypeError, match="increments must be integers"): + global_counter_add_many( + metadata, + { + ("usage", "input"): 3, + ("usage", "output"): "invalid", + }, + ) + + assert COUNTER_TREE_KEY not in metadata + + def test_counters_are_isolated_by_key_path(): """Sibling and nested keys, plus the root, hold independent counters.""" metadata: dict = {}