Repository: python-arq/arq (language: python)

Issue:
Workers stuck in terminating state when 'job_completion_wait' is specified and 'handle_sig_wait_for_completion' is triggered
Currently the workers will be stuck in a termination state when the handle_sig_wait_for_completion is activated, causing ineffective k8s pod lifecycles.

This appears to be happening since the '_sleep_until_tasks_complete' function is checking len() of self.tasks, but the deletion of done task ids is nested under the if self.allow_pick_jobs. reference: https://github.com/samuelcolvin/arq/blob/main/arq/worker.py#L388

Have tried a couple things:
1. tabbing over the task deletion out from the if self.allow_pick_jobs check
2. having '_sleep_until_tasks_complete' read from self.job_tasks, which are being deleted during finally block in 'run_job'. reference: https://github.com/samuelcolvin/arq/blob/main/arq/worker.py#L583

For both attempted fixes, have been met with asyncio.exceptions.CancelledError errors.


A fix that seems to be working, iterate tasks during '_sleep_until_tasks_complete' while loop and search for not done.

Proposed fix:

Existing code:


    async def _sleep_until_tasks_complete(self) -> None:  
        """
        Sleeps until all tasks are done. Used together with asyncio.wait_for()
        """
        while len(self.tasks):
            await asyncio.sleep(0.1)



Recommended fix:


    async def _sleep_until_tasks_complete(self) -> None:  
        """
        Sleeps until all tasks are done. Used together with asyncio.wait_for()
        """
        while len([task for task in self.tasks.values() if not task.done()]):
            await asyncio.sleep(0.1)


Is this a change you would support if a PR is made, or is there a better way around this issue?
Workers stuck in terminating state when 'job_completion_wait' is specified and 'handle_sig_wait_for_completion' is triggered
Currently the workers will be stuck in a termination state when the handle_sig_wait_for_completion is activated, causing ineffective k8s pod lifecycles.

This appears to be happening since the '_sleep_until_tasks_complete' function is checking len() of self.tasks, but the deletion of done task ids is nested under the if self.allow_pick_jobs. reference: https://github.com/samuelcolvin/arq/blob/main/arq/worker.py#L388

Have tried a couple things:
1. tabbing over the task deletion out from the if self.allow_pick_jobs check
2. having '_sleep_until_tasks_complete' read from self.job_tasks, which are being deleted during finally block in 'run_job'. reference: https://github.com/samuelcolvin/arq/blob/main/arq/worker.py#L583

For both attempted fixes, have been met with asyncio.exceptions.CancelledError errors.


A fix that seems to be working, iterate tasks during '_sleep_until_tasks_complete' while loop and search for not done.

Proposed fix:

Existing code:


    async def _sleep_until_tasks_complete(self) -> None:  
        """
        Sleeps until all tasks are done. Used together with asyncio.wait_for()
        """
        while len(self.tasks):
            await asyncio.sleep(0.1)



Recommended fix:


    async def _sleep_until_tasks_complete(self) -> None:  
        """
        Sleeps until all tasks are done. Used together with asyncio.wait_for()
        """
        while len([task for task in self.tasks.values() if not task.done()]):
            await asyncio.sleep(0.1)


Is this a change you would support if a PR is made, or is there a better way around this issue?
Workers stuck in terminating state when 'job_completion_wait' is specified and 'handle_sig_wait_for_completion' is triggered
Currently the workers will be stuck in a termination state when the handle_sig_wait_for_completion is activated, causing ineffective k8s pod lifecycles.

This appears to be happening since the '_sleep_until_tasks_complete' function is checking len() of self.tasks, but the deletion of done task ids is nested under the if self.allow_pick_jobs. reference: https://github.com/samuelcolvin/arq/blob/main/arq/worker.py#L388

Have tried a couple things:
1. tabbing over the task deletion out from the if self.allow_pick_jobs check
2. having '_sleep_until_tasks_complete' read from self.job_tasks, which are being deleted during finally block in 'run_job'. reference: https://github.com/samuelcolvin/arq/blob/main/arq/worker.py#L583

For both attempted fixes, have been met with asyncio.exceptions.CancelledError errors.


A fix that seems to be working, iterate tasks during '_sleep_until_tasks_complete' while loop and search for not done.

Proposed fix:

Existing code:


    async def _sleep_until_tasks_complete(self) -> None:  
        """
        Sleeps until all tasks are done. Used together with asyncio.wait_for()
        """
        while len(self.tasks):
            await asyncio.sleep(0.1)



Recommended fix:


    async def _sleep_until_tasks_complete(self) -> None:  
        """
        Sleeps until all tasks are done. Used together with asyncio.wait_for()
        """
        while len([task for task in self.tasks.values() if not task.done()]):
            await asyncio.sleep(0.1)


Is this a change you would support if a PR is made, or is there a better way around this issue?


Discussion / hints:
@JonasKs nice work on getting this feature implemented! 
It would be nice to not have to wait for the configured job_completion_wait time in order for the workers to shutdown, if you had any ideas.
Hey @mernmic! Thanks for the good bug report!

I was unsure how I missed that, but I actually used your suggested implementation when I was looking at my [logs in my k8s lab](https://github.com/JonasKs/arq/blob/450207732208f43dc19863b0add59df682bd4ee0/arq/worker.py#L794), but acted on `self.tasks()` since it was used in the other signal handler.

I'm sorry I didn't catch this. PR is very welcome - we deployed this to QA yesterday, so I assume I'll see your findings when I get to work in a few hours..  
(PR world preferably come together with a test that proves it's a bug, so we don't implement it again. I'll look into helping with this ASAP.)
@JonasKs nice work on getting this feature implemented! 
It would be nice to not have to wait for the configured job_completion_wait time in order for the workers to shutdown, if you had any ideas.
Hey @mernmic! Thanks for the good bug report!

I was unsure how I missed that, but I actually used your suggested implementation when I was looking at my [logs in my k8s lab](https://github.com/JonasKs/arq/blob/450207732208f43dc19863b0add59df682bd4ee0/arq/worker.py#L794), but acted on `self.tasks()` since it was used in the other signal handler.

I'm sorry I didn't catch this. PR is very welcome - we deployed this to QA yesterday, so I assume I'll see your findings when I get to work in a few hours..  
(PR world preferably come together with a test that proves it's a bug, so we don't implement it again. I'll look into helping with this ASAP.)
@JonasKs nice work on getting this feature implemented! 
It would be nice to not have to wait for the configured job_completion_wait time in order for the workers to shutdown, if you had any ideas.
Hey @mernmic! Thanks for the good bug report!

I was unsure how I missed that, but I actually used your suggested implementation when I was looking at my [logs in my k8s lab](https://github.com/JonasKs/arq/blob/450207732208f43dc19863b0add59df682bd4ee0/arq/worker.py#L794), but acted on `self.tasks()` since it was used in the other signal handler.

I'm sorry I didn't catch this. PR is very welcome - we deployed this to QA yesterday, so I assume I'll see your findings when I get to work in a few hours..  
(PR world preferably come together with a test that proves it's a bug, so we don't implement it again. I'll look into helping with this ASAP.)

Fix patch:
diff --git a/arq/worker.py b/arq/worker.py
index b19a3f78..81afd5b7 100644
--- a/arq/worker.py
+++ b/arq/worker.py
@@ -382,14 +382,14 @@ async def _poll_iteration(self) -> None:
 
             await self.start_jobs(job_ids)
 
-            if self.allow_abort_jobs:
-                await self._cancel_aborted_jobs()
-
-            for job_id, t in list(self.tasks.items()):
-                if t.done():
-                    del self.tasks[job_id]
-                    # required to make sure errors in run_job get propagated
-                    t.result()
+        if self.allow_abort_jobs:
+            await self._cancel_aborted_jobs()
+
+        for job_id, t in list(self.tasks.items()):
+            if t.done():
+                del self.tasks[job_id]
+                # required to make sure errors in run_job get propagated
+                t.result()
 
         await self.heart_beat()
 


Test patch:
diff --git a/tests/test_worker.py b/tests/test_worker.py
index 28a9fda0..aa56085b 100644
--- a/tests/test_worker.py
+++ b/tests/test_worker.py
@@ -116,19 +116,23 @@ async def test_worker_signal_completes_job_before_shutting_down(caplog, arq_redi
     async def sleep_job(ctx, time):
         await asyncio.sleep(time)
 
-    await arq_redis.enqueue_job('sleep_job', 0.2, _job_id='short_sleep')  # should be cancelled
+    await arq_redis.enqueue_job('sleep_job', 0.2, _job_id='short_sleep')  # should be completed
     await arq_redis.enqueue_job('sleep_job', 5, _job_id='long_sleep')  # should be cancelled
     worker = worker(
         functions=[func(sleep_job, name='sleep_job', max_tries=1)],
-        job_completion_wait=0.4,
+        job_completion_wait=0.5,
         job_timeout=10,
     )
     assert worker.jobs_complete == 0
     asyncio.create_task(worker.main())
     await asyncio.sleep(0.1)
     worker.handle_sig_wait_for_completion(signal.SIGINT)
+    assert len(worker.tasks) == 2  # should be two tasks when sigint is sent
     assert worker.allow_pick_jobs is False
-    await asyncio.sleep(0.5)
+    await asyncio.sleep(0.3)
+    assert len(worker.tasks) == 1  # slept a bit, first job should now be complete and self.tasks should be updated
+    await asyncio.sleep(0.3)
+    assert len(worker.tasks) == 0  # slept longer than `job_completion_wait`, task should be cancelled and updated
     logs = [rec.message for rec in caplog.records]
     assert 'shutdown on SIGINT ◆ 0 jobs complete ◆ 0 failed ◆ 0 retries ◆ 2 to be completed' in logs
     assert 'shutdown on SIGINT, wait complete ◆ 1 jobs complete ◆ 0 failed ◆ 0 retries ◆ 1 ongoing to cancel' in logs