fix(router.py): create a semaphore for each deployment with rpm

run semaphore logic for each deployment with rpm
This commit is contained in:
Krrish Dholakia 2024-04-12 18:03:23 -07:00
parent 87c621d726
commit 5f1fcaad6d

View file

@ -487,7 +487,9 @@ class Router:
} }
) )
rpm_semaphore = self.semaphore rpm_semaphore = self._get_client(
deployment=deployment, kwargs=kwargs, client_type="rpm_client"
)
if rpm_semaphore is not None and isinstance( if rpm_semaphore is not None and isinstance(
rpm_semaphore, asyncio.Semaphore rpm_semaphore, asyncio.Semaphore
@ -1723,16 +1725,15 @@ class Router:
model_name = litellm_params.get("model") model_name = litellm_params.get("model")
model_id = model["model_info"]["id"] model_id = model["model_info"]["id"]
# ### IF RPM SET - initialize a semaphore ### # ### IF RPM SET - initialize a semaphore ###
# rpm = litellm_params.get("rpm", None) rpm = litellm_params.get("rpm", None)
# print(f"rpm: {rpm}") if rpm:
# if rpm: semaphore = asyncio.Semaphore(rpm)
# semaphore = asyncio.Semaphore(rpm) cache_key = f"{model_id}_rpm_client"
# cache_key = f"{model_id}_rpm_client" self.cache.set_cache(
# self.cache.set_cache( key=cache_key,
# key=cache_key, value=semaphore,
# value=semaphore, local_only=True,
# local_only=True, )
# )
# print("STORES SEMAPHORE IN CACHE") # print("STORES SEMAPHORE IN CACHE")