feature: readiness probe checks cache state

Signed-off-by: Haoyu Sun <[email protected]>

ols/app/endpoints/health.py

@@ -66,6 +66,13 @@ def index_is_ready() -> bool:
     return True
 
 
+def cache_is_ready() -> bool:
+    """Check if the cache is ready."""
+    if config.conversation_cache is None:
+        return False
+    return config.conversation_cache.ready()
+
+
 get_readiness_responses: dict[int | str, dict[str, Any]] = {
     200: {
         "description": "Service is ready",
@@ -97,6 +104,15 @@ def readiness_probe_get_method() -> ReadinessResponse:
                 "cause": "LLM is not ready",
             },
         )
+    if not cache_is_ready():
+        raise HTTPException(
+            status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
+            detail={
+                "response": "Service is not ready",
+                "cause": "Cache is not ready",
+            },
+        )
+
     return ReadinessResponse(ready=True, reason="service is ready")
 
 

ols/src/cache/cache.py

@@ -64,3 +64,11 @@ def insert_or_append(
             conversation_id: Conversation ID unique for given user.
             cache_entry: The value to store.
         """
+
+    @abstractmethod
+    def ready(self) -> bool:
+        """Check if the cache is ready.
+
+        Returns:
+            True if the cache is ready, False otherwise.
+        """

ols/src/cache/in_memory_cache.py

@@ -83,3 +83,13 @@ def insert_or_append(
                 old_value.append(value)
                 self.cache[key] = old_value
             self.deque.appendleft(key)
+
+    def ready(self) -> bool:
+        """Check if the cache is ready.
+
+           In memory cache is always ready.
+
+        Returns:
+            True if the cache is ready, False otherwise.
+        """
+        return True

ols/src/cache/postgres_cache.py

@@ -164,6 +164,16 @@ def insert_or_append(
                 logger.error("PostgresCache.insert_or_append: %s", e)
                 raise CacheError("PostgresCache.insert_or_append", e) from e
 
+    def ready(self) -> bool:
+        """Check if the cache is ready.
+
+        Postgres cache checks if the connection is alive.
+
+        Returns:
+            True if the cache is ready, False otherwise.
+        """
+        return self.conn is not None and self.conn.closed == 0
+
     @staticmethod
     def _select(
         cursor: psycopg2.extensions.cursor, user_id: str, conversation_id: str

ols/src/cache/redis_cache.py

@@ -115,3 +115,19 @@ def insert_or_append(
                 )
             else:
                 self.redis_client.set(key, json.dumps([cache_entry.to_dict()]))
+
+    def ready(self):
+        """Check if the cache is ready.
+
+           Redis cache checks if the connection is alive.
+
+        Returns:
+            True if the cache is ready, False otherwise.
+        """
+        if self.redis_client is None:
+            return False
+        try:
+            self.redis_client.ping()
+        except ConnectionError:
+            return False
+        return True

tests/mock_classes/mock_redis_client.py

@@ -41,3 +41,7 @@ def set(self, key, value, *args, **kwargs):
         assert isinstance(value, (str, bytes, int, float))
 
         self.cache[key] = value
+
+    def ping(self):
+        """Return PONG."""
+        return "PONG"

tests/unit/app/endpoints/test_health.py

@@ -1,7 +1,7 @@
 """Unit tests for health endpoints handlers."""
 
 import time
-from unittest.mock import patch
+from unittest.mock import Mock, patch
 
 import pytest
 from fastapi import HTTPException
@@ -14,7 +14,18 @@
     llm_is_ready,
     readiness_probe_get_method,
 )
+from ols.app.models.config import InMemoryCacheConfig
 from ols.app.models.models import LivenessResponse, ReadinessResponse
+from ols.src.cache.in_memory_cache import InMemoryCache
+
+
+def mock_cache():
+    """Fixture with constructed and initialized in memory cache object."""
+    mc = InMemoryCacheConfig({"max_entries": "10"})
+    mc_kls = Mock(InMemoryCache)
+    c = mc_kls(mc)
+    c.initialize_cache(mc)
+    return c
 
 
 class MockedLLM:
@@ -53,6 +64,7 @@ def test_readiness_probe_llm_check__llm_unexpected_response(mocked_load_llm):
     assert not llm_is_ready()
 
 
+@patch("ols.config._conversation_cache", create=True, new=mock_cache())
 @patch("ols.app.endpoints.health.llm_is_ready_persistent_state", new=False)
 @patch("ols.app.endpoints.health.load_llm")
 def test_readiness_probe_llm_check__state_cache(mocked_load_llm):
@@ -69,6 +81,7 @@ def test_readiness_probe_llm_check__state_cache(mocked_load_llm):
     assert mocked_load_llm.call_count == 1
 
 
+@patch("ols.config._conversation_cache", create=True, new=mock_cache())
 @patch("ols.app.endpoints.health.llm_is_ready_persistent_state", new=False)
 @patch("ols.app.endpoints.health.load_llm")
 def test_readiness_probe_llm_check__state_cache_not_expired(mocked_load_llm):
@@ -91,6 +104,7 @@ def test_readiness_probe_llm_check__state_cache_not_expired(mocked_load_llm):
         config.ols_config.expire_llm_is_ready_persistent_state = -1
 
 
+@patch("ols.config._conversation_cache", create=True, new=mock_cache())
 @patch("ols.app.endpoints.health.llm_is_ready_persistent_state", new=False)
 @patch("ols.app.endpoints.health.load_llm")
 def test_readiness_probe_llm_check__state_cache_expired(mocked_load_llm):
@@ -123,6 +137,7 @@ def test_readiness_probe_llm_check__llm_raise(mocked_load_llm):
     assert not llm_is_ready()
 
 
+@patch("ols.config._conversation_cache", create=True, new=mock_cache())
 @patch("ols.app.endpoints.health.llm_is_ready_persistent_state", new=False)
 @patch("ols.app.endpoints.health.load_llm")
 def test_readiness_probe_get_method_service_is_ready(mocked_load_llm):
@@ -133,6 +148,7 @@ def test_readiness_probe_get_method_service_is_ready(mocked_load_llm):
     assert response == ReadinessResponse(ready=True, reason="service is ready")
 
 
+@patch("ols.config._conversation_cache", create=True, new=mock_cache())
 @patch("ols.app.endpoints.health.llm_is_ready_persistent_state", new=True)
 def test_readiness_probe_get_method_index_is_ready():
     """Test the readiness_probe function when index is loaded."""
@@ -163,6 +179,21 @@ def test_readiness_probe_get_method_index_not_ready():
         readiness_probe_get_method()
 
 
+def test_readiness_probe_get_method_cache_not_ready():
+    """Test the readiness_probe function when cache is not ready."""
+    # simulate that the cache is not ready
+    with patch("ols.config._conversation_cache", new=None) as mocked_cache:
+        with pytest.raises(HTTPException, match="Service is not ready"):
+            readiness_probe_get_method()
+
+    with patch(
+        "ols.config._conversation_cache", create=True, new=mock_cache()
+    ) as mocked_cache:
+        mocked_cache.ready.return_value = False
+        with pytest.raises(HTTPException, match="Service is not ready"):
+            readiness_probe_get_method()
+
+
 def test_liveness_probe_get_method():
     """Test the liveness_probe function."""
     # the tested function returns constant right now

tests/unit/cache/test_in_memory_cache.py

@@ -84,6 +84,11 @@ def test_get_nonexistent_user(cache):
     assert cache.get("ffffffff-ffff-ffff-ffff-ffffffffffff", conversation_id) is None
 
 
+def test_ready(cache):
+    """Test if in memory cache always report ready."""
+    assert cache.ready()
+
+
 improper_user_uuids = [
     None,
     "",

tests/unit/cache/test_postgres_cache.py

@@ -217,3 +217,21 @@ def test_insert_or_append_operation_on_exception(mock_connect):
     # error must be raised during cache operation
     with pytest.raises(CacheError, match="PLSQL error"):
         cache.insert_or_append(user_id, conversation_id, history)
+
+
+@patch("psycopg2.connect")
+def test_ready(mock_connect):
+    """Test the Cache.ready operation."""
+    # initialize Postgres cache
+    config = PostgresConfig()
+    cache = PostgresCache(config)
+
+    # mock the connection state 0 - open
+    cache.conn.closed = 0
+    # cache is ready
+    assert cache.ready()
+
+    # mock the connection state 1 - closed
+    cache.conn.closed = 1
+    # cache is not ready
+    assert not cache.ready()

tests/unit/cache/test_redis_cache.py

@@ -1,8 +1,9 @@
 """Unit tests for RedisCache class."""
 
-from unittest.mock import patch
+from unittest.mock import MagicMock, patch
 
 import pytest
+from redis.exceptions import ConnectionError  # noqa: A004
 
 from ols import constants
 from ols.app.models.config import RedisConfig
@@ -208,3 +209,16 @@ def test_initialize_redis_retry_settings():
         assert cache.redis_client.kwargs["retry_on_timeout"] is True
         assert cache.redis_client.kwargs["retry_on_error"] is not None
         assert cache.redis_client.kwargs["retry"]._retries == 10
+
+
+def test_ready():
+    """Test the ready method."""
+    with patch("redis.StrictRedis", new=MockRedisClient):
+        config = RedisConfig({})
+        cache = RedisCache(config)
+        assert cache.ready()
+
+        cache.redis_client.ping = MagicMock(
+            side_effect=ConnectionError("Redis connection error")
+        )
+        assert not cache.ready()