From fefbf0e2566806cebcacd03046b1b316d47f370c Mon Sep 17 00:00:00 2001 From: Enoch Tang Date: Fri, 13 Nov 2026 15:48:29 -0500 Subject: [PATCH] feat(worker): Expose occupancy for Prometheus scraping Read LAUNCHPAD_WORKER_PROMETHEUS_PORT and pass it to PushTaskWorker so the worker serves taskworker_worker_occupancy for KEDA autoscaling. Unset by default, so nothing changes until ops sets it. Pull mode ignores the port, since the client only publishes occupancy from push workers. Refs STREAM-2108 Co-Authored-By: Claude Opus 5.5 --- src/launchpad/worker/config.py | 14 ++++- tests/unit/test_worker_config.py | 89 ++++++++++++++++++++++++++++++++ 2 files changed, 102 insertions(+), 1 deletion(-) diff --git a/src/launchpad/worker/config.py b/src/launchpad/worker/config.py index dfe3c8b3..39e4f1a7 100644 --- a/src/launchpad/worker/config.py +++ b/src/launchpad/worker/config.py @@ -30,6 +30,7 @@ class WorkerConfig: push_mode: bool push_timeout_sec: int wait_for_child_exit: bool + prometheus_port: int | None def get_worker_config() -> WorkerConfig: @@ -70,6 +71,12 @@ def get_worker_config() -> WorkerConfig: f"LAUNCHPAD_WORKER_PUSH_TIMEOUT_SEC must be a valid integer, got: {os.getenv('LAUNCHPAD_WORKER_PUSH_TIMEOUT_SEC', '5')}" ) + prometheus_port_str = os.getenv("LAUNCHPAD_WORKER_PROMETHEUS_PORT", "").strip() + try: + prometheus_port = int(prometheus_port_str) if prometheus_port_str else None + except ValueError: + raise ValueError(f"LAUNCHPAD_WORKER_PROMETHEUS_PORT must be a valid integer, got: {prometheus_port_str}") + return WorkerConfig( rpc_hosts=rpc_hosts, rpc_host_service=rpc_host_service, @@ -80,6 +87,7 @@ def get_worker_config() -> WorkerConfig: push_mode=push_mode, push_timeout_sec=push_timeout_sec, wait_for_child_exit=os.getenv("LAUNCHPAD_WORKER_WAIT_FOR_CHILD_EXIT", "false") == "true", + prometheus_port=prometheus_port, ) @@ -93,7 +101,7 @@ 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}, wait_for_child_exit={config.wait_for_child_exit}, " - f"health_check_file_path={config.health_check_file_path})" + f"health_check_file_path={config.health_check_file_path}, prometheus_port={config.prometheus_port})" ) if config.push_mode: assert config.rpc_host_service, ( @@ -113,9 +121,13 @@ def run_worker( health_check_file_path=config.health_check_file_path, pod_name=config.pod_name, push_task_timeout=config.push_timeout_sec, + prometheus_port=config.prometheus_port, ) else: assert config.rpc_hosts, "LAUNCHPAD_WORKER_RPC_HOST environment variable is required when push mode is disabled" + if config.prometheus_port is not None: + # taskbroker-client only publishes occupancy from PushTaskWorker. + logger.warning(f"prometheus_port={config.prometheus_port} is ignored in pull mode") worker = TaskWorker( app_module="launchpad.worker.app:app", broker_hosts=config.rpc_hosts, diff --git a/tests/unit/test_worker_config.py b/tests/unit/test_worker_config.py index 7bd1e9a7..550b85e4 100644 --- a/tests/unit/test_worker_config.py +++ b/tests/unit/test_worker_config.py @@ -9,6 +9,7 @@ DEFAULT_MAX_CHILD_TASK_COUNT, WorkerConfig, get_worker_config, + run_worker, ) @@ -32,6 +33,7 @@ def test_valid_config(self): push_mode=False, push_timeout_sec=5, wait_for_child_exit=False, + prometheus_port=None, ) def test_custom_health_check_file_path(self): @@ -125,3 +127,90 @@ def test_wait_for_child_exit_enabled(self): ): config = get_worker_config() assert config.wait_for_child_exit is True + + def test_prometheus_port(self): + with patch.dict( + os.environ, + { + "LAUNCHPAD_WORKER_RPC_HOST_SERVICE": "localhost:50051", + "LAUNCHPAD_WORKER_CONCURRENCY": "2", + "LAUNCHPAD_WORKER_PROMETHEUS_PORT": "9100", + }, + clear=True, + ): + config = get_worker_config() + assert config.prometheus_port == 9100 + + def test_invalid_prometheus_port(self): + with patch.dict( + os.environ, + { + "LAUNCHPAD_WORKER_RPC_HOST_SERVICE": "localhost:50051", + "LAUNCHPAD_WORKER_CONCURRENCY": "2", + "LAUNCHPAD_WORKER_PROMETHEUS_PORT": "not-a-port", + }, + clear=True, + ): + with pytest.raises(ValueError, match="LAUNCHPAD_WORKER_PROMETHEUS_PORT must be a valid integer"): + get_worker_config() + + +class TestRunWorker: + @pytest.fixture(autouse=True) + def _skip_bootup(self): + with ( + patch("launchpad.worker.config.initialize_sentry_sdk"), + patch("launchpad.worker.config.init_options"), + ): + yield + + def test_push_mode_passes_prometheus_port(self): + env = { + "LAUNCHPAD_WORKER_RPC_HOST_SERVICE": "task-launchpad-push-broker-grpc:50051", + "LAUNCHPAD_WORKER_CONCURRENCY": "2", + "LAUNCHPAD_WORKER_PUSH_MODE": "true", + "LAUNCHPAD_WORKER_PROMETHEUS_PORT": "9100", + } + with ( + patch.dict(os.environ, env, clear=True), + patch("launchpad.worker.config.PushTaskWorker") as push_worker, + ): + push_worker.return_value.start.return_value = 0 + with pytest.raises(SystemExit): + run_worker(processing_pool_name="launchpad-push") + + kwargs = push_worker.call_args.kwargs + assert kwargs["prometheus_port"] == 9100 + assert kwargs["processing_pool_name"] == "launchpad-push" + + def test_push_mode_prometheus_disabled_by_default(self): + env = { + "LAUNCHPAD_WORKER_RPC_HOST_SERVICE": "task-launchpad-push-broker-grpc:50051", + "LAUNCHPAD_WORKER_CONCURRENCY": "2", + "LAUNCHPAD_WORKER_PUSH_MODE": "true", + } + with ( + patch.dict(os.environ, env, clear=True), + patch("launchpad.worker.config.PushTaskWorker") as push_worker, + ): + push_worker.return_value.start.return_value = 0 + with pytest.raises(SystemExit): + run_worker() + + assert push_worker.call_args.kwargs["prometheus_port"] is None + + def test_pull_mode_ignores_prometheus_port(self): + env = { + "LAUNCHPAD_WORKER_RPC_HOST": "task-launchpad-broker-grpc-0:50051", + "LAUNCHPAD_WORKER_CONCURRENCY": "2", + "LAUNCHPAD_WORKER_PROMETHEUS_PORT": "9100", + } + with ( + patch.dict(os.environ, env, clear=True), + patch("launchpad.worker.config.TaskWorker") as pull_worker, + ): + pull_worker.return_value.start.return_value = 0 + with pytest.raises(SystemExit): + run_worker() + + assert "prometheus_port" not in pull_worker.call_args.kwargs