Graceful exit with Python multiprocessing

import os, signal, logging, sys
from time import sleep, time
from random import random, randint
from multiprocessing import JoinableQueue, Event, Process
from queue import Empty
from typing import Optional

logger = logging.getLogger(__name__)

# Make sure to log process/thread name
logging.basicConfig(level=logging.INFO, format="%(processName)s %(relativeCreated)d %(message)s")

print("Python version:", sys.version.split(" ")[0])

def do_task(sleep_sec: int) -> int:
    "Dummy task"
    sleep(sleep_sec)
    logger.info("Slept for %.3fs", sleep_sec)
    return randint(0, 10)


def worker(q: JoinableQueue, stop_event: Event):
    logger.info("Starting worker...")
    while 1:
        if stop_event.is_set():
            logger.info("Worker exiting because of stop_event")
            break
        # We set a timeout so we loop past "stop_event" even if the queue is empty
        try:
            args = q.get(timeout=.01)
        except Empty:
            # Run next iteration of loop
            continue
        
        # Exit if end of queue
        if args is None:
            logger.info("Worker exiting because of None on queue")
            q.task_done()
            break
        
        # Do the task
        try:
            do_task(*args)
        # Will also catch KeyboardInterrupt
        except:
            logger.exception("Failed to process args %s", args)
            # Can implement some kind of retry handling here
        finally:
            q.task_done()

def main(
        n_workers: int = 2,
        n_tasks: int = 10,
        max_queue_size: int = 1,
        grace_period: int = 2,
        kill_period: int = 30,
        interrupt: Optional[int] = None
    ) -> None:
    """
    Run a process pool of workers.
    
    Args:
        n_workers: Start this many processes
        n_tasks: Launch this many tasks
        max_queue_size: If queue exceeds this size, block when putting items on the queue
        grace_period: Send SIGINT to processes if they don't exit within this time after SIGINT/SIGTERM
        kill_period: Send SIGKILL to processes if they don't exit after this many seconds
        interrupt: If given, send signal SIGTERM to itself after queueing this many tasks
    """
    # The JoinableQueue has an internal counter that increments when an item is put on the queue and
    # decrements when q.task_done() is called. This allows us to wait until it's empty using .join()
    q = JoinableQueue(maxsize=max_queue_size)
    # This is a process-safe version of the "panic" variable shown above
    stop_event = Event()
    
    def handler(signalname):
        """
        Python 3.9 has `signal.strsignal(signalnum)` so this closure would not be needed.
        Also, 3.8 includes `signal.valid_signals()` that can be used to create a mapping for the same purpose.
        """
        def f(signal_received, frame):
            raise KeyboardInterrupt(f"{signalname} received")
        return f

    # This will be inherited by the child process if it is forked (not spawned)
    signal.signal(signal.SIGINT, handler("SIGINT"))
    signal.signal(signal.SIGTERM, handler("SIGTERM"))
    
    procs = []

    for i in range(n_workers):
        # Make it a daemon process so it is definitely terminated when this process exits,
        # might be overkill but is a nice feature. See
        # https://docs.python.org/3.8/library/multiprocessing.html#multiprocessing.Process.daemon
        p = Process(name=f"Worker-{i:02d}", daemon=True, target=worker, args=(q, stop_event))
        procs.append(p)
        p.start()
    try:
        # Put tasks on queue
        for i_task in range(n_tasks):
            # For demonstration purposes
            if interrupt and i_task == interrupt:
                os.kill(os.getpid(), signal.SIGTERM)
            s = random()
            logger.info("Put %.3f on queue", s)
            q.put([s])
        # Put exit tasks on queue
        for i in range(n_workers):
            q.put(None)
        # Wait until all tasks are processed
        q.join()
    except KeyboardInterrupt:
        logger.warning("Caught KeyboardInterrupt! Setting stop event...")
    finally:
        stop_event.set()
        t = time()
        # Send SIGINT if process doesn't exit quickly enough, and kill it as last resort
        # .is_alive() also implicitly joins the process (good practice in linux)
        while alive_procs := [p for p in procs if p.is_alive()]:
            if time() > t + grace_period:
                for p in alive_procs:
                    os.kill(p.pid, signal.SIGINT)
                    logger.warning("Sending SIGINT to %s", p)
            elif time() > t + kill_period:
                for p in alive_procs:
                    logger.warning("Sending SIGKILL to %s", p)
                    # Queues and other inter-process communication primitives can break when
                    # process is killed, but we don't care here
                    p.kill()
            sleep(.01)

        sleep(.1)
        for p in procs:
            logger.info("Process status: %s", p)

main()

main(interrupt=4)

main(interrupt=3, grace_period=.01)

main(interrupt=3, kill_period=.01)