From 018292163907036824d592f12e74c07ce978a96e Mon Sep 17 00:00:00 2001 From: Jamie <2119834+jamieQ@users.noreply.github.com> Date: Thu, 12 Nov 2026 17:30:56 -0600 Subject: [PATCH] feat(worker): Add opt-in wait for child exit during recycling --- pyproject.toml | 2 +- src/launchpad/worker/config.py | 7 ++++++- tests/unit/test_worker_config.py | 14 ++++++++++++++ uv.lock | 8 ++++---- 4 files changed, 25 insertions(+), 6 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index e2c13153..11909e48 100644 --- a/pyproject.toml +++ b/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", ] diff --git a/src/launchpad/worker/config.py b/src/launchpad/worker/config.py index 73e464f8..dfe3c8b3 100644 --- a/src/launchpad/worker/config.py +++ b/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, diff --git a/tests/unit/test_worker_config.py b/tests/unit/test_worker_config.py index f03859f9..7bd1e9a7 100644 --- a/tests/unit/test_worker_config.py +++ b/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 diff --git a/uv.lock b/uv.lock index 1e759437..155cab5a 100644 --- a/uv.lock +++ b/uv.lock @@ -907,7 +907,7 @@ requires-dist = [ { name = "sentry-options", specifier = "==1.2.10" }, { name = "sentry-sdk", specifier = ">=2.74.0" }, { name = "sortedcontainers", specifier = ">=2.4.0" }, - { name = "taskbroker-client", specifier = ">=0.20.32,<1" }, + { name = "taskbroker-client", specifier = ">=0.20.35,<1" }, { name = "typing-extensions", specifier = ">=4.15.0" }, { name = "zipfile-zstd", specifier = "==0.0.4" }, ] @@ -2097,7 +2097,7 @@ wheels = [ [[package]] name = "taskbroker-client" -version = "0.20.32" +version = "0.20.35" source = { registry = "https://sfw.security.sentry.io/pypi/simple" } dependencies = [ { name = "confluent-kafka" }, @@ -2113,9 +2113,9 @@ dependencies = [ { name = "sentry-sdk", extra = ["http2"] }, { name = "zstandard" }, ] -sdist = { url = "https://sfw.security.sentry.io/pypi/packages/19/7c/8c7d0654efd0ff578dd94d5421d44c7488fb2236092861af1ed7b5dd08a4/taskbroker_client-0.20.32.tar.gz", hash = "sha256:74cbc58e706c4053dab20cdd7200a94388c9eaf748fd526e8e61dfab88bfbb04", size = 55084, upload-time = "2026-10-28T20:32:22.158Z" } +sdist = { url = "https://sfw.security.sentry.io/pypi/packages/d9/da/c2c57c0a6980fc31fd0fb00c82d13c3b605e43356496629a3c43f432726a/taskbroker_client-0.20.35.tar.gz", hash = "sha256:323dfe9dfcdde9035a8037400feb3e131da19e292063611b7a4b5246bfed7088", size = 55663, upload-time = "2026-11-12T21:01:20.799Z" } wheels = [ - { url = "https://sfw.security.sentry.io/pypi/packages/16/87/5f5c39a3d5d2cfd503f3410630a3b8d74b5460dade8674d6e03c40d1a09f/taskbroker_client-0.20.32-py3-none-any.whl", hash = "sha256:f519142f82ac157d99c051bf6afb8e22d5b9e327d7a894efcde1b98a267009e4", size = 65337, upload-time = "2026-10-28T20:32:21.118Z" }, + { url = "https://sfw.security.sentry.io/pypi/packages/b0/61/3a6d0103156b2527b258750f6595790c194f59ba04f953b42e7542b196a6/taskbroker_client-0.20.35-py3-none-any.whl", hash = "sha256:ef5e4eb6944e7d3f431d90db54aac0197c41c7bb44c5b5d14f20f7a642f29cd8", size = 65892, upload-time = "2026-11-12T21:01:19.673Z" }, ] [[package]]