From 3cc19c56bab2fa42b93ff39da42624ba0585a127 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Tue, 20 Jan 2026 18:46:26 +0530 Subject: [PATCH] Revert "feat: Add Redis-based migration lock with bug fixes (#19261)" This reverts commit 98e87c3e67ab36948020f769da6eea87fd3ae2c3. --- .../litellm_proxy_extras/utils.py | 250 ++---------------- litellm/proxy/db/prisma_client.py | 13 +- .../test_litellm_proxy_extras_utils.py | 214 +-------------- 3 files changed, 25 insertions(+), 452 deletions(-) diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index 4e95267d8ab..7ffbe95be13 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/litellm-proxy-extras/litellm_proxy_extras/utils.py @@ -11,11 +11,6 @@ from typing import Optional from litellm_proxy_extras._logging import logger -try: - from litellm.caching.redis_cache import RedisCache -except ImportError: - RedisCache = None # type: ignore - def str_to_bool(value: Optional[str]) -> bool: if value is None: @@ -23,154 +18,6 @@ def str_to_bool(value: Optional[str]) -> bool: return value.lower() in ("true", "1", "t", "y", "yes") -class MigrationLockManager: - """Redis-based lock manager for database migrations. - - Prevents concurrent Prisma migrations in multi-pod deployments by using - a distributed lock. Only one pod can hold the lock and run migrations at a time. - """ - - MIGRATION_LOCK_KEY = "migration_lock" - LOCK_TTL_SECONDS = 300 # 5 minutes TTL - - def __init__(self, redis_cache: Optional["RedisCache"] = None): - """Initialize the migration lock manager. - - Args: - redis_cache: Optional RedisCache instance for distributed locking. - If None, migrations run without lock protection (single instance mode). - """ - self.redis_cache = redis_cache - self.lock_acquired = False - self.pod_id = f"pod_{os.getpid()}_{int(time.time())}" - - def _get_redis_lock_key(self) -> str: - """Get Redis lock key for migration.""" - return f"migration_lock:{self.MIGRATION_LOCK_KEY}" - - def acquire_lock(self) -> bool: - """Acquire migration lock using Redis SET NX. - - Returns: - bool: True if lock acquired, False otherwise. - """ - if self.redis_cache is None: - logger.warning( - "Redis cache is not available, running migration without lock protection" - ) - self.lock_acquired = True - return True - - try: - lock_key = self._get_redis_lock_key() - - # FIX: Use native Redis SET NX instead of RedisCache.set_cache() - # Original bug: set_cache() doesn't support nx parameter and returns None - # Fixed: Use redis_client.set() directly which returns True/False - acquired = self.redis_cache.redis_client.set( - name=lock_key, - value=self.pod_id, - nx=True, # Only set if key doesn't exist - ex=self.LOCK_TTL_SECONDS, # Set expiration time - ) - - if acquired: - self.lock_acquired = True - logger.info(f"Migration lock acquired by pod {self.pod_id}") - return True - else: - logger.info("Migration lock is already held by another pod") - return False - - except Exception as e: - logger.warning(f"Failed to acquire migration lock: {e}") - return False - - def wait_for_lock_release( - self, check_interval: int = 5, max_wait: int = 300 - ) -> bool: - """Wait for another process to release the lock. - - Args: - check_interval: Seconds to wait between lock acquisition attempts. - max_wait: Maximum seconds to wait for lock release. - - Returns: - bool: True if lock acquired after waiting, False if timeout. - """ - if self.redis_cache is None: - logger.warning("Redis cache is not available, cannot wait for lock") - return False - - logger.info(f"Waiting for migration lock to be released (max {max_wait}s)...") - start_time = time.time() - - while time.time() - start_time < max_wait: - # Try to acquire lock using the public acquire_lock method - if self.acquire_lock(): - logger.info( - f"Migration lock acquired after waiting by pod {self.pod_id}" - ) - return True - - time.sleep(check_interval) - - logger.warning(f"Failed to acquire migration lock within {max_wait} seconds") - return False - - def release_lock(self): - """Release migration lock atomically using Lua script. - - FIX: Use Lua script for atomic compare-and-delete to prevent race conditions. - Original bug: Non-atomic GET then DELETE allows another pod to acquire lock - between the GET and DELETE operations. - """ - if not self.lock_acquired or self.redis_cache is None: - return - - try: - lock_key = self._get_redis_lock_key() - - # FIX: Use Lua script for atomic compare-and-delete - # This prevents race condition where: - # 1. Pod A reads lock value (sees its own pod_id) - # 2. Lock TTL expires - # 3. Pod B acquires lock - # 4. Pod A deletes lock (deletes Pod B's lock!) - lua_script = """ - if redis.call("GET", KEYS[1]) == ARGV[1] then - return redis.call("DEL", KEYS[1]) - else - return 0 - end - """ - - result = self.redis_cache.redis_client.eval( - lua_script, - 1, # Number of keys - lock_key, # KEYS[1] - self.pod_id, # ARGV[1] - ) - - if result == 1: - logger.info(f"Migration lock released by pod {self.pod_id}") - else: - logger.warning(f"Pod {self.pod_id} cannot release lock (not owner)") - - except Exception as e: - logger.warning(f"Failed to release migration lock: {e}") - finally: - self.lock_acquired = False - - def __enter__(self): - """Context manager entry - acquire lock when entering with statement.""" - self.acquire_lock() - return self - - def __exit__(self, exc_type, exc_val, exc_tb): - """Context manager exit - release lock when exiting with statement.""" - self.release_lock() - def _get_prisma_env() -> dict: """Get environment variables for Prisma, handling offline mode if configured.""" @@ -178,9 +25,7 @@ def _get_prisma_env() -> dict: if str_to_bool(os.getenv("PRISMA_OFFLINE_MODE")): # These env vars prevent Prisma from attempting downloads prisma_env["NPM_CONFIG_PREFER_OFFLINE"] = "true" - prisma_env["NPM_CONFIG_CACHE"] = os.getenv( - "NPM_CONFIG_CACHE", "/app/.cache/npm" - ) + prisma_env["NPM_CONFIG_CACHE"] = os.getenv("NPM_CONFIG_CACHE", "/app/.cache/npm") return prisma_env @@ -189,28 +34,29 @@ def _get_prisma_command() -> str: if str_to_bool(os.getenv("PRISMA_OFFLINE_MODE")): # Primary location where Prisma Python package installs the CLI default_cli_path = "/app/.cache/prisma-python/binaries/node_modules/.bin/prisma" - + # Check if custom path is provided (for flexibility) custom_cli_path = os.getenv("PRISMA_CLI_PATH") if custom_cli_path and os.path.exists(custom_cli_path): logger.info(f"Using custom Prisma CLI at {custom_cli_path}") return custom_cli_path - + # Check the default location if os.path.exists(default_cli_path): logger.info(f"Using cached Prisma CLI at {default_cli_path}") return default_cli_path - + # If not found, log warning and fall back logger.warning( f"Prisma CLI not found at {default_cli_path}. " "Falling back to Python wrapper (may attempt downloads)" ) - + # Fall back to the Python wrapper (will work in online mode) return "prisma" + class ProxyExtrasDBManager: @staticmethod def _get_prisma_dir() -> str: @@ -273,7 +119,7 @@ class ProxyExtrasDBManager: stdout=open(migration_file, "w"), check=True, timeout=30, - env=prisma_env, + env=prisma_env ) # 3. Mark the migration as applied since it represents current state @@ -288,7 +134,7 @@ class ProxyExtrasDBManager: ], check=True, timeout=30, - env=prisma_env, + env=prisma_env ) return True @@ -313,20 +159,14 @@ class ProxyExtrasDBManager: @staticmethod def _roll_back_migration(migration_name: str): """Mark a specific migration as rolled back""" - # Set up environment for offline mode if configured + # Set up environment for offline mode if configured prisma_env = _get_prisma_env() subprocess.run( - [ - _get_prisma_command(), - "migrate", - "resolve", - "--rolled-back", - migration_name, - ], + [_get_prisma_command(), "migrate", "resolve", "--rolled-back", migration_name], timeout=60, check=True, capture_output=True, - env=prisma_env, + env=prisma_env ) @staticmethod @@ -338,7 +178,7 @@ class ProxyExtrasDBManager: timeout=60, check=True, capture_output=True, - env=prisma_env, + env=prisma_env ) @staticmethod @@ -408,7 +248,7 @@ class ProxyExtrasDBManager: if not database_url: logger.error("DATABASE_URL not set") return - + diff_dir = ( Path(migrations_dir) / "migrations" @@ -443,7 +283,7 @@ class ProxyExtrasDBManager: check=True, timeout=60, stdout=f, - env=_get_prisma_env(), + env=_get_prisma_env() ) except subprocess.CalledProcessError as e: logger.warning(f"Failed to generate migration diff: {e.stderr}") @@ -473,7 +313,7 @@ class ProxyExtrasDBManager: check=True, capture_output=True, text=True, - env=_get_prisma_env(), + env=_get_prisma_env() ) logger.info(f"prisma db execute stdout: {result.stdout}") logger.info("✅ Migration diff applied successfully") @@ -491,18 +331,12 @@ class ProxyExtrasDBManager: try: logger.info(f"Resolving migration: {migration_name}") subprocess.run( - [ - _get_prisma_command(), - "migrate", - "resolve", - "--applied", - migration_name, - ], + [_get_prisma_command(), "migrate", "resolve", "--applied", migration_name], timeout=60, check=True, capture_output=True, text=True, - env=_get_prisma_env(), + env=_get_prisma_env() ) logger.debug(f"Resolved migration: {migration_name}") except subprocess.CalledProcessError as e: @@ -512,57 +346,19 @@ class ProxyExtrasDBManager: ) @staticmethod - def setup_database( - use_migrate: bool = False, redis_cache: Optional["RedisCache"] = None - ) -> bool: + def setup_database(use_migrate: bool = False) -> bool: """ - Set up the database using either prisma migrate or prisma db push. - Uses migrations from litellm-proxy-extras package. - In multi-instance environment, use redis lock to prevent concurrent execution. + Set up the database using either prisma migrate or prisma db push + Uses migrations from litellm-proxy-extras package Args: + schema_path (str): Path to the Prisma schema file use_migrate (bool): Whether to use prisma migrate instead of db push - redis_cache: Redis cache instance for distributed locking Returns: bool: True if setup was successful, False otherwise """ schema_path = ProxyExtrasDBManager._get_prisma_dir() + "/schema.prisma" - - database_url = os.getenv("DATABASE_URL") - if not database_url: - logger.error("DATABASE_URL environment variable is not set") - return False - - # Use MigrationLockManager to prevent concurrent migration execution - with MigrationLockManager(redis_cache) as lock_manager: - # Lock is already acquired in __enter__, check if it was successful - if not lock_manager.lock_acquired: - # Cannot acquire lock, another process is running migration - logger.info( - "Another pod is running migration, waiting for completion..." - ) - - # Wait for other process to complete migration - if not lock_manager.wait_for_lock_release(): - logger.error("Failed to acquire migration lock after waiting") - return False - - # Successfully acquired lock, proceed with migration - logger.info("Acquired migration lock, proceeding with migration") - return ProxyExtrasDBManager._execute_migration(use_migrate, schema_path) - - @staticmethod - def _execute_migration(use_migrate: bool, schema_path: str) -> bool: - """Execute the actual migration. - - Args: - use_migrate: Whether to use prisma migrate instead of db push - schema_path: Path to the Prisma schema file - - Returns: - bool: True if migration was successful, False otherwise - """ for attempt in range(4): original_dir = os.getcwd() migrations_dir = ProxyExtrasDBManager._get_prisma_dir() @@ -579,7 +375,7 @@ class ProxyExtrasDBManager: check=True, capture_output=True, text=True, - env=_get_prisma_env(), + env=_get_prisma_env() ) logger.info(f"prisma migrate deploy stdout: {result.stdout}") @@ -617,7 +413,7 @@ class ProxyExtrasDBManager: check=True, capture_output=True, text=True, - env=_get_prisma_env(), + env=_get_prisma_env() ) logger.info( f"✅ Migration {failed_migration} marked as rolled back... retrying" diff --git a/litellm/proxy/db/prisma_client.py b/litellm/proxy/db/prisma_client.py index 65fed92340b..c9c0cfe8f68 100644 --- a/litellm/proxy/db/prisma_client.py +++ b/litellm/proxy/db/prisma_client.py @@ -383,18 +383,7 @@ class PrismaManager: prisma_dir = PrismaManager._get_prisma_dir() - # Import redis_usage_cache for distributed locking - try: - from litellm.proxy.proxy_server import redis_usage_cache - except ImportError: - verbose_proxy_logger.warning( - "Could not import redis_usage_cache, migrations will run without distributed locking" - ) - redis_usage_cache = None - - return ProxyExtrasDBManager.setup_database( - use_migrate=use_migrate, redis_cache=redis_usage_cache - ) + return ProxyExtrasDBManager.setup_database(use_migrate=use_migrate) else: # Use prisma db push with increased timeout subprocess.run( diff --git a/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py b/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py index cff3a88ab6c..5714cd5c487 100644 --- a/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py +++ b/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py @@ -1,13 +1,11 @@ import os import sys -import time -from unittest.mock import Mock, patch sys.path.insert( 0, os.path.abspath("../..") ) # Adds the parent directory to the system path -from litellm_proxy_extras.utils import ProxyExtrasDBManager, MigrationLockManager +from litellm_proxy_extras.utils import ProxyExtrasDBManager def test_custom_prisma_dir(monkeypatch): @@ -127,213 +125,3 @@ class TestErrorClassificationPriority: error_message = "connection timeout" assert ProxyExtrasDBManager._is_permission_error(error_message) is False assert ProxyExtrasDBManager._is_idempotent_error(error_message) is False - - -class TestMigrationLockManager: - """Test cases for MigrationLockManager""" - - def test_acquire_lock_success(self): - """Test successful lock acquisition using Redis SET NX""" - mock_redis = Mock() - # FIX VERIFICATION: redis_client.set() returns True for successful SET NX - mock_redis.redis_client.set.return_value = True - - lock_manager = MigrationLockManager(mock_redis) - result = lock_manager.acquire_lock() - - assert result is True - assert lock_manager.lock_acquired is True - # Verify SET NX was called with correct parameters - mock_redis.redis_client.set.assert_called_once() - call_args = mock_redis.redis_client.set.call_args - assert call_args[1]["nx"] is True # SET NX parameter - assert call_args[1]["ex"] == 300 # TTL - - def test_acquire_lock_failure(self): - """Test lock acquisition failure when lock is already held""" - mock_redis = Mock() - # FIX VERIFICATION: redis_client.set() returns False when key exists - mock_redis.redis_client.set.return_value = False - - lock_manager = MigrationLockManager(mock_redis) - result = lock_manager.acquire_lock() - - assert result is False - assert lock_manager.lock_acquired is False - - def test_acquire_lock_without_redis(self): - """Test lock acquisition without Redis (graceful fallback)""" - lock_manager = MigrationLockManager(redis_cache=None) - result = lock_manager.acquire_lock() - - assert result is True - assert lock_manager.lock_acquired is True - - def test_release_lock_success(self): - """Test successful lock release using Lua script""" - mock_redis = Mock() - # FIX VERIFICATION: Lua script returns 1 for successful delete - mock_redis.redis_client.eval.return_value = 1 - - lock_manager = MigrationLockManager(mock_redis) - lock_manager.pod_id = "pod_123_456" - lock_manager.lock_acquired = True - - lock_manager.release_lock() - - # Verify Lua script was called for atomic compare-and-delete - mock_redis.redis_client.eval.assert_called_once() - call_args = mock_redis.redis_client.eval.call_args - # Verify Lua script contains atomic compare-and-delete logic - lua_script = call_args[0][0] - assert "GET" in lua_script - assert "DEL" in lua_script - assert lock_manager.lock_acquired is False - - def test_release_lock_wrong_owner(self): - """Test releasing lock when not the owner""" - mock_redis = Mock() - # FIX VERIFICATION: Lua script returns 0 when ownership check fails - mock_redis.redis_client.eval.return_value = 0 - - lock_manager = MigrationLockManager(mock_redis) - lock_manager.pod_id = "pod_123_456" - lock_manager.lock_acquired = True - - lock_manager.release_lock() - - mock_redis.redis_client.eval.assert_called_once() - assert lock_manager.lock_acquired is False - - def test_context_manager(self): - """Test MigrationLockManager as context manager""" - mock_redis = Mock() - mock_redis.redis_client.set.return_value = True - mock_redis.redis_client.eval.return_value = 1 - - lock_manager = MigrationLockManager(mock_redis) - lock_manager.pod_id = "pod_123_456" - - with lock_manager: - assert lock_manager.lock_acquired is True - - # Should call release_lock when exiting context - mock_redis.redis_client.eval.assert_called_once() - - def test_wait_for_lock_release_success(self): - """Test waiting for lock release and acquiring it""" - mock_redis = Mock() - # First call fails, second call succeeds - mock_redis.redis_client.set.side_effect = [False, True] - - lock_manager = MigrationLockManager(mock_redis) - result = lock_manager.wait_for_lock_release(check_interval=0.1, max_wait=1) - - assert result is True - assert lock_manager.lock_acquired is True - assert mock_redis.redis_client.set.call_count == 2 - - def test_wait_for_lock_release_timeout(self): - """Test timeout when waiting for lock release""" - mock_redis = Mock() - # Always fails to acquire lock - mock_redis.redis_client.set.return_value = False - - lock_manager = MigrationLockManager(mock_redis) - start_time = time.time() - result = lock_manager.wait_for_lock_release(check_interval=0.1, max_wait=0.5) - end_time = time.time() - - assert result is False - assert lock_manager.lock_acquired is False - assert end_time - start_time >= 0.5 # Should wait at least max_wait time - assert mock_redis.redis_client.set.call_count > 1 # Multiple attempts - - -class TestProxyExtrasDBManagerMigrationLock: - """Test cases for ProxyExtrasDBManager with migration locking""" - - @patch("litellm_proxy_extras.utils.ProxyExtrasDBManager._execute_migration") - def test_setup_database_with_redis_lock_success( - self, mock_execute_migration, monkeypatch - ): - """Test successful database setup with Redis lock""" - monkeypatch.setenv("DATABASE_URL", "postgresql://test:test@localhost/test") - mock_execute_migration.return_value = True - - # Mock Redis cache - mock_redis = Mock() - mock_redis.redis_client.set.return_value = True # Lock acquired - mock_redis.redis_client.eval.return_value = 1 # Lock released - - result = ProxyExtrasDBManager.setup_database( - use_migrate=True, redis_cache=mock_redis - ) - - assert result is True - mock_execute_migration.assert_called_once() - # Verify lock was acquired - mock_redis.redis_client.set.assert_called_once() - - @patch("litellm_proxy_extras.utils.ProxyExtrasDBManager._execute_migration") - def test_setup_database_with_redis_lock_wait_and_acquire( - self, mock_execute_migration, monkeypatch - ): - """Test database setup when lock is held, then acquired after waiting""" - monkeypatch.setenv("DATABASE_URL", "postgresql://test:test@localhost/test") - mock_execute_migration.return_value = True - - # Mock Redis cache - first call fails, second call succeeds - mock_redis = Mock() - mock_redis.redis_client.set.side_effect = [False, True] - mock_redis.redis_client.eval.return_value = 1 - - result = ProxyExtrasDBManager.setup_database( - use_migrate=True, redis_cache=mock_redis - ) - - assert result is True - mock_execute_migration.assert_called_once() - # Should have tried to acquire lock twice - assert mock_redis.redis_client.set.call_count == 2 - - @patch("litellm_proxy_extras.utils.ProxyExtrasDBManager._execute_migration") - def test_setup_database_without_redis(self, mock_execute_migration, monkeypatch): - """Test database setup without Redis cache (single instance mode)""" - monkeypatch.setenv("DATABASE_URL", "postgresql://test:test@localhost/test") - mock_execute_migration.return_value = True - - result = ProxyExtrasDBManager.setup_database(use_migrate=True, redis_cache=None) - - assert result is True - mock_execute_migration.assert_called_once() - - def test_setup_database_no_database_url(self): - """Test database setup without DATABASE_URL""" - with patch.dict(os.environ, {}, clear=True): - result = ProxyExtrasDBManager.setup_database( - use_migrate=True, redis_cache=None - ) - - assert result is False - - @patch("litellm_proxy_extras.utils.ProxyExtrasDBManager._execute_migration") - def test_setup_database_lock_timeout(self, mock_execute_migration, monkeypatch): - """Test database setup when lock acquisition times out""" - monkeypatch.setenv("DATABASE_URL", "postgresql://test:test@localhost/test") - - # Mock Redis cache - always fails to acquire lock - mock_redis = Mock() - mock_redis.redis_client.set.return_value = False - - # Patch wait_for_lock_release to simulate timeout - with patch.object(MigrationLockManager, "wait_for_lock_release") as mock_wait: - mock_wait.return_value = False # Simulate timeout - - result = ProxyExtrasDBManager.setup_database( - use_migrate=True, redis_cache=mock_redis - ) - - assert result is False - mock_execute_migration.assert_not_called() # Should not run migration - mock_wait.assert_called_once()