feat(BA-3002): Replace time() with Redis TIME command in ValkeyScheduleClient (#6695)

Co-authored-by: Claude <[email protected]>

Author: HyeockJinKim

changes/6695.feature.md

@@ -0,0 +1 @@
+Replace time() with Redis TIME command in ValkeyScheduleClient

src/ai/backend/common/clients/valkey_client/valkey_schedule/client.py

@@ -1,6 +1,5 @@
 from collections.abc import Mapping, Sequence
 from dataclasses import dataclass
-from time import time
 from typing import Optional, Self, cast
 from uuid import UUID
 
@@ -130,9 +129,20 @@ def _get_route_health_key(self, route_id: str) -> str:
         """
         return f"route:health:{route_id}"
 
-    def _is_health_status_valid(self, status: str, timestamp_str: str) -> bool:
+    async def _get_redis_time(self) -> int:
         """
-        Check if health status is healthy and timestamp is not stale.
+        Get current Unix timestamp from Redis server using TIME command.
+        This ensures consistent timestamps across distributed systems.
+
+        :return: Current Unix timestamp in seconds
+        """
+        result = await self._client.client.time()
+        seconds_bytes, _ = result
+        return int(seconds_bytes)
+
+    async def _validate_health_status(self, status: str, timestamp_str: str) -> bool:
+        """
+        Validate health status by checking if it's healthy and timestamp is not stale.
 
         :param status: The status string ("1" for healthy, "0" for unhealthy)
         :param timestamp_str: The timestamp string value from Redis
@@ -142,7 +152,7 @@ def _is_health_status_valid(self, status: str, timestamp_str: str) -> bool:
             return False
         try:
             timestamp = int(timestamp_str)
-            current_time = int(time())
+            current_time = await self._get_redis_time()
             return (current_time - timestamp) <= MAX_HEALTH_STALENESS_SEC
         except (ValueError, TypeError):
             return False
@@ -323,10 +333,10 @@ async def get_route_health_status(self, route_id: str) -> Optional[HealthStatus]
         data = {k.decode(): v.decode() for k, v in result.items()}
 
         # Parse boolean values using validation helper (checks both status and staleness)
-        readiness = self._is_health_status_valid(
+        readiness = await self._validate_health_status(
             data.get("readiness", "0"), data.get("last_readiness", "0")
         )
-        liveness = self._is_health_status_valid(
+        liveness = await self._validate_health_status(
             data.get("liveness", "0"), data.get("last_liveness", "0")
         )
         last_check = int(data["last_check"]) if "last_check" in data else 0
@@ -345,7 +355,7 @@ async def initialize_routes_health_status_batch(self, route_ids: list[str]) -> N
         if not route_ids:
             return
 
-        current_time = str(int(time()))
+        current_time = str(await self._get_redis_time())
         batch = Batch(is_atomic=False)
 
         for route_id in route_ids:
@@ -373,7 +383,7 @@ async def update_route_readiness(self, route_id: str, readiness: bool) -> None:
         key = self._get_route_health_key(route_id)
         data: Mapping[str | bytes, str | bytes] = {
             "readiness": "1" if readiness else "0",
-            "last_readiness": str(int(time())),
+            "last_readiness": str(await self._get_redis_time()),
         }
 
         batch = Batch(is_atomic=False)
@@ -393,7 +403,7 @@ async def update_route_liveness(self, route_id: str, liveness: bool) -> None:
         key = self._get_route_health_key(route_id)
         data: Mapping[str | bytes, str | bytes] = {
             "liveness": "1" if liveness else "0",
-            "last_liveness": str(int(time())),
+            "last_liveness": str(await self._get_redis_time()),
         }
 
         batch = Batch(is_atomic=False)
@@ -415,7 +425,7 @@ async def check_route_health_status(
         if not route_ids:
             return {}
 
-        current_time = str(int(time()))
+        current_time = str(await self._get_redis_time())
         batch = Batch(is_atomic=False)
 
         # Single batch: update last_check, refresh TTL, and get all data
@@ -447,13 +457,15 @@ async def check_route_health_status(
             # Parse existing data
             data = {k.decode(): v.decode() for k, v in result.items()}
             # Parse boolean values using validation helper (checks both status and staleness)
+            readiness = await self._validate_health_status(
+                data.get("readiness", "0"), data.get("last_readiness", "0")
+            )
+            liveness = await self._validate_health_status(
+                data.get("liveness", "0"), data.get("last_liveness", "0")
+            )
             health_statuses[route_id] = HealthStatus(
-                readiness=self._is_health_status_valid(
-                    data.get("readiness", "0"), data.get("last_readiness", "0")
-                ),
-                liveness=self._is_health_status_valid(
-                    data.get("liveness", "0"), data.get("last_liveness", "0")
-                ),
+                readiness=readiness,
+                liveness=liveness,
                 last_check=int(data["last_check"]) if "last_check" in data else 0,
             )
 
@@ -470,7 +482,7 @@ async def update_routes_readiness_batch(self, route_readiness: Mapping[str, bool
         if not route_readiness:
             return
 
-        current_time = str(int(time()))
+        current_time = str(await self._get_redis_time())
         batch = Batch(is_atomic=False)
         for route_id, readiness in route_readiness.items():
             key = self._get_route_health_key(route_id)

tests/common/clients/valkey_client/test_valkey_schedule_client.py

@@ -61,36 +61,36 @@ async def _set_stale_health_data(self, client: ValkeyScheduleClient, route_id: s
         )
 
     @pytest.mark.asyncio
-    async def test_is_health_status_valid_with_valid_timestamp(
+    async def test_validate_health_status_with_valid_timestamp(
         self, valkey_schedule_client: ValkeyScheduleClient
     ) -> None:
         """Test validation with healthy status and fresh timestamp"""
         fresh_timestamp = str(int(time()))
-        assert valkey_schedule_client._is_health_status_valid("1", fresh_timestamp) is True
+        assert await valkey_schedule_client._validate_health_status("1", fresh_timestamp) is True
 
     @pytest.mark.asyncio
-    async def test_is_health_status_valid_with_stale_timestamp(
+    async def test_validate_health_status_with_stale_timestamp(
         self, valkey_schedule_client: ValkeyScheduleClient
     ) -> None:
         """Test validation with healthy status but stale timestamp"""
         stale_timestamp = str(int(time()) - MAX_HEALTH_STALENESS_SEC - 10)
-        assert valkey_schedule_client._is_health_status_valid("1", stale_timestamp) is False
+        assert await valkey_schedule_client._validate_health_status("1", stale_timestamp) is False
 
     @pytest.mark.asyncio
-    async def test_is_health_status_valid_with_unhealthy_status(
+    async def test_validate_health_status_with_unhealthy_status(
         self, valkey_schedule_client: ValkeyScheduleClient
     ) -> None:
         """Test validation with unhealthy status regardless of timestamp"""
         fresh_timestamp = str(int(time()))
-        assert valkey_schedule_client._is_health_status_valid("0", fresh_timestamp) is False
+        assert await valkey_schedule_client._validate_health_status("0", fresh_timestamp) is False
 
     @pytest.mark.asyncio
-    async def test_is_health_status_valid_with_invalid_timestamp(
+    async def test_validate_health_status_with_invalid_timestamp(
         self, valkey_schedule_client: ValkeyScheduleClient
     ) -> None:
         """Test validation with invalid timestamp formats"""
-        assert valkey_schedule_client._is_health_status_valid("1", "invalid") is False
-        assert valkey_schedule_client._is_health_status_valid("1", "") is False
+        assert await valkey_schedule_client._validate_health_status("1", "invalid") is False
+        assert await valkey_schedule_client._validate_health_status("1", "") is False
 
     @pytest.mark.asyncio
     async def test_initialize_routes_health_status_batch(