mirror of
https://github.com/BerriAI/litellm.git
synced 2025-04-26 03:04:13 +00:00
(Infra/DB) - Allow running older litellm version when out of sync with current state of DB (#8695)
* fix check migration * clean up should_update_prisma_schema * update test * db_migration_disable_update_check * Check container logs for expected message * db_migration_disable_update_check * test_check_migration_out_of_sync * test_should_update_prisma_schema * db_migration_disable_update_check * pip install aiohttp
This commit is contained in:
parent
300d7825f5
commit
55b938dd6e
7 changed files with 168 additions and 22 deletions
|
@ -1112,6 +1112,23 @@ jobs:
|
||||||
working_directory: ~/project
|
working_directory: ~/project
|
||||||
steps:
|
steps:
|
||||||
- checkout
|
- checkout
|
||||||
|
- 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
|
||||||
- run:
|
- run:
|
||||||
name: Build Docker image
|
name: Build Docker image
|
||||||
command: |
|
command: |
|
||||||
|
@ -1119,29 +1136,48 @@ jobs:
|
||||||
- run:
|
- run:
|
||||||
name: Run Docker container
|
name: Run Docker container
|
||||||
command: |
|
command: |
|
||||||
docker run --name my-app \
|
docker run -d \
|
||||||
-p 4000:4000 \
|
-p 4000:4000 \
|
||||||
-e DATABASE_URL=$PROXY_DATABASE_URL \
|
-e DATABASE_URL=$PROXY_DATABASE_URL \
|
||||||
-e DISABLE_SCHEMA_UPDATE="True" \
|
-e DISABLE_SCHEMA_UPDATE="True" \
|
||||||
-v $(pwd)/litellm/proxy/example_config_yaml/bad_schema.prisma:/app/schema.prisma \
|
-v $(pwd)/litellm/proxy/example_config_yaml/bad_schema.prisma:/app/schema.prisma \
|
||||||
-v $(pwd)/litellm/proxy/example_config_yaml/bad_schema.prisma:/app/litellm/proxy/schema.prisma \
|
-v $(pwd)/litellm/proxy/example_config_yaml/bad_schema.prisma:/app/litellm/proxy/schema.prisma \
|
||||||
-v $(pwd)/litellm/proxy/example_config_yaml/disable_schema_update.yaml:/app/config.yaml \
|
-v $(pwd)/litellm/proxy/example_config_yaml/disable_schema_update.yaml:/app/config.yaml \
|
||||||
|
--name my-app \
|
||||||
myapp:latest \
|
myapp:latest \
|
||||||
--config /app/config.yaml \
|
--config /app/config.yaml \
|
||||||
--port 4000 > docker_output.log 2>&1 || true
|
--port 4000
|
||||||
- run:
|
- run:
|
||||||
name: Display Docker logs
|
name: Install curl and dockerize
|
||||||
command: cat docker_output.log
|
|
||||||
- run:
|
|
||||||
name: Check for expected error
|
|
||||||
command: |
|
command: |
|
||||||
if grep -q "prisma schema out of sync with db. Consider running these sql_commands to sync the two" docker_output.log; then
|
sudo apt-get update
|
||||||
echo "Expected error found. Test passed."
|
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: Wait for container to be ready
|
||||||
|
command: dockerize -wait http://localhost:4000 -timeout 1m
|
||||||
|
- run:
|
||||||
|
name: Check container logs for expected message
|
||||||
|
command: |
|
||||||
|
echo "=== Printing Full Container Startup Logs ==="
|
||||||
|
docker logs my-app
|
||||||
|
echo "=== End of Full Container Startup Logs ==="
|
||||||
|
|
||||||
|
if docker logs my-app 2>&1 | grep -q "prisma schema out of sync with db. Consider running these sql_commands to sync the two"; then
|
||||||
|
echo "Expected message found in logs. Test passed."
|
||||||
else
|
else
|
||||||
echo "Expected error not found. Test failed."
|
echo "Expected message not found in logs. Test failed."
|
||||||
cat docker_output.log
|
|
||||||
exit 1
|
exit 1
|
||||||
fi
|
fi
|
||||||
|
- run:
|
||||||
|
name: Run Basic Proxy Startup Tests (Health Readiness and Chat Completion)
|
||||||
|
command: |
|
||||||
|
python -m pytest -vv tests/basic_proxy_startup_tests -x --junitxml=test-results/junit-2.xml --durations=5
|
||||||
|
no_output_timeout: 120m
|
||||||
|
|
||||||
|
|
||||||
build_and_test:
|
build_and_test:
|
||||||
machine:
|
machine:
|
||||||
|
|
|
@ -4,6 +4,8 @@ import os
|
||||||
import subprocess
|
import subprocess
|
||||||
from typing import List, Optional, Tuple
|
from typing import List, Optional, Tuple
|
||||||
|
|
||||||
|
from litellm._logging import verbose_logger
|
||||||
|
|
||||||
|
|
||||||
def extract_sql_commands(diff_output: str) -> List[str]:
|
def extract_sql_commands(diff_output: str) -> List[str]:
|
||||||
"""
|
"""
|
||||||
|
@ -52,6 +54,7 @@ def check_prisma_schema_diff_helper(db_url: str) -> Tuple[bool, List[str]]:
|
||||||
subprocess.CalledProcessError: If the Prisma command fails.
|
subprocess.CalledProcessError: If the Prisma command fails.
|
||||||
Exception: For any other errors during execution.
|
Exception: For any other errors during execution.
|
||||||
"""
|
"""
|
||||||
|
verbose_logger.debug("Checking for Prisma schema diff...") # noqa: T201
|
||||||
try:
|
try:
|
||||||
result = subprocess.run(
|
result = subprocess.run(
|
||||||
[
|
[
|
||||||
|
@ -94,8 +97,8 @@ def check_prisma_schema_diff(db_url: Optional[str] = None) -> None:
|
||||||
raise Exception("DATABASE_URL not set")
|
raise Exception("DATABASE_URL not set")
|
||||||
has_diff, message = check_prisma_schema_diff_helper(db_url)
|
has_diff, message = check_prisma_schema_diff_helper(db_url)
|
||||||
if has_diff:
|
if has_diff:
|
||||||
raise Exception(
|
verbose_logger.exception(
|
||||||
"prisma schema out of sync with db. Consider running these sql_commands to sync the two - {}".format(
|
"🚨🚨🚨 prisma schema out of sync with db. Consider running these sql_commands to sync the two - {}".format(
|
||||||
message
|
message
|
||||||
)
|
)
|
||||||
)
|
)
|
||||||
|
|
|
@ -7,7 +7,7 @@ import os
|
||||||
import urllib
|
import urllib
|
||||||
import urllib.parse
|
import urllib.parse
|
||||||
from datetime import datetime, timedelta
|
from datetime import datetime, timedelta
|
||||||
from typing import Any, Optional
|
from typing import Any, Optional, Union
|
||||||
|
|
||||||
from litellm.secret_managers.main import str_to_bool
|
from litellm.secret_managers.main import str_to_bool
|
||||||
|
|
||||||
|
@ -112,12 +112,28 @@ class PrismaWrapper:
|
||||||
return original_attr
|
return original_attr
|
||||||
|
|
||||||
|
|
||||||
def should_update_schema(disable_prisma_schema_update: Optional[bool]):
|
def should_update_prisma_schema(
|
||||||
|
disable_updates: Optional[Union[bool, str]] = None
|
||||||
|
) -> bool:
|
||||||
"""
|
"""
|
||||||
This function is used to determine if the Prisma schema should be updated.
|
Determines if Prisma Schema updates should be applied during startup.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
disable_updates: Controls whether schema updates are disabled.
|
||||||
|
Accepts boolean or string ('true'/'false'). Defaults to checking DISABLE_SCHEMA_UPDATE env var.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
bool: True if schema updates should be applied, False if updates are disabled.
|
||||||
|
|
||||||
|
Examples:
|
||||||
|
>>> should_update_prisma_schema() # Checks DISABLE_SCHEMA_UPDATE env var
|
||||||
|
>>> should_update_prisma_schema(True) # Explicitly disable updates
|
||||||
|
>>> should_update_prisma_schema("false") # Enable updates using string
|
||||||
"""
|
"""
|
||||||
if disable_prisma_schema_update is None:
|
if disable_updates is None:
|
||||||
disable_prisma_schema_update = str_to_bool(os.getenv("DISABLE_SCHEMA_UPDATE"))
|
disable_updates = os.getenv("DISABLE_SCHEMA_UPDATE", "false")
|
||||||
if disable_prisma_schema_update is True:
|
|
||||||
return False
|
if isinstance(disable_updates, str):
|
||||||
return True
|
disable_updates = str_to_bool(disable_updates)
|
||||||
|
|
||||||
|
return not bool(disable_updates)
|
||||||
|
|
|
@ -4,6 +4,11 @@ model_list:
|
||||||
model: openai/fake
|
model: openai/fake
|
||||||
api_key: fake-key
|
api_key: fake-key
|
||||||
api_base: https://exampleopenaiendpoint-production.up.railway.app/
|
api_base: https://exampleopenaiendpoint-production.up.railway.app/
|
||||||
|
- model_name: gpt-4
|
||||||
|
litellm_params:
|
||||||
|
model: openai/gpt-4
|
||||||
|
api_key: fake-key
|
||||||
|
api_base: https://exampleopenaiendpoint-production.up.railway.app/
|
||||||
|
|
||||||
litellm_settings:
|
litellm_settings:
|
||||||
callbacks: ["gcs_bucket"]
|
callbacks: ["gcs_bucket"]
|
||||||
|
|
|
@ -607,10 +607,10 @@ def run_server( # noqa: PLR0915
|
||||||
|
|
||||||
if is_prisma_runnable:
|
if is_prisma_runnable:
|
||||||
from litellm.proxy.db.check_migration import check_prisma_schema_diff
|
from litellm.proxy.db.check_migration import check_prisma_schema_diff
|
||||||
from litellm.proxy.db.prisma_client import should_update_schema
|
from litellm.proxy.db.prisma_client import should_update_prisma_schema
|
||||||
|
|
||||||
if (
|
if (
|
||||||
should_update_schema(
|
should_update_prisma_schema(
|
||||||
general_settings.get("disable_prisma_schema_update")
|
general_settings.get("disable_prisma_schema_update")
|
||||||
)
|
)
|
||||||
is False
|
is False
|
||||||
|
|
51
tests/litellm/proxy/db/test_check_migration.py
Normal file
51
tests/litellm/proxy/db/test_check_migration.py
Normal file
|
@ -0,0 +1,51 @@
|
||||||
|
import json
|
||||||
|
import os
|
||||||
|
import sys
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
from fastapi.testclient import TestClient
|
||||||
|
|
||||||
|
sys.path.insert(
|
||||||
|
0, os.path.abspath("../../../..")
|
||||||
|
) # Adds the parent directory to the system path
|
||||||
|
|
||||||
|
|
||||||
|
import json
|
||||||
|
import os
|
||||||
|
import sys
|
||||||
|
import time
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
from fastapi.testclient import TestClient
|
||||||
|
|
||||||
|
import litellm
|
||||||
|
|
||||||
|
|
||||||
|
def test_check_migration_out_of_sync(mocker):
|
||||||
|
"""
|
||||||
|
Test that the check_prisma_schema_diff function
|
||||||
|
- 🚨 [IMPORTANT] Does NOT Raise an Exception when the Prisma schema is out of sync with the database.
|
||||||
|
- logs an error when the Prisma schema is out of sync with the database.
|
||||||
|
"""
|
||||||
|
# Mock the logger BEFORE importing the function
|
||||||
|
mock_logger = mocker.patch("litellm._logging.verbose_logger")
|
||||||
|
|
||||||
|
# Import the function after mocking the logger
|
||||||
|
from litellm.proxy.db.check_migration import check_prisma_schema_diff
|
||||||
|
|
||||||
|
# Mock the helper function to simulate out-of-sync state
|
||||||
|
mock_diff_helper = mocker.patch(
|
||||||
|
"litellm.proxy.db.check_migration.check_prisma_schema_diff_helper",
|
||||||
|
return_value=(True, ["ALTER TABLE users ADD COLUMN new_field TEXT;"]),
|
||||||
|
)
|
||||||
|
|
||||||
|
# Run the function - it should not raise an error
|
||||||
|
try:
|
||||||
|
check_prisma_schema_diff(db_url="mock_url")
|
||||||
|
except Exception as e:
|
||||||
|
pytest.fail(f"check_prisma_schema_diff raised an unexpected exception: {e}")
|
||||||
|
|
||||||
|
# Verify the logger was called with the expected message
|
||||||
|
mock_logger.exception.assert_called_once()
|
||||||
|
actual_message = mock_logger.exception.call_args[0][0]
|
||||||
|
assert "prisma schema out of sync with db" in actual_message
|
35
tests/litellm/proxy/db/test_prisma_client.py
Normal file
35
tests/litellm/proxy/db/test_prisma_client.py
Normal file
|
@ -0,0 +1,35 @@
|
||||||
|
import json
|
||||||
|
import os
|
||||||
|
import sys
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
from fastapi.testclient import TestClient
|
||||||
|
|
||||||
|
sys.path.insert(
|
||||||
|
0, os.path.abspath("../../../..")
|
||||||
|
) # Adds the parent directory to the system path
|
||||||
|
|
||||||
|
|
||||||
|
from litellm.proxy.db.prisma_client import should_update_prisma_schema
|
||||||
|
|
||||||
|
|
||||||
|
def test_should_update_prisma_schema(monkeypatch):
|
||||||
|
# CASE 1: Environment variable behavior
|
||||||
|
# When DISABLE_SCHEMA_UPDATE is not set -> should update
|
||||||
|
monkeypatch.setenv("DISABLE_SCHEMA_UPDATE", None)
|
||||||
|
assert should_update_prisma_schema() == True
|
||||||
|
|
||||||
|
# When DISABLE_SCHEMA_UPDATE="true" -> should not update
|
||||||
|
monkeypatch.setenv("DISABLE_SCHEMA_UPDATE", "true")
|
||||||
|
assert should_update_prisma_schema() == False
|
||||||
|
|
||||||
|
# When DISABLE_SCHEMA_UPDATE="false" -> should update
|
||||||
|
monkeypatch.setenv("DISABLE_SCHEMA_UPDATE", "false")
|
||||||
|
assert should_update_prisma_schema() == True
|
||||||
|
|
||||||
|
# CASE 2: Explicit parameter behavior (overrides env var)
|
||||||
|
monkeypatch.setenv("DISABLE_SCHEMA_UPDATE", None)
|
||||||
|
assert should_update_prisma_schema(True) == False # Param True -> should not update
|
||||||
|
|
||||||
|
monkeypatch.setenv("DISABLE_SCHEMA_UPDATE", None) # Set env var opposite to param
|
||||||
|
assert should_update_prisma_schema(False) == True # Param False -> should update
|
Loading…
Add table
Add a link
Reference in a new issue