Merge pull request #1403 from BerriAI/litellm_latency_routing_updates

fix(lowest_latency.py): add back tpm/rpm checks, configurable time window support, improved latency tracking
This commit is contained in:
Krish Dholakia 2024-01-10 21:34:05 +05:30 • committed by GitHub
commit 9e97227625
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
6 changed files with 436 additions and 188 deletions

View file

@ -77,7 +77,65 @@ print(response)
Router provides 4 strategies for routing your calls across multiple deployments:
<Tabs>
<TabItem value="simple-shuffle" label="Weighted Pick">
<TabItem value="latency-based" label="Latency-Based">
Picks the deployment with the lowest response time.
It caches, and updates the response times for deployments based on when a request was sent and received from a deployment.
[**How to test**](https://github.com/BerriAI/litellm/blob/main/litellm/tests/test_lowest_latency_routing.py)
```python
from litellm import Router
import asyncio
model_list = [{ ... }]
# init router
router = Router(model_list=model_list, routing_strategy="latency-based-routing") # 👈 set routing strategy
## CALL 1+2
tasks = []
response = None
final_response = None
for _ in range(2):
tasks.append(router.acompletion(model=model, messages=messages))
response = await asyncio.gather(*tasks)
if response is not None:
## CALL 3
await asyncio.sleep(1) # let the cache update happen
picked_deployment = router.lowestlatency_logger.get_available_deployments(
model_group=model, healthy_deployments=router.healthy_deployments
)
final_response = await router.acompletion(model=model, messages=messages)
print(f"min deployment id: {picked_deployment}")
print(f"model id: {final_response._hidden_params['model_id']}")
assert (
final_response._hidden_params["model_id"]
== picked_deployment["model_info"]["id"]
)
```
### Set Time Window
Set time window for how far back to consider when averaging latency for a deployment.
**In Router**
```python
router = Router(..., routing_strategy_args={"ttl": 10})
```
**In Proxy**
```yaml
router_settings:
routing_strategy_args: {"ttl": 10}
```
</TabItem>
<TabItem value="simple-shuffle" label="(Default) Weighted Pick">
**Default** Picks a deployment based on the provided **Requests per minute (rpm) or Tokens per minute (tpm)**
@ -235,58 +293,7 @@ asyncio.run(router_acompletion())
```
</TabItem>
<TabItem value="latency-based" label="Latency-Based">
Picks the deployment with the lowest response time.
It caches, and updates the response times for deployments based on when a request was sent and received from a deployment.
[**How to test**](https://github.com/BerriAI/litellm/blob/main/litellm/tests/test_lowest_latency_routing.py)
```python
from litellm import Router
import asyncio
model_list = [{ # list of model deployments
"model_name": "gpt-3.5-turbo", # model alias
"litellm_params": { # params for litellm completion/embedding call
"model": "azure/chatgpt-v-2", # actual model name
"api_key": os.getenv("AZURE_API_KEY"),
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
}
}, {
"model_name": "gpt-3.5-turbo",
"litellm_params": { # params for litellm completion/embedding call
"model": "azure/chatgpt-functioncalling",
"api_key": os.getenv("AZURE_API_KEY"),
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
}
}, {
"model_name": "gpt-3.5-turbo",
"litellm_params": { # params for litellm completion/embedding call
"model": "gpt-3.5-turbo",
"api_key": os.getenv("OPENAI_API_KEY"),
}
}]
# init router
router = Router(model_list=model_list, routing_strategy="latency-based-routing")
async def router_acompletion():
response = await router.acompletion(
model="gpt-3.5-turbo",
messages=[{"role": "user", "content": "Hey, how's it going?"}]
)
print(response)
return response
asyncio.run(router_acompletion())
```
</TabItem>
</Tabs>
## Basic Reliability
@ -608,4 +615,4 @@ def __init__(
"latency-based-routing",
] = "simple-shuffle",
):
```
```

View file

@ -105,7 +105,7 @@ class Router:
"usage-based-routing",
"latency-based-routing",
] = "simple-shuffle",
routing_strategy_args: dict = {}, # just for latency-based routing
routing_strategy_args: dict = {}, # just for latency-based routing
) -> None:
self.set_verbose = set_verbose
self.deployment_names: List = (
@ -218,7 +218,9 @@ class Router:
litellm.callbacks.append(self.lowesttpm_logger) # type: ignore
elif routing_strategy == "latency-based-routing":
self.lowestlatency_logger = LowestLatencyLoggingHandler(
router_cache=self.cache, model_list=self.model_list, routing_args=routing_strategy_args
router_cache=self.cache,
model_list=self.model_list,
routing_args=routing_strategy_args,
)
if isinstance(litellm.callbacks, list):
litellm.callbacks.append(self.lowestlatency_logger) # type: ignore
@ -1428,9 +1430,8 @@ class Router:
http_client=httpx.AsyncClient(
transport=AsyncCustomHTTPTransport(),
limits=httpx.Limits(
max_connections=1000,
max_keepalive_connections=100
)
max_connections=1000, max_keepalive_connections=100
),
), # type: ignore
)
self.cache.set_cache(
@ -1450,9 +1451,8 @@ class Router:
http_client=httpx.Client(
transport=CustomHTTPTransport(),
limits=httpx.Limits(
max_connections=1000,
max_keepalive_connections=100
)
max_connections=1000, max_keepalive_connections=100
),
), # type: ignore
)
self.cache.set_cache(
@ -1472,10 +1472,9 @@ class Router:
max_retries=max_retries,
http_client=httpx.AsyncClient(
limits=httpx.Limits(
max_connections=1000,
max_keepalive_connections=100
max_connections=1000, max_keepalive_connections=100
)
)
),
)
self.cache.set_cache(
key=cache_key,
@ -1493,10 +1492,9 @@ class Router:
max_retries=max_retries,
http_client=httpx.Client(
limits=httpx.Limits(
max_connections=1000,
max_keepalive_connections=100
max_connections=1000, max_keepalive_connections=100
)
)
),
)
self.cache.set_cache(
key=cache_key,

View file

@ -2,13 +2,32 @@
# picks based on response time (for streaming, this is time to first token)
from pydantic import BaseModel, Extra, Field, root_validator
import dotenv, os, requests, random
from typing import Optional
from typing import Optional, Union, List, Dict
from datetime import datetime, timedelta
dotenv.load_dotenv() # Loading env variables using dotenv
import traceback
from litellm.caching import DualCache
from litellm.integrations.custom_logger import CustomLogger
from litellm import ModelResponse
from litellm import token_counter
class LiteLLMBase(BaseModel):
"""
Implements default functions, all pydantic objects should have.
"""
def json(self, **kwargs):
try:
return self.model_dump() # noqa
except:
# if using pydantic v1
return self.dict()
class RoutingArgs(LiteLLMBase):
ttl: int = 1 * 60 * 60 # 1 hour
class LiteLLMBase(BaseModel):
"""
@ -30,7 +49,9 @@ class LowestLatencyLoggingHandler(CustomLogger):
logged_success: int = 0
logged_failure: int = 0
def __init__(self, router_cache: DualCache, model_list: list, routing_args: dict={}):
def __init__(
self, router_cache: DualCache, model_list: list, routing_args: dict = {}
):
self.router_cache = router_cache
self.model_list = model_list
self.routing_args = RoutingArgs(**routing_args)
@ -51,25 +72,64 @@ class LowestLatencyLoggingHandler(CustomLogger):
if model_group is None or id is None:
return
response_ms = end_time - start_time
# ------------
# Setup values
# ------------
latency_key = f"{model_group}_latency_map"
"""
{
{model_group}_map: {
id: {
"latency": [..]
f"{date:hour:minute}" : {"tpm": 34, "rpm": 3}
}
}
}
"""
latency_key = f"{model_group}_map"
current_date = datetime.now().strftime("%Y-%m-%d")
current_hour = datetime.now().strftime("%H")
current_minute = datetime.now().strftime("%M")
precise_minute = f"{current_date}-{current_hour}-{current_minute}"
response_ms: timedelta = end_time - start_time
final_value = response_ms
total_tokens = 0
if isinstance(response_obj, ModelResponse):
completion_tokens = response_obj.usage.completion_tokens
total_tokens = response_obj.usage.total_tokens
final_value = float(completion_tokens / response_ms.total_seconds())
# ------------
# Update usage
# ------------
## Latency
request_count_dict = self.router_cache.get_cache(key=latency_key) or {}
if id in request_count_dict and isinstance(request_count_dict[id], list):
request_count_dict[id] = request_count_dict[id].append(response_ms)
else:
request_count_dict[id] = [response_ms]
self.router_cache.set_cache(key=latency_key, value=request_count_dict, ttl=self.routing_args.ttl) # reset map within window
if id not in request_count_dict:
request_count_dict[id] = {}
## Latency
request_count_dict[id].setdefault("latency", []).append(final_value)
if precise_minute not in request_count_dict[id]:
request_count_dict[id][precise_minute] = {}
## TPM
request_count_dict[id][precise_minute]["tpm"] = (
request_count_dict[id][precise_minute].get("tpm", 0) + total_tokens
)
## RPM
request_count_dict[id][precise_minute]["rpm"] = (
request_count_dict[id][precise_minute].get("rpm", 0) + 1
)
self.router_cache.set_cache(
key=latency_key, value=request_count_dict, ttl=self.routing_args.ttl
) # reset map within window
### TESTING ###
if self.test_flag:
@ -94,26 +154,65 @@ class LowestLatencyLoggingHandler(CustomLogger):
if model_group is None or id is None:
return
response_ms = end_time - start_time
# ------------
# Setup values
# ------------
latency_key = f"{model_group}_latency_map"
"""
{
{model_group}_map: {
id: {
"latency": [..]
f"{date:hour:minute}" : {"tpm": 34, "rpm": 3}
}
}
}
"""
latency_key = f"{model_group}_map"
current_date = datetime.now().strftime("%Y-%m-%d")
current_hour = datetime.now().strftime("%H")
current_minute = datetime.now().strftime("%M")
precise_minute = f"{current_date}-{current_hour}-{current_minute}"
response_ms: timedelta = end_time - start_time
final_value = response_ms
total_tokens = 0
if isinstance(response_obj, ModelResponse):
completion_tokens = response_obj.usage.completion_tokens
total_tokens = response_obj.usage.total_tokens
final_value = float(completion_tokens / response_ms.total_seconds())
# ------------
# Update usage
# ------------
## Latency
request_count_dict = self.router_cache.get_cache(key=latency_key) or {}
if id in request_count_dict and isinstance(request_count_dict[id], list):
request_count_dict[id] = request_count_dict[id] + [response_ms]
else:
request_count_dict[id] = [response_ms]
self.router_cache.set_cache(key=latency_key, value=request_count_dict, ttl=self.routing_args.ttl) # reset map within window
if id not in request_count_dict:
request_count_dict[id] = {}
## Latency
request_count_dict[id].setdefault("latency", []).append(final_value)
if precise_minute not in request_count_dict[id]:
request_count_dict[id][precise_minute] = {}
## TPM
request_count_dict[id][precise_minute]["tpm"] = (
request_count_dict[id][precise_minute].get("tpm", 0) + total_tokens
)
## RPM
request_count_dict[id][precise_minute]["rpm"] = (
request_count_dict[id][precise_minute].get("rpm", 0) + 1
)
self.router_cache.set_cache(
key=latency_key, value=request_count_dict, ttl=self.routing_args.ttl
) # reset map within window
### TESTING ###
if self.test_flag:
self.logged_success += 1
@ -121,12 +220,18 @@ class LowestLatencyLoggingHandler(CustomLogger):
traceback.print_exc()
pass
def get_available_deployments(self, model_group: str, healthy_deployments: list):
def get_available_deployments(
self,
model_group: str,
healthy_deployments: list,
messages: Optional[List[Dict[str, str]]] = None,
input: Optional[Union[str, List]] = None,
):
"""
Returns a deployment with the lowest latency
"""
# get list of potential deployments
latency_key = f"{model_group}_latency_map"
latency_key = f"{model_group}_map"
request_count_dict = self.router_cache.get_cache(key=latency_key) or {}
@ -134,6 +239,12 @@ class LowestLatencyLoggingHandler(CustomLogger):
# Find lowest used model
# ----------------------
lowest_latency = float("inf")
current_date = datetime.now().strftime("%Y-%m-%d")
current_hour = datetime.now().strftime("%H")
current_minute = datetime.now().strftime("%M")
precise_minute = f"{current_date}-{current_hour}-{current_minute}"
deployment = None
if request_count_dict is None: # base case
@ -143,9 +254,17 @@ class LowestLatencyLoggingHandler(CustomLogger):
for d in healthy_deployments:
## if healthy deployment not yet used
if d["model_info"]["id"] not in all_deployments:
all_deployments[d["model_info"]["id"]] = [0]
all_deployments[d["model_info"]["id"]] = {
"latency": [0],
precise_minute: {"tpm": 0, "rpm": 0},
}
for item, item_latency in all_deployments.items():
try:
input_tokens = token_counter(messages=messages, text=input)
except:
input_tokens = 0
for item, item_map in all_deployments.items():
## get the item from model list
_deployment = None
for m in healthy_deployments:
@ -154,18 +273,38 @@ class LowestLatencyLoggingHandler(CustomLogger):
if _deployment is None:
continue # skip to next one
# get average latency
total = 0.0
_deployment_tpm = (
_deployment.get("tpm", None)
or _deployment.get("litellm_params", {}).get("tpm", None)
or _deployment.get("model_info", {}).get("tpm", None)
or float("inf")
)
_deployment_rpm = (
_deployment.get("rpm", None)
or _deployment.get("litellm_params", {}).get("rpm", None)
or _deployment.get("model_info", {}).get("rpm", None)
or float("inf")
)
item_latency = item_map.get("latency", [])
item_rpm = item_map.get(precise_minute, {}).get("rpm", 0)
item_tpm = item_map.get(precise_minute, {}).get("tpm", 0)
# get average latency
total: float = 0.0
for _call_latency in item_latency:
if isinstance(_call_latency, timedelta):
total += float(_call_latency.total_seconds())
elif isinstance(_call_latency, float):
if isinstance(_call_latency, float):
total += _call_latency
item_latency = total/len(item_latency)
item_latency = total / len(item_latency)
if item_latency == 0:
deployment = _deployment
break
elif (
item_tpm + input_tokens > _deployment_tpm
or item_rpm + 1 > _deployment_rpm
): # if user passed in tpm / rpm in the model_list
continue
elif item_latency < lowest_latency:
lowest_latency = item_latency
deployment = _deployment

View file

@ -29,20 +29,6 @@ def logger_fn(user_model_dict):
pass
# normal call
def test_completion_custom_provider_model_name():
try:
response = completion_with_retries(
model="together_ai/togethercomputer/llama-2-70b-chat",
messages=messages,
logger_fn=logger_fn,
)
# Add any assertions here to check the response
print(response)
except Exception as e:
pytest.fail(f"Error occurred: {e}")
# completion with num retries + impact on exception mapping
def test_completion_with_num_retries():
try:
@ -75,7 +61,3 @@ def test_completion_with_0_num_retries():
except Exception as e:
print("exception", e)
pass
# Call the test function
test_completion_with_0_num_retries()

View file

@ -48,8 +48,11 @@ def test_latency_updated():
start_time=start_time,
end_time=end_time,
)
latency_key = f"{model_group}_latency_map"
assert end_time - start_time == test_cache.get_cache(key=latency_key)[deployment_id][0]
latency_key = f"{model_group}_map"
assert (
end_time - start_time
== test_cache.get_cache(key=latency_key)[deployment_id]["latency"][0]
)
# test_tpm_rpm_updated()
@ -92,6 +95,45 @@ def test_latency_updated_custom_ttl():
assert test_cache.get_cache(key=latency_key) is None
def test_latency_updated_custom_ttl():
"""
Invalidate the cached request.
Test that the cache is empty
"""
test_cache = DualCache()
model_list = []
cache_time = 3
lowest_latency_logger = LowestLatencyLoggingHandler(
router_cache=test_cache, model_list=model_list, routing_args={"ttl": cache_time}
)
model_group = "gpt-3.5-turbo"
deployment_id = "1234"
kwargs = {
"litellm_params": {
"metadata": {
"model_group": "gpt-3.5-turbo",
"deployment": "azure/chatgpt-v-2",
},
"model_info": {"id": deployment_id},
}
}
start_time = time.time()
response_obj = {"usage": {"total_tokens": 50}}
time.sleep(5)
end_time = time.time()
lowest_latency_logger.log_success_event(
response_obj=response_obj,
kwargs=kwargs,
start_time=start_time,
end_time=end_time,
)
latency_key = f"{model_group}_map"
assert isinstance(test_cache.get_cache(key=latency_key), dict)
time.sleep(cache_time)
assert test_cache.get_cache(key=latency_key) is None
def test_get_available_deployments():
test_cache = DualCache()
model_list = [
@ -170,6 +212,90 @@ def test_get_available_deployments():
# test_get_available_deployments()
def test_get_available_endpoints_tpm_rpm_check():
"""
Pass in list of 2 valid models
Update cache with 1 model clearly being at tpm/rpm limit
assert that only the valid model is returned
"""
test_cache = DualCache()
model_list = [
{
"model_name": "gpt-3.5-turbo",
"litellm_params": {"model": "azure/chatgpt-v-2"},
"model_info": {"id": "1234", "rpm": 10},
},
{
"model_name": "gpt-3.5-turbo",
"litellm_params": {"model": "azure/chatgpt-v-2"},
"model_info": {"id": "5678", "rpm": 3},
},
]
lowest_latency_logger = LowestLatencyLoggingHandler(
router_cache=test_cache, model_list=model_list
)
model_group = "gpt-3.5-turbo"
## DEPLOYMENT 1 ##
deployment_id = "1234"
kwargs = {
"litellm_params": {
"metadata": {
"model_group": "gpt-3.5-turbo",
"deployment": "azure/chatgpt-v-2",
},
"model_info": {"id": deployment_id},
}
}
for _ in range(3):
start_time = time.time()
response_obj = {"usage": {"total_tokens": 50}}
time.sleep(0.05)
end_time = time.time()
lowest_latency_logger.log_success_event(
response_obj=response_obj,
kwargs=kwargs,
start_time=start_time,
end_time=end_time,
)
## DEPLOYMENT 2 ##
deployment_id = "5678"
kwargs = {
"litellm_params": {
"metadata": {
"model_group": "gpt-3.5-turbo",
"deployment": "azure/chatgpt-v-2",
},
"model_info": {"id": deployment_id},
}
}
for _ in range(3):
start_time = time.time()
response_obj = {"usage": {"total_tokens": 20}}
time.sleep(2)
end_time = time.time()
lowest_latency_logger.log_success_event(
response_obj=response_obj,
kwargs=kwargs,
start_time=start_time,
end_time=end_time,
)
## CHECK WHAT'S SELECTED ##
print(
lowest_latency_logger.get_available_deployments(
model_group=model_group, healthy_deployments=model_list
)
)
assert (
lowest_latency_logger.get_available_deployments(
model_group=model_group, healthy_deployments=model_list
)["model_info"]["id"]
== "1234"
)
def test_router_get_available_deployments():
"""
Test if routers 'get_available_deployments' returns the fastest deployment
@ -250,9 +376,6 @@ def test_router_get_available_deployments():
assert router.get_available_deployment(model="azure-model")["model_info"]["id"] == 2
# test_get_available_deployments()
# test_router_get_available_deployments()

View file

@ -70,86 +70,85 @@ for pr in repo.get_pulls():
print(f"The pull request number for branch {branch_name} is: {pr_number}")
def test_add_new_key():
max_retries = 3
retry_delay = 1 # seconds
# def test_add_new_key():
# max_retries = 3
# retry_delay = 10 # seconds
for retry in range(max_retries + 1):
try:
# Your test data
test_data = {
"models": ["gpt-3.5-turbo", "gpt-4", "claude-2", "azure-model"],
"aliases": {"mistral-7b": "gpt-3.5-turbo"},
"duration": "20m",
}
print("testing proxy server")
# for retry in range(max_retries + 1):
# try:
# # Your test data
# test_data = {
# "models": ["gpt-3.5-turbo", "gpt-4", "claude-2", "azure-model"],
# "aliases": {"mistral-7b": "gpt-3.5-turbo"},
# "duration": "20m",
# }
# print("testing proxy server")
# Your bearer token
token = os.getenv("PROXY_MASTER_KEY")
headers = {"Authorization": f"Bearer {token}"}
# # Your bearer token
# token = os.getenv("PROXY_MASTER_KEY")
# headers = {"Authorization": f"Bearer {token}"}
endpoint = f"https://litellm-litellm-pr-{pr_number}.up.railway.app"
# endpoint = f"https://litellm-litellm-pr-{pr_number}.up.railway.app"
# Make a request to the staging endpoint
response = requests.post(
endpoint + "/key/generate", json=test_data, headers=headers
)
# # Make a request to the staging endpoint
# response = requests.post(
# endpoint + "/key/generate", json=test_data, headers=headers
# )
print(f"response: {response.text}")
# print(f"response: {response.text}")
if response.status_code == 200:
result = response.json()
break # Successful response, exit the loop
elif response.status_code == 503 and retry < max_retries:
print(
f"Retrying in {retry_delay} seconds... (Retry {retry + 1}/{max_retries})"
)
time.sleep(retry_delay)
else:
assert False, f"Unexpected response status code: {response.status_code}"
# if response.status_code == 200:
# result = response.json()
# break # Successful response, exit the loop
# elif response.status_code == 503 and retry < max_retries:
# print(
# f"Retrying in {retry_delay} seconds... (Retry {retry + 1}/{max_retries})"
# )
# time.sleep(retry_delay)
# else:
# assert False, f"Unexpected response status code: {response.status_code}"
except Exception as e:
print(traceback.format_exc())
pytest.fail(f"An error occurred {e}")
# except Exception as e:
# print(traceback.format_exc())
# pytest.fail(f"An error occurred {e}")
def test_update_new_key():
try:
# Your test data
test_data = {
"models": ["gpt-3.5-turbo", "gpt-4", "claude-2", "azure-model"],
"aliases": {"mistral-7b": "gpt-3.5-turbo"},
"duration": "20m",
}
print("testing proxy server")
# Your bearer token
token = os.getenv("PROXY_MASTER_KEY")
headers = {"Authorization": f"Bearer {token}"}
# def test_update_new_key():
# try:
# # Your test data
# test_data = {
# "models": ["gpt-3.5-turbo", "gpt-4", "claude-2", "azure-model"],
# "aliases": {"mistral-7b": "gpt-3.5-turbo"},
# "duration": "20m",
# }
# print("testing proxy server")
# # Your bearer token
# token = os.getenv("PROXY_MASTER_KEY")
# headers = {"Authorization": f"Bearer {token}"}
endpoint = f"https://litellm-litellm-pr-{pr_number}.up.railway.app"
# endpoint = f"https://litellm-litellm-pr-{pr_number}.up.railway.app"
# Make a request to the staging endpoint
response = requests.post(
endpoint + "/key/generate", json=test_data, headers=headers
)
assert response.status_code == 200
result = response.json()
assert result["key"].startswith("sk-")
# # Make a request to the staging endpoint
# response = requests.post(
# endpoint + "/key/generate", json=test_data, headers=headers
# )
# assert response.status_code == 200
# result = response.json()
# assert result["key"].startswith("sk-")
def _post_data():
json_data = {"models": ["bedrock-models"], "key": result["key"]}
response = requests.post(
endpoint + "/key/generate", json=json_data, headers=headers
)
print(f"response text: {response.text}")
assert response.status_code == 200
return response
_post_data()
print(f"Received response: {result}")
except Exception as e:
pytest.fail(f"LiteLLM Proxy test failed. Exception: {str(e)}")
# def _post_data():
# json_data = {"models": ["bedrock-models"], "key": result["key"]}
# response = requests.post(
# endpoint + "/key/generate", json=json_data, headers=headers
# )
# print(f"response text: {response.text}")
# assert response.status_code == 200
# return response
# _post_data()
# print(f"Received response: {result}")
# except Exception as e:
# pytest.fail(f"LiteLLM Proxy test failed. Exception: {str(e)}")
# def test_add_new_key_max_parallel_limit():
# try: