Skip to content

Fix unsynchronized call counter in CallLimitLanguageModel - #285

Open
jaisinha77777 wants to merge 1 commit into
google-deepmind:mainfrom
jaisinha77777:fix-call-limit-race-condition
Open

Fix unsynchronized call counter in CallLimitLanguageModel#285
jaisinha77777 wants to merge 1 commit into
google-deepmind:mainfrom
jaisinha77777:fix-call-limit-race-condition

Conversation

@jaisinha77777

Copy link
Copy Markdown

Summary

`CallLimitLanguageModel` (used to hard-cap the number of calls made to an underlying language model, e.g. to bound API spend) checked and incremented its counter with no synchronization:

if self._calls >= self._max_calls:
  return ...
self._calls += 1

A single wrapped model instance is normally shared across an entire simulation, and `EntityAgent` dispatches component calls concurrently through a `ThreadPoolExecutor` (see `entity_agent.py`'s `parallel_call`), so `sample_text`/`sample_choice` on this shared instance can genuinely be called from many threads at once. That makes the check-then-increment an unsynchronized read-modify-write on shared state — the same class of bug I found (and fixed, in a merged PR) in `Measurements.close()`, except there this codebase's own convention is normally to guard exactly this pattern with a lock (`Measurements`, `ProfilerContext`, `AssociativeMemoryBank`, etc. all do). This class was the one outlier that didn't.

Being upfront about what I could and couldn't demonstrate: I stress-tested the original code extensively (hundreds of trials, hundreds of threads each, against a dummy model) and was not able to force a standard CPython build to actually exceed the configured limit — the GIL happens to make this specific narrow check-then-increment hard to interleave in practice on this build. So I want to be precise about the claim: this is not a bug I observed causing a limit overrun in current CPython. It's a genuine synchronization bug by the threading model's own contract (nothing guarantees += on an attribute is atomic across threads), and it's a real, not-hidden-by-anything bug on free-threaded (no-GIL) CPython builds, which PEP 703 made an officially supported build target starting with 3.13. Since the entire purpose of max_calls is to hard-cap spend/rate-limit exposure, an unenforced limit under a future Python build silently defeats that.

Fix

Adds a lock around the check-and-increment only. The actual (potentially slow) call to the underlying model stays outside the lock, so concurrent calls still execute in parallel once past the counter check — this doesn't serialize LLM calls, just the counter.

Test plan

New `concordia/language_model/call_limit_wrapper_test.py` (previously no coverage at all):

  • Calls pass through under the limit; `sample_text` returns `''` and `sample_choice` returns the first response once the limit is reached, without invoking the underlying model.
  • `sample_text` and `sample_choice` share a single budget.
  • A concurrency test: 200 threads racing a budget of 37 must produce exactly 37 underlying calls, not "close to" 37 — this is deterministic and passes reliably with the lock in place.

`python -m pytest concordia/language_model/ -q` → 14 passed.
`python -m pyink --check` on both changed/new files → clean.

CallLimitLanguageModel.sample_text()/sample_choice() checked and
incremented self._calls with no locking:

  if self._calls >= self._max_calls:
    return ...
  self._calls += 1

A single model instance wrapped this way is normally shared across an
entire simulation, and EntityAgent dispatches component calls
concurrently through a ThreadPoolExecutor, so this check-then-increment
runs from many threads against the same counter. That's an
unsynchronized read-modify-write on shared state -- exactly the pattern
this codebase otherwise guards with a lock everywhere else it appears
(Measurements, ProfilerContext, AssociativeMemoryBank, and others all
follow this convention; this class was the outlier).

In extensive local stress-testing (hundreds of trials, hundreds of
threads each) I could not force a standard CPython build to actually
exceed the configured limit -- the GIL happens to make this particular
narrow check-then-increment hard to interleave in practice. That's not
a correctness guarantee: nothing in the threading model promises this
is atomic, and it is a real bug on free-threaded (no-GIL) CPython
builds (PEP 703, officially supported since 3.13), which is where the
ecosystem is headed. Since max_calls exists specifically to hard-cap
API spend/rate limit exposure, silently under-enforcing it under any
future Python build defeats the point of the class.

Adds a lock around the check-and-increment only, leaving the actual
(slow) underlying model call outside the lock so concurrent calls
still execute in parallel once past the counter check.

Also adds call_limit_wrapper_test.py, which had no coverage: the
under-limit passthrough, empty-string/first-choice behavior once the
limit is reached, that sample_text and sample_choice share one budget,
and a concurrency test asserting the call count is exact (not just
"close") across 200 threads racing a budget of 37.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant