Skip to content

Commit 6901e2d

Browse files
SuperMarioYLcodelionclaude
authored
fix: read MCTS params per-request instead of mutating shared global (#304) (#314)
* fix: read MCTS params per-request instead of mutating shared global (#304) proxy() wrote the request's mcts_depth/mcts_exploration/mcts_simulations into the module-level server_config on every request, and execute_single_approach read them back out of that same global. Under Flask's threaded mode (or gunicorn/uWSGI), concurrent requests overwrite each other's MCTS parameters between the write and the read, silently corrupting results. The write also leaked one request's params into every subsequent request. Resolve the params from the per-request request_config (with the server defaults as fallback) and stop writing them into the shared global, so each request uses its own values with no cross-request interference. Add regression tests covering the per-request read, the CLI-default fallback, and an end-to-end proxy check that the global is no longer mutated by a request. Fixes #304 * Bump version to 0.3.17 Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> --------- Co-authored-by: Asankhaya Sharma <codelion@users.noreply.github.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
1 parent 159d340 commit 6901e2d

4 files changed

Lines changed: 181 additions & 8 deletions

File tree

‎optillm/__init__.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
# Version information
2-
__version__ = "0.3.16"
2+
__version__ = "0.3.17"
33

44
import os as _os
55

‎optillm/server.py‎

Lines changed: 17 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -416,8 +416,15 @@ def execute_single_approach(approach, system_prompt, initial_query, client, mode
416416
# since the full token count is already in the response
417417
return response, 0
418418
elif approach == 'mcts':
419-
return chat_with_mcts(system_prompt, initial_query, client, model, server_config['mcts_simulations'],
420-
server_config['mcts_exploration'], server_config['mcts_depth'], request_config, request_id)
419+
# Read the MCTS params from the per-request config (falling back to
420+
# the server defaults) instead of the shared global server_config.
421+
# Reading them off the global races with concurrent requests that
422+
# overwrite it between the write in proxy() and this read (issue #304).
423+
mcts_config = request_config or {}
424+
return chat_with_mcts(system_prompt, initial_query, client, model,
425+
mcts_config.get('mcts_simulations', server_config['mcts_simulations']),
426+
mcts_config.get('mcts_exploration', server_config['mcts_exploration']),
427+
mcts_config.get('mcts_depth', server_config['mcts_depth']), request_config, request_id)
421428
elif approach == 'bon':
422429
return best_of_n_sampling(system_prompt, initial_query, client, model, server_config['best_of_n'], request_config, request_id)
423430
elif approach == 'moa':
@@ -744,9 +751,14 @@ def proxy():
744751

745752
optillm_approach = data.get('optillm_approach', server_config['approach'])
746753
logger.debug(data)
747-
server_config['mcts_depth'] = data.get('mcts_depth', server_config['mcts_depth'])
748-
server_config['mcts_exploration'] = data.get('mcts_exploration', server_config['mcts_exploration'])
749-
server_config['mcts_simulations'] = data.get('mcts_simulations', server_config['mcts_simulations'])
754+
# MCTS params are resolved per-request in execute_single_approach: any
755+
# client-sent mcts_* values are already in request_config (copied from the
756+
# request body above) and the server defaults are used as the fallback there.
757+
# We deliberately do NOT write them back into the shared global server_config:
758+
# concurrent requests would overwrite each other's values between the write
759+
# here and the read in execute_single_approach, silently corrupting results
760+
# under any threaded WSGI server (and the write also leaked one request's
761+
# params into every later request). See issue #304.
750762

751763
system_prompt, initial_query, message_optillm_approach = parse_conversation(messages)
752764

‎pyproject.toml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
44

55
[project]
66
name = "optillm"
7-
version = "0.3.16"
7+
version = "0.3.17"
88
description = "An optimizing inference proxy for LLMs."
99
readme = "README.md"
1010
license = "Apache-2.0"

‎tests/test_approaches.py‎

Lines changed: 162 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -207,12 +207,173 @@ def test_approaches_handle_bad_responses():
207207
print("✅ approaches handle bad responses gracefully")
208208

209209

210+
def test_mcts_params_are_request_scoped():
211+
"""MCTS params must be read from the per-request config, not the shared
212+
global ``server_config`` (issue #304).
213+
214+
Before the fix, ``proxy`` wrote the request's ``mcts_*`` params into the
215+
module-level ``server_config`` and ``execute_single_approach`` read them
216+
back out of it. Under Flask's threaded mode (or gunicorn/uWSGI), two
217+
concurrent requests race on that global: request B overwrites request A's
218+
params between A's write and A's read, silently corrupting A's MCTS run.
219+
220+
This pins the fix: a per-request ``mcts_*`` value must flow through to
221+
``chat_with_mcts`` untouched, and the shared global must not be mutated as a
222+
side effect of running a request.
223+
"""
224+
import optillm.server as server
225+
226+
captured = {}
227+
228+
def fake_chat_with_mcts(system_prompt, initial_query, client, model,
229+
num_simulations, exploration_weight, depth,
230+
request_config=None, request_id=None):
231+
captured['num_simulations'] = num_simulations
232+
captured['exploration_weight'] = exploration_weight
233+
captured['depth'] = depth
234+
return "mock mcts response", 0
235+
236+
original_mcts = server.chat_with_mcts
237+
saved_defaults = {k: server.server_config[k]
238+
for k in ('mcts_simulations', 'mcts_exploration', 'mcts_depth')}
239+
try:
240+
server.chat_with_mcts = fake_chat_with_mcts
241+
# Pin the global defaults to values distinct from the per-request ones,
242+
# so a read from the global is observably different from a per-request read.
243+
server.server_config['mcts_simulations'] = 2
244+
server.server_config['mcts_exploration'] = 0.2
245+
server.server_config['mcts_depth'] = 1
246+
247+
# Case 1: a per-request override must reach chat_with_mcts verbatim.
248+
per_request = {'mcts_simulations': 99, 'mcts_exploration': 0.9, 'mcts_depth': 7}
249+
server.execute_single_approach(
250+
'mcts', "You are a helpful assistant.", "What is 2 + 2?",
251+
MockClient(), "mock-model",
252+
request_config=dict(per_request), request_id="req-304")
253+
254+
assert captured.get('num_simulations') == 99, (
255+
f"expected per-request mcts_simulations=99, got {captured.get('num_simulations')} "
256+
"— params were read from the shared global instead of request_config (issue #304)")
257+
assert captured.get('exploration_weight') == 0.9, (
258+
f"expected per-request mcts_exploration=0.9, got {captured.get('exploration_weight')}")
259+
assert captured.get('depth') == 7, (
260+
f"expected per-request mcts_depth=7, got {captured.get('depth')}")
261+
262+
# Running a request must not mutate the shared global — that mutation is
263+
# exactly what let concurrent requests corrupt each other.
264+
assert server.server_config['mcts_simulations'] == 2, (
265+
"execute_single_approach must not write per-request MCTS params back "
266+
"into the global server_config (issue #304)")
267+
268+
# Case 2: with no per-request override (e.g. a user who set --simulations
269+
# globally via the CLI), the server defaults must still be used.
270+
captured.clear()
271+
server.execute_single_approach(
272+
'mcts', "You are a helpful assistant.", "What is 2 + 2?",
273+
MockClient(), "mock-model",
274+
request_config={}, request_id="req-304b")
275+
276+
assert captured.get('num_simulations') == 2, (
277+
f"expected the server-default mcts_simulations=2 when no per-request "
278+
f"value is sent, got {captured.get('num_simulations')}")
279+
assert captured.get('exploration_weight') == 0.2, (
280+
f"expected the server-default mcts_exploration=0.2, got {captured.get('exploration_weight')}")
281+
assert captured.get('depth') == 1, (
282+
f"expected the server-default mcts_depth=1, got {captured.get('depth')}")
283+
finally:
284+
server.chat_with_mcts = original_mcts
285+
server.server_config.update(saved_defaults)
286+
287+
print("✅ mcts params are request-scoped (issue #304)")
288+
289+
290+
def test_proxy_does_not_mutate_global_mcts_config():
291+
"""End-to-end guard for the #304 fix at the proxy layer.
292+
293+
Drives a real request through the Flask ``/v1/chat/completions`` handler with
294+
per-request MCTS params and asserts the shared global ``server_config`` is
295+
NOT mutated as a side effect (the mutation is what let concurrent requests
296+
corrupt each other), while the per-request values still reach ``chat_with_mcts``.
297+
"""
298+
import optillm.server as server
299+
300+
class _FakeResp:
301+
def model_dump(self):
302+
return {"choices": [{"index": 0,
303+
"message": {"role": "assistant", "content": "ok"},
304+
"finish_reason": "stop"}],
305+
"usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}}
306+
307+
class _FakeClient:
308+
def __init__(self):
309+
self.chat = type("_Chat", (), {
310+
"completions": type("_Completions", (), {"create": lambda self, **k: _FakeResp()})()
311+
})()
312+
313+
captured = {}
314+
315+
def fake_chat_with_mcts(system_prompt, initial_query, client, model,
316+
num_simulations, exploration_weight, depth,
317+
request_config=None, request_id=None):
318+
captured['num_simulations'] = num_simulations
319+
captured['exploration_weight'] = exploration_weight
320+
captured['depth'] = depth
321+
return "mock mcts response", 5
322+
323+
original_get_config = server.get_config
324+
original_mcts = server.chat_with_mcts
325+
saved_defaults = {k: server.server_config[k]
326+
for k in ('mcts_simulations', 'mcts_exploration', 'mcts_depth')}
327+
try:
328+
server.get_config = lambda: (_FakeClient(), 'optillm')
329+
server.chat_with_mcts = fake_chat_with_mcts
330+
server.server_config['mcts_simulations'] = 2
331+
server.server_config['mcts_exploration'] = 0.2
332+
server.server_config['mcts_depth'] = 1
333+
334+
server.app.config['TESTING'] = True
335+
client = server.app.test_client()
336+
resp = client.post(
337+
'/v1/chat/completions',
338+
json={"model": "gpt-4o-mini",
339+
"messages": [{"role": "user", "content": "hello"}],
340+
"optillm_approach": "mcts",
341+
"mcts_simulations": 99, "mcts_exploration": 0.9, "mcts_depth": 7},
342+
headers={"Authorization": "Bearer optillm"})
343+
344+
assert resp.status_code == 200, f"expected 200, got {resp.status_code}: {resp.get_data(as_text=True)[:200]}"
345+
346+
# Per-request params must have reached the MCTS approach.
347+
assert captured.get('num_simulations') == 99, (
348+
f"per-request mcts_simulations should reach chat_with_mcts, got {captured.get('num_simulations')}")
349+
assert captured.get('depth') == 7, (
350+
f"per-request mcts_depth should reach chat_with_mcts, got {captured.get('depth')}")
351+
352+
# ...but the shared global must be untouched (issue #304). On the buggy
353+
# code proxy() wrote these back into server_config, so this would be 99/0.9/7.
354+
assert server.server_config['mcts_simulations'] == 2, (
355+
f"proxy must not write per-request mcts_simulations into the global "
356+
f"server_config, found {server.server_config['mcts_simulations']} (issue #304)")
357+
assert server.server_config['mcts_exploration'] == 0.2, (
358+
f"proxy leaked mcts_exploration into the global: {server.server_config['mcts_exploration']}")
359+
assert server.server_config['mcts_depth'] == 1, (
360+
f"proxy leaked mcts_depth into the global: {server.server_config['mcts_depth']}")
361+
finally:
362+
server.get_config = original_get_config
363+
server.chat_with_mcts = original_mcts
364+
server.server_config.update(saved_defaults)
365+
366+
print("✅ proxy does not mutate global mcts config (issue #304)")
367+
368+
210369
if __name__ == "__main__":
211370
print("Running approach tests...")
212371

213372
test_approach_imports()
214373
test_basic_approach_calls()
215374
test_approach_parameters()
216375
test_approaches_handle_bad_responses()
376+
test_mcts_params_are_request_scoped()
377+
test_proxy_does_not_mutate_global_mcts_config()
217378

218-
print("\nAll tests completed!")
379+
print("\nAll tests completed!")

0 commit comments

Comments
 (0)