diff --git a/litellm/constants.py b/litellm/constants.py index a845b1a49ae..0657d89529f 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -467,6 +467,9 @@ MAX_TIME_TO_CLEAR_QUEUE: Final = float(os.getenv("MAX_TIME_TO_CLEAR_QUEUE", 5.0) LOGGING_WORKER_AGGRESSIVE_CLEAR_COOLDOWN_SECONDS: Final = float( os.getenv("LOGGING_WORKER_AGGRESSIVE_CLEAR_COOLDOWN_SECONDS", 0.5) ) # Cooldown time in seconds before allowing another aggressive clear (default: 0.5s) +LOGGING_EXECUTOR_MAX_THREADS: Final = get_env_int("LOGGING_EXECUTOR_MAX_THREADS", 100) +LOGGING_EXECUTOR_MAX_PENDING_TASKS: Final = get_env_int("LOGGING_EXECUTOR_MAX_PENDING_TASKS", 10_000) +LOGGING_EXECUTOR_DROPPED_TASK_LOG_INTERVAL_SECONDS: Final = 30.0 DD_TRACER_STREAMING_CHUNK_YIELD_RESOURCE: Final = os.getenv( "DD_TRACER_STREAMING_CHUNK_YIELD_RESOURCE", "streaming.chunk.yield" ) diff --git a/litellm/litellm_core_utils/thread_pool_executor.py b/litellm/litellm_core_utils/thread_pool_executor.py index 881a91400df..f989f20247f 100644 --- a/litellm/litellm_core_utils/thread_pool_executor.py +++ b/litellm/litellm_core_utils/thread_pool_executor.py @@ -1,6 +1,82 @@ -from concurrent.futures import ThreadPoolExecutor -from typing import Final +import logging +import threading +import time +from collections.abc import Callable +from concurrent.futures import Future, ThreadPoolExecutor +from typing import Final, ParamSpec, TypeVar -MAX_THREADS: Final = 100 -# Create a ThreadPoolExecutor -executor: Final = ThreadPoolExecutor(max_workers=MAX_THREADS) +from litellm._logging import verbose_logger +from litellm.constants import ( + LOGGING_EXECUTOR_DROPPED_TASK_LOG_INTERVAL_SECONDS, + LOGGING_EXECUTOR_MAX_PENDING_TASKS, + LOGGING_EXECUTOR_MAX_THREADS, +) + +MAX_THREADS: Final = LOGGING_EXECUTOR_MAX_THREADS + +_P = ParamSpec("_P") +_T = TypeVar("_T") + + +class BoundedLoggingThreadPoolExecutor(ThreadPoolExecutor): + """ThreadPoolExecutor with a cap on queued-plus-running tasks. + + The default ThreadPoolExecutor work queue is unbounded, and every queued + logging task pins its request/response payload in memory, so a sustained + burst of sync callbacks slower than request arrival grows memory without + bound. Logging is best-effort: once the cap is reached, new submissions + are dropped with a rate-limited warning instead of queueing forever. + """ + + def __init__( + self, + max_workers: int, + max_pending_tasks: int, + drop_log_interval_seconds: float = LOGGING_EXECUTOR_DROPPED_TASK_LOG_INTERVAL_SECONDS, + logger: logging.Logger = verbose_logger, + ) -> None: + super().__init__(max_workers=max_workers, thread_name_prefix="litellm-logging") + self._max_pending_tasks: Final = max_pending_tasks + self._drop_log_interval_seconds: Final = drop_log_interval_seconds + self._logger: Final = logger + self._pending_slots: Final = threading.Semaphore(max_pending_tasks) + self._drop_lock: Final = threading.Lock() + self._dropped_since_last_log = 0 + self._last_drop_log_time = 0.0 + + def submit(self, fn: Callable[_P, _T], /, *args: _P.args, **kwargs: _P.kwargs) -> Future[_T]: + if not self._pending_slots.acquire(blocking=False): + self._record_drop() + dropped_future: Final[Future[_T]] = Future() + dropped_future.cancel() + return dropped_future + try: + future: Final = super().submit(fn, *args, **kwargs) + except BaseException: + self._pending_slots.release() + raise + future.add_done_callback(lambda _: self._pending_slots.release()) + return future + + def _record_drop(self) -> None: + with self._drop_lock: + self._dropped_since_last_log += 1 + now: Final = time.monotonic() + if now - self._last_drop_log_time < self._drop_log_interval_seconds: + return + dropped_count: Final = self._dropped_since_last_log + self._dropped_since_last_log = 0 + self._last_drop_log_time = now + + self._logger.warning( + "litellm logging executor backlog is full (max_pending_tasks=%s); dropped %s logging task(s) " + "since the last warning. Set LOGGING_EXECUTOR_MAX_PENDING_TASKS to raise the cap.", + self._max_pending_tasks, + dropped_count, + ) + + +executor: Final = BoundedLoggingThreadPoolExecutor( + max_workers=MAX_THREADS, + max_pending_tasks=LOGGING_EXECUTOR_MAX_PENDING_TASKS, +) diff --git a/tests/test_litellm/litellm_core_utils/test_thread_pool_executor.py b/tests/test_litellm/litellm_core_utils/test_thread_pool_executor.py new file mode 100644 index 00000000000..e81de277eaa --- /dev/null +++ b/tests/test_litellm/litellm_core_utils/test_thread_pool_executor.py @@ -0,0 +1,128 @@ +import logging +import threading +import time +from typing import Final + +from litellm._logging import verbose_logger +from litellm.constants import LOGGING_EXECUTOR_MAX_PENDING_TASKS +from litellm.litellm_core_utils.thread_pool_executor import ( + BoundedLoggingThreadPoolExecutor, + executor, +) + + +def test_submit_drops_tasks_when_backlog_is_full(): + release: Final = threading.Event() + started: Final = threading.Event() + ran_first: Final = threading.Event() + ran_second: Final = threading.Event() + ran_dropped: Final = threading.Event() + + def blocking_task(ran: threading.Event) -> None: + ran.set() + started.set() + release.wait(timeout=10) + + pool: Final = BoundedLoggingThreadPoolExecutor(max_workers=1, max_pending_tasks=2) + try: + first: Final = pool.submit(blocking_task, ran_first) + assert started.wait(timeout=10) + second: Final = pool.submit(blocking_task, ran_second) + dropped: Final = pool.submit(blocking_task, ran_dropped) + + assert dropped.cancelled() + assert not first.cancelled() + assert not second.cancelled() + + release.set() + first.result(timeout=10) + second.result(timeout=10) + assert ran_first.is_set() + assert ran_second.is_set() + assert not ran_dropped.is_set() + finally: + release.set() + pool.shutdown(wait=True) + + +def test_submit_releases_slots_after_completion(): + pool: Final = BoundedLoggingThreadPoolExecutor(max_workers=1, max_pending_tasks=1) + + def submit_and_wait() -> str: + future: Final = pool.submit(lambda: "ok") + assert not future.cancelled() + return future.result(timeout=10) + + try: + results: Final = tuple(submit_and_wait() for _ in range(5)) + assert results == ("ok",) * 5 + finally: + pool.shutdown(wait=True) + + +def test_drop_warning_is_rate_limited(caplog): + release: Final = threading.Event() + started: Final = threading.Event() + + def blocking_task() -> None: + started.set() + release.wait(timeout=10) + + drop_logger: Final = logging.getLogger("test_bounded_logging_executor") + pool: Final = BoundedLoggingThreadPoolExecutor( + max_workers=1, + max_pending_tasks=1, + drop_log_interval_seconds=60.0, + logger=drop_logger, + ) + try: + pool.submit(blocking_task) + assert started.wait(timeout=10) + + with caplog.at_level(logging.WARNING, logger=drop_logger.name): + assert pool.submit(time.sleep, 0).cancelled() + assert pool.submit(time.sleep, 0).cancelled() + assert pool.submit(time.sleep, 0).cancelled() + + warnings: Final = tuple(record for record in caplog.records if record.name == drop_logger.name) + assert len(warnings) == 1 + assert warnings[0].args == (1, 1) + finally: + release.set() + pool.shutdown(wait=True) + + +def test_each_drop_warning_counts_only_drops_since_the_last_one(caplog): + release: Final = threading.Event() + started: Final = threading.Event() + + def blocking_task() -> None: + started.set() + release.wait(timeout=10) + + drop_logger: Final = logging.getLogger("test_bounded_logging_executor_every_drop") + pool: Final = BoundedLoggingThreadPoolExecutor( + max_workers=1, + max_pending_tasks=1, + drop_log_interval_seconds=0.0, + logger=drop_logger, + ) + try: + pool.submit(blocking_task) + assert started.wait(timeout=10) + + with caplog.at_level(logging.WARNING, logger=drop_logger.name): + assert pool.submit(time.sleep, 0).cancelled() + assert pool.submit(time.sleep, 0).cancelled() + + warnings: Final = tuple(record for record in caplog.records if record.name == drop_logger.name) + assert tuple(record.args for record in warnings) == ((1, 1), (1, 1)) + finally: + release.set() + pool.shutdown(wait=True) + + +def test_global_executor_is_bounded(): + assert isinstance(executor, BoundedLoggingThreadPoolExecutor) + assert executor._max_pending_tasks == LOGGING_EXECUTOR_MAX_PENDING_TASKS + assert executor._logger is verbose_logger