From b9132968b2a20060dc4b5d47f2e9f7a6a238232b Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Sat, 23 Aug 2025 13:13:23 -0700 Subject: [PATCH] [Perf] Improvements for Async Success Handler (Logging Callbacks) - Approx +130 RPS (#13905) * [Performance] Reduce Significant CPU overhead from litellm_logging.py (#13895) * fix: litellm.configured_cold_storage_logger * fix Session Management - Non-OpenAI Models docs * ruff fix * test fix * create LoggingWorker * add GLOBAL_LOGGING_WORKER for async task handling * fix logging tests * add conftest * fix conftest * test fix location of encode bedrock runtime modelid arn * fix conftest.py * tuning LoggingWorker * conftest.py * fix conftest batches/ * test_async_chat_azure * event_loop * test_bedrock_streaming_passthrough_test2 * fix GLOBAL_LOGGING_WORKER * logging worker * add flush for global logging worker * Revert "fix GLOBAL_LOGGING_WORKER" This reverts commit d254f508f48935652f054777652938ad71976cce. * fix conftest clear_queue * fix conftest clear_queue * setup_and_teardown for llm translation * docs AWS_REGION * test_async_chat_azure * change test DIR * run ci/cd again * use 1 job for litellm_router_unit_testing * fix space * fix litellm_router_unit_testing * test_aaarouter_dynamic_cooldown_message_retry_time * litellm_router_unit_testing * conftest.py clearing qu * fixes litellm_router_unit_testing * fixes clear_queue * fix router_unit_tests * remove conftest * add back conftest for router * fix event loop test * test fix * fixes for LoggingWorker * ruff fix --- .circleci/config.yml | 59 +++++++- litellm/litellm_core_utils/logging_worker.py | 131 ++++++++++++++++++ litellm/utils.py | 18 ++- tests/batches_tests/conftest.py | 22 +++ tests/image_gen_tests/conftest.py | 22 +++ tests/llm_responses_api_testing/conftest.py | 11 ++ tests/llm_translation/conftest.py | 11 +- .../test_bedrock_completion.py | 8 +- tests/local_testing/conftest.py | 19 +++ tests/local_testing/test_router.py | 2 +- tests/logging_callback_tests/conftest.py | 25 ++++ .../test_custom_callback_router.py | 12 +- tests/mcp_tests/conftest.py | 64 +++++++++ tests/router_unit_tests/conftest.py | 82 +++++++++++ .../test_router_endpoints.py | 2 +- tests/test_litellm/conftest.py | 14 ++ tests/unified_google_tests/conftest.py | 64 +++++++++ 17 files changed, 549 insertions(+), 17 deletions(-) create mode 100644 litellm/litellm_core_utils/logging_worker.py create mode 100644 tests/batches_tests/conftest.py create mode 100644 tests/image_gen_tests/conftest.py rename tests/{local_testing => logging_callback_tests}/test_custom_callback_router.py (98%) create mode 100644 tests/mcp_tests/conftest.py create mode 100644 tests/router_unit_tests/conftest.py create mode 100644 tests/unified_google_tests/conftest.py diff --git a/.circleci/config.yml b/.circleci/config.yml index e5b1a2d0bf3..e49dca6faae 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -439,6 +439,7 @@ jobs: paths: - auth_ui_unit_tests_coverage.xml - auth_ui_unit_tests_coverage + litellm_router_testing: # Runs all tests with the "router" keyword docker: - image: cimg/python:3.11 @@ -469,7 +470,55 @@ jobs: command: | pwd ls - python -m pytest tests/local_testing tests/router_unit_tests --cov=litellm --cov-report=xml -vv -k "router" -x -v --junitxml=test-results/junit.xml --durations=5 + python -m pytest tests/local_testing --cov=litellm --cov-report=xml -vv -k "router" -x -v --junitxml=test-results/junit.xml --durations=5 + no_output_timeout: 120m + - run: + name: Rename the coverage files + command: | + mv coverage.xml litellm_router_coverage.xml + mv .coverage litellm_router_coverage + # Store test results + - store_test_results: + path: test-results + + - persist_to_workspace: + root: . + paths: + - litellm_router_coverage.xml + - litellm_router_coverage + + litellm_router_unit_testing: # Runs all tests with the "router" keyword + docker: + - image: cimg/python:3.11 + auth: + username: ${DOCKERHUB_USERNAME} + password: ${DOCKERHUB_PASSWORD} + working_directory: ~/project + + steps: + - checkout + - setup_google_dns + - run: + name: Install Dependencies + command: | + python -m pip install --upgrade pip + python -m pip install -r requirements.txt + pip install "pytest==7.3.1" + pip install "respx==0.22.0" + pip install "pytest-cov==5.0.0" + pip install "pytest-retry==1.6.3" + pip install "pytest-asyncio==0.21.1" + pip install semantic_router --no-deps + pip install aurelio_sdk --no-deps + pip install "pytest-xdist==3.6.1" + # Run pytest and generate JUnit XML report + - setup_litellm_enterprise_pip + - run: + name: Run tests + command: | + pwd + ls + python -m pytest -vv tests/router_unit_tests --cov=litellm --cov-report=xml -x -s -v --junitxml=test-results/junit.xml --durations=5 no_output_timeout: 120m - run: name: Rename the coverage files @@ -2982,6 +3031,12 @@ workflows: only: - main - /litellm_.*/ + - litellm_router_unit_testing: + filters: + branches: + only: + - main + - /litellm_.*/ - check_code_and_doc_quality: filters: branches: @@ -3128,6 +3183,7 @@ workflows: - image_gen_testing - logging_testing - litellm_router_testing + - litellm_router_unit_testing - caching_unit_tests - litellm_proxy_unit_testing - litellm_security_tests @@ -3187,6 +3243,7 @@ workflows: - image_gen_testing - logging_testing - litellm_router_testing + - litellm_router_unit_testing - caching_unit_tests - langfuse_logging_unit_tests - litellm_assistants_api_testing diff --git a/litellm/litellm_core_utils/logging_worker.py b/litellm/litellm_core_utils/logging_worker.py new file mode 100644 index 00000000000..faf8413bd57 --- /dev/null +++ b/litellm/litellm_core_utils/logging_worker.py @@ -0,0 +1,131 @@ +import asyncio +import contextlib +from typing import Coroutine, Optional + +from litellm._logging import verbose_logger + + +class LoggingWorker: + """ + A simple, async logging worker that processes log coroutines in the background. + Designed to be best-effort with bounded queues to prevent backpressure. + + This leads to a +200 RPS performance improvement when using LiteLLM Python SDK or Proxy Server. + - Use this to queue coroutine tasks that are not critical to the main flow of the application. e.g Success/Error callbacks, logging, etc. + """ + LOGGING_WORKER_MAX_QUEUE_SIZE = 50_000 + LOGGING_WORKER_MAX_TIME_PER_COROUTINE = 20.0 + + MAX_ITERATIONS_TO_CLEAR_QUEUE = 200 + MAX_TIME_TO_CLEAR_QUEUE = 5.0 + + def __init__( + self, + timeout: float = LOGGING_WORKER_MAX_TIME_PER_COROUTINE, + max_queue_size: int = LOGGING_WORKER_MAX_QUEUE_SIZE, + ): + self.timeout = timeout + self.max_queue_size = max_queue_size + self._queue: Optional[asyncio.Queue] = None + self._worker_task: Optional[asyncio.Task] = None + + def _ensure_queue(self) -> None: + """Initialize the queue if it doesn't exist.""" + if self._queue is None: + self._queue = asyncio.Queue(maxsize=self.max_queue_size) + + def start(self) -> None: + """Start the logging worker. Idempotent - safe to call multiple times.""" + self._ensure_queue() + if self._worker_task is None or self._worker_task.done(): + self._worker_task = asyncio.create_task(self._worker_loop()) + + async def _worker_loop(self) -> None: + """Main worker loop that processes log coroutines sequentially.""" + try: + if self._queue is None: + return + + while True: + # Process one coroutine at a time to keep event loop load predictable + coroutine = await self._queue.get() + try: + await asyncio.wait_for(coroutine, timeout=self.timeout) + except Exception as e: + verbose_logger.exception(f"LoggingWorker error: {e}") + pass + finally: + self._queue.task_done() + + except asyncio.CancelledError as e: + verbose_logger.exception(f"LoggingWorker cancelled: {e}") + pass + + def enqueue(self, coroutine: Coroutine) -> None: + """ + Add a coroutine to the logging queue. + Hot path: never blocks, drops logs if queue is full. + """ + if self._queue is None: + return + + try: + self._queue.put_nowait(coroutine) + except asyncio.QueueFull as e: + verbose_logger.exception(f"LoggingWorker queue is full: {e}") + # Drop logs on overload to protect request throughput + pass + + def ensure_initialized_and_enqueue(self, async_coroutine: Coroutine): + """ + Ensure the logging worker is initialized and enqueue the coroutine. + """ + self.start() + self.enqueue(async_coroutine) + + async def stop(self) -> None: + """Stop the logging worker and clean up resources.""" + if self._worker_task: + self._worker_task.cancel() + with contextlib.suppress(Exception): + await self._worker_task + self._worker_task = None + + async def flush(self) -> None: + """Flush the logging queue.""" + if self._queue is None: + return + while not self._queue.empty(): + await self._queue.join() + + async def clear_queue(self): + """ + Clear the queue with a maximum time limit. + """ + if self._queue is None: + return + + start_time = asyncio.get_event_loop().time() + + for _ in range(self.MAX_ITERATIONS_TO_CLEAR_QUEUE): + # Check if we've exceeded the maximum time + if asyncio.get_event_loop().time() - start_time >= self.MAX_TIME_TO_CLEAR_QUEUE: + verbose_logger.warning(f"clear_queue exceeded max_time of {self.MAX_TIME_TO_CLEAR_QUEUE}s, stopping early") + break + + try: + coroutine = self._queue.get_nowait() + # Await the coroutine to properly execute and avoid "never awaited" warnings + try: + await asyncio.wait_for(coroutine, timeout=self.timeout) + except Exception: + # Suppress errors during cleanup + pass + self._queue.task_done() # If you're using join() elsewhere + except asyncio.QueueEmpty: + break + + +# Global instance for backward compatibility +GLOBAL_LOGGING_WORKER = LoggingWorker() + diff --git a/litellm/utils.py b/litellm/utils.py index 2821720e06d..aa3c00735ec 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -833,10 +833,22 @@ async def _client_async_logging_helper( print_verbose( f"Async Wrapper: Completed Call, calling async_success_handler: {logging_obj.async_success_handler}" ) - # check if user does not want this to be logged - asyncio.create_task( - logging_obj.async_success_handler(result, start_time, end_time) + ################################################ + # Async Logging Worker + ################################################ + from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER + GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue( + async_coroutine = logging_obj.async_success_handler( + result=result, + start_time=start_time, + end_time=end_time + ) ) + + + ################################################ + # Sync Logging Worker + ################################################ logging_obj.handle_sync_success_callbacks_for_async_calls( result=result, start_time=start_time, diff --git a/tests/batches_tests/conftest.py b/tests/batches_tests/conftest.py new file mode 100644 index 00000000000..e504ee5ac0a --- /dev/null +++ b/tests/batches_tests/conftest.py @@ -0,0 +1,22 @@ +# conftest.py + +import importlib +import os +import sys + +import pytest + +sys.path.insert( + 0, os.path.abspath("../..") +) # Adds the parent directory to the system path +import litellm +import asyncio + +@pytest.fixture(scope="session") +def event_loop(): + try: + loop = asyncio.get_running_loop() + except RuntimeError: + loop = asyncio.new_event_loop() + yield loop + loop.close() diff --git a/tests/image_gen_tests/conftest.py b/tests/image_gen_tests/conftest.py new file mode 100644 index 00000000000..fe04da2c66e --- /dev/null +++ b/tests/image_gen_tests/conftest.py @@ -0,0 +1,22 @@ + +import importlib +import os +import sys +import asyncio +import pytest + +sys.path.insert( + 0, os.path.abspath("../..") +) # Adds the parent directory to the system path +import litellm + +import asyncio + +@pytest.fixture(scope="session") +def event_loop(): + try: + loop = asyncio.get_running_loop() + except RuntimeError: + loop = asyncio.new_event_loop() + yield loop + loop.close() \ No newline at end of file diff --git a/tests/llm_responses_api_testing/conftest.py b/tests/llm_responses_api_testing/conftest.py index b3561d8a626..197dce9a020 100644 --- a/tests/llm_responses_api_testing/conftest.py +++ b/tests/llm_responses_api_testing/conftest.py @@ -11,6 +11,17 @@ sys.path.insert( ) # Adds the parent directory to the system path import litellm +import asyncio + +@pytest.fixture(scope="session") +def event_loop(): + try: + loop = asyncio.get_running_loop() + except RuntimeError: + loop = asyncio.new_event_loop() + yield loop + loop.close() + @pytest.fixture(scope="function", autouse=True) def setup_and_teardown(): diff --git a/tests/llm_translation/conftest.py b/tests/llm_translation/conftest.py index 38262eb590d..97edb4c023c 100644 --- a/tests/llm_translation/conftest.py +++ b/tests/llm_translation/conftest.py @@ -11,11 +11,14 @@ sys.path.insert( ) # Adds the parent directory to the system path import litellm +import asyncio @pytest.fixture(scope="session") def event_loop(): - """Create an instance of the default event loop for each test session.""" - loop = asyncio.get_event_loop_policy().new_event_loop() + try: + loop = asyncio.get_running_loop() + except RuntimeError: + loop = asyncio.new_event_loop() yield loop loop.close() @@ -28,6 +31,10 @@ def setup_and_teardown(event_loop): # Add event_loop as a dependency import litellm from litellm import Router + from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER + # flush all logs + asyncio.run(GLOBAL_LOGGING_WORKER.clear_queue()) + importlib.reload(litellm) # Set the event loop from the fixture diff --git a/tests/llm_translation/test_bedrock_completion.py b/tests/llm_translation/test_bedrock_completion.py index 4bab551d9fb..c246094c81c 100644 --- a/tests/llm_translation/test_bedrock_completion.py +++ b/tests/llm_translation/test_bedrock_completion.py @@ -3308,7 +3308,7 @@ async def test_bedrock_converse__streaming_passthrough(monkeypatch): @pytest.mark.asyncio -async def test_bedrock_streaming_passthrough(monkeypatch): +async def test_bedrock_streaming_passthrough_test2(monkeypatch): import litellm import time import asyncio @@ -3350,7 +3350,7 @@ async def test_bedrock_streaming_passthrough(monkeypatch): async for chunk in response: print(chunk) - await asyncio.sleep(1) + await asyncio.sleep(5) mock_callback.assert_called_once() # check standard logging payload created @@ -3360,7 +3360,7 @@ async def test_bedrock_streaming_passthrough(monkeypatch): @pytest.mark.asyncio -async def test_bedrock_streaming_passthrough(monkeypatch): +async def test_bedrock_streaming_passthrough_test1(monkeypatch): import litellm import time import asyncio @@ -3402,7 +3402,7 @@ async def test_bedrock_streaming_passthrough(monkeypatch): async for chunk in response: print(chunk) - await asyncio.sleep(1) + await asyncio.sleep(5) mock_callback.assert_called_once() # check standard logging payload created diff --git a/tests/local_testing/conftest.py b/tests/local_testing/conftest.py index b3561d8a626..7290f3e75ff 100644 --- a/tests/local_testing/conftest.py +++ b/tests/local_testing/conftest.py @@ -11,6 +11,19 @@ sys.path.insert( ) # Adds the parent directory to the system path import litellm +import asyncio + +@pytest.fixture(scope="session") +def event_loop(): + try: + loop = asyncio.get_running_loop() + except RuntimeError: + loop = asyncio.new_event_loop() + yield loop + loop.close() + + + @pytest.fixture(scope="function", autouse=True) def setup_and_teardown(): @@ -24,6 +37,12 @@ def setup_and_teardown(): import litellm from litellm import Router + import asyncio + + from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER + # flush all logs + asyncio.run(GLOBAL_LOGGING_WORKER.clear_queue()) + importlib.reload(litellm) diff --git a/tests/local_testing/test_router.py b/tests/local_testing/test_router.py index f719f0b6061..10634ba3a73 100644 --- a/tests/local_testing/test_router.py +++ b/tests/local_testing/test_router.py @@ -2328,7 +2328,7 @@ async def test_aaarouter_dynamic_cooldown_message_retry_time(sync_mode): except litellm.RateLimitError: pass - await asyncio.sleep(2) + await asyncio.sleep(5) if sync_mode: cooldown_deployments = _get_cooldown_deployments( diff --git a/tests/logging_callback_tests/conftest.py b/tests/logging_callback_tests/conftest.py index eca0bc431a5..e47df872d3f 100644 --- a/tests/logging_callback_tests/conftest.py +++ b/tests/logging_callback_tests/conftest.py @@ -10,7 +10,16 @@ sys.path.insert( 0, os.path.abspath("../..") ) # Adds the parent directory to the system path import litellm +import asyncio +@pytest.fixture(scope="session") +def event_loop(): + try: + loop = asyncio.get_running_loop() + except RuntimeError: + loop = asyncio.new_event_loop() + yield loop + loop.close() @pytest.fixture(scope="function", autouse=True) def setup_and_teardown(): @@ -24,8 +33,23 @@ def setup_and_teardown(): import litellm from litellm import Router + import asyncio + + from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER + # flush all logs + asyncio.run(GLOBAL_LOGGING_WORKER.clear_queue()) + importlib.reload(litellm) + + try: + if hasattr(litellm, "proxy") and hasattr(litellm.proxy, "proxy_server"): + import litellm.proxy.proxy_server + + importlib.reload(litellm.proxy.proxy_server) + except Exception as e: + print(f"Error reloading litellm.proxy.proxy_server: {e}") + import asyncio loop = asyncio.get_event_loop_policy().new_event_loop() @@ -39,6 +63,7 @@ def setup_and_teardown(): asyncio.set_event_loop(None) # Remove the reference to the loop + def pytest_collection_modifyitems(config, items): # Separate tests in 'test_amazing_proxy_custom_logger.py' and other tests custom_logger_tests = [ diff --git a/tests/local_testing/test_custom_callback_router.py b/tests/logging_callback_tests/test_custom_callback_router.py similarity index 98% rename from tests/local_testing/test_custom_callback_router.py rename to tests/logging_callback_tests/test_custom_callback_router.py index 17b37dca5ca..e8fc1ac8676 100644 --- a/tests/local_testing/test_custom_callback_router.py +++ b/tests/logging_callback_tests/test_custom_callback_router.py @@ -34,7 +34,7 @@ from litellm.integrations.custom_logger import CustomLogger ## 5. Azure OpenAI acompletion + streaming call with fallbacks ## 6. Azure OpenAI aembedding call with fallbacks -# Test interfaces +## Test interfaces ## 1. router.completion() + router.embeddings() ## 2. proxy.completions + proxy.embeddings @@ -265,8 +265,9 @@ class CompletionCustomHandler( async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): try: + print("CompletionCustomHandler.async_log_success_event, kwargs: ", kwargs) self.states.append("async_success") - print("in async success, kwargs: ", kwargs) + print("############### CompletionCustomHandler async success, kwargs: ", kwargs) ## START TIME assert isinstance(start_time, datetime) ## END TIME @@ -409,7 +410,8 @@ async def test_async_chat_azure(): model="gpt-3.5-turbo", messages=[{"role": "user", "content": "Hi 👋 - i'm openai"}], ) - await asyncio.sleep(2) + print("got response, sleeping 5 seconds....") + await asyncio.sleep(5) assert len(customHandler_completion_azure_router.errors) == 0 assert ( len(customHandler_completion_azure_router.states) == 3 @@ -427,7 +429,7 @@ async def test_async_chat_azure(): async for chunk in response: print(f"async azure router chunk: {chunk}") continue - await asyncio.sleep(2) + await asyncio.sleep(5) print(f"customHandler.states: {customHandler_streaming_azure_router.states}") assert len(customHandler_streaming_azure_router.errors) == 0 assert ( @@ -459,7 +461,7 @@ async def test_async_chat_azure(): print(f"response in router3 acompletion: {response}") except Exception: pass - await asyncio.sleep(1) + await asyncio.sleep(5) print(f"customHandler.states: {customHandler_failure.states}") assert len(customHandler_failure.errors) == 0 assert len(customHandler_failure.states) == 3 # pre, post, failure diff --git a/tests/mcp_tests/conftest.py b/tests/mcp_tests/conftest.py new file mode 100644 index 00000000000..f74a3569c19 --- /dev/null +++ b/tests/mcp_tests/conftest.py @@ -0,0 +1,64 @@ +# conftest.py + +import importlib +import os +import sys + +import pytest + +sys.path.insert( + 0, os.path.abspath("../..") +) # Adds the parent directory to the system path +import litellm +import asyncio + +@pytest.fixture(scope="session") +def event_loop(): + try: + loop = asyncio.get_running_loop() + except RuntimeError: + loop = asyncio.new_event_loop() + yield loop + loop.close() + + +@pytest.fixture(scope="function", autouse=True) +def setup_and_teardown(): + """ + This fixture reloads litellm before every function. To speed up testing by removing callbacks being chained. + """ + curr_dir = os.getcwd() # Get the current working directory + sys.path.insert( + 0, os.path.abspath("../..") + ) # Adds the project directory to the system path + + import litellm + from litellm import Router + + importlib.reload(litellm) + import asyncio + + loop = asyncio.get_event_loop_policy().new_event_loop() + asyncio.set_event_loop(loop) + print(litellm) + # from litellm import Router, completion, aembedding, acompletion, embedding + yield + + # Teardown code (executes after the yield point) + loop.close() # Close the loop created earlier + asyncio.set_event_loop(None) # Remove the reference to the loop + + +def pytest_collection_modifyitems(config, items): + # Separate tests in 'test_amazing_proxy_custom_logger.py' and other tests + custom_logger_tests = [ + item for item in items if "custom_logger" in item.parent.name + ] + other_tests = [item for item in items if "custom_logger" not in item.parent.name] + + # Sort tests based on their names + custom_logger_tests.sort(key=lambda x: x.name) + other_tests.sort(key=lambda x: x.name) + + # Reorder the items list + items[:] = custom_logger_tests + other_tests diff --git a/tests/router_unit_tests/conftest.py b/tests/router_unit_tests/conftest.py new file mode 100644 index 00000000000..7290f3e75ff --- /dev/null +++ b/tests/router_unit_tests/conftest.py @@ -0,0 +1,82 @@ +# conftest.py + +import importlib +import os +import sys + +import pytest + +sys.path.insert( + 0, os.path.abspath("../..") +) # Adds the parent directory to the system path +import litellm + +import asyncio + +@pytest.fixture(scope="session") +def event_loop(): + try: + loop = asyncio.get_running_loop() + except RuntimeError: + loop = asyncio.new_event_loop() + yield loop + loop.close() + + + + +@pytest.fixture(scope="function", autouse=True) +def setup_and_teardown(): + """ + This fixture reloads litellm before every function. To speed up testing by removing callbacks being chained. + """ + curr_dir = os.getcwd() # Get the current working directory + sys.path.insert( + 0, os.path.abspath("../..") + ) # Adds the project directory to the system path + + import litellm + from litellm import Router + import asyncio + + from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER + # flush all logs + asyncio.run(GLOBAL_LOGGING_WORKER.clear_queue()) + + + importlib.reload(litellm) + + try: + if hasattr(litellm, "proxy") and hasattr(litellm.proxy, "proxy_server"): + import litellm.proxy.proxy_server + + importlib.reload(litellm.proxy.proxy_server) + except Exception as e: + print(f"Error reloading litellm.proxy.proxy_server: {e}") + + import asyncio + + loop = asyncio.get_event_loop_policy().new_event_loop() + asyncio.set_event_loop(loop) + print(litellm) + # from litellm import Router, completion, aembedding, acompletion, embedding + yield + + # Teardown code (executes after the yield point) + loop.close() # Close the loop created earlier + asyncio.set_event_loop(None) # Remove the reference to the loop + + +def pytest_collection_modifyitems(config, items): + # Separate tests in 'test_amazing_proxy_custom_logger.py' and other tests + custom_logger_tests = [ + item for item in items if "custom_logger" in item.parent.name + ] + other_tests = [item for item in items if "custom_logger" not in item.parent.name] + + # Sort tests based on their names + custom_logger_tests.sort(key=lambda x: x.name) + other_tests.sort(key=lambda x: x.name) + + # Reorder the items list + items[:] = custom_logger_tests + other_tests diff --git a/tests/router_unit_tests/test_router_endpoints.py b/tests/router_unit_tests/test_router_endpoints.py index 0b60d5b26fa..87af688461d 100644 --- a/tests/router_unit_tests/test_router_endpoints.py +++ b/tests/router_unit_tests/test_router_endpoints.py @@ -14,7 +14,7 @@ sys.path.insert( from litellm import Router, CustomLogger from litellm.types.utils import StandardLoggingPayload -# Get the current directory of the file being run +## Get the current directory of the file being run pwd = os.path.dirname(os.path.realpath(__file__)) print(pwd) diff --git a/tests/test_litellm/conftest.py b/tests/test_litellm/conftest.py index a88148f9d11..ac8a00d850c 100644 --- a/tests/test_litellm/conftest.py +++ b/tests/test_litellm/conftest.py @@ -9,9 +9,23 @@ import pytest sys.path.insert( 0, os.path.abspath("../..") ) # Adds the parent directory to the system path +import asyncio + import litellm +@pytest.fixture(scope="session") +def event_loop(): + try: + loop = asyncio.get_running_loop() + except RuntimeError: + loop = asyncio.new_event_loop() + yield loop + loop.close() + + + + @pytest.fixture(scope="function", autouse=True) def setup_and_teardown(): """ diff --git a/tests/unified_google_tests/conftest.py b/tests/unified_google_tests/conftest.py new file mode 100644 index 00000000000..f74a3569c19 --- /dev/null +++ b/tests/unified_google_tests/conftest.py @@ -0,0 +1,64 @@ +# conftest.py + +import importlib +import os +import sys + +import pytest + +sys.path.insert( + 0, os.path.abspath("../..") +) # Adds the parent directory to the system path +import litellm +import asyncio + +@pytest.fixture(scope="session") +def event_loop(): + try: + loop = asyncio.get_running_loop() + except RuntimeError: + loop = asyncio.new_event_loop() + yield loop + loop.close() + + +@pytest.fixture(scope="function", autouse=True) +def setup_and_teardown(): + """ + This fixture reloads litellm before every function. To speed up testing by removing callbacks being chained. + """ + curr_dir = os.getcwd() # Get the current working directory + sys.path.insert( + 0, os.path.abspath("../..") + ) # Adds the project directory to the system path + + import litellm + from litellm import Router + + importlib.reload(litellm) + import asyncio + + loop = asyncio.get_event_loop_policy().new_event_loop() + asyncio.set_event_loop(loop) + print(litellm) + # from litellm import Router, completion, aembedding, acompletion, embedding + yield + + # Teardown code (executes after the yield point) + loop.close() # Close the loop created earlier + asyncio.set_event_loop(None) # Remove the reference to the loop + + +def pytest_collection_modifyitems(config, items): + # Separate tests in 'test_amazing_proxy_custom_logger.py' and other tests + custom_logger_tests = [ + item for item in items if "custom_logger" in item.parent.name + ] + other_tests = [item for item in items if "custom_logger" not in item.parent.name] + + # Sort tests based on their names + custom_logger_tests.sort(key=lambda x: x.name) + other_tests.sort(key=lambda x: x.name) + + # Reorder the items list + items[:] = custom_logger_tests + other_tests