From b5bfcca5175196f8bf02bca86a644fc69a5dfd15 Mon Sep 17 00:00:00 2001 From: Alexsander Hamir Date: Tue, 10 Feb 2026 17:29:35 -0800 Subject: [PATCH] Fix OOM issues across integrations --- .semgrep/rules/python/unbounded-memory.yml | 7 +++---- cookbook/nova_sonic_realtime.py | 6 +++++- docs/my-website/docs/proxy/config_settings.md | 1 + litellm/constants.py | 4 ++++ litellm/integrations/gcs_bucket/gcs_bucket.py | 14 ++++++++++---- .../db/db_transaction_queue/base_update_queue.py | 8 ++++++-- .../daily_spend_update_queue.py | 3 ++- .../db/db_transaction_queue/spend_update_queue.py | 5 ++++- 8 files changed, 35 insertions(+), 13 deletions(-) diff --git a/.semgrep/rules/python/unbounded-memory.yml b/.semgrep/rules/python/unbounded-memory.yml index 18b18b7944a..811ef689344 100644 --- a/.semgrep/rules/python/unbounded-memory.yml +++ b/.semgrep/rules/python/unbounded-memory.yml @@ -1,15 +1,14 @@ # Unbounded memory growth – data structures without a clear max limit -# Can lead to OOM under load or with unbounded input. +# Can lead to OOM under load. rules: - # asyncio.Queue() with no maxsize is unbounded – bad pattern for integrations (log queues, etc.) - id: unbounded-asyncio-queue message: asyncio.Queue() with no maxsize can grow unbounded. Use asyncio.Queue(maxsize=N) for integrations (e.g. log queues). - severity: WARNING + severity: ERROR languages: [python] pattern-either: - pattern: asyncio.Queue() - pattern: asyncio.Queue(maxsize=0) metadata: category: correctness - cwe: "CWE-400: Uncontrolled Resource Consumption" + cwe: "CWE-400: Uncontrolled Resource Consumption" \ No newline at end of file diff --git a/cookbook/nova_sonic_realtime.py b/cookbook/nova_sonic_realtime.py index 0ea0badfb01..c7a73c1d00f 100644 --- a/cookbook/nova_sonic_realtime.py +++ b/cookbook/nova_sonic_realtime.py @@ -16,10 +16,14 @@ Usage: import asyncio import base64 import json +import os import pyaudio import websockets from typing import Optional +# Bounded queue size for audio chunks (configurable via env to avoid unbounded memory) +AUDIO_QUEUE_MAXSIZE = int(os.getenv("LITELLM_ASYNCIO_QUEUE_MAXSIZE", 10_000)) + # Audio configuration (matching Nova Sonic requirements) INPUT_SAMPLE_RATE = 16000 # Nova Sonic expects 16kHz input OUTPUT_SAMPLE_RATE = 24000 # Nova Sonic outputs 24kHz @@ -40,7 +44,7 @@ class RealtimeClient: self.api_key = api_key self.ws: Optional[websockets.WebSocketClientProtocol] = None self.is_active = False - self.audio_queue = asyncio.Queue() + self.audio_queue = asyncio.Queue(maxsize=AUDIO_QUEUE_MAXSIZE) self.pyaudio = pyaudio.PyAudio() self.input_stream = None self.output_stream = None diff --git a/docs/my-website/docs/proxy/config_settings.md b/docs/my-website/docs/proxy/config_settings.md index 5b6c6669b91..52504faede8 100644 --- a/docs/my-website/docs/proxy/config_settings.md +++ b/docs/my-website/docs/proxy/config_settings.md @@ -784,6 +784,7 @@ router_settings: | LITELLM_USER_AGENT | Custom user agent string for LiteLLM API requests. Used for partner telemetry attribution | LITELLM_PRINT_STANDARD_LOGGING_PAYLOAD | If true, prints the standard logging payload to the console - useful for debugging | LITELM_ENVIRONMENT | Environment for LiteLLM Instance. This is currently only logged to DeepEval to determine the environment for DeepEval integration. +| LITELLM_ASYNCIO_QUEUE_MAXSIZE | Maximum size for asyncio queues (e.g. log queues, spend update queues, and cookbook examples such as realtime audio in `nova_sonic_realtime.py`). Bounds in-memory growth to prevent OOM. Default is 1000. | LOGFIRE_TOKEN | Token for Logfire logging service | LOGFIRE_BASE_URL | Base URL for Logfire logging service (useful for self hosted deployments) | LOGGING_WORKER_CONCURRENCY | Maximum number of concurrent coroutine slots for the logging worker on the asyncio event loop. Default is 100. Setting too high will flood the event loop with logging tasks which will lower the overall latency of the requests. diff --git a/litellm/constants.py b/litellm/constants.py index 180315ace0e..2bad2e212b6 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -213,6 +213,10 @@ REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY = "litellm_daily_agent_spend_update_bu REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY = "litellm_daily_tag_spend_update_buffer" MAX_REDIS_BUFFER_DEQUEUE_COUNT = int(os.getenv("MAX_REDIS_BUFFER_DEQUEUE_COUNT", 100)) MAX_SIZE_IN_MEMORY_QUEUE = int(os.getenv("MAX_SIZE_IN_MEMORY_QUEUE", 2000)) +# Bounds asyncio.Queue() instances (log queues, spend update queues, etc.) to prevent unbounded memory growth +LITELLM_ASYNCIO_QUEUE_MAXSIZE = int( + os.getenv("LITELLM_ASYNCIO_QUEUE_MAXSIZE", 1000) +) MAX_IN_MEMORY_QUEUE_FLUSH_COUNT = int( os.getenv("MAX_IN_MEMORY_QUEUE_FLUSH_COUNT", 1000) ) diff --git a/litellm/integrations/gcs_bucket/gcs_bucket.py b/litellm/integrations/gcs_bucket/gcs_bucket.py index 3cb62905531..0f1ba4a4093 100644 --- a/litellm/integrations/gcs_bucket/gcs_bucket.py +++ b/litellm/integrations/gcs_bucket/gcs_bucket.py @@ -9,6 +9,7 @@ from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple from urllib.parse import quote from litellm._logging import verbose_logger +from litellm.constants import LITELLM_ASYNCIO_QUEUE_MAXSIZE from litellm.integrations.additional_logging_utils import AdditionalLoggingUtils from litellm.integrations.gcs_bucket.gcs_bucket_base import GCSBucketBase from litellm.proxy._types import CommonProxyErrors @@ -41,7 +42,9 @@ class GCSBucketLogger(GCSBucketBase, AdditionalLoggingUtils): batch_size=self.batch_size, flush_interval=self.flush_interval, ) - self.log_queue: asyncio.Queue[GCSLogQueueItem] = asyncio.Queue() # type: ignore[assignment] + self.log_queue: asyncio.Queue[GCSLogQueueItem] = asyncio.Queue( # type: ignore[assignment] + maxsize=LITELLM_ASYNCIO_QUEUE_MAXSIZE + ) asyncio.create_task(self.periodic_flush()) AdditionalLoggingUtils.__init__(self) @@ -69,6 +72,9 @@ class GCSBucketLogger(GCSBucketBase, AdditionalLoggingUtils): ) if logging_payload is None: raise ValueError("standard_logging_object not found in kwargs") + # When queue is at maxsize, flush immediately to make room (no blocking, no data dropped) + if self.log_queue.full(): + await self.flush_queue() await self.log_queue.put( GCSLogQueueItem( payload=logging_payload, kwargs=kwargs, response_obj=response_obj @@ -91,9 +97,9 @@ class GCSBucketLogger(GCSBucketBase, AdditionalLoggingUtils): ) if logging_payload is None: raise ValueError("standard_logging_object not found in kwargs") - # Add to logging queue - this will be flushed periodically - # Use asyncio.Queue.put() for thread-safe concurrent access - # If queue is full, this will block until space is available (backpressure) + # When queue is at maxsize, flush immediately to make room (no blocking, no data dropped) + if self.log_queue.full(): + await self.flush_queue() await self.log_queue.put( GCSLogQueueItem( payload=logging_payload, kwargs=kwargs, response_obj=response_obj diff --git a/litellm/proxy/db/db_transaction_queue/base_update_queue.py b/litellm/proxy/db/db_transaction_queue/base_update_queue.py index 202829b78b6..a5ec1c3eaf4 100644 --- a/litellm/proxy/db/db_transaction_queue/base_update_queue.py +++ b/litellm/proxy/db/db_transaction_queue/base_update_queue.py @@ -10,14 +10,18 @@ from litellm._service_logger import ServiceLogging service_logger_obj = ( ServiceLogging() ) # used for tracking metrics for In memory buffer, redis buffer, pod lock manager -from litellm.constants import MAX_IN_MEMORY_QUEUE_FLUSH_COUNT, MAX_SIZE_IN_MEMORY_QUEUE +from litellm.constants import ( + LITELLM_ASYNCIO_QUEUE_MAXSIZE, + MAX_IN_MEMORY_QUEUE_FLUSH_COUNT, + MAX_SIZE_IN_MEMORY_QUEUE, +) class BaseUpdateQueue: """Base class for in memory buffer for database transactions""" def __init__(self): - self.update_queue = asyncio.Queue() + self.update_queue = asyncio.Queue(maxsize=LITELLM_ASYNCIO_QUEUE_MAXSIZE) self.MAX_SIZE_IN_MEMORY_QUEUE = MAX_SIZE_IN_MEMORY_QUEUE async def add_update(self, update): diff --git a/litellm/proxy/db/db_transaction_queue/daily_spend_update_queue.py b/litellm/proxy/db/db_transaction_queue/daily_spend_update_queue.py index c3074e641b2..5ba8fb13596 100644 --- a/litellm/proxy/db/db_transaction_queue/daily_spend_update_queue.py +++ b/litellm/proxy/db/db_transaction_queue/daily_spend_update_queue.py @@ -3,6 +3,7 @@ from copy import deepcopy from typing import Dict, List, Optional from litellm._logging import verbose_proxy_logger +from litellm.constants import LITELLM_ASYNCIO_QUEUE_MAXSIZE from litellm.proxy._types import BaseDailySpendTransaction from litellm.proxy.db.db_transaction_queue.base_update_queue import ( BaseUpdateQueue, @@ -54,7 +55,7 @@ class DailySpendUpdateQueue(BaseUpdateQueue): def __init__(self): super().__init__() self.update_queue: asyncio.Queue[Dict[str, BaseDailySpendTransaction]] = ( - asyncio.Queue() + asyncio.Queue(maxsize=LITELLM_ASYNCIO_QUEUE_MAXSIZE) ) async def add_update(self, update: Dict[str, BaseDailySpendTransaction]): diff --git a/litellm/proxy/db/db_transaction_queue/spend_update_queue.py b/litellm/proxy/db/db_transaction_queue/spend_update_queue.py index 9b0449bb9ab..c96564252d0 100644 --- a/litellm/proxy/db/db_transaction_queue/spend_update_queue.py +++ b/litellm/proxy/db/db_transaction_queue/spend_update_queue.py @@ -2,6 +2,7 @@ import asyncio from typing import Dict, List, Optional from litellm._logging import verbose_proxy_logger +from litellm.constants import LITELLM_ASYNCIO_QUEUE_MAXSIZE from litellm.proxy._types import ( DBSpendUpdateTransactions, Litellm_EntityType, @@ -21,7 +22,9 @@ class SpendUpdateQueue(BaseUpdateQueue): def __init__(self): super().__init__() - self.update_queue: asyncio.Queue[SpendUpdateQueueItem] = asyncio.Queue() + self.update_queue: asyncio.Queue[SpendUpdateQueueItem] = asyncio.Queue( + maxsize=LITELLM_ASYNCIO_QUEUE_MAXSIZE + ) async def flush_and_get_aggregated_db_spend_update_transactions( self,