mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-15 04:19:08 +00:00
1587faeeb9
* test(unit): wait on the work itself instead of fixed sleeps in 21 flaky files Replace sleeps that raced background logging, cache writes, refreshes and thread work with either awaiting that work directly (logging worker flush, the batch logger's flush_queue, the tasks the call spawned, thread joins) or an Event the test's own fake sets when it runs. Each wait keeps a 5s hang guard that never decides pass or fail. Tests that need a fake clock, real sockets, or measure latency are left unchanged, and no production code moves * test(unit): await the tasks each completion spawns instead of yielding once The logging tests yielded once with asyncio.sleep(0) so the task that enqueues the success log could run before the worker flush. That relied on the enqueue happening before the task's first await. Snapshot the running tasks around each acompletion call and await the ones it spawned, so the guarantee holds however many steps sit before the enqueue
270 lines
8.8 KiB
Python
270 lines
8.8 KiB
Python
import asyncio
|
|
import json
|
|
import threading
|
|
from datetime import datetime
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
# Adds the grandparent directory to sys.path to allow importing project modules
|
|
|
|
import pytest
|
|
|
|
import litellm
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_mlflow_logging_functionality():
|
|
"""Test that inputs, outputs and tags are properly logged in MLflow traces."""
|
|
|
|
# Mock MLflow client and dependencies
|
|
mock_client = MagicMock()
|
|
mock_span = MagicMock()
|
|
mock_span.parent_id = None # Simulate root trace
|
|
mock_span.request_id = "test_trace_id"
|
|
trace_started = threading.Event()
|
|
|
|
def start_trace(*args: object, **kwargs: object) -> MagicMock:
|
|
trace_started.set()
|
|
return mock_span
|
|
|
|
mock_client.start_trace.side_effect = start_trace
|
|
|
|
# Mock all MLflow-related imports to avoid requiring MLflow as a dependency
|
|
mock_mlflow_tracking = MagicMock()
|
|
mock_mlflow_tracking.MlflowClient = MagicMock(return_value=mock_client)
|
|
|
|
mock_mlflow_entities = MagicMock()
|
|
mock_mlflow_entities.SpanStatusCode.OK = "OK"
|
|
mock_mlflow_entities.SpanStatusCode.ERROR = "ERROR"
|
|
mock_mlflow_entities.SpanType.LLM = "LLM"
|
|
|
|
mock_mlflow = MagicMock()
|
|
mock_mlflow.get_current_active_span.return_value = None
|
|
|
|
with patch.dict(
|
|
"sys.modules",
|
|
{
|
|
"mlflow": mock_mlflow,
|
|
"mlflow.tracking": mock_mlflow_tracking,
|
|
"mlflow.entities": mock_mlflow_entities,
|
|
"mlflow.tracing.utils": MagicMock(),
|
|
},
|
|
):
|
|
# Now we can safely import MlflowLogger
|
|
from litellm.integrations.mlflow import MlflowLogger
|
|
|
|
# Create MlflowLogger instance
|
|
mlflow_logger = MlflowLogger()
|
|
litellm.callbacks = [mlflow_logger]
|
|
|
|
# Test completion with request_tags and prediction parameter
|
|
test_prediction = {"type": "content", "content": "This is a predicted output"}
|
|
await litellm.acompletion(
|
|
model="gpt-3.5-turbo",
|
|
messages=[{"role": "user", "content": "test message"}],
|
|
prediction=test_prediction,
|
|
mock_response="test response",
|
|
metadata={
|
|
"tags": [
|
|
"tag1",
|
|
"tag2",
|
|
"production",
|
|
"jobID:214590dsff09fds",
|
|
"taskName:run_page_classification",
|
|
]
|
|
},
|
|
)
|
|
|
|
assert await asyncio.to_thread(trace_started.wait, 5)
|
|
|
|
# Verify start_trace was called with tags parameter
|
|
assert mock_client.start_trace.called, "start_trace should have been called"
|
|
|
|
# Get the call arguments
|
|
call_args = mock_client.start_trace.call_args
|
|
assert call_args is not None, "start_trace call args should not be None"
|
|
|
|
# Check that tags parameter was included and properly transformed
|
|
tags_param = call_args.kwargs.get("tags", {})
|
|
expected_tags = {
|
|
"tag1": "",
|
|
"tag2": "",
|
|
"production": "",
|
|
"jobID": "214590dsff09fds",
|
|
"taskName": "run_page_classification",
|
|
}
|
|
assert (
|
|
tags_param == expected_tags
|
|
), f"Expected tags {expected_tags}, got {tags_param}"
|
|
|
|
# Check that prediction parameter was included in inputs
|
|
inputs_param = call_args.kwargs.get("inputs", {})
|
|
assert (
|
|
"prediction" in inputs_param
|
|
), "Prediction should be included in span inputs"
|
|
assert (
|
|
inputs_param["prediction"] == test_prediction
|
|
), f"Expected prediction {test_prediction}, got {inputs_param['prediction']}"
|
|
|
|
|
|
def test_mlflow_token_usage_attribute_structure():
|
|
"""Ensure token usage attributes are formatted with mlflow.chat.tokenUsage."""
|
|
|
|
mock_mlflow_tracking = MagicMock()
|
|
mock_mlflow_tracking.MlflowClient = MagicMock()
|
|
|
|
with patch.dict(
|
|
"sys.modules",
|
|
{
|
|
"mlflow": MagicMock(),
|
|
"mlflow.tracking": mock_mlflow_tracking,
|
|
"mlflow.tracing.utils": MagicMock(),
|
|
},
|
|
):
|
|
from litellm.integrations.mlflow import MlflowLogger
|
|
|
|
mlflow_logger = MlflowLogger()
|
|
|
|
attrs = mlflow_logger._extract_attributes( # type: ignore
|
|
{
|
|
"litellm_call_id": "123",
|
|
"call_type": "completion",
|
|
"model": "gpt-3.5-turbo",
|
|
"standard_logging_object": {
|
|
"prompt_tokens": 5,
|
|
"completion_tokens": 7,
|
|
"total_tokens": 12,
|
|
},
|
|
}
|
|
)
|
|
|
|
assert attrs["mlflow.chat.tokenUsage"] == {
|
|
"input_tokens": 5,
|
|
"output_tokens": 7,
|
|
"total_tokens": 12,
|
|
}
|
|
|
|
|
|
def _mock_mlflow_modules():
|
|
mock_tracking = MagicMock()
|
|
mock_tracking.MlflowClient = MagicMock()
|
|
|
|
class DummySpanEvent:
|
|
def __init__(self, name, attributes):
|
|
self.name = name
|
|
self.attributes = attributes
|
|
|
|
mock_entities = MagicMock()
|
|
mock_entities.SpanStatusCode.OK = "OK"
|
|
mock_entities.SpanEvent = DummySpanEvent
|
|
|
|
return {
|
|
"mlflow": MagicMock(),
|
|
"mlflow.tracking": mock_tracking,
|
|
"mlflow.entities": mock_entities,
|
|
"mlflow.tracing.utils": MagicMock(),
|
|
}
|
|
|
|
|
|
def test_mlflow_stream_handler_uses_async_complete_response():
|
|
modules = _mock_mlflow_modules()
|
|
with patch.dict("sys.modules", modules):
|
|
from litellm.integrations.mlflow import MlflowLogger
|
|
|
|
mlflow_logger = MlflowLogger()
|
|
mlflow_logger._start_span_or_trace = MagicMock(return_value="mock_span")
|
|
mlflow_logger._end_span_or_trace = MagicMock()
|
|
mlflow_logger._extract_and_set_chat_attributes = MagicMock()
|
|
|
|
class DummyDelta:
|
|
def model_dump(self, exclude_none=True):
|
|
return {"content": "chunk"}
|
|
|
|
response_obj = MagicMock()
|
|
response_obj.choices = [MagicMock(delta=DummyDelta())]
|
|
|
|
final_response = MagicMock()
|
|
kwargs = {
|
|
"litellm_call_id": "abc123",
|
|
"async_complete_streaming_response": final_response,
|
|
}
|
|
|
|
mlflow_logger._handle_stream_event(
|
|
kwargs=kwargs,
|
|
response_obj=response_obj,
|
|
start_time=datetime.utcnow(),
|
|
end_time=datetime.utcnow(),
|
|
)
|
|
|
|
mlflow_logger._end_span_or_trace.assert_called_once()
|
|
assert (
|
|
mlflow_logger._end_span_or_trace.call_args.kwargs["outputs"]
|
|
is final_response
|
|
)
|
|
assert "abc123" not in mlflow_logger._stream_id_to_span
|
|
|
|
|
|
def test_mlflow_stream_handler_pops_span_when_end_raises():
|
|
modules = _mock_mlflow_modules()
|
|
with patch.dict("sys.modules", modules):
|
|
from litellm.integrations.mlflow import MlflowLogger
|
|
|
|
mlflow_logger = MlflowLogger()
|
|
mlflow_logger._start_span_or_trace = MagicMock(return_value="mock_span")
|
|
mlflow_logger._end_span_or_trace = MagicMock(
|
|
side_effect=TypeError("unexpected keyword argument 'trace_id'")
|
|
)
|
|
mlflow_logger._extract_and_set_chat_attributes = MagicMock()
|
|
|
|
response_obj = MagicMock()
|
|
response_obj.choices = []
|
|
|
|
kwargs = {
|
|
"litellm_call_id": "leak123",
|
|
"complete_streaming_response": MagicMock(),
|
|
}
|
|
|
|
with pytest.raises(TypeError):
|
|
mlflow_logger._handle_stream_event(
|
|
kwargs=kwargs,
|
|
response_obj=response_obj,
|
|
start_time=datetime.utcnow(),
|
|
end_time=datetime.utcnow(),
|
|
)
|
|
|
|
assert "leak123" not in mlflow_logger._stream_id_to_span
|
|
|
|
|
|
class _Mlflow2StyleClient:
|
|
"""Mimics the mlflow 2.x client signatures, which have no trace_id kwarg."""
|
|
|
|
def __init__(self):
|
|
self.ended_traces = []
|
|
self.ended_spans = []
|
|
|
|
def end_trace(self, request_id, outputs=None, attributes=None, status="OK", end_time_ns=None):
|
|
self.ended_traces.append(request_id)
|
|
|
|
def end_span(self, request_id, span_id, outputs=None, attributes=None, status="OK", end_time_ns=None):
|
|
self.ended_spans.append((request_id, span_id))
|
|
|
|
|
|
def test_mlflow_end_span_or_trace_works_with_mlflow_2x_client():
|
|
modules = _mock_mlflow_modules()
|
|
with patch.dict("sys.modules", modules):
|
|
from litellm.integrations.mlflow import MlflowLogger
|
|
|
|
mlflow_logger = MlflowLogger()
|
|
client = _Mlflow2StyleClient()
|
|
mlflow_logger._client = client
|
|
|
|
root_span = MagicMock(parent_id=None, request_id="req-1")
|
|
mlflow_logger._end_span_or_trace(
|
|
span=root_span, outputs="out", end_time_ns=1, status="OK"
|
|
)
|
|
assert client.ended_traces == ["req-1"]
|
|
|
|
child_span = MagicMock(parent_id="parent-1", request_id="req-2", span_id="span-2")
|
|
mlflow_logger._end_span_or_trace(
|
|
span=child_span, outputs="out", end_time_ns=1, status="OK"
|
|
)
|
|
assert client.ended_spans == [("req-2", "span-2")]
|