fix: Add timeout to prepare()'s PersistentTaskGroup (#514)

* Add `async_timeout` when calling `TaskGroup`
* Change `prepare()`'s `TaskGroup` to `PersistentTaskGroup`

changes/514.fix

@@ -0,0 +1 @@
+Fix `prepare()` not running when `start_session()` call hangs without raising Exception

src/ai/backend/manager/scheduler/dispatcher.py

@@ -17,6 +17,7 @@
 )
 
 import aiotools
+import async_timeout
 import sqlalchemy as sa
 from dateutil.tz import tzutc
 from sqlalchemy.engine.row import Row
@@ -781,7 +782,10 @@ async def _mark_session_preparing() -> Sequence[PendingSession]:
 
                 scheduled_sessions = await execute_with_retry(_mark_session_preparing)
                 log.debug("prepare(): preparing {} session(s)", len(scheduled_sessions))
-                async with aiotools.TaskGroup() as tg:
+                async with (
+                    async_timeout.timeout(delay=50.0),
+                    aiotools.PersistentTaskGroup() as tg,
+                ):
                     for scheduled_session in scheduled_sessions:
                         await self.registry.event_producer.produce_event(
                             SessionPreparingEvent(
@@ -800,6 +804,8 @@ async def _mark_session_preparing() -> Sequence[PendingSession]:
                          "maybe another prepare() call is still running")
                 raise asyncio.CancelledError()
             raise
+        except asyncio.TimeoutError:
+            log.warn('prepare(): timeout while executing start_session()')
 
     async def start_session(
         self,