mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-13 23:11:40 +00:00
b16cfd7de9
* test: point router/completion/triton tests at the local fake OpenAI endpoint The shared Railway-hosted mock (exampleopenaiendpoint-production.up.railway.app) takes down unrelated CI jobs whenever it is unreachable. #30695 moved the mounted proxy configs onto a job-local fake server but left these in-Python api_base literals pointing at the dead host, so litellm_router_testing, local_testing_part1, local_testing_part2 and llm_translation_testing still fail with a 404 "Application not found" when Railway is down Resolve the api_base from FAKE_OPENAI_API_BASE (default http://127.0.0.1:8190) through a shared helper, auto-start the canned server from the local_testing and llm_translation conftests when nothing is already serving, and extend the server with a Triton embeddings route and a slow-endpoint delay so the triton and latency-timeout tests run fully offline. The deliberately broken fallback URL is left as-is so fallback handling still has a failing upstream * fix: ignore non-loopback FAKE_OPENAI_API_BASE so the local mock is used in CI * fix: drop 0.0.0.0 from loopback hosts, an unreliable client connect target * fix(tests): keep fake OpenAI mock alive across xdist workers ensure_fake_openai_endpoint registered atexit on the worker that spawned the subprocess, so under -n 4 the first worker to drain its queue would terminate the shared mock while siblings were still hitting it. Detach the child via start_new_session and drop the per-worker teardown; reuse on /health handles re-runs and CI containers clean up themselves
115 lines
3.3 KiB
Python
115 lines
3.3 KiB
Python
import asyncio
|
|
import os
|
|
import sys
|
|
import time
|
|
|
|
from dotenv import load_dotenv
|
|
|
|
load_dotenv()
|
|
|
|
sys.path.insert(
|
|
0, os.path.abspath("../..")
|
|
) # Adds the parent directory to the system path
|
|
from typing import Dict, List, Optional, Union
|
|
|
|
import pytest
|
|
|
|
import litellm
|
|
from litellm import Router
|
|
from litellm.router import CustomRoutingStrategyBase
|
|
|
|
from tests.fake_openai_endpoint import FAKE_OPENAI_API_BASE
|
|
|
|
|
|
def _create_router():
|
|
return Router(
|
|
model_list=[
|
|
{
|
|
"model_name": "azure-model",
|
|
"litellm_params": {
|
|
"model": "openai/very-special-endpoint",
|
|
"api_base": FAKE_OPENAI_API_BASE,
|
|
"api_key": "fake-key",
|
|
},
|
|
"model_info": {"id": "very-special-endpoint"},
|
|
},
|
|
{
|
|
"model_name": "azure-model",
|
|
"litellm_params": {
|
|
"model": "openai/fast-endpoint",
|
|
"api_base": FAKE_OPENAI_API_BASE,
|
|
"api_key": "fake-key",
|
|
},
|
|
"model_info": {"id": "fast-endpoint"},
|
|
},
|
|
],
|
|
set_verbose=True,
|
|
debug_level="DEBUG",
|
|
)
|
|
|
|
|
|
class CustomRoutingStrategy(CustomRoutingStrategyBase):
|
|
def __init__(self, router_instance: Router):
|
|
self._router = router_instance
|
|
|
|
async def async_get_available_deployment(
|
|
self,
|
|
model: str,
|
|
messages: Optional[List[Dict[str, str]]] = None,
|
|
input: Optional[Union[str, List]] = None,
|
|
specific_deployment: Optional[bool] = False,
|
|
request_kwargs: Optional[Dict] = None,
|
|
):
|
|
print("In CUSTOM async get available deployment")
|
|
model_list = self._router.model_list
|
|
print("router model list=", model_list)
|
|
for model in model_list:
|
|
if isinstance(model, dict):
|
|
if model["litellm_params"]["model"] == "openai/very-special-endpoint":
|
|
return model
|
|
pass
|
|
|
|
def get_available_deployment(
|
|
self,
|
|
model: str,
|
|
messages: Optional[List[Dict[str, str]]] = None,
|
|
input: Optional[Union[str, List]] = None,
|
|
specific_deployment: Optional[bool] = False,
|
|
request_kwargs: Optional[Dict] = None,
|
|
):
|
|
pass
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_custom_routing():
|
|
litellm.set_verbose = True
|
|
|
|
router = _create_router()
|
|
router.set_custom_routing_strategy(CustomRoutingStrategy(router))
|
|
|
|
# make 4 requests
|
|
for _ in range(4):
|
|
try:
|
|
response = await router.acompletion(
|
|
model="azure-model", messages=[{"role": "user", "content": "hello"}]
|
|
)
|
|
print(response)
|
|
except Exception as e:
|
|
print("got exception", e)
|
|
|
|
await asyncio.sleep(1)
|
|
print("done sending initial requests to collect latency")
|
|
|
|
deployments = {}
|
|
# make 10 requests
|
|
for _ in range(10):
|
|
response = await router.acompletion(
|
|
model="azure-model", messages=[{"role": "user", "content": "hello"}]
|
|
)
|
|
print(response)
|
|
_picked_model_id = response._hidden_params["model_id"]
|
|
if _picked_model_id not in deployments:
|
|
deployments[_picked_model_id] = 1
|
|
else:
|
|
deployments[_picked_model_id] += 1
|
|
print("deployments", deployments)
|