mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-22 00:31:44 +00:00
refactor: Clean up GCS bucket logging code
- Remove redundant comments that don't add value - Remove maxsize queue limit (queue is now unbounded) - Restore original verbose logging patterns from before branch - Simplify code while maintaining functionality
This commit is contained in:
Alexsander Hamir
parent
29c6b8ca52
commit
8b464dc78a
1 changed files with 4 additions and 70 deletions
|
|
@ -28,7 +28,6 @@ class GCSBucketLogger(GCSBucketBase, AdditionalLoggingUtils):
|
|||
|
||||
super().__init__(bucket_name=bucket_name)
|
||||
|
||||
# Init Batch logging settings
|
||||
self.batch_size = int(os.getenv("GCS_BATCH_SIZE", GCS_DEFAULT_BATCH_SIZE))
|
||||
self.flush_interval = int(
|
||||
os.getenv("GCS_FLUSH_INTERVAL", GCS_DEFAULT_FLUSH_INTERVAL_SECONDS)
|
||||
|
|
@ -39,16 +38,10 @@ class GCSBucketLogger(GCSBucketBase, AdditionalLoggingUtils):
|
|||
batch_size=self.batch_size,
|
||||
flush_interval=self.flush_interval,
|
||||
)
|
||||
# Override log_queue with bounded asyncio.Queue for thread-safe concurrent access
|
||||
# Bounded queue prevents OOM when GCS is slower than log accumulation
|
||||
# Default maxsize is 10x batch_size to allow buffering during slow periods
|
||||
queue_maxsize = int(os.getenv("GCS_QUEUE_MAXSIZE", str(self.batch_size * 10)))
|
||||
self.log_queue: asyncio.Queue[GCSLogQueueItem] = asyncio.Queue(maxsize=queue_maxsize) # type: ignore[assignment]
|
||||
self.log_queue: asyncio.Queue[GCSLogQueueItem] = asyncio.Queue() # type: ignore[assignment]
|
||||
asyncio.create_task(self.periodic_flush())
|
||||
AdditionalLoggingUtils.__init__(self)
|
||||
|
||||
print(f"GCS Bucket Logger initialized: bucket_name={bucket_name or 'from env'}, batch_size={self.batch_size}, flush_interval={self.flush_interval}s, queue_maxsize={queue_maxsize}")
|
||||
|
||||
if premium_user is not True:
|
||||
raise ValueError(
|
||||
f"GCS Bucket logging is a premium feature. Please upgrade to use it. {CommonProxyErrors.not_premium_user.value}"
|
||||
|
|
@ -73,25 +66,13 @@ 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)
|
||||
await self.log_queue.put(
|
||||
GCSLogQueueItem(
|
||||
payload=logging_payload, kwargs=kwargs, response_obj=response_obj
|
||||
)
|
||||
)
|
||||
queue_size = self.log_queue.qsize()
|
||||
queue_maxsize = self.log_queue.maxsize if hasattr(self.log_queue, 'maxsize') else None
|
||||
# Warn if queue is getting full (>80% capacity)
|
||||
if queue_maxsize and queue_size > queue_maxsize * 0.8:
|
||||
print(f"GCS Bucket: Success event queued. Queue size: {queue_size}/{queue_maxsize} (WARNING: queue nearly full)")
|
||||
verbose_logger.warning(f"GCS Bucket queue is {queue_size}/{queue_maxsize} full, processing may be slow")
|
||||
else:
|
||||
print(f"GCS Bucket: Success event queued. Queue size: {queue_size}")
|
||||
|
||||
except Exception as e:
|
||||
print(f"GCS Bucket: Error queueing success event: {str(e)}")
|
||||
verbose_logger.exception(f"GCS Bucket logging error: {str(e)}")
|
||||
|
||||
async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time):
|
||||
|
|
@ -115,17 +96,8 @@ class GCSBucketLogger(GCSBucketBase, AdditionalLoggingUtils):
|
|||
payload=logging_payload, kwargs=kwargs, response_obj=response_obj
|
||||
)
|
||||
)
|
||||
queue_size = self.log_queue.qsize()
|
||||
queue_maxsize = self.log_queue.maxsize if hasattr(self.log_queue, 'maxsize') else None
|
||||
# Warn if queue is getting full (>80% capacity)
|
||||
if queue_maxsize and queue_size > queue_maxsize * 0.8:
|
||||
print(f"GCS Bucket: Failure event queued. Queue size: {queue_size}/{queue_maxsize} (WARNING: queue nearly full)")
|
||||
verbose_logger.warning(f"GCS Bucket queue is {queue_size}/{queue_maxsize} full, processing may be slow")
|
||||
else:
|
||||
print(f"GCS Bucket: Failure event queued. Queue size: {queue_size}")
|
||||
|
||||
except Exception as e:
|
||||
print(f"GCS Bucket: Error queueing failure event: {str(e)}")
|
||||
verbose_logger.exception(f"GCS Bucket logging error: {str(e)}")
|
||||
|
||||
def _drain_queue_batch(self) -> List[GCSLogQueueItem]:
|
||||
|
|
@ -164,13 +136,9 @@ class GCSBucketLogger(GCSBucketBase, AdditionalLoggingUtils):
|
|||
"""
|
||||
standard_callback_dynamic_params = kwargs.get("standard_callback_dynamic_params", None) or {}
|
||||
|
||||
# Extract bucket name (dynamic or default)
|
||||
bucket_name = standard_callback_dynamic_params.get("gcs_bucket_name", None) or self.BUCKET_NAME or "default"
|
||||
|
||||
# Extract service account path (dynamic or default)
|
||||
path_service_account = standard_callback_dynamic_params.get("gcs_path_service_account", None) or self.path_service_account_json or "default"
|
||||
|
||||
# Create unique key from both values
|
||||
return f"{bucket_name}|{path_service_account}"
|
||||
|
||||
def _sanitize_config_key(self, config_key: str) -> str:
|
||||
|
|
@ -180,7 +148,6 @@ class GCSBucketLogger(GCSBucketBase, AdditionalLoggingUtils):
|
|||
|
||||
Returns a short hash prefix for safe logging.
|
||||
"""
|
||||
# Use hash to avoid exposing sensitive information in logs
|
||||
hash_obj = hashlib.sha256(config_key.encode('utf-8'))
|
||||
return f"config-{hash_obj.hexdigest()[:8]}"
|
||||
|
||||
|
|
@ -207,7 +174,6 @@ class GCSBucketLogger(GCSBucketBase, AdditionalLoggingUtils):
|
|||
lines = []
|
||||
for item in items:
|
||||
logging_payload = item["payload"]
|
||||
# Serialize each payload as a JSON line
|
||||
json_line = json.dumps(logging_payload, default=str, ensure_ascii=False)
|
||||
lines.append(json_line)
|
||||
return "\n".join(lines)
|
||||
|
|
@ -222,7 +188,6 @@ class GCSBucketLogger(GCSBucketBase, AdditionalLoggingUtils):
|
|||
if not items:
|
||||
return (0, 0)
|
||||
|
||||
# Use first item's kwargs to get the config (all items in group have same config)
|
||||
first_kwargs = items[0]["kwargs"]
|
||||
|
||||
try:
|
||||
|
|
@ -236,36 +201,27 @@ class GCSBucketLogger(GCSBucketBase, AdditionalLoggingUtils):
|
|||
)
|
||||
bucket_name = gcs_logging_config["bucket_name"]
|
||||
|
||||
# Generate batch object name with timestamp and unique ID
|
||||
current_date = self._get_object_date_from_datetime(datetime.now(timezone.utc))
|
||||
batch_id = f"{int(time.time() * 1000)}-{uuid.uuid4().hex[:8]}"
|
||||
object_name = self._generate_batch_object_name(current_date, batch_id)
|
||||
|
||||
# Combine all payloads into NDJSON format
|
||||
combined_payload = self._combine_payloads_to_ndjson(items)
|
||||
|
||||
# Upload single batched object
|
||||
await self._log_json_data_on_gcs(
|
||||
headers=headers,
|
||||
bucket_name=bucket_name,
|
||||
object_name=object_name,
|
||||
logging_payload=combined_payload, # Pass as string (NDJSON)
|
||||
logging_payload=combined_payload,
|
||||
)
|
||||
|
||||
success_count = len(items)
|
||||
error_count = 0
|
||||
sanitized_key = self._sanitize_config_key(config_key)
|
||||
print(f"GCS Bucket: Successfully sent batch of {success_count} logs to bucket '{bucket_name}', object '{object_name}' (config: {sanitized_key})")
|
||||
return (success_count, error_count)
|
||||
|
||||
except Exception as e:
|
||||
# If batch upload fails, count all items as errors
|
||||
success_count = 0
|
||||
error_count = len(items)
|
||||
sanitized_key = self._sanitize_config_key(config_key)
|
||||
print(f"GCS Bucket: Error sending batch to GCS bucket (config: {sanitized_key}): {str(e)}")
|
||||
verbose_logger.exception(
|
||||
f"GCS Bucket error logging batch payload to GCS bucket (config: {sanitized_key}): {str(e)}"
|
||||
f"GCS Bucket error logging batch payload to GCS bucket: {str(e)}"
|
||||
)
|
||||
return (success_count, error_count)
|
||||
|
||||
|
|
@ -290,28 +246,17 @@ class GCSBucketLogger(GCSBucketBase, AdditionalLoggingUtils):
|
|||
|
||||
if not items_to_process:
|
||||
return
|
||||
|
||||
print(f"GCS Bucket: Starting batch send. Processing {len(items_to_process)} items")
|
||||
|
||||
# Group items by GCS config (bucket, credentials) to handle mixed configs
|
||||
# This ensures logs go to the correct bucket with correct credentials
|
||||
grouped_items = self._group_items_by_config(items_to_process)
|
||||
|
||||
if len(grouped_items) > 1:
|
||||
sanitized_keys = [self._sanitize_config_key(key) for key in grouped_items.keys()]
|
||||
print(f"GCS Bucket: Items grouped into {len(grouped_items)} config groups: {sanitized_keys}")
|
||||
|
||||
total_success = 0
|
||||
total_errors = 0
|
||||
|
||||
# Process each config group separately
|
||||
for config_key, group_items in grouped_items.items():
|
||||
success_count, error_count = await self._send_grouped_batch(group_items, config_key)
|
||||
total_success += success_count
|
||||
total_errors += error_count
|
||||
|
||||
print(f"GCS Bucket: Batch send completed. Success: {total_success}, Errors: {total_errors}, Remaining queue size: {self.log_queue.qsize()}")
|
||||
|
||||
def _get_object_name(
|
||||
self, kwargs: Dict, logging_payload: StandardLoggingPayload, response_obj: Any
|
||||
) -> str:
|
||||
|
|
@ -352,7 +297,6 @@ class GCSBucketLogger(GCSBucketBase, AdditionalLoggingUtils):
|
|||
"start_time_utc is required for getting a payload from GCS Bucket"
|
||||
)
|
||||
|
||||
# Try current day, next day, and previous day
|
||||
dates_to_try = [
|
||||
start_time_utc,
|
||||
start_time_utc + timedelta(days=1),
|
||||
|
|
@ -399,26 +343,16 @@ class GCSBucketLogger(GCSBucketBase, AdditionalLoggingUtils):
|
|||
async def flush_queue(self):
|
||||
"""
|
||||
Override flush_queue to work with asyncio.Queue.
|
||||
|
||||
No lock needed: asyncio.Queue.get_nowait() is atomic, and async_send_batch()
|
||||
drains the queue completely, so concurrent flushes just compete for items safely.
|
||||
No qsize() check needed: async_send_batch() handles empty queues gracefully.
|
||||
"""
|
||||
await self.async_send_batch()
|
||||
# Note: async_send_batch() already drains the queue and handles empty case
|
||||
self.last_flush_time = time.time()
|
||||
|
||||
async def periodic_flush(self):
|
||||
"""
|
||||
Override periodic_flush to add queue size observability.
|
||||
Logs the GCS queue size before each flush operation.
|
||||
Override periodic_flush to work with asyncio.Queue.
|
||||
"""
|
||||
while True:
|
||||
await asyncio.sleep(self.flush_interval)
|
||||
queue_size = self.log_queue.qsize()
|
||||
print(
|
||||
f"GCS Bucket queue status: {queue_size} logs queued, batch_size={self.batch_size}, flush_interval={self.flush_interval}s"
|
||||
)
|
||||
verbose_logger.debug(
|
||||
f"GCS Bucket periodic flush after {self.flush_interval} seconds"
|
||||
)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue