fix: cancel scheduler tasks on shutdown

Otherwise the currently running tasks will never exit (before they
actually complete), which means the process can't be properly shut down
(only with SIGKILL).

Ideally, we let tasks know that they are about to shutdown and give them
some time to do so; but in the lack of the mechanism, it's better to
cancel than linger forever.

Signed-off-by: Ihar Hrachyshka <[email protected]>

Author: booxter

llama_stack/providers/utils/scheduler.py

@@ -157,10 +157,14 @@ def _run_loop(self) -> None:
         asyncio.set_event_loop(self._loop)
         self._loop.run_forever()
 
-        # When stopping the loop, give tasks a chance to finish
+        # TODO: When stopping the loop, give tasks a chance to finish
         # TODO: should we explicitly inform jobs of pending stoppage?
+
+        # cancel all tasks
         for task in asyncio.all_tasks(self._loop):
-            self._loop.run_until_complete(task)
+            if not task.done():
+                task.cancel()
+
         self._loop.close()
 
     async def shutdown(self) -> None:

tests/unit/providers/utils/test_scheduler.py

@@ -52,6 +52,9 @@ async def job_handler(on_log, on_status, on_artifact):
     assert sched.get_jobs("unknown") == []
     assert sched.get_jobs(job_type) == [sched.get_job(job_id)]
 
+    # give the job handler a chance to run
+    await asyncio.sleep(0)
+
     # now shut the scheduler down and make sure the job ran
     await sched.shutdown()
 
@@ -112,6 +115,9 @@ async def successful_job_handler(on_log, on_status, on_artifact):
     job_id = "test_job_id2"
     sched.schedule(job_type, job_id, successful_job_handler)
 
+    # give the job handler a chance to run
+    await asyncio.sleep(0)
+
     await sched.shutdown()
 
     assert called