mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-11 08:38:24 +00:00
refactor(logging_worker): build the aggressive-clear run set immutably
Final tuple instead of an appended list, with the re-enqueue in a _requeue helper, and Final annotations on the overflow regression test Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
parent
a61f8b8036
commit
54c0fa922f
2 changed files with 19 additions and 10 deletions
|
|
@ -423,13 +423,11 @@ class LoggingWorker:
|
|||
return
|
||||
|
||||
extracted_tasks: Final = self._extract_tasks_from_queue()
|
||||
runs: Final = [self._process_single_task(task) for task in extracted_tasks]
|
||||
if new_task is not None:
|
||||
try:
|
||||
self._queue.put_nowait(new_task)
|
||||
except asyncio.QueueFull:
|
||||
self._track_dequeued(new_task)
|
||||
runs.append(self._process_single_task(new_task, counted=False))
|
||||
uncounted: Final = None if new_task is None or self._requeue(new_task) else new_task
|
||||
runs: Final = (
|
||||
*(self._process_single_task(task) for task in extracted_tasks),
|
||||
*(() if uncounted is None else (self._process_single_task(uncounted, counted=False),)),
|
||||
)
|
||||
if runs:
|
||||
await asyncio.gather(*runs)
|
||||
except Exception as e:
|
||||
|
|
@ -438,6 +436,17 @@ class LoggingWorker:
|
|||
# Always reset the flag even if an error occurs
|
||||
self._aggressive_clear_in_progress = False
|
||||
|
||||
def _requeue(self, task: LoggingTask) -> bool:
|
||||
"""Put a rejected task back on the queue once the clear made room; track it as dequeued otherwise."""
|
||||
if self._queue is None:
|
||||
return False
|
||||
try:
|
||||
self._queue.put_nowait(task)
|
||||
except asyncio.QueueFull:
|
||||
self._track_dequeued(task)
|
||||
return False
|
||||
return True
|
||||
|
||||
async def _process_single_task(self, task: LoggingTask, *, counted: bool = True) -> None:
|
||||
"""Process a single task; ``counted`` is False for a task the full queue never accepted."""
|
||||
if self._queue is None:
|
||||
|
|
|
|||
|
|
@ -167,9 +167,9 @@ async def test_overflow_callback_does_not_undercount_the_unfinished_queue() -> N
|
|||
queue.task_done() it never contributed, otherwise flush() returns while a real callback is still
|
||||
running and the shutdown drain cancels it with budget to spare.
|
||||
"""
|
||||
worker = LoggingWorker(timeout=5.0, max_queue_size=2, concurrency=1)
|
||||
gates = {name: asyncio.Event() for name in ("first", "second", "third")}
|
||||
completed: list[str] = []
|
||||
worker: Final = LoggingWorker(timeout=5.0, max_queue_size=2, concurrency=1)
|
||||
gates: Final = {name: asyncio.Event() for name in ("first", "second", "third")}
|
||||
completed: Final[list[str]] = []
|
||||
|
||||
async def gated(name: str) -> None:
|
||||
await gates[name].wait()
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue