feat(worker): Add opt-in wait for child exit during recycling

pyproject.toml

@@ -45,7 +45,7 @@ dependencies = [
     "sentry-options==1.2.10",
     "sentry-sdk>=2.74.0",
     "sortedcontainers>=2.4.0",
-    "taskbroker-client>=0.20.32,<1",
+    "taskbroker-client>=0.20.35,<1",
     "typing-extensions>=4.15.0",
     "zipfile-zstd==0.0.4",
 ]

src/launchpad/worker/config.py

@@ -29,6 +29,7 @@ class WorkerConfig:
     pod_name: str
     push_mode: bool
     push_timeout_sec: int
+    wait_for_child_exit: bool
 
 
 def get_worker_config() -> WorkerConfig:
@@ -78,6 +79,7 @@ def get_worker_config() -> WorkerConfig:
         pod_name=pod_name,
         push_mode=push_mode,
         push_timeout_sec=push_timeout_sec,
+        wait_for_child_exit=os.getenv("LAUNCHPAD_WORKER_WAIT_FOR_CHILD_EXIT", "false") == "true",
     )
 
 
@@ -90,7 +92,8 @@ def run_worker(
 
     logger.info(
         f"Starting TaskWorker (rpc_hosts={config.rpc_hosts}, concurrency={config.concurrency}, "
-        f"max_child_task_count={config.max_child_task_count}, health_check_file_path={config.health_check_file_path})"
+        f"max_child_task_count={config.max_child_task_count}, wait_for_child_exit={config.wait_for_child_exit}, "
+        f"health_check_file_path={config.health_check_file_path})"
     )
     if config.push_mode:
         assert config.rpc_host_service, (
@@ -101,6 +104,7 @@ def run_worker(
             broker_service=config.rpc_host_service,
             max_child_task_count=config.max_child_task_count,
             concurrency=config.concurrency,
+            wait_for_child_exit=config.wait_for_child_exit,
             child_tasks_queue_maxsize=1,
             result_queue_maxsize=1,
             rebalance_after=16,
@@ -117,6 +121,7 @@ def run_worker(
             broker_hosts=config.rpc_hosts,
             max_child_task_count=config.max_child_task_count,
             concurrency=config.concurrency,
+            wait_for_child_exit=config.wait_for_child_exit,
             child_tasks_queue_maxsize=1,
             result_queue_maxsize=1,
             rebalance_after=16,

tests/unit/test_worker_config.py

@@ -31,6 +31,7 @@ def test_valid_config(self):
                 pod_name="",
                 push_mode=False,
                 push_timeout_sec=5,
+                wait_for_child_exit=False,
             )
 
     def test_custom_health_check_file_path(self):
@@ -111,3 +112,16 @@ def test_rpc_host_service(self):
         ):
             config = get_worker_config()
             assert config.rpc_host_service == "localhost:50051"
+
+    def test_wait_for_child_exit_enabled(self):
+        with patch.dict(
+            os.environ,
+            {
+                "LAUNCHPAD_WORKER_RPC_HOST": "localhost:50051",
+                "LAUNCHPAD_WORKER_CONCURRENCY": "2",
+                "LAUNCHPAD_WORKER_WAIT_FOR_CHILD_EXIT": "true",
+            },
+            clear=True,
+        ):
+            config = get_worker_config()
+            assert config.wait_for_child_exit is True