From a9e8a36f89c38fdc528fb1c4b537abd0a28ef570 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Wed, 16 Apr 2025 09:47:10 -0700 Subject: [PATCH] [Bug Fix] Azure Blob Storage fixes (#10059) * Simple fix for #9339 - upgrade the underlying library and cache the azure storage client (#9965) * fix - use constants for caching azure storage client --------- Co-authored-by: Adrian Lyjak --- litellm/constants.py | 3 ++ .../azure_storage/azure_storage.py | 44 ++++++++++++++----- litellm/llms/custom_httpx/http_handler.py | 2 +- requirements.txt | 4 +- 4 files changed, 38 insertions(+), 15 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index 0f35520b2e..f25c7c71e4 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -21,6 +21,9 @@ DEFAULT_MAX_TOKENS = 256 # used when providers need a default MAX_SIZE_PER_ITEM_IN_MEMORY_CACHE_IN_KB = 1024 # 1MB = 1024KB SINGLE_DEPLOYMENT_TRAFFIC_FAILURE_THRESHOLD = 1000 # Minimum number of requests to consider "reasonable traffic". Used for single-deployment cooldown logic. +########## Networking constants ############################################################## +_DEFAULT_TTL_FOR_HTTPX_CLIENTS = 3600 # 1 hour, re-use the same httpx client for 1 hour + ########### v2 Architecture constants for managing writing updates to the database ########### REDIS_UPDATE_BUFFER_KEY = "litellm_spend_update_buffer" REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY = "litellm_daily_spend_update_buffer" diff --git a/litellm/integrations/azure_storage/azure_storage.py b/litellm/integrations/azure_storage/azure_storage.py index 24a3f5749a..6ffb1e542f 100644 --- a/litellm/integrations/azure_storage/azure_storage.py +++ b/litellm/integrations/azure_storage/azure_storage.py @@ -1,12 +1,13 @@ import asyncio import json import os +import time import uuid from datetime import datetime, timedelta from typing import List, Optional from litellm._logging import verbose_logger -from litellm.constants import AZURE_STORAGE_MSFT_VERSION +from litellm.constants import _DEFAULT_TTL_FOR_HTTPX_CLIENTS, AZURE_STORAGE_MSFT_VERSION from litellm.integrations.custom_batch_logger import CustomBatchLogger from litellm.llms.azure.common_utils import get_azure_ad_token_from_entra_id from litellm.llms.custom_httpx.http_handler import ( @@ -48,14 +49,17 @@ 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[ - str - ] = None # the Azure AD token to use for Azure Storage API requests - self.token_expiry: Optional[ - datetime - ] = None # the expiry time of the currentAzure AD token + self.azure_auth_token: Optional[str] = ( + None # the Azure AD token to use for Azure Storage API requests + ) + self.token_expiry: Optional[datetime] = ( + None # the expiry time of the currentAzure AD token + ) asyncio.create_task(self.periodic_flush()) self.flush_lock = asyncio.Lock() @@ -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() + _DEFAULT_TTL_FOR_HTTPX_CLIENTS + 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 diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index f1aa5627dc..627dd8c9f9 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/litellm/llms/custom_httpx/http_handler.py @@ -8,6 +8,7 @@ import httpx from httpx import USE_CLIENT_DEFAULT, AsyncHTTPTransport, HTTPTransport import litellm +from litellm.constants import _DEFAULT_TTL_FOR_HTTPX_CLIENTS from litellm.litellm_core_utils.logging_utils import track_llm_api_timing from litellm.types.llms.custom_http import * @@ -31,7 +32,6 @@ headers = { # https://www.python-httpx.org/advanced/timeouts _DEFAULT_TIMEOUT = httpx.Timeout(timeout=5.0, connect=5.0) -_DEFAULT_TTL_FOR_HTTPX_CLIENTS = 3600 # 1 hour, re-use the same httpx client for 1 hour def mask_sensitive_info(error_message): diff --git a/requirements.txt b/requirements.txt index dca68b6f93..ebc525bafa 100644 --- a/requirements.txt +++ b/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 -#### \ No newline at end of file +####