feat(taskworker): Emit worker occupancy metric for autoscaling

clients/python/pyproject.toml

@@ -12,6 +12,7 @@ dependencies = [
     "grpcio>=1.67.1",
     "grpcio-health-checking>=1.67.1",
     "msgpack>=1.0.0",
+    "prometheus_client>=0.20",
     "protobuf>=5.28.3",
     "redis>=3.4.1",
     "zstandard>=0.18.0",

clients/python/src/examples/cli.py

@@ -22,12 +22,18 @@ def main() -> None:
     help="The number of tasks to generate",
     default=1,
 )
-def spawn(count: int = 1) -> None:
+@click.option(
+    "--sleep-seconds",
+    help="How long each task sleeps. Use a larger value to make occupancy observable.",
+    default=0.1,
+    type=float,
+)
+def spawn(count: int = 1, sleep_seconds: float = 0.1) -> None:
     from examples.tasks import timed_task
 
     click.echo(f"Spawning {count} tasks")
     for _ in range(0, count):
-        timed_task.delay(sleep_seconds=0.1)
+        timed_task.delay(sleep_seconds=sleep_seconds)
     click.echo("Complete")
 
 
@@ -85,8 +91,19 @@ def scheduler() -> None:
     default=50052,
     type=int,
 )
+@click.option(
+    "--prometheus-port",
+    help="Expose prometheus metrics on this port for scraping. Unset = disabled.",
+    default=None,
+    type=int,
+)
 def worker(
-    rpc_host: str, concurrency: int, push_mode: bool, batch_push_mode: bool, grpc_port: int
+    rpc_host: str,
+    concurrency: int,
+    push_mode: bool,
+    batch_push_mode: bool,
+    grpc_port: int,
+    prometheus_port: int | None,
 ) -> None:
     from taskbroker_client.worker import BatchPushTaskWorker, PushTaskWorker, TaskWorker
 
@@ -104,6 +121,7 @@ def worker(
             process_type="forkserver",
             grpc_port=grpc_port,
             update_in_batches=True,
+            prometheus_port=prometheus_port,
         )
     elif push_mode:
         worker = PushTaskWorker(
@@ -118,6 +136,7 @@ def worker(
             process_type="forkserver",
             grpc_port=grpc_port,
             push_task_timeout=5,
+            prometheus_port=prometheus_port,
         )
     else:
         worker = TaskWorker(

clients/python/src/taskbroker_client/worker/worker.py

@@ -14,6 +14,7 @@
 from typing import TYPE_CHECKING, Any, Callable, List
 
 import grpc
+import prometheus_client
 from grpc_health.v1 import health, health_pb2, health_pb2_grpc
 from sentry_protos.taskbroker.v1 import taskbroker_pb2_grpc
 from sentry_protos.taskbroker.v1.taskbroker_pb2 import (
@@ -55,6 +56,27 @@
 WORKER_SERVICE_NAME = "sentry_protos.taskbroker.v1.WorkerService"
 
 
+class WorkerPrometheusMetrics:
+    """
+    Owns the Prometheus registry, server, and metrics we expose for scraping.
+    """
+
+    def __init__(
+        self, port: int, registry: prometheus_client.CollectorRegistry | None = None
+    ) -> None:
+        self.registry = registry or prometheus_client.CollectorRegistry()
+
+        self.occupancy = prometheus_client.Gauge(
+            "taskworker_worker_occupancy",
+            "Fraction of worker child slots currently executing a task (busy / concurrency).",
+            ["processing_pool"],
+            registry=self.registry,
+        )
+
+        prometheus_client.start_http_server(port, registry=self.registry)
+        logger.info("taskworker.worker.prometheus_server_started", extra={"port": port})
+
+
 class WorkerServicer(taskbroker_pb2_grpc.WorkerServiceServicer):
     """
     gRPC servicer that receives task activations pushed from the broker
@@ -145,6 +167,7 @@ def __init__(
         update_in_batches: bool = False,
         skip_awaiting_futures: bool = True,
         warmup_timeout: float = DEFAULT_WORKER_WARMUP_TIMEOUT_SEC,
+        prometheus_port: int | None = None,
         future_checking_frequency: float = 0.1,
     ) -> None:
         app = import_app(app_module)
@@ -171,6 +194,7 @@ def __init__(
             process_type=process_type,
             update_in_batches=update_in_batches,
             skip_awaiting_futures=skip_awaiting_futures,
+            prometheus_port=prometheus_port,
             future_checking_frequency=future_checking_frequency,
         )
 
@@ -781,6 +805,7 @@ def __init__(
         process_type: str = "spawn",
         update_in_batches: bool = False,
         skip_awaiting_futures: bool = True,
+        prometheus_port: int | None = None,
         future_checking_frequency: float = 0.1,
     ) -> None:
         self._concurrency = concurrency
@@ -809,6 +834,9 @@ def __init__(
         self._children: list[BaseProcess] = []
         self._shutdown_event = self._mp_context.Event()
         self._ready_counter = self._mp_context.Value("i", 0)
+        self._busy_counter = self._mp_context.Value("i", 0)
+        self._prometheus_port = prometheus_port
+        self._prom: WorkerPrometheusMetrics | None = None
         self._result_thread: threading.Thread | None = None
         self._metrics_thread: threading.Thread | None = None
         self._spawn_children_thread: threading.Thread | None = None
@@ -842,6 +870,8 @@ def start_metrics_thread(self) -> None:
         """
         Start a thread that emits metrics on an interval.
         """
+        if self._prometheus_port is not None and self._prom is None:
+            self._prom = WorkerPrometheusMetrics(self._prometheus_port)
 
         def metrics_thread() -> None:
             tags = {
@@ -851,6 +881,18 @@ def metrics_thread() -> None:
 
             while True:
                 try:
+                    busy = max(0, min(self._busy_counter.value, self._concurrency))
+                    occupancy = busy / self._concurrency if self._concurrency else 0.0
+                    self._metrics.gauge(
+                        "taskworker.worker.occupancy",
+                        occupancy,
+                        tags=tags,
+                    )
+                    if self._prom is not None:
+                        self._prom.occupancy.labels(processing_pool=self._processing_pool_name).set(
+                            occupancy
+                        )
+
                     # 'qsize' is not implemented on all platforms, such as macOS
                     self._metrics.gauge(
                         "taskworker.child_tasks.size",
@@ -943,6 +985,7 @@ def spawn_children_thread() -> None:
                             self._skip_awaiting_futures,
                             self._future_checking_frequency,
                             self._ready_counter,
+                            self._busy_counter,
                         ),
                     )
                     process.start()

clients/python/src/taskbroker_client/worker/workerchild.py

@@ -166,6 +166,20 @@ def _log_task_retry_exhausted(
     )
 
 
+def _adjust_busy(counter: "Synchronized[int] | None", delta: int) -> None:
+    """
+    Adjust the shared count of children currently executing a task.
+
+    The parent worker pool divides this by concurrency to emit occupancy, the
+    autoscaling signal. A child that is hard-killed (e.g. OOM) mid-task leaks
+    its increment; the parent clamps occupancy to [0, 1] to bound the drift.
+    """
+    if counter is None:
+        return
+    with counter.get_lock():
+        counter.value += delta
+
+
 def child_process(
     app_module: str,
     child_tasks: queue.Queue[InflightTaskActivation],
@@ -177,6 +191,7 @@ def child_process(
     skip_awaiting_futures: bool,
     future_checking_frequency: float,
     ready_counter: "Synchronized[int] | None" = None,
+    busy_counter: "Synchronized[int] | None" = None,
 ) -> None:
     """
     The entrypoint for spawned worker children.
@@ -438,6 +453,7 @@ def check_task_future_completion(
             next_state = TASK_ACTIVATION_STATUS_FAILURE
             # Use time.time() so we can measure against activation.received_at
             execution_start_time = time.time()
+            _adjust_busy(busy_counter, 1)
             try:
                 with timeout_alarm(inflight.activation.processing_deadline_duration, handle_alarm):
                     _execute_activation(task_func, inflight.activation, app.context_hooks)
@@ -509,6 +525,8 @@ def check_task_future_completion(
                     and next_state != TASK_ACTIVATION_STATUS_RETRY
                 ):
                     _log_task_failed(inflight.activation, err, processing_pool_name)
+            finally:
+                _adjust_busy(busy_counter, -1)
 
             clear_current_task()
             processed_task_count += 1

clients/python/tests/worker/test_worker.py

@@ -919,6 +919,32 @@ def test_child_process_increments_ready_counter() -> None:
     assert ready_counter.value == 1
 
 
+def test_child_process_busy_counter_returns_to_zero() -> None:
+    todo: queue.Queue[InflightTaskActivation] = queue.Queue()
+    processed: queue.Queue[ProcessingResult] = queue.Queue()
+    shutdown = Event()
+    ctx = get_context("fork")
+    busy_counter = ctx.Value("i", 0)
+
+    todo.put(SIMPLE_TASK)
+    child_process(
+        "examples.app:app",
+        todo,
+        processed,
+        shutdown,
+        max_task_count=1,
+        processing_pool_name="test",
+        process_type="fork",
+        skip_awaiting_futures=False,
+        future_checking_frequency=0.1,
+        busy_counter=busy_counter,
+    )
+
+    # Incremented while executing, decremented in the finally afterwards, so the
+    # slot is released back to idle once the task is done.
+    assert busy_counter.value == 0
+
+
 def test_child_process_remove_start_time_kwargs() -> None:
     activation = InflightTaskActivation(
         host="localhost:50051",