feat: Set expiry to Bgtask metadata id sets (#5736)

Author: fregataa

changes/5736.feature.md

@@ -0,0 +1 @@
+Set expiry to set records of Bgtask metadata ids

src/ai/backend/common/bgtask/bgtask.py

@@ -460,7 +460,11 @@ async def _heartbeat_loop(self) -> None:
                 for task_id, bg_task in self._ongoing_tasks.items():
                     if not bg_task.done():
                         alive_task_ids.append(task_id)
-                await self._valkey_client.heartbeat(alive_task_ids)
+                await self._valkey_client.heartbeat(
+                    alive_task_ids,
+                    server_id=self._server_id,
+                    tags=self._tags,
+                )
             except Exception as e:
                 log.exception("Exception in heartbeat loop: {}", e)
             await asyncio.sleep(_HEARTBEAT_INTERVAL)

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

@@ -118,9 +118,11 @@ async def register_task(self, metadata: BackgroundTaskMetadata) -> None:
         for tag in metadata.tags:
             tag_key = self._get_tag_key(tag)
             batch.sadd(tag_key, [metadata.task_id.hex])
+            batch.expire(tag_key, TASK_METADATA_TTL)
 
         server_key = self._get_server_key(metadata.server_id)
         batch.sadd(server_key, [metadata.task_id.hex])
+        batch.expire(server_key, TASK_METADATA_TTL)
         await self._client.client.exec(batch, raise_on_error=True)
 
     @valkey_decorator()
@@ -240,14 +242,28 @@ def _resolve_script_result(
         return result_type, metadata
 
     @valkey_decorator()
-    async def heartbeat(self, task_ids: Collection[TaskID]) -> None:
+    async def heartbeat(
+        self,
+        task_ids: Collection[TaskID],
+        server_id: Optional[str],
+        tags: Collection[str],
+    ) -> None:
         """
         Extend TTL to 24 hours for active tasks. Non-existent tasks are ignored.
         """
-        if not task_ids:
+        keys: list[str] = []
+
+        task_keys = [self._get_task_key(task_id) for task_id in task_ids]
+        keys.extend(task_keys)
+        if server_id is not None:
+            server_key = self._get_server_key(server_id)
+            keys.append(server_key)
+        tag_keys = [self._get_tag_key(tag) for tag in tags]
+        keys.extend(tag_keys)
+
+        if not keys:
             return
         batch = self._create_batch()
-        keys = [self._get_task_key(task_id) for task_id in task_ids]
         for key in keys:
             batch.expire(key, TASK_METADATA_TTL)
         await self._client.client.exec(batch, raise_on_error=True)

src/ai/backend/common/data/bgtask/defs.py

@@ -1,2 +1,2 @@
-TASK_METADATA_TTL = 7200  # 2 hours
-TASK_TTL_THRESHOLD = 6900  # 1 hour 55 minutes
+TASK_METADATA_TTL = 60 * 60 * 24  # 24 hours
+TASK_TTL_THRESHOLD = TASK_METADATA_TTL - 60 * 5  # 5 minutes margin