Python Asyncio Graceful Shutdown (Interrupt Sleep)

December 10, 2020

Solution 1: Shutdown Flag

This will enable graceful shutdown, but it will not interrupt ayncio.sleep (wait until sleep end before shutdown).

import asyncio
import signal
import os

shutdown = False

def stop(signum, frame):
    global shutdown
    print('STOP', signum)
    shutdown = True
    print('STOP2')

signal.signal(signal.SIGTERM, stop)

async def run():
    while not shutdown:
        print('.', end='', flush=True)
        await asyncio.sleep(10)

    print('run-END')

async def main():
    await run()

    print('main-END')

if __name__ == '__main__':
    print('pid', os.getpid())
    asyncio.run(main())

Shutdown

kill [PID]

Solution 2: Task.cancel

  • Call Task.cancel to stop asyncio task (will interrupt sleep for immediate shudown)
  • Need to use loop.add_signal_handler to listen for signal, else Task.cancel would not interrupt sleep immediately
  • Cons: sleep will raise asyncio.CancelledError. If not handled properly, the task might end adruptly.
import asyncio
import signal
import os
import functools

shutdown = False

#def stop(signum, frame):
def stop(signame, loop):
    global shutdown
    print('STOP', signame)
    shutdown = True
    # loop = asyncio.get_event_loop()
    # loop.close()


    tasks = asyncio.all_tasks()
    for _task in tasks:
        print('cancel task')
        _task.cancel()


    print('STOP2')

#signal.signal(signal.SIGTERM, stop)

async def run():
    while not shutdown:
        print('.', end='', flush=True)
        try:
            await asyncio.sleep(10)
        except asyncio.CancelledError as e:
            print('run', 'CancelledError', flush=True)
            # raise e

    print('run-END')

async def main():
    loop = asyncio.get_running_loop()

    for signame in {'SIGINT', 'SIGTERM'}:
        loop.add_signal_handler(
            getattr(signal, signame),
            functools.partial(stop, signame, loop))

    task = asyncio.create_task(run())
    try:
        await asyncio.gather(task)
    except asyncio.CancelledError as e:
        print('main', 'cancelledError')
    print('main-END')


if __name__ == '__main__':
    print('pid', os.getpid())
    asyncio.run(main())

Solution 3: Interrupt Sleep

  • We create a special sleep function which are cancellable tasks.
  • Pros:
    • Don’t have to handle asyncio.CancelledError for every sleep call.
    • Asyncio.Task won’t end adruptly.
class Sleep:
    def __init__(self):
        self.tasks = set()

    async def sleep(self, delay, result=None):
        task = asyncio.create_task(asyncio.sleep(delay, result))
        self.tasks.add(task)
        try:
            return await task
        except asyncio.CancelledError:
            return result
        finally:
            self.tasks.remove(task)

    def cancel_all(self):
        for _task in self.tasks:
            _task.cancel()
        # self.tasks = set()
light = Sleep()

def stop(signame, loop):
    global shutdown
    print('STOP', signame)
    shutdown = True

    light.cancel_all()
    

    print('STOP2')

async def run():
    while not shutdown:
        print('.', end='', flush=True)
        # wait asyncio.sleep(10)
        await light.sleep(10)


    print('run-END')

async def main():
    loop = asyncio.get_running_loop()

    for signame in {'SIGINT', 'SIGTERM'}:
        loop.add_signal_handler(
            getattr(signal, signame),
            functools.partial(stop, signame, loop))

    task = asyncio.create_task(run())
    await asyncio.gather(task)
    print('main-END')


if __name__ == '__main__':
    print('pid', os.getpid())
    asyncio.run(main())
This work is licensed under a
Creative Commons Attribution-NonCommercial 4.0 International License.