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
This commit is contained in:
Krish Dholakia
2025-07-21 20:17:56 -07:00
committed by GitHub
parent ff22aed1ea
commit 00677afefe
5 changed files with 37 additions and 30 deletions
@@ -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}
+1
View File
@@ -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)
)
+6 -19
View File
@@ -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
+9 -3
View File
@@ -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:
@@ -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