Simple fix for #9339 - upgrade the underlying library and cache the azure storage client (#9965)

litellm/integrations/azure_storage/azure_storage.py

@@ -1,6 +1,7 @@
 import asyncio
 import json
 import os
+import time
 import uuid
 from datetime import datetime, timedelta
 from typing import List, Optional
@@ -48,6 +49,9 @@ class AzureBlobStorageLogger(CustomBatchLogger):
                     "Missing required environment variable: AZURE_STORAGE_FILE_SYSTEM"
                 )
             self.azure_storage_file_system: str = _azure_storage_file_system
+            self._service_client = None
+            # Time that the azure service client expires, in order to reset the connection pool and keep it fresh
+            self._service_client_timeout: Optional[float] = None
 
             # Internal variables used for Token based authentication
             self.azure_auth_token: Optional[
@@ -324,6 +328,25 @@ class AzureBlobStorageLogger(CustomBatchLogger):
                 f"AzureBlobStorageLogger is only available for premium users. {CommonProxyErrors.not_premium_user}"
             )
 
+    async def get_service_client(self):
+        from azure.storage.filedatalake.aio import DataLakeServiceClient
+
+        # expire old clients to recover from connection issues
+        if (
+            self._service_client_timeout
+            and self._service_client
+            and self._service_client_timeout > time.time()
+        ):
+            await self._service_client.close()
+            self._service_client = None
+        if not self._service_client:
+            self._service_client = DataLakeServiceClient(
+                account_url=f"https://{self.azure_storage_account_name}.dfs.core.windows.net",
+                credential=self.azure_storage_account_key,
+            )
+            self._service_client_timeout = time.time() + 3600
+        return self._service_client
+
     async def upload_to_azure_data_lake_with_azure_account_key(
         self, payload: StandardLoggingPayload
     ):
@@ -332,13 +355,10 @@ class AzureBlobStorageLogger(CustomBatchLogger):
 
         This is used when Azure Storage Account Key is set - Azure Storage Account Key does not work directly with Azure Rest API
         """
-        from azure.storage.filedatalake.aio import DataLakeServiceClient
 
         # Create an async service client
-        service_client = DataLakeServiceClient(
-            account_url=f"https://{self.azure_storage_account_name}.dfs.core.windows.net",
-            credential=self.azure_storage_account_key,
-        )
+
+        service_client = await self.get_service_client()
         # Get file system client
         file_system_client = service_client.get_file_system_client(
             file_system=self.azure_storage_file_system

requirements.txt

@@ -29,7 +29,7 @@ python-multipart==0.0.18 # admin UI
 Pillow==11.0.0
 azure-ai-contentsafety==1.0.0 # for azure content safety
 azure-identity==1.16.1 # for azure content safety
-azure-storage-file-datalake==12.15.0 # for azure buck storage logging
+azure-storage-file-datalake==12.20.0 # for azure buck storage logging
 opentelemetry-api==1.25.0
 opentelemetry-sdk==1.25.0
 opentelemetry-exporter-otlp==1.25.0
@@ -51,4 +51,4 @@ tenacity==8.2.3  # for retrying requests, when litellm.num_retries set
 pydantic==2.10.2 # proxy + openai req.
 jsonschema==4.22.0 # validating json schema
 websockets==13.1.0 # for realtime API
-####
+####