Graceful Shutdown Issue with Taskiq in K8s

[Bohdan-Ilchyshyn]

I am using Taskiq with Kubernetes and KEDA scaling, and I am encountering issues with implementing a graceful shutdown process for my workers.

Problem Description

When a pod receives a SIGTERM signal:

• I need to stop the worker from receiving new tasks.
• The worker should wait for currently running tasks to complete before shutting down.

I have configured the following settings:

shutdown-timeout = 70
wait-tasks-timeout = 60

However, this does not work as expected because wait-tasks-timeout is only used when the listener receives the QUEUE_DONE message. The QUEUE_DONE message is sent exclusively in the Listener.prefetcher.

Expected Behavior

Upon receiving a SIGTERM signal, the worker should:

• Immediately stop accepting new tasks.
• Wait for all in-progress tasks to complete within the wait-tasks-timeout duration.

Questions

• How can I properly implement this behavior for a graceful shutdown?
• Are there existing hooks or configurations I should use to address this issue?
• Should QUEUE_DONE messages be triggered elsewhere to support this scenario?

[s3rius]

This sounds like a serious bug. I'll look into it.

[Bohdan-Ilchyshyn]

Oh, thank you so much for the quick fix!

[Bohdan-Ilchyshyn]

I started working on my own solution yesterday, but since you've already completed yours, I’ll wait for the new version and then ask a few questions. Perhaps I don’t fully understand why the following things can't be done:

• Why can't we move the task registry into init?
• Why can't we start two tasks and store them in init?
• Why can't we create a shutdown method to terminate the tasks?

For example, I was thinking of doing it this way:

Receiver

    def __init__(...):
        ...
        self._allow_pick_jobs: bool = False
        self._tasks: Set[asyncio.Task[Any]] = set()

        self.prefetcher_task: Task | None = None
        self.runner_task: Task | None = None

    def _task_cb(self, task: "asyncio.Task[Any]") -> None:
        self._tasks.discard(task)
        if self.sem is not None:
            self.sem.release()

    async def _sleep_until_tasks_complete(self) -> None:
        while len(self._tasks) != 0:
            await asyncio.sleep(0.1)

    async def wait_for_tasks_to_complete(self) -> None:
        if self.prefetcher_task is not None:
            self.prefetcher_task.cancel()

        with contextlib.suppress(asyncio.TimeoutError):
            await asyncio.wait_for(
                self._sleep_until_tasks_complete(),
                self.wait_tasks_timeout,
            )

        for task in self._tasks:
            if not task.done():
                task.cancel()
                
        if self.runner_task is not None:
            self.runner_task.cancel()
                
    async def close(self) -> None:
        await asyncio.gather(*self._tasks)
                
    async def listen(self) -> None:  # pragma: no cover
        """
        This function iterates over tasks asynchronously.

        It uses listen() method of an AsyncBroker
        to get new messages from queues.
        """
        if self.run_startup:
            await self.broker.startup()

        logger.info("Listening started.")
        queue: "asyncio.Queue[Union[bytes, AckableMessage]]" = asyncio.Queue()

        def _cancel(*_: Any):
            asyncio.create_task(self.wait_for_tasks_to_complete())

        self.prefetcher_task = asyncio.create_task(self.prefetcher(queue))
        self.runner_task = asyncio.create_task(self.runner(queue))

        signal.signal(signal.SIGTERM, _cancel)

        if self.on_exit is not None:
            self.on_exit(self)

run.py

    pool = ThreadPoolExecutor(args.max_threadpool_threads)

    try:
        logger.debug("Initialize receiver.")
        receiver = receiver_type(
            broker=broker,
            executor=pool,
            validate_params=not args.no_parse,
            max_async_tasks=args.max_async_tasks,
            max_prefetch=args.max_prefetch,
            propagate_exceptions=not args.no_propagate_errors,
            ack_type=args.ack_type,
            max_tasks_to_execute=args.max_tasks_per_child,
            wait_tasks_timeout=args.wait_tasks_timeout,
            **receiver_kwargs,  # type: ignore
        )
        receiver_task = loop.create_task(receiver.listen())
        loop.run_until_complete(receiver_task)
    except (KeyboardInterrupt, asyncio.CancelledError):
        loop.run_until_complete(receiver.close())
        logger.warning("Worker process interrupted.")
        loop.run_until_complete(shutdown_broker(broker, args.shutdown_timeout))
    finally:
        pool.shutdown(wait=True)

But your solution with the signal is much better! :)

[s3rius]

@Bohdan-Ilchyshyn, thanks for your interest, and I really appreciate your help with this. I also think the Receiver class and the subprocess manager can both be refactored to make them simpler and nicer.

Why can't we move the task registry into init?

That's a good question. I don't see why we can't.

Why can't we start two tasks and store them in init?

Because before prefetch and listen, we need to run startup actions, which can be async. You can't run async tasks in __init__.

Why can't we create a shutdown method to terminate the tasks?

I guess we can, actually, and it might be better. But it could be tricky to handle signals in those cases.

To be honest, this part is the core of the whole system, which is why I've stuck with the original solution since it works. I’m looking forward to seeing any new implementation, but right now I won’t be doing it myself since the current one works. However, I'd be happy if anyone finds a way to improve it. That would be really helpful for future features.

Also, another key issue here is that this part is poorly tested. The receiver is at least covered by unit tests, but the process manager is really hard to test. Maybe a new implementation will make this part easier to test automatically.

So, to wrap it up: I'll stick with my solution for now, but if you find a way to refactor it without affecting the project's functionality, I'll be happy to help with the PRs to get it into the release.