From dfd4520161ff9ab0ec8a2dd0bc8d187aa9498c58 Mon Sep 17 00:00:00 2001 From: Amin Vakil Date: Wed, 12 Aug 2026 15:04:07 +0330 Subject: [PATCH 1/2] fix(worker): keep pull worker healthy while queue is full --- .../src/taskbroker_client/worker/client.py | 3 ++ .../src/taskbroker_client/worker/worker.py | 1 + clients/python/tests/worker/test_worker.py | 29 +++++++++++++++++++ 3 files changed, 33 insertions(+) diff --git a/clients/python/src/taskbroker_client/worker/client.py b/clients/python/src/taskbroker_client/worker/client.py index 5e8e1ba9..8a6f1da9 100644 --- a/clients/python/src/taskbroker_client/worker/client.py +++ b/clients/python/src/taskbroker_client/worker/client.py @@ -328,6 +328,9 @@ def _emit_health_check(self) -> None: ) self._timestamp_since_touch = cur_time + def emit_health_check(self) -> None: + self._emit_health_check() + def _connect_to_host(self, host: str) -> ConsumerServiceStub: logger.info("taskworker.client.connect", extra={"host": host}) channel = grpc.insecure_channel(host, options=self._grpc_options) diff --git a/clients/python/src/taskbroker_client/worker/worker.py b/clients/python/src/taskbroker_client/worker/worker.py index 30f7c298..58e2e55d 100644 --- a/clients/python/src/taskbroker_client/worker/worker.py +++ b/clients/python/src/taskbroker_client/worker/worker.py @@ -656,6 +656,7 @@ def signal_handler(*args: Any) -> None: def run_once(self) -> None: """Access point for tests to run a single worker loop""" + self.client.emit_health_check() self._add_task() def _add_task(self) -> bool: diff --git a/clients/python/tests/worker/test_worker.py b/clients/python/tests/worker/test_worker.py index 556e8023..940f1e21 100644 --- a/clients/python/tests/worker/test_worker.py +++ b/clients/python/tests/worker/test_worker.py @@ -723,6 +723,35 @@ def test_constructor_push_mode(self) -> None: self.assertEqual(taskworker._grpc_port, 50099) +def test_pull_worker_health_check_touches_while_full(tmp_path: Path) -> None: + health_check_path = tmp_path / "health" + taskworker = TaskWorker( + app_module="examples.app:app", + broker_hosts=["127.0.0.1:50051"], + max_child_task_count=100, + process_type="fork", + health_check_file_path=str(health_check_path), + health_check_sec_per_touch=1, + ) + + with ( + mock.patch.object(taskworker.worker_pool, "is_worker_full", return_value=True), + mock.patch.object(taskworker.client, "get_task") as mock_get_task, + mock.patch("taskbroker_client.worker.worker.time.sleep"), + mock.patch("taskbroker_client.worker.client.time") as mock_time, + ): + mock_time.time.return_value = 1 + taskworker.run_once() + assert health_check_path.exists() + + health_check_path.unlink() + mock_time.time.return_value = 3 + taskworker.run_once() + + assert health_check_path.exists() + mock_get_task.assert_not_called() + + def test_push_worker_health_check_touches_while_idle(tmp_path: Path) -> None: taskworker = PushTaskWorker( app_module="examples.app:app", From d0f16e9d21c12e7cf9e08e5afcfd5bedf80dbb6c Mon Sep 17 00:00:00 2001 From: Amin Vakil Date: Fri, 14 Aug 2026 20:51:20 +0330 Subject: [PATCH 2/2] Simplify new health check emitting --- clients/python/src/taskbroker_client/worker/client.py | 9 +++------ clients/python/src/taskbroker_client/worker/worker.py | 2 +- 2 files changed, 4 insertions(+), 7 deletions(-) diff --git a/clients/python/src/taskbroker_client/worker/client.py b/clients/python/src/taskbroker_client/worker/client.py index 8a6f1da9..717d0379 100644 --- a/clients/python/src/taskbroker_client/worker/client.py +++ b/clients/python/src/taskbroker_client/worker/client.py @@ -309,7 +309,7 @@ def __init__( self._timestamp_since_touch_lock = threading.Lock() self._timestamp_since_touch = 0.0 - def _emit_health_check(self) -> None: + def emit_health_check(self) -> None: if self._health_check_settings is None: return @@ -328,9 +328,6 @@ def _emit_health_check(self) -> None: ) self._timestamp_since_touch = cur_time - def emit_health_check(self) -> None: - self._emit_health_check() - def _connect_to_host(self, host: str) -> ConsumerServiceStub: logger.info("taskworker.client.connect", extra={"host": host}) channel = grpc.insecure_channel(host, options=self._grpc_options) @@ -406,7 +403,7 @@ def get_task(self, namespace: str | None = None) -> InflightTaskActivation | Non If a namespace is provided, only tasks for that namespace will be fetched. This will return None if there are no tasks to fetch. """ - self._emit_health_check() + self.emit_health_check() request = GetTaskRequest(application=self._application, namespace=namespace) try: @@ -458,7 +455,7 @@ def update_task( The return value is the next task that should be executed. """ - self._emit_health_check() + self.emit_health_check() if fetch_next_task is not None: fetch_next_task.application = self._application diff --git a/clients/python/src/taskbroker_client/worker/worker.py b/clients/python/src/taskbroker_client/worker/worker.py index 58e2e55d..2db3a147 100644 --- a/clients/python/src/taskbroker_client/worker/worker.py +++ b/clients/python/src/taskbroker_client/worker/worker.py @@ -656,7 +656,6 @@ def signal_handler(*args: Any) -> None: def run_once(self) -> None: """Access point for tests to run a single worker loop""" - self.client.emit_health_check() self._add_task() def _add_task(self) -> bool: @@ -664,6 +663,7 @@ def _add_task(self) -> bool: Add a task to child tasks queue. Returns False if no new task was fetched. """ if self.worker_pool.is_worker_full(): + self.client.emit_health_check() self._metrics.incr( "taskworker.worker.add_tasks.child_tasks_full", tags={"processing_pool": self._processing_pool_name},