Skip to content

Commit 8a1ee06

Browse files
authored
fix(core): make OGXAsLibraryClient thread-safe (#5773)
# What does this PR do? Rework the sync-on-async client. The previous implementation called `loop.run_until_complete()` on a shared loop instance inside `.request()` method. When multiple threads called `request()` concurrently, this raised `RuntimeError: This event loop is already running`. The fix uses a single dedicated daemon thread running `loop.run_forever()`, with all coroutines submitted with a `asyncio.run_coroutine_threadsafe()` call. ## Known limitations and tradeoffs - **Not fully thread-safe**: per @mattf's comment on #5752, global state in the async client may still cause races under concurrent load. This PR fixes the event loop collision specifically. - Uses a bounded `threading.Queue` with busy-poll fallback (`asyncio.sleep(0.01)`) when full. My judgement call is: This is good enough for LLM token rates. - `future.cancel()` on the streaming coroutine is best-effort; a coroutine blocked on network I/O won't respond immediately. - Decided against the use of [janus](https://github.com/aio-libs/janus) to avoid extra dependencies. Closes #5752 ## Test Plan Added unit tests covering: - Concurrent requests from multiple threads. - Init failure cleanup (background thread stops on failed initialization). Launch just them with: ``` pytest tests/unit/ -vs -k TestOGXAsLibraryClientSyncOnAsync ```
1 parent d56fc4a commit 8a1ee06

2 files changed

Lines changed: 211 additions & 42 deletions

File tree

src/ogx/core/library_client.py

Lines changed: 137 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -5,11 +5,15 @@
55
# the root directory of this source tree.
66

77
import asyncio
8+
import atexit
9+
import concurrent.futures
810
import inspect
911
import json
1012
import logging # allow-direct-logging
1113
import os
14+
import queue
1215
import sys
16+
import threading
1317
import typing
1418
from collections.abc import AsyncGenerator, Generator
1519
from enum import Enum
@@ -54,6 +58,12 @@
5458

5559
T = TypeVar("T")
5660

61+
_INIT_TIMEOUT: float = 60.0
62+
_SHUTDOWN_TIMEOUT: float = 10.0
63+
_CLEANUP_TIMEOUT: float = 5.0
64+
_HANG_GUARD_TIMEOUT: float = 600.0
65+
_STREAM_HEARTBEAT_INTERVAL: float = 1.0
66+
5767

5868
def convert_pydantic_to_json_value(value: Any) -> Any:
5969
"""Recursively convert Pydantic models, enums, and nested structures to JSON-serializable values.
@@ -156,7 +166,11 @@ def __init__(self, response: FastAPIResponse) -> None:
156166

157167

158168
class OGXAsLibraryClient(OgxClient):
159-
"""Synchronous client that runs a OGX distribution in-process as a library."""
169+
"""Synchronous client that runs a OGX distribution in-process as a library.
170+
171+
This is a sync-on-async implementation wrapping `AsyncOGXAsLibraryClient` class,
172+
starting a daemon loop thread, which will be shut down when the main thread exits.
173+
"""
160174

161175
def __init__(
162176
self,
@@ -170,24 +184,39 @@ def __init__(
170184
config_path_or_distro_name, custom_provider_registry, provider_data, skip_logger_removal
171185
)
172186
self.provider_data = provider_data
187+
self._shutdown_lock: threading.Lock = threading.Lock()
188+
self._shutdown = False
173189

190+
# stick with one loop and run it in a dedicated daemon thread
174191
self.loop = asyncio.new_event_loop()
192+
self.loop_thread = threading.Thread(
193+
target=self._run_event_loop, daemon=True, name="ogx-lib-sync-client-event-loop"
194+
)
195+
self.loop_thread.start()
175196

176-
# use a new event loop to avoid interfering with the main event loop
177-
loop = asyncio.new_event_loop()
178-
asyncio.set_event_loop(loop)
179197
try:
180-
loop.run_until_complete(self.async_client.initialize())
198+
future = asyncio.run_coroutine_threadsafe(self.async_client.initialize(), self.loop)
199+
future.result(timeout=_INIT_TIMEOUT) # Block until initialization completes + timeout if hangs
200+
except Exception:
201+
self.loop.call_soon_threadsafe(self.loop.stop)
202+
self.loop_thread.join(timeout=_CLEANUP_TIMEOUT)
203+
raise
204+
205+
atexit.register(self.shutdown) # Safety net: if the user forgets to shutdown properly
206+
207+
def _run_event_loop(self) -> None:
208+
"""Runs forever in the background thread."""
209+
asyncio.set_event_loop(self.loop)
210+
try:
211+
self.loop.run_forever()
181212
finally:
182-
asyncio.set_event_loop(None)
213+
self.loop.close() # Close the loop when the thread is instructed to stop
183214

184215
def initialize(self) -> None:
185-
"""
186-
Deprecated method for backward compatibility.
187-
"""
216+
"""Deprecated method for backward compatibility."""
188217
pass
189218

190-
def shutdown(self) -> None:
219+
def shutdown(self, timeout: float = _SHUTDOWN_TIMEOUT) -> None:
191220
"""Shutdown the client and release all resources.
192221
193222
This method should be called when you're done using the client to properly
@@ -197,18 +226,40 @@ def shutdown(self) -> None:
197226
198227
This method is idempotent and can be called multiple times safely.
199228
229+
Args:
230+
timeout: Maximum seconds to wait for graceful shutdown before forcing close.
231+
232+
**IMPORTANT!** `shutdown()` is not safe to call concurrently with requests!
233+
Use the client as a context manager to assure proper shutdown.
234+
200235
Example:
201-
client = OGXAsLibraryClient("starter")
202-
# ... use the client ...
203-
client.shutdown()
236+
with OGXAsLibraryClient("starter") as client:
237+
# ... use the client ...
204238
"""
205-
loop = self.loop
206-
asyncio.set_event_loop(loop)
239+
# Guard against calling shutdown before init finishes, or multiple times
240+
with self._shutdown_lock:
241+
if self._shutdown:
242+
return
243+
self._shutdown = True
244+
if not self.loop.is_running():
245+
return
246+
247+
future = asyncio.run_coroutine_threadsafe(self.async_client.shutdown(), self.loop)
207248
try:
208-
loop.run_until_complete(self.async_client.shutdown())
209-
finally:
210-
loop.close()
211-
asyncio.set_event_loop(None)
249+
future.result(timeout=timeout)
250+
except concurrent.futures.TimeoutError:
251+
logger.warning("Async client shutdown timed out", timeout=timeout)
252+
future.cancel()
253+
except Exception as e:
254+
logger.warning("Unexpected error during async client shutdown", exception=e)
255+
256+
# Safely instruct the background loop to stop
257+
self.loop.call_soon_threadsafe(self.loop.stop)
258+
259+
# Wait for the thread to actually exit
260+
self.loop_thread.join(timeout=timeout)
261+
if self.loop_thread.is_alive():
262+
logger.error("Background event loop thread failed to join (zombie thread)")
212263

213264
def __enter__(self) -> "OGXAsLibraryClient":
214265
"""Enter the context manager.
@@ -227,33 +278,77 @@ def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
227278
self.shutdown()
228279

229280
def request(self, *args: Any, **kwargs: Any) -> Any:
230-
loop = self.loop
231-
asyncio.set_event_loop(loop)
232-
281+
# Route streaming vs non-streaming
233282
if kwargs.get("stream"):
283+
return self._stream_request(*args, **kwargs)
234284

235-
def sync_generator() -> Generator[Any, None, None]:
236-
try:
237-
async_stream = loop.run_until_complete(self.async_client.request(*args, **kwargs))
238-
while True:
239-
chunk = loop.run_until_complete(async_stream.__anext__())
240-
yield chunk
241-
except StopAsyncIteration:
242-
pass
243-
finally:
244-
pending = asyncio.all_tasks(loop)
245-
if pending:
246-
loop.run_until_complete(asyncio.gather(*pending, return_exceptions=True))
247-
248-
return sync_generator()
249-
else:
285+
coro = self.async_client.request(*args, **kwargs)
286+
future = asyncio.run_coroutine_threadsafe(coro, self.loop)
287+
# the giant timeout here is to prevent it from hanging forever:
288+
return future.result(timeout=_HANG_GUARD_TIMEOUT)
289+
290+
def _stream_request(self, *args: Any, **kwargs: Any) -> Generator[Any, None, None]:
291+
"""Thread-safe synchronous generator wrapper around an async generator."""
292+
# 32 chunks of buffering. LLM token rate makes OOM from unbounded queue unlikely
293+
# but a bound prevents runaway memory if the consumer stalls.
294+
q: queue.Queue = queue.Queue(maxsize=32)
295+
296+
async def _consume() -> None:
297+
async_gen = None
250298
try:
251-
result = loop.run_until_complete(self.async_client.request(*args, **kwargs))
299+
async_gen = await self.async_client.request(*args, **kwargs)
300+
async for chunk in async_gen:
301+
while True:
302+
try:
303+
q.put_nowait(("chunk", chunk))
304+
break
305+
except queue.Full:
306+
await asyncio.sleep(0.01)
307+
308+
while True:
309+
try:
310+
q.put_nowait(("done", None))
311+
break
312+
except queue.Full:
313+
await asyncio.sleep(0.01)
314+
315+
except asyncio.CancelledError:
316+
pass
317+
except Exception as err:
318+
while True:
319+
try:
320+
q.put_nowait(("error", err))
321+
break
322+
except queue.Full:
323+
await asyncio.sleep(0.01)
252324
finally:
253-
pending = asyncio.all_tasks(loop)
254-
if pending:
255-
loop.run_until_complete(asyncio.gather(*pending, return_exceptions=True))
256-
return result
325+
if async_gen is not None:
326+
await async_gen.aclose()
327+
328+
future = asyncio.run_coroutine_threadsafe(_consume(), self.loop)
329+
330+
try:
331+
while True:
332+
try:
333+
# Timeout prevents the sync thread from hanging forever if the loop dies or shutdown is called.
334+
msg_type, payload = q.get(timeout=1.0)
335+
except queue.Empty as err:
336+
with self._shutdown_lock:
337+
if self._shutdown:
338+
raise RuntimeError("Client was shut down during streaming") from err
339+
340+
if not self.loop.is_running():
341+
raise RuntimeError("Event loop crashed during streaming") from err
342+
continue
343+
344+
if msg_type == "chunk":
345+
yield payload
346+
elif msg_type == "error":
347+
raise payload
348+
elif msg_type == "done":
349+
break
350+
finally:
351+
future.cancel()
257352

258353

259354
class AsyncOGXAsLibraryClient(AsyncOgxClient):

tests/unit/distribution/test_library_client_initialization.py

Lines changed: 74 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,9 @@
1111
and ready to use immediately after construction.
1212
"""
1313

14+
import threading
15+
import time
16+
1417
import pytest
1518

1619
from ogx.core.library_client import (
@@ -523,3 +526,74 @@ def mock_initialize_route_impls(impls):
523526
client = OGXAsLibraryClient("ci-tests")
524527
assert hasattr(client, "__enter__")
525528
assert hasattr(client, "__exit__")
529+
530+
531+
class TestOGXAsLibraryClientSyncOnAsync:
532+
def test_sync_client_concurrent_requests(self, monkeypatch):
533+
"""Test that multiple threads can make requests concurrently without
534+
RuntimeError: This event loop is already running.
535+
"""
536+
mock_impls = {}
537+
mock_route_impls = RouteImpls({})
538+
539+
class MockStack:
540+
def __init__(self, config, custom_provider_registry=None):
541+
self.impls = mock_impls
542+
543+
async def initialize(self):
544+
pass
545+
546+
async def shutdown(self):
547+
pass
548+
549+
monkeypatch.setattr("ogx.core.library_client.Stack", MockStack)
550+
monkeypatch.setattr("ogx.core.library_client.initialize_route_impls", lambda impls: mock_route_impls)
551+
552+
errors = []
553+
554+
# Mock request to return something
555+
async def mock_request(*args, **kwargs):
556+
return "ok"
557+
558+
with OGXAsLibraryClient("ci-tests") as client:
559+
client.async_client.request = mock_request
560+
561+
def make_request():
562+
try:
563+
client.request(...)
564+
except Exception as e:
565+
errors.append(e)
566+
567+
threads = [threading.Thread(target=make_request) for _ in range(10)]
568+
for t in threads:
569+
t.start()
570+
for t in threads:
571+
t.join()
572+
573+
assert not errors, f"Concurrent requests failed: {errors}"
574+
575+
def test_sync_client_cleanup_on_init_failure(self, monkeypatch):
576+
"""Test that background thread is cleaned up if initialization fails."""
577+
threads_before = set(threading.enumerate())
578+
579+
class MockStack:
580+
def __init__(self, config, custom_provider_registry=None):
581+
self.impls = {}
582+
583+
async def initialize(self):
584+
raise RuntimeError("Init failed intentionally!")
585+
586+
monkeypatch.setattr("ogx.core.library_client.Stack", MockStack)
587+
monkeypatch.setattr("ogx.core.library_client.initialize_route_impls", lambda impls: None)
588+
589+
with pytest.raises(RuntimeError, match="Init failed"):
590+
OGXAsLibraryClient("ci-tests")
591+
592+
deadline = time.monotonic() + 1.0
593+
while time.monotonic() < deadline:
594+
new_threads = set(threading.enumerate()) - threads_before
595+
if not any("ogx-lib-sync-client-event-loop" in t.name for t in new_threads):
596+
break
597+
time.sleep(0.01)
598+
else:
599+
pytest.fail("Background event loop thread was not cleaned up after init failure")

0 commit comments

Comments
 (0)