From 02ab3cb73d93d71a7bc73ec4feb397100978e11d Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Thu, 11 Jul 2024 08:47:16 -0700 Subject: [PATCH] test- otel span recording --- .circleci/config.yml | 98 ++++++++++++++++ litellm/integrations/opentelemetry.py | 6 + litellm/proxy/common_utils/debug_utils.py | 16 +++ .../example_config_yaml/otel_test_config.yaml | 11 ++ tests/otel_tests/test_otel.py | 111 ++++++++++++++++++ 5 files changed, 242 insertions(+) create mode 100644 litellm/proxy/example_config_yaml/otel_test_config.yaml create mode 100644 tests/otel_tests/test_otel.py diff --git a/.circleci/config.yml b/.circleci/config.yml index 40d498d6e..58a3cfd6c 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -249,6 +249,104 @@ jobs: # Store test results - store_test_results: path: test-results + proxy_log_to_otel_tests: + machine: + image: ubuntu-2204:2023.10.1 + resource_class: xlarge + working_directory: ~/project + steps: + - checkout + - run: + name: Install Docker CLI (In case it's not already installed) + command: | + sudo apt-get update + sudo apt-get install -y docker-ce docker-ce-cli containerd.io + - run: + name: Install Python 3.9 + command: | + curl https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh --output miniconda.sh + bash miniconda.sh -b -p $HOME/miniconda + export PATH="$HOME/miniconda/bin:$PATH" + conda init bash + source ~/.bashrc + conda create -n myenv python=3.9 -y + conda activate myenv + python --version + - run: + name: Install Dependencies + command: | + pip install "pytest==7.3.1" + pip install "pytest-asyncio==0.21.1" + pip install aiohttp + pip install openai + python -m pip install --upgrade pip + python -m pip install -r .circleci/requirements.txt + pip install "pytest==7.3.1" + pip install "pytest-mock==3.12.0" + pip install "pytest-asyncio==0.21.1" + pip install mypy + pip install "google-generativeai==0.3.2" + pip install "google-cloud-aiplatform==1.43.0" + pip install pyarrow + pip install "boto3==1.34.34" + pip install "aioboto3==12.3.0" + pip install numpydoc + pip install prisma + pip install fastapi + pip install jsonschema + pip install "httpx==0.24.1" + pip install "gunicorn==21.2.0" + pip install "anyio==3.7.1" + pip install "aiodynamo==23.10.1" + pip install "asyncio==3.4.3" + pip install "PyGithub==1.59.1" + - run: + name: Build Docker image + command: docker build -t my-app:latest -f Dockerfile.database . + - run: + name: Run Docker container + command: | + docker run -d \ + -p 4000:4000 \ + -e DATABASE_URL=$PROXY_DATABASE_URL \ + -e REDIS_HOST=$REDIS_HOST \ + -e REDIS_PASSWORD=$REDIS_PASSWORD \ + -e REDIS_PORT=$REDIS_PORT \ + -e OPENAI_API_KEY=$OPENAI_API_KEY \ + -e LITELLM_LICENSE=$LITELLM_LICENSE \ + -e OTEL_EXPORTER="in_memory" \ + --name my-app \ + -v $(pwd)/litellm/proxy/example_config_yaml/otel_test_config.yaml.yaml:/app/config.yaml \ + my-app:latest \ + --config /app/config.yaml \ + --port 4000 \ + --detailed_debug \ + - run: + name: Install curl and dockerize + command: | + sudo apt-get update + sudo apt-get install -y curl + sudo wget https://github.com/jwilder/dockerize/releases/download/v0.6.1/dockerize-linux-amd64-v0.6.1.tar.gz + sudo tar -C /usr/local/bin -xzvf dockerize-linux-amd64-v0.6.1.tar.gz + sudo rm dockerize-linux-amd64-v0.6.1.tar.gz + - run: + name: Start outputting logs + command: docker logs -f my-app + background: true + - run: + name: Wait for app to be ready + command: dockerize -wait http://localhost:4000 -timeout 5m + - run: + name: Run tests + command: | + pwd + ls + python -m pytest -vv tests/otel_tests/test_otel.py -x --junitxml=test-results/junit.xml --durations=5 + no_output_timeout: 120m + + # Store test results + - store_test_results: + path: test-results publish_to_pypi: docker: diff --git a/litellm/integrations/opentelemetry.py b/litellm/integrations/opentelemetry.py index c15161fc7..4ed561116 100644 --- a/litellm/integrations/opentelemetry.py +++ b/litellm/integrations/opentelemetry.py @@ -52,6 +52,12 @@ class OpenTelemetryConfig: OTEL_HEADERS gets sent as headers = {"x-honeycomb-team": "B85YgLm96******"} """ + from opentelemetry.sdk.trace.export.in_memory_span_exporter import ( + InMemorySpanExporter, + ) + + if os.getenv("OTEL_EXPORTER") == "in_memory": + return cls(exporter=InMemorySpanExporter()) return cls( exporter=os.getenv("OTEL_EXPORTER", "console"), endpoint=os.getenv("OTEL_ENDPOINT"), diff --git a/litellm/proxy/common_utils/debug_utils.py b/litellm/proxy/common_utils/debug_utils.py index dc77958a6..6c96967aa 100644 --- a/litellm/proxy/common_utils/debug_utils.py +++ b/litellm/proxy/common_utils/debug_utils.py @@ -25,3 +25,19 @@ if os.environ.get("LITELLM_PROFILE", "false").lower() == "true": result.append(f"{stat.traceback.format()}: {stat.size / 1024} KiB") return {"top_50_memory_usage": result} + + +@router.get("/otel-spans", include_in_schema=False) +async def get_otel_spans(): + from litellm.integrations.opentelemetry import OpenTelemetry + from litellm.proxy.proxy_server import open_telemetry_logger + + open_telemetry_logger: OpenTelemetry = open_telemetry_logger + otel_exporter = open_telemetry_logger.OTEL_EXPORTER + recorded_spans = otel_exporter.get_finished_spans() + + print("Spans: ", recorded_spans) # noqa + + # these are otel spans - get the span name + span_names = [span.name for span in recorded_spans] + return {"otel_spans": span_names} diff --git a/litellm/proxy/example_config_yaml/otel_test_config.yaml b/litellm/proxy/example_config_yaml/otel_test_config.yaml new file mode 100644 index 000000000..2e2537443 --- /dev/null +++ b/litellm/proxy/example_config_yaml/otel_test_config.yaml @@ -0,0 +1,11 @@ +model_list: + - model_name: fake-openai-endpoint + litellm_params: + model: openai/fake + api_key: fake-key + api_base: https://exampleopenaiendpoint-production.up.railway.app/ + +litellm_settings: + cache: true + callbacks: ["otel"] + diff --git a/tests/otel_tests/test_otel.py b/tests/otel_tests/test_otel.py new file mode 100644 index 000000000..a07e93e3f --- /dev/null +++ b/tests/otel_tests/test_otel.py @@ -0,0 +1,111 @@ +# What this tests ? +## Tests /chat/completions by generating a key and then making a chat completions request +import pytest +import asyncio +import aiohttp, openai +from openai import OpenAI, AsyncOpenAI +from typing import Optional, List, Union + + +async def generate_key( + session, + models=[ + "gpt-4", + "text-embedding-ada-002", + "dall-e-2", + "fake-openai-endpoint", + "mistral-embed", + ], +): + url = "http://0.0.0.0:4000/key/generate" + headers = {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"} + data = { + "models": models, + "duration": None, + } + + async with session.post(url, headers=headers, json=data) as response: + status = response.status + response_text = await response.text() + + print(response_text) + print() + + if status != 200: + raise Exception(f"Request did not return a 200 status code: {status}") + + return await response.json() + + +async def chat_completion(session, key, model: Union[str, List] = "gpt-4"): + url = "http://0.0.0.0:4000/chat/completions" + headers = { + "Authorization": f"Bearer {key}", + "Content-Type": "application/json", + } + data = { + "model": model, + "messages": [ + {"role": "user", "content": "Hello!"}, + ], + } + + async with session.post(url, headers=headers, json=data) as response: + status = response.status + response_text = await response.text() + + print(response_text) + print() + + if status != 200: + raise Exception(f"Request did not return a 200 status code: {status}") + + return await response.json() + + +async def get_otel_spans(session, key): + url = "http://0.0.0.0:4000/otel-spans" + headers = { + "Authorization": f"Bearer {key}", + "Content-Type": "application/json", + } + + async with session.get(url, headers=headers) as response: + status = response.status + response_text = await response.text() + + print(response_text) + print() + + if status != 200: + raise Exception(f"Request did not return a 200 status code: {status}") + + return await response.json() + + +@pytest.mark.asyncio +async def test_chat_completion_check_otel_spans(): + """ + - Create key + Make chat completion call + - Create user + make chat completion call + """ + async with aiohttp.ClientSession() as session: + key_gen = await generate_key(session=session) + key = key_gen["key"] + await chat_completion(session=session, key=key, model="fake-openai-endpoint") + + otel_spans = await get_otel_spans(session=session, key=key) + print("otel_spans: ", otel_spans) + + all_otel_spans = otel_spans["otel_spans"] + + assert len(all_otel_spans) == 5 + + # 'postgres', 'redis', 'raw_gen_ai_request', 'litellm_request', 'Received Proxy Server Request' in the span + assert "postgres" in all_otel_spans + assert "redis" in all_otel_spans + assert "raw_gen_ai_request" in all_otel_spans + assert "litellm_request" in all_otel_spans + assert "Received Proxy Server Request" in all_otel_spans