Python - Cancel task in asyncio?

Question

I have written code for async pool below. in __aexit__ i'm cancelling the _worker tasks after the tasks get finished. But when i run the code, the worker tasks are not getting cancelled and the code is running forever. This what the task looks like: <Task pending coro=<AsyncPool._worker() running at \async_pool.py:17> wait_for=<Future cancelled>> . The asyncio.wait_for is getting cancelled but not the worker tasks.

class AsyncPool:
    def __init__(self,coroutine,no_of_workers,timeout):
        self._loop           = asyncio.get_event_loop()
        self._queue          = asyncio.Queue()
        self._no_of_workers  = no_of_workers
        self._coroutine      = coroutine
        self._timeout        = timeout
        self._workers        = None

    async def _worker(self): 
        while True:
            try:
                ret = False
                queue_item           = await self._queue.get()
                ret = True
                result               = await asyncio.wait_for(self._coroutine(queue_item), timeout = self._timeout,loop= self._loop)
            except Exception as e:
                print(e)
            finally:
                if ret:
                    self._queue.task_done()


    async def push_to_queue(self,item):
        self._queue.put_nowait(item)
    
    async def __aenter__(self):
        assert self._workers == None
        self._workers = [asyncio.create_task(self._worker()) for _ in range(self._no_of_workers)]
        return self
    
    async def __aexit__(self,type,value,traceback):
        await self._queue.join()

        for worker in self._workers:
            worker.cancel()

        await asyncio.gather(*self._workers, loop=self._loop, return_exceptions =True)

To use the Asyncpool:

async def something(item):
    print("got", item)
    await asyncio.sleep(item)
 
async def main():
    async with AsyncPool(something, 5, 2) as pool:
        for i in range(10):
            await pool.push_to_queue(i)
 
asyncio.run(main())

The Output in my terminal:

Answer 1

When you have an asyncio task created and then cancelled, you still have the task alive that need to be "reclaimed". So you want to await worker for it. However, once you await such a cancelled task, as it will never give you back the expected return value, the asyncio.CancelledError will be raised and you need to catch it somewhere.

Because of this behavior, I don't think you should gather them but to await for each of the cancelled tasks, as they are supposed to return right away:

async def __aexit__(self,type,value,traceback):
    await self._queue.join()

    for worker in self._workers:
        worker.cancel()
    for worker in self._workers:
        try:
           await worker
        except asyncio.CancelledError:
           print("worker cancelled:", worker)

Answer 2

The problem is that your except Exception exception clause also catches cancellation, and ignores it. To add to the confusion, print(e) just prints an empty line in case of a CancelledError , which is where the empty lines in the output come from. (Changing it to print(type(e)) shows what's going on.)

To correct the issue, change except Exception to something more specific, like except asyncio.TimeoutError . This change is not needed in Python 3.8 where asyncio.CancelledError no longer derives from Exception , but from BaseException , so except Exception doesn't catch it.

Answer 3

This appears to work. The event is a counting timer and when it expires it cancels the tasks.

import asyncio
from datetime import datetime as dt
from datetime import timedelta as td
import random
import time

class Program:
    def __init__(self):
        self.duration_in_seconds = 20
        self.program_start = dt.now()
        self.event_has_expired = False
        self.canceled_success = False

        

    async def on_start(self):
        print("On Start Event Start! Applying Overrides!!!")
        await asyncio.sleep(random.randint(3, 9))


    async def on_end(self):
        print("On End Releasing All Overrides!")
        await asyncio.sleep(random.randint(3, 9))
        

    async def get_sensor_readings(self):
        print("getting sensor readings!!!")
        await asyncio.sleep(random.randint(3, 9))   

 
    async def evauluate_data(self):
        print("checking data!!!")
        await asyncio.sleep(random.randint(3, 9))   


    async def check_time(self):
        if (dt.now() - self.program_start > td(seconds = self.duration_in_seconds)):
            self.event_has_expired = True
            print("Event is DONE!!!")

        else:
            print("Event is not done! ",dt.now() - self.program_start)



    async def main(self):
        # script starts, do only once self.on_start()
        await self.on_start()
        print("On Start Done!")

        while not self.canceled_success:

            readings = asyncio.ensure_future(self.get_sensor_readings())
            analysis = asyncio.ensure_future(self.evauluate_data())
            checker = asyncio.ensure_future(self.check_time())
            
            if not self.event_has_expired:
                await readings   
                await analysis           
                await checker
                
            else:
                # close other tasks before final shutdown
                readings.cancel()
                analysis.cancel()
                checker.cancel()
                self.canceled_success = True
                print("cancelled hit!")


        # script ends, do only once self.on_end() when even is done
        await self.on_end()
        print('Done Deal!')


async def main():
    program = Program()
    await program.main()