feat(worker): Expose occupancy for Prometheus scraping

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,

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