mirror of
https://github.com/BerriAI/litellm.git
synced 2026-08-28 05:25:59 +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
392 lines
12 KiB
Python
392 lines
12 KiB
Python
import json
|
|
import os
|
|
import sys
|
|
import traceback
|
|
|
|
from dotenv import load_dotenv
|
|
|
|
load_dotenv()
|
|
import io
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
|
|
sys.path.insert(
|
|
0, os.path.abspath("../..")
|
|
) # Adds the parent directory to the system path
|
|
import pytest
|
|
import litellm
|
|
|
|
import pytest
|
|
from litellm.llms.triton.embedding.transformation import TritonEmbeddingConfig
|
|
import litellm
|
|
|
|
from tests.fake_openai_endpoint import FAKE_OPENAI_API_BASE
|
|
|
|
|
|
def test_split_embedding_by_shape_passes():
|
|
try:
|
|
data = [
|
|
{
|
|
"shape": [2, 3],
|
|
"data": [1, 2, 3, 4, 5, 6],
|
|
}
|
|
]
|
|
split_output_data = TritonEmbeddingConfig.split_embedding_by_shape(
|
|
data[0]["data"], data[0]["shape"]
|
|
)
|
|
assert split_output_data == [[1, 2, 3], [4, 5, 6]]
|
|
except Exception as e:
|
|
pytest.fail(f"An exception occured: {e}")
|
|
|
|
|
|
def test_split_embedding_by_shape_fails_with_shape_value_error():
|
|
data = [
|
|
{
|
|
"shape": [2],
|
|
"data": [1, 2, 3, 4, 5, 6],
|
|
}
|
|
]
|
|
with pytest.raises(ValueError):
|
|
TritonEmbeddingConfig.split_embedding_by_shape(
|
|
data[0]["data"], data[0]["shape"]
|
|
)
|
|
|
|
|
|
def test_triton_embedding_response_sets_usage_with_token_counter():
|
|
config = TritonEmbeddingConfig()
|
|
mock_http_response = MagicMock()
|
|
mock_http_response.status_code = 200
|
|
mock_http_response.json.return_value = {
|
|
"model_name": "gte-base-en-v1",
|
|
"outputs": [
|
|
{
|
|
"name": "embedding",
|
|
"shape": [1, 2],
|
|
"data": [0.1, 0.2],
|
|
}
|
|
],
|
|
}
|
|
model_response = litellm.EmbeddingResponse()
|
|
request_data = {
|
|
"inputs": [
|
|
{
|
|
"name": "input_text",
|
|
"shape": [1],
|
|
"datatype": "BYTES",
|
|
"data": ["hello from triton"],
|
|
}
|
|
]
|
|
}
|
|
|
|
with patch(
|
|
"litellm.llms.triton.embedding.transformation.token_counter",
|
|
return_value=7,
|
|
):
|
|
transformed = config.transform_embedding_response(
|
|
model="triton/gte-base-en-v1",
|
|
raw_response=mock_http_response,
|
|
model_response=model_response,
|
|
logging_obj=MagicMock(),
|
|
request_data=request_data,
|
|
)
|
|
|
|
assert transformed.usage is not None
|
|
assert transformed.usage.prompt_tokens == 7
|
|
assert transformed.usage.completion_tokens == 0
|
|
assert transformed.usage.total_tokens == 7
|
|
|
|
|
|
def test_triton_embedding_response_sets_usage_with_word_count_fallback():
|
|
config = TritonEmbeddingConfig()
|
|
mock_http_response = MagicMock()
|
|
mock_http_response.status_code = 200
|
|
mock_http_response.json.return_value = {
|
|
"model_name": "gte-base-en-v1",
|
|
"outputs": [
|
|
{
|
|
"name": "embedding",
|
|
"shape": [1, 2],
|
|
"data": [0.1, 0.2],
|
|
}
|
|
],
|
|
}
|
|
model_response = litellm.EmbeddingResponse()
|
|
request_data = {
|
|
"inputs": [
|
|
{
|
|
"name": "input_text",
|
|
"shape": [1],
|
|
"datatype": "BYTES",
|
|
"data": ["hello from triton"],
|
|
}
|
|
]
|
|
}
|
|
|
|
with patch(
|
|
"litellm.llms.triton.embedding.transformation.token_counter",
|
|
side_effect=Exception("tokenizer error"),
|
|
):
|
|
transformed = config.transform_embedding_response(
|
|
model="triton/gte-base-en-v1",
|
|
raw_response=mock_http_response,
|
|
model_response=model_response,
|
|
logging_obj=MagicMock(),
|
|
request_data=request_data,
|
|
)
|
|
|
|
assert transformed.usage is not None
|
|
assert transformed.usage.prompt_tokens == 3
|
|
assert transformed.usage.completion_tokens == 0
|
|
assert transformed.usage.total_tokens == 3
|
|
|
|
|
|
def test_triton_embedding_batch_usage_sums_per_input_token_counts():
|
|
"""Batch inputs must not be joined before token counting (avoids extra newline tokens)."""
|
|
config = TritonEmbeddingConfig()
|
|
mock_http_response = MagicMock()
|
|
mock_http_response.status_code = 200
|
|
mock_http_response.json.return_value = {
|
|
"model_name": "gte-base-en-v1",
|
|
"outputs": [
|
|
{
|
|
"name": "embedding",
|
|
"shape": [2, 2],
|
|
"data": [0.1, 0.2, 0.3, 0.4],
|
|
}
|
|
],
|
|
}
|
|
model_response = litellm.EmbeddingResponse()
|
|
request_data = {
|
|
"inputs": [
|
|
{
|
|
"name": "input_text",
|
|
"shape": [2],
|
|
"datatype": "BYTES",
|
|
"data": ["first input", "second input"],
|
|
}
|
|
]
|
|
}
|
|
|
|
with patch(
|
|
"litellm.llms.triton.embedding.transformation.token_counter",
|
|
side_effect=[5, 7],
|
|
):
|
|
transformed = config.transform_embedding_response(
|
|
model="triton/gte-base-en-v1",
|
|
raw_response=mock_http_response,
|
|
model_response=model_response,
|
|
logging_obj=MagicMock(),
|
|
request_data=request_data,
|
|
)
|
|
|
|
assert transformed.usage is not None
|
|
assert transformed.usage.prompt_tokens == 12
|
|
assert transformed.usage.total_tokens == 12
|
|
|
|
|
|
@pytest.mark.parametrize("stream", [True, False])
|
|
def test_completion_triton_generate_api(stream):
|
|
try:
|
|
mock_response = MagicMock()
|
|
if stream:
|
|
|
|
def mock_iter_lines():
|
|
mock_output = "".join(
|
|
[
|
|
'data: {"model_name":"ensemble","model_version":"1","sequence_end":false,"sequence_id":0,"sequence_start":false,"text_output":"'
|
|
+ t
|
|
+ '"}\n\n'
|
|
for t in ["I", " am", " an", " AI", " assistant"]
|
|
]
|
|
)
|
|
for out in mock_output.split("\n"):
|
|
yield out
|
|
|
|
mock_response.iter_lines = mock_iter_lines
|
|
else:
|
|
|
|
def return_val():
|
|
return {
|
|
"text_output": "I am an AI assistant",
|
|
}
|
|
|
|
mock_response.json = return_val
|
|
mock_response.status_code = 200
|
|
|
|
with patch(
|
|
"litellm.llms.custom_httpx.http_handler.HTTPHandler.post",
|
|
return_value=mock_response,
|
|
) as mock_post:
|
|
response = litellm.completion(
|
|
model="triton/llama-3-8b-instruct",
|
|
messages=[{"role": "user", "content": "who are u?"}],
|
|
max_tokens=10,
|
|
timeout=5,
|
|
api_base="http://localhost:8000/generate",
|
|
stream=stream,
|
|
)
|
|
|
|
# Verify the call was made
|
|
mock_post.assert_called_once()
|
|
|
|
# Get the arguments passed to the post request
|
|
print("call args", mock_post.call_args)
|
|
call_kwargs = mock_post.call_args.kwargs # Access kwargs directly
|
|
|
|
# Verify URL
|
|
if stream:
|
|
assert call_kwargs["url"] == "http://localhost:8000/generate_stream"
|
|
else:
|
|
assert call_kwargs["url"] == "http://localhost:8000/generate"
|
|
|
|
# Parse the request data from the JSON string
|
|
request_data = json.loads(call_kwargs["data"])
|
|
|
|
# Verify request data
|
|
assert request_data["text_input"] == "who are u?"
|
|
assert request_data["parameters"]["max_tokens"] == 10
|
|
|
|
# Verify response
|
|
if stream:
|
|
tokens = ["I", " am", " an", " AI", " assistant", None]
|
|
idx = 0
|
|
for chunk in response:
|
|
assert chunk.choices[0].delta.content == tokens[idx]
|
|
idx += 1
|
|
assert idx == len(tokens)
|
|
else:
|
|
assert response.choices[0].message.content == "I am an AI assistant"
|
|
|
|
except Exception as e:
|
|
print("exception", e)
|
|
import traceback
|
|
|
|
traceback.print_exc()
|
|
pytest.fail(f"Error occurred: {e}")
|
|
|
|
|
|
def test_completion_triton_infer_api():
|
|
litellm.set_verbose = True
|
|
try:
|
|
mock_response = MagicMock()
|
|
|
|
def return_val():
|
|
return {
|
|
"model_name": "basketgpt",
|
|
"model_version": "2",
|
|
"outputs": [
|
|
{
|
|
"name": "text_output",
|
|
"datatype": "BYTES",
|
|
"shape": [1],
|
|
"data": [
|
|
"0004900005024 0004900006774 0004900005024 0004900005027 0004900005026 0004900005025 0004900005027 0004900005024 0004900006774 0004900005027"
|
|
],
|
|
},
|
|
{
|
|
"name": "debug_probs",
|
|
"datatype": "FP32",
|
|
"shape": [0],
|
|
"data": [],
|
|
},
|
|
{
|
|
"name": "debug_tokens",
|
|
"datatype": "BYTES",
|
|
"shape": [0],
|
|
"data": [],
|
|
},
|
|
],
|
|
}
|
|
|
|
mock_response.json = return_val
|
|
mock_response.status_code = 200
|
|
|
|
with patch(
|
|
"litellm.llms.custom_httpx.http_handler.HTTPHandler.post",
|
|
return_value=mock_response,
|
|
) as mock_post:
|
|
response = litellm.completion(
|
|
model="triton/llama-3-8b-instruct",
|
|
messages=[
|
|
{
|
|
"role": "user",
|
|
"content": "0004900005025 0004900005026 0004900005027",
|
|
}
|
|
],
|
|
api_base="http://localhost:8000/infer",
|
|
)
|
|
|
|
print("litellm response", response.model_dump_json(indent=4))
|
|
|
|
# Verify the call was made
|
|
mock_post.assert_called_once()
|
|
|
|
# Get the arguments passed to the post request
|
|
call_kwargs = mock_post.call_args.kwargs
|
|
|
|
# Verify URL
|
|
assert call_kwargs["url"] == "http://localhost:8000/infer"
|
|
|
|
# Parse the request data from the JSON string
|
|
request_data = json.loads(call_kwargs["data"])
|
|
|
|
# Verify request matches expected Triton format
|
|
assert request_data["inputs"][0]["name"] == "text_input"
|
|
assert request_data["inputs"][0]["shape"] == [1]
|
|
assert request_data["inputs"][0]["datatype"] == "BYTES"
|
|
assert request_data["inputs"][0]["data"] == [
|
|
"0004900005025 0004900005026 0004900005027"
|
|
]
|
|
|
|
assert request_data["inputs"][1]["shape"] == [1]
|
|
assert request_data["inputs"][1]["datatype"] == "INT32"
|
|
assert request_data["inputs"][1]["data"] == [20]
|
|
|
|
# Verify response format matches expected completion format
|
|
assert (
|
|
response.choices[0].message.content
|
|
== "0004900005024 0004900006774 0004900005024 0004900005027 0004900005026 0004900005025 0004900005027 0004900005024 0004900006774 0004900005027"
|
|
)
|
|
assert response.choices[0].finish_reason == "stop"
|
|
assert response.choices[0].index == 0
|
|
assert response.object == "chat.completion"
|
|
|
|
except Exception as e:
|
|
print("exception", e)
|
|
traceback.print_exc()
|
|
pytest.fail(f"Error occurred: {e}")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_triton_embeddings():
|
|
try:
|
|
litellm.set_verbose = True
|
|
response = await litellm.aembedding(
|
|
model="triton/my-triton-model",
|
|
api_base=f"{FAKE_OPENAI_API_BASE}/triton/embeddings",
|
|
input=["good morning from litellm"],
|
|
)
|
|
print(f"response: {response}")
|
|
|
|
# stubbed endpoint is setup to return this
|
|
assert response.data[0]["embedding"] == [0.1, 0.2]
|
|
except Exception as e:
|
|
pytest.fail(f"Error occurred: {e}")
|
|
|
|
|
|
def test_triton_generate_raw_request():
|
|
from litellm.utils import return_raw_request
|
|
from litellm.types.utils import CallTypes
|
|
|
|
try:
|
|
kwargs = {
|
|
"model": "triton/llama-3-8b-instruct",
|
|
"messages": [{"role": "user", "content": "who are u?"}],
|
|
"api_base": "http://localhost:8000/generate",
|
|
}
|
|
raw_request = return_raw_request(endpoint=CallTypes.completion, kwargs=kwargs)
|
|
print("raw_request", raw_request)
|
|
assert raw_request is not None
|
|
assert "bad_words" not in json.dumps(raw_request["raw_request_body"])
|
|
assert "stop_words" not in json.dumps(raw_request["raw_request_body"])
|
|
except Exception as e:
|
|
pytest.fail(f"Error occurred: {e}")
|