(e2e testing) - add tests for using litellm /team/ updates in multi-instance deployments with Redis (#8440)

* add team block/unblock test

* test_team_blocking_behavior_multi_instance

* proxy_multi_instance_tests

* test - Run Docker container 2

.circleci/config.yml

@@ -1518,6 +1518,117 @@ jobs:
       - store_test_results:
           path: test-results
 
+  proxy_multi_instance_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
+            python -m pip install --upgrade pip
+            python -m pip install -r requirements.txt
+            pip install "pytest==7.3.1"
+            pip install "pytest-retry==1.6.3"
+            pip install "pytest-mock==3.12.0"
+            pip install "pytest-asyncio==0.21.1"
+      - run:
+          name: Build Docker image
+          command: docker build -t my-app:latest -f ./docker/Dockerfile.database .
+      - run:
+          name: Run Docker container 1
+          # intentionally give bad redis credentials here
+          # the OTEL test - should get this as a trace
+          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 LITELLM_MASTER_KEY="sk-1234" \
+              -e LITELLM_LICENSE=$LITELLM_LICENSE \
+              -e USE_DDTRACE=True \
+              -e DD_API_KEY=$DD_API_KEY \
+              -e DD_SITE=$DD_SITE \
+              --name my-app \
+              -v $(pwd)/litellm/proxy/example_config_yaml/multi_instance_simple_config.yaml:/app/config.yaml \
+              my-app:latest \
+              --config /app/config.yaml \
+              --port 4000 \
+              --detailed_debug \
+      - run:
+          name: Run Docker container 2
+          command: |
+            docker run -d \
+              -p 4001:4001 \
+              -e DATABASE_URL=$PROXY_DATABASE_URL \
+              -e REDIS_HOST=$REDIS_HOST \
+              -e REDIS_PASSWORD=$REDIS_PASSWORD \
+              -e REDIS_PORT=$REDIS_PORT \
+              -e LITELLM_MASTER_KEY="sk-1234" \
+              -e LITELLM_LICENSE=$LITELLM_LICENSE \
+              -e USE_DDTRACE=True \
+              -e DD_API_KEY=$DD_API_KEY \
+              -e DD_SITE=$DD_SITE \
+              --name my-app-2 \
+              -v $(pwd)/litellm/proxy/example_config_yaml/multi_instance_simple_config.yaml:/app/config.yaml \
+              my-app:latest \
+              --config /app/config.yaml \
+              --port 4001 \
+              --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 instance 1 to be ready
+          command: dockerize -wait http://localhost:4000 -timeout 5m
+      - run:
+          name: Wait for instance 2 to be ready
+          command: dockerize -wait http://localhost:4001 -timeout 5m
+      - run:
+          name: Run tests
+          command: |
+            pwd
+            ls
+            python -m pytest -vv tests/multi_instance_e2e_tests -x --junitxml=test-results/junit.xml --durations=5
+          no_output_timeout:
+            120m
+            # Clean up first container
+      # Store test results
+      - store_test_results:
+          path: test-results
+
   proxy_store_model_in_db_tests:
     machine:
       image: ubuntu-2204:2023.10.1
@@ -2173,6 +2284,12 @@ workflows:
               only:
                 - main
                 - /litellm_.*/
+      - proxy_multi_instance_tests:
+          filters:
+            branches:
+              only:
+                - main
+                - /litellm_.*/
       - proxy_store_model_in_db_tests:
           filters:
             branches:
@@ -2304,6 +2421,7 @@ workflows:
             - installing_litellm_on_python
             - installing_litellm_on_python_3_13
             - proxy_logging_guardrails_model_info_tests
+            - proxy_multi_instance_tests
             - proxy_store_model_in_db_tests
             - proxy_build_from_pip_tests
             - proxy_pass_through_endpoint_tests

litellm/proxy/example_config_yaml/multi_instance_simple_config.yaml

@@ -0,0 +1,12 @@
+model_list:
+  - model_name: fake-openai-endpoint
+    litellm_params:
+      model: openai/my-fake-model
+      api_key: my-fake-key
+      api_base: https://exampleopenaiendpoint-production.up.railway.app/
+
+litellm_settings:
+  cache: True
+  cache_params:
+    type: redis
+

litellm/proxy/proxy_config.yaml

@@ -4,6 +4,9 @@ model_list:
       model: openai/my-fake-model
       api_key: my-fake-key
       api_base: https://exampleopenaiendpoint-production.up.railway.app/
-general_settings:
+
-  store_model_in_db: true
+litellm_settings:
+  cache: True
+  cache_params:
+    type: redis
 

tests/multi_instance_e2e_tests/test_update_team_e2e.py

@@ -0,0 +1,195 @@
+import pytest
+import asyncio
+import aiohttp
+import json
+from httpx import AsyncClient
+from typing import Any, Optional
+
+
+# =====================================================================
+# NEW HELPER FUNCTIONS FOR TEAM BLOCKING TESTS
+# =====================================================================
+async def generate_team_key(
+    session,
+    team_id: str,
+    max_budget: Optional[float] = None,
+):
+    """Helper function to generate a key for a specific team"""
+    url = "http://0.0.0.0:4000/key/generate"
+    headers = {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"}
+    data: dict[str, Any] = {"team_id": team_id}
+    if max_budget is not None:
+        data["max_budget"] = max_budget
+    async with session.post(url, headers=headers, json=data) as response:
+        return await response.json()
+
+
+async def update_team_block_status(session, team_id: str, blocked: bool, port: int):
+    """Helper to update a team's 'blocked' status on a given instance port."""
+    url = f"http://0.0.0.0:{port}/team/update"
+    headers = {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"}
+    data = {"team_id": team_id, "blocked": blocked}
+    async with session.post(url, headers=headers, json=data) as response:
+        return await response.json()
+
+
+async def get_team_info(session, team_id: str, port: int):
+    """Helper to retrieve team info from a specific instance port."""
+    url = f"http://0.0.0.0:{port}/team/info"
+    headers = {"Authorization": "Bearer sk-1234"}
+    async with session.get(
+        url, headers=headers, params={"team_id": team_id}
+    ) as response:
+        data = await response.json()
+        return data["team_info"]
+
+
+async def chat_completion_on_port(
+    session, key: str, model: str, port: int, prompt: Optional[str] = None
+):
+    """
+    Helper function to make a chat completion request on a specified instance port.
+    Accepts an optional prompt string.
+    """
+    from openai import AsyncOpenAI
+    import uuid
+
+    if prompt is None:
+        prompt = f"Say hello! {uuid.uuid4()}" * 100
+    client = AsyncOpenAI(api_key=key, base_url=f"http://0.0.0.0:{port}/v1")
+    response = await client.chat.completions.create(
+        model=model,
+        messages=[{"role": "user", "content": prompt}],
+    )
+    return response
+
+
+# =====================================================================
+# NEW END‑TO‑END TEST FOR TEAM BLOCKING ACROSS MULTI‑INSTANCE SETUP
+# =====================================================================
+
+
+@pytest.mark.asyncio
+async def test_team_blocking_behavior_multi_instance():
+    """
+    Test team blocking scenario across multi-instance setup:
+
+    1. Create a new team on port 4000.
+    2. Verify (via team/info on port 4001) that the team is not blocked.
+    3. Create a key for that team.
+    4. Make a chat completion request (via instance on port 4000) and verify that it works.
+    6. Update the team to set 'blocked': True via the update endpoint on port 4001.
+    --- Sleep for 61 seconds --- the in-memory team obj ttl is 60 seconds
+    7. Verify (via team/info on port 4000) that the team is now blocked.
+    8. Make a chat completion request (using instance on port 4000) with a new prompt; expect it to be blocked.
+    9. Repeat the chat completion request with another new prompt; expect it to be blocked.
+    10. Confirm via team/info endpoints on both ports that the team remains blocked.
+    """
+    async with aiohttp.ClientSession() as session:
+        headers = {
+            "Authorization": "Bearer sk-1234",
+            "Content-Type": "application/json",
+        }
+
+        # 1. Create a new team on instance (port 4000)
+        url_new_team = "http://0.0.0.0:4000/team/new"
+        team_data = {}
+        async with session.post(
+            url_new_team, headers=headers, json=team_data
+        ) as response:
+            assert response.status == 200, "Failed to create team"
+            team_resp = await response.json()
+        team_id = team_resp["team_id"]
+
+        # 2. Verify via team/info on port 4001 that team is not blocked.
+        team_info_4001 = await get_team_info(session, team_id, port=4001)
+        assert "blocked" in team_info_4001, "Team info missing 'blocked' field"
+        assert (
+            team_info_4001["blocked"] is False
+        ), "Team should not be blocked initially"
+
+        # 3. Create a key for the team using the existing helper.
+        key_gen = await generate_team_key(session=session, team_id=team_id)
+        key = key_gen["key"]
+
+        # 4. Make a chat completion request on port 4000 and verify it works.
+        response = await chat_completion_on_port(
+            session,
+            key=key,
+            model="fake-openai-endpoint",
+            port=4000,
+            prompt="Non-cached prompt 1",
+        )
+        assert (
+            response is not None
+        ), "Chat completion should succeed when team is not blocked"
+
+        # 5. Update the team to set 'blocked': True on instance port 4001.
+        await update_team_block_status(session, team_id, blocked=True, port=4001)
+        print("sleeping for 61 seconds")
+        await asyncio.sleep(61)
+
+        # 6. Verify via team/info on port 4000 that the team is blocked.
+        team_info_4000 = await get_team_info(session, team_id, port=4000)
+        assert "blocked" in team_info_4000, "Team info missing 'blocked' field"
+        print(
+            "Team info on port 4000: ",
+            json.dumps(team_info_4000, indent=4, default=str),
+        )
+        assert team_info_4000["blocked"] is True, "Team should be blocked after update"
+        # 7.  Verify via team/info on port 4001 that the team is blocked.
+        team_info_4001 = await get_team_info(session, team_id, port=4001)
+        assert "blocked" in team_info_4001, "Team info missing 'blocked' field"
+        assert team_info_4001["blocked"] is True, "Team should be blocked after update"
+
+        # 8. Make a chat completion request on port 4000 with a new prompt; expect it to be blocked.
+        with pytest.raises(Exception) as excinfo:
+            await chat_completion_on_port(
+                session,
+                key=key,
+                model="fake-openai-endpoint",
+                port=4001,
+                prompt="Non-cached prompt 2",
+            )
+        error_msg = str(excinfo.value)
+        assert (
+            "blocked" in error_msg.lower()
+        ), f"Expected error indicating team blocked, got: {error_msg}"
+
+        # 9. Make a chat completion request on port 4000 with a new prompt; expect it to be blocked.
+        with pytest.raises(Exception) as excinfo:
+            await chat_completion_on_port(
+                session,
+                key=key,
+                model="fake-openai-endpoint",
+                port=4000,
+                prompt="Non-cached prompt 2",
+            )
+        error_msg = str(excinfo.value)
+        assert (
+            "blocked" in error_msg.lower()
+        ), f"Expected error indicating team blocked, got: {error_msg}"
+
+        # 9. Repeat the chat completion request with another new prompt; expect it to be blocked.
+        with pytest.raises(Exception) as excinfo_second:
+            await chat_completion_on_port(
+                session,
+                key=key,
+                model="fake-openai-endpoint",
+                port=4000,
+                prompt="Non-cached prompt 3",
+            )
+        error_msg_second = str(excinfo_second.value)
+        assert (
+            "blocked" in error_msg_second.lower()
+        ), f"Expected error indicating team blocked, got: {error_msg_second}"
+
+        # 10. Final verification: check team info on both ports indicates the team is blocked.
+        final_team_info_4000 = await get_team_info(session, team_id, port=4000)
+        final_team_info_4001 = await get_team_info(session, team_id, port=4001)
+        assert (
+            final_team_info_4000.get("blocked") is True
+        ), "Team on port 4000 should be blocked"
+        assert (
+            final_team_info_4001.get("blocked") is True
+        ), "Team on port 4001 should be blocked"