fix(lowest_tpm_rpm_v2.py): update key to use model name

This commit is contained in:
Krrish Dholakia
2025-03-18 16:19:47 -07:00
parent 00b338cb9c
commit 9bf6028f14
+10 -5
View File
@@ -151,7 +151,9 @@ class LowestTPMLoggingHandler_v2(CustomLogger):
dt = get_utc_datetime()
current_minute = dt.strftime("%H-%M")
model_id = deployment.get("model_info", {}).get("id")
rpm_key = f"{model_id}:rpm:{current_minute}"
deployment_name = deployment.get("litellm_params", {}).get("model")
rpm_key = f"{model_id}:{deployment_name}:rpm:{current_minute}"
local_result = await self.router_cache.async_get_cache(
key=rpm_key, local_only=True
) # check local result first
@@ -228,6 +230,7 @@ class LowestTPMLoggingHandler_v2(CustomLogger):
if standard_logging_object is None:
raise ValueError("standard_logging_object not passed in.")
model_group = standard_logging_object.get("model_group")
model = standard_logging_object.get("model")
id = standard_logging_object.get("model_id")
if model_group is None or id is None:
return
@@ -244,7 +247,7 @@ class LowestTPMLoggingHandler_v2(CustomLogger):
"%H-%M"
) # use the same timezone regardless of system clock
tpm_key = f"{id}:tpm:{current_minute}"
tpm_key = f"{id}:{model}:tpm:{current_minute}"
# ------------
# Update usage
# ------------
@@ -276,6 +279,7 @@ class LowestTPMLoggingHandler_v2(CustomLogger):
if standard_logging_object is None:
raise ValueError("standard_logging_object not passed in.")
model_group = standard_logging_object.get("model_group")
model = standard_logging_object.get("model")
id = standard_logging_object.get("model_id")
if model_group is None or id is None:
return
@@ -290,7 +294,7 @@ class LowestTPMLoggingHandler_v2(CustomLogger):
"%H-%M"
) # use the same timezone regardless of system clock
tpm_key = f"{id}:tpm:{current_minute}"
tpm_key = f"{id}:{model}:tpm:{current_minute}"
# ------------
# Update usage
# ------------
@@ -458,8 +462,9 @@ class LowestTPMLoggingHandler_v2(CustomLogger):
id = m.get("model_info", {}).get(
"id"
) # a deployment should always have an 'id'. this is set in router.py
tpm_key = "{}:tpm:{}".format(id, current_minute)
rpm_key = "{}:rpm:{}".format(id, current_minute)
deployment_name = m.get("model_name")
tpm_key = "{}:{}:tpm:{}".format(id, deployment_name, current_minute)
rpm_key = "{}:{}:rpm:{}".format(id, deployment_name, current_minute)
tpm_keys.append(tpm_key)
rpm_keys.append(rpm_key)