From 00677afefed474ad2e754ea93dcb3eee5798fbff Mon Sep 17 00:00:00 2001 From: Krish Dholakia Date: Mon, 21 Jul 2025 20:17:56 -0700 Subject: [PATCH] Litellm batch cost tracking debug (#12782) * feat(proxy_server.py): support batch polling interval allows admin to control batch polling interval (default is 3600s) easier debugging * fix(proxy_settings_endpoint.py): ensure value is actually set before updating env var --- .../proxy/common_utils/check_batch_cost.py | 27 ++++++++++++++----- litellm/constants.py | 1 + litellm/proxy/_new_secret_config.yaml | 25 +++++------------ litellm/proxy/proxy_server.py | 12 ++++++--- .../proxy_setting_endpoints.py | 2 +- 5 files changed, 37 insertions(+), 30 deletions(-) diff --git a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py index d8b8efeef4..6edd198cd8 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -83,15 +83,25 @@ class CheckBatchCost: ) continue - response = await self.llm_router.aretrieve_batch( - model=model_id, - batch_id=batch_id, - litellm_metadata={ - "user_api_key_user_id": job.created_by or "default-user-id", - "batch_ignore_default_logging": True, - }, + verbose_proxy_logger.info( + f"Querying model ID: {model_id} for cost and usage of batch ID: {batch_id}" ) + try: + response = await self.llm_router.aretrieve_batch( + model=model_id, + batch_id=batch_id, + litellm_metadata={ + "user_api_key_user_id": job.created_by or "default-user-id", + "batch_ignore_default_logging": True, + }, + ) + except Exception as e: + verbose_proxy_logger.info( + f"Skipping job {unified_object_id} because of error querying model ID: {model_id} for cost and usage of batch ID: {batch_id}: {e}" + ) + continue + ## RETRIEVE THE BATCH JOB OUTPUT FILE managed_files_obj = cast( Optional[_PROXY_LiteLLMManagedFiles], @@ -102,6 +112,9 @@ class CheckBatchCost: and response.output_file_id is not None and managed_files_obj is not None ): + verbose_proxy_logger.info( + f"Batch ID: {batch_id} is complete, tracking cost and usage" + ) # track cost model_file_id_mapping = { response.output_file_id: {model_id: response.output_file_id} diff --git a/litellm/constants.py b/litellm/constants.py index a0dcd80c07..970b16617f 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -819,6 +819,7 @@ DEFAULT_CRON_JOB_LOCK_TTL_SECONDS = int( PROXY_BUDGET_RESCHEDULER_MIN_TIME = int( os.getenv("PROXY_BUDGET_RESCHEDULER_MIN_TIME", 597) ) +PROXY_BATCH_POLLING_INTERVAL = int(os.getenv("PROXY_BATCH_POLLING_INTERVAL", 3600)) PROXY_BUDGET_RESCHEDULER_MAX_TIME = int( os.getenv("PROXY_BUDGET_RESCHEDULER_MAX_TIME", 605) ) diff --git a/litellm/proxy/_new_secret_config.yaml b/litellm/proxy/_new_secret_config.yaml index 054491ba41..00d3ef1adf 100644 --- a/litellm/proxy/_new_secret_config.yaml +++ b/litellm/proxy/_new_secret_config.yaml @@ -1,24 +1,11 @@ model_list: - - model_name: gpt-3.5-turbo-allow + - model_name: gpt-4o-mini-batch litellm_params: - model: gpt-3.5-turbo + model: azure/gpt-4o-mini + api_key: os.environ/AZURE_API_KEY_HIDDEN + api_base: os.environ/AZURE_API_BASE_HIDDEN model_info: version: 2 - - model_name: gpt-3.5-turbo-disallow - litellm_params: - model: gpt-3.5-turbo - model_info: - version: 2 - - model_name: openai/* - litellm_params: - model: openai/* - - model_name: openai/gpt-4o-mini - litellm_params: - model: openai/gpt-4o-mini -litellm_settings: - model_group_alias: {"gpt-3.5-turbo-custom": "gpt-3.5-turbo-disallow"} - model_group_settings: - forward_client_headers_to_llm_api: - - "gpt-3.5-turbo-allow" - - "openai/*" +general_settings: + proxy_batch_polling_interval: 10 \ No newline at end of file diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 50a044d983..c2bf2f4608 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -132,6 +132,7 @@ from litellm.constants import ( DEFAULT_MODEL_CREATED_AT_TIME, LITELLM_PROXY_ADMIN_NAME, PROMETHEUS_FALLBACK_STATS_SEND_TIME_HOURS, + PROXY_BATCH_POLLING_INTERVAL, PROXY_BATCH_WRITE_AT, PROXY_BUDGET_RESCHEDULER_MAX_TIME, PROXY_BUDGET_RESCHEDULER_MIN_TIME, @@ -538,7 +539,7 @@ async def proxy_shutdown_event(): @asynccontextmanager async def proxy_startup_event(app: FastAPI): - global prisma_client, master_key, use_background_health_checks, llm_router, llm_model_list, general_settings, proxy_budget_rescheduler_min_time, proxy_budget_rescheduler_max_time, litellm_proxy_admin_name, db_writer_client, store_model_in_db, premium_user, _license_check + global prisma_client, master_key, use_background_health_checks, llm_router, llm_model_list, general_settings, proxy_budget_rescheduler_min_time, proxy_budget_rescheduler_max_time, litellm_proxy_admin_name, db_writer_client, store_model_in_db, premium_user, _license_check, proxy_batch_polling_interval import json init_verbose_loggers() @@ -940,6 +941,7 @@ litellm_proxy_admin_name = LITELLM_PROXY_ADMIN_NAME ui_access_mode: Union[Literal["admin", "all"], Dict] = "all" proxy_budget_rescheduler_min_time = PROXY_BUDGET_RESCHEDULER_MIN_TIME proxy_budget_rescheduler_max_time = PROXY_BUDGET_RESCHEDULER_MAX_TIME +proxy_batch_polling_interval = PROXY_BATCH_POLLING_INTERVAL proxy_batch_write_at = PROXY_BATCH_WRITE_AT litellm_master_key_hash = None disable_spend_logs = False @@ -1697,7 +1699,7 @@ class ProxyConfig: """ Load config values into proxy global state """ - global master_key, user_config_file_path, otel_logging, user_custom_auth, user_custom_auth_path, user_custom_key_generate, user_custom_sso, user_custom_ui_sso_sign_in_handler, use_background_health_checks, health_check_interval, use_queue, proxy_budget_rescheduler_max_time, proxy_budget_rescheduler_min_time, ui_access_mode, litellm_master_key_hash, proxy_batch_write_at, disable_spend_logs, prompt_injection_detection_obj, redis_usage_cache, store_model_in_db, premium_user, open_telemetry_logger, health_check_details, callback_settings + global master_key, user_config_file_path, otel_logging, user_custom_auth, user_custom_auth_path, user_custom_key_generate, user_custom_sso, user_custom_ui_sso_sign_in_handler, use_background_health_checks, health_check_interval, use_queue, proxy_budget_rescheduler_max_time, proxy_budget_rescheduler_min_time, ui_access_mode, litellm_master_key_hash, proxy_batch_write_at, disable_spend_logs, prompt_injection_detection_obj, redis_usage_cache, store_model_in_db, premium_user, open_telemetry_logger, health_check_details, callback_settings, proxy_batch_polling_interval config: dict = await self.get_config(config_file_path=config_file_path) @@ -2040,6 +2042,10 @@ class ProxyConfig: proxy_budget_rescheduler_max_time = general_settings.get( "proxy_budget_rescheduler_max_time", proxy_budget_rescheduler_max_time ) + ## BATCH POLLING INTERVAL ## + proxy_batch_polling_interval = general_settings.get( + "proxy_batch_polling_interval", proxy_batch_polling_interval + ) ## BATCH WRITER ## proxy_batch_write_at = general_settings.get( "proxy_batch_write_at", proxy_batch_write_at @@ -3564,7 +3570,7 @@ class ProxyStartupEvent: scheduler.add_job( check_batch_cost_job.check_batch_cost, "interval", - seconds=3600, # these can run infrequently, as batch jobs take time to complete + seconds=proxy_batch_polling_interval, # these can run infrequently, as batch jobs take time to complete ) except Exception: diff --git a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py index 415ec923c4..38002df386 100644 --- a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py +++ b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py @@ -480,7 +480,7 @@ async def update_sso_settings(sso_config: SSOConfig): elif field_name == "ui_access_mode" and value is not None: config["general_settings"]["ui_access_mode"] = value - elif field_name in env_var_mapping and value is not None: + elif field_name in env_var_mapping and value is not None and len(value) > 0: env_var_name = env_var_mapping[field_name] # Update in config config["environment_variables"][env_var_name] = value