SystemExit handling in TaskGroups

[vxgmichel]

This a follow-up for the discussion in PR #95571.

Bug report

A Task exception was never retrieved warning is shown when a SystemExit is raised in a task of a TaskGroup, e.g:

import asyncio

async def system_exit():
    raise SystemExit

async def main():
    async with asyncio.TaskGroup() as tg:
        tg.create_task(system_exit())

if __name__ == "__main__":
    asyncio.run(main())

Task exception was never retrieved
future: <Task finished name='Task-1' coro=<main() done, defined at /home/vinmic/debug.py:6> exception=SystemExit()>
Traceback (most recent call last):
[...]
  File "/home/vinmic/debug.py", line 4, in system_exit
    raise SystemExit
SystemExit

After a bit of investigation, it turns out that the SystemExit exception from the subtask raises up to the runner (without the main task being finished):

cpython/Lib/asyncio/runners.py

Line 118 in ee21110

return self._loop.run_until_complete(task)

The runner then proceeds to teardown its context (Runner.close()) and cancel-join the remaining tasks here:

cpython/Lib/asyncio/runners.py

Lines 202 to 205 in ee21110

	for task in to_cancel:
	task.cancel()
	loop.run_until_complete(tasks.gather(*to_cancel, return_exceptions=True))

But the SystemExit leaks again out of run_until_complete without gather being finished. The main task is now done with a SystemExit as its exception but it's never been accessed so the warning is prompted when it is garbage collected.

Note that the documentation for the TaskGroup says:

Two base exceptions are treated specially: If any task fails with KeyboardInterrupt or SystemExit, the task group still cancels the remaining tasks and waits for them, but then the initial KeyboardInterrupt or SystemExit is re-raised instead of ExceptionGroup or BaseExceptionGroup.

This seems to imply that raising a SystemExit in a task group (or in asyncio in general) is a sound way to shutdown an asyncio application.

Possible fix

Here's a patch that solves the issue, although it might not be right way to address the underlying problem.

diff --git a/Lib/asyncio/runners.py b/Lib/asyncio/runners.py
index 1b89236599..1cd0214730 100644
--- a/Lib/asyncio/runners.py
+++ b/Lib/asyncio/runners.py
@@ -115,7 +115,14 @@ def run(self, coro, *, context=None):
 
         self._interrupt_count = 0
         try:
-            return self._loop.run_until_complete(task)
+            while True:
+                try:
+                    return self._loop.run_until_complete(task)
+                except BaseException:
+                    if not task.done():
+                        task.cancel()
+                        continue
+                    raise
         except exceptions.CancelledError:
             if self._interrupt_count > 0:
                 uncancel = getattr(task, "uncancel", None)

It's a bit naive, but it kind of makes sense: when the run of the main task finishes without the main task being actually done (because a base exception such as KeyboardError or SystemExit bubbled up from another task), we cancel it and run it again. The while-loop is needed because it might take as many tries as the depth of the nested task groups, e.g:

import asyncio

async def level3():
    raise SystemExit

async def level2():
    async with asyncio.TaskGroup() as tg:
        tg.create_task(level3())

async def level1():
    async with asyncio.TaskGroup() as tg:
        tg.create_task(level2())

async def main():
    async with asyncio.TaskGroup() as tg:
        tg.create_task(level1())

if __name__ == "__main__":
    asyncio.run(main())

Your environment

• CPython versions tested on: 3.11 and forward (tested on the latest main 62251c3)
• Operating system and architecture: Linux

[vxgmichel]

Note that this bug might also lead to async cleanup issues. Consider the following example:

import asyncio

async def system_exit():
    raise SystemExit

async def level1():
    try:
        async with asyncio.TaskGroup() as tg:
            tg.create_task(system_exit())
    finally:
        await asyncio.sleep(0)
        print("Level1 cleanup")

async def main():
    try:
        async with asyncio.TaskGroup() as tg:
            tg.create_task(level1())
    finally:
        await asyncio.sleep(0)
        print("Main cleanup")

if __name__ == "__main__":
    asyncio.run(main())

In this case the level 1 cleanup does run, but the main cleanup does not.

[YvesDup]

Occasionaly I am working on this issue #93122 and I get the same traceback messages with Task exception was never retrieved in asyncio.run when SystemExit or KeybordInterrupt are raised (This is not my initial main problem).

@vxgmichel It's look like your first case described above.

To eliminate these warning messages, I wrote a fix in the _cancel_all_tasks function of file runnners.py. For reminder, this function is called on the close method from the __exit__ method of Runner class. The fix adds a try bloc around the call of asyncio.run_until_complete(gather(*....), and in the dedicated except bloc on SystemExit or KeybordInterrupt, there are:

• a check of cancelled tasks,
• a comsuption of complete task with an exception (no more message 'Task exception ...'),
• a new call to the _cancel_all_tasks function, if tasks are still pending.

    to_cancel = tasks.all_tasks(loop)
    if not to_cancel:
        return

    for task in to_cancel:
        task.cancel()

    try:
        loop.run_until_complete(tasks.gather(*to_cancel, return_exceptions=True))
    except (SystemExit, KeyboardInterrupt):
        # Checks if all pending tasks are cancelled, done or
        # still pending.
        _pending_tasks = False
        for task in to_cancel:
            if task.cancelled():
                continue
            elif task.done():
                # task raises this exception, consumes it
                # now elimintate a warning message 'Task exception ...'
                task.exception()
            else:
                _pending_tasks = True
        # cancels last pending tasks.
        if _pending_tasks:
            _cancel_all_tasks(loop)
        raise

    for task in to_cancel:
        if task.cancelled():
            continue
        if task.exception() is not None:
            loop.call_exception_handler({
                'message': 'unhandled exception during asyncio.run() shutdown',
                'exception': task.exception(),
                'task': task,
            })

@vxgmichel I hope this fix will be useful for your case.
Thank you for your feedback.