From 6d4c2dc109e169e40990310f2989f296b1a6355c Mon Sep 17 00:00:00 2001 From: Enoch Tang Date: Mon, 29 Jun 2026 14:23:31 -0400 Subject: [PATCH 1/4] Expose occupancy metric and prom server in taskworker --- clients/python/pyproject.toml | 1 + .../src/taskbroker_client/worker/worker.py | 43 +++++++++++++++++++ .../taskbroker_client/worker/workerchild.py | 18 ++++++++ clients/python/tests/worker/test_worker.py | 25 +++++++++++ 4 files changed, 87 insertions(+) diff --git a/clients/python/pyproject.toml b/clients/python/pyproject.toml index ccc42a90..0c0b0067 100644 --- a/clients/python/pyproject.toml +++ b/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", diff --git a/clients/python/src/taskbroker_client/worker/worker.py b/clients/python/src/taskbroker_client/worker/worker.py index 690b6514..3c34cac4 100644 --- a/clients/python/src/taskbroker_client/worker/worker.py +++ b/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, ) -> None: app = import_app(app_module) @@ -170,6 +193,7 @@ def __init__( process_type=process_type, update_in_batches=update_in_batches, skip_awaiting_futures=skip_awaiting_futures, + prometheus_port=prometheus_port, ) logger.info("Running in PUSH mode") @@ -777,6 +801,7 @@ def __init__( process_type: str = "spawn", update_in_batches: bool = False, skip_awaiting_futures: bool = True, + prometheus_port: int | None = None, ) -> None: self._concurrency = concurrency self._processing_pool_name = processing_pool_name or "unknown" @@ -803,6 +828,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 @@ -836,6 +864,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 = { @@ -857,6 +887,18 @@ def metrics_thread() -> None: float(self._processed_tasks.qsize()), tags=tags, ) + + 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 + ) except Exception as e: logger.debug( "taskworker.worker.queue_gauges.error", @@ -936,6 +978,7 @@ def spawn_children_thread() -> None: self._process_type, self._skip_awaiting_futures, self._ready_counter, + self._busy_counter, ), ) process.start() diff --git a/clients/python/src/taskbroker_client/worker/workerchild.py b/clients/python/src/taskbroker_client/worker/workerchild.py index 5395920c..2f952c1d 100644 --- a/clients/python/src/taskbroker_client/worker/workerchild.py +++ b/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], @@ -176,6 +190,7 @@ def child_process( process_type: str, skip_awaiting_futures: bool, ready_counter: "Synchronized[int] | None" = None, + busy_counter: "Synchronized[int] | None" = None, ) -> None: """ The entrypoint for spawned worker children. @@ -429,6 +444,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) @@ -500,6 +516,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 diff --git a/clients/python/tests/worker/test_worker.py b/clients/python/tests/worker/test_worker.py index bcf6186a..68e2ee04 100644 --- a/clients/python/tests/worker/test_worker.py +++ b/clients/python/tests/worker/test_worker.py @@ -917,6 +917,31 @@ 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, + 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", From bcf91cdb5a1ec10cc2f26f3b89f18861187210b6 Mon Sep 17 00:00:00 2001 From: Enoch Tang Date: Mon, 29 Jun 2026 15:48:48 -0400 Subject: [PATCH 2/4] fix --- .../scripts/occupancy_prometheus_demo.py | 107 ++++++++++++++++++ clients/python/src/examples/cli.py | 25 +++- .../src/taskbroker_client/worker/worker.py | 24 ++-- uv.lock | 10 ++ 4 files changed, 151 insertions(+), 15 deletions(-) create mode 100644 clients/python/scripts/occupancy_prometheus_demo.py diff --git a/clients/python/scripts/occupancy_prometheus_demo.py b/clients/python/scripts/occupancy_prometheus_demo.py new file mode 100644 index 00000000..6f0ed804 --- /dev/null +++ b/clients/python/scripts/occupancy_prometheus_demo.py @@ -0,0 +1,107 @@ +""" +Local demo: watch worker occupancy on the Prometheus endpoint. + +Spins up a real TaskWorkerProcessingPool (no broker needed), feeds it slow +`examples.timed` tasks to saturate the child slots, and exposes occupancy on a +Prometheus /metrics endpoint. Use it to confirm the scrape endpoint works and +that occupancy tracks real load. + + python scripts/occupancy_prometheus_demo.py --port 9100 --concurrency 4 + +Then in another terminal: + + watch -n1 'curl -s localhost:9100/metrics | grep taskworker_worker_occupancy' + +You should see occupancy climb toward 1.0 while tasks are fed, then fall to 0 +once the feeder stops. +""" + +from __future__ import annotations + +import argparse +import multiprocessing as mp +import threading +import time + +import msgpack +from sentry_protos.taskbroker.v1.taskbroker_pb2 import TaskActivation + +from taskbroker_client.types import InflightTaskActivation +from taskbroker_client.worker.worker import TaskWorkerProcessingPool + + +def make_task(task_id: int, sleep_seconds: float) -> InflightTaskActivation: + return InflightTaskActivation( + host="localhost:0", + receive_timestamp=time.monotonic(), + activation=TaskActivation( + id=str(task_id), + taskname="examples.timed", + namespace="examples", + parameters_bytes=msgpack.packb( + {"args": [sleep_seconds], "kwargs": {}}, use_bin_type=True + ), + processing_deadline_duration=60, + ), + ) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--port", type=int, default=9100, help="Prometheus port") + parser.add_argument("--concurrency", type=int, default=4, help="child processes") + parser.add_argument("--task-seconds", type=float, default=2.0, help="task duration") + parser.add_argument("--feed-seconds", type=float, default=20.0, help="how long to feed load") + parser.add_argument("--drain-seconds", type=float, default=15.0, help="observe drain after") + args = parser.parse_args() + + pool = TaskWorkerProcessingPool( + app_module="examples.app:app", + mp_context=mp.get_context("fork"), + # No broker in this demo: results are simply dropped. + send_result_fn=lambda results, is_draining: None, + concurrency=args.concurrency, + processing_pool_name="local-demo", + prometheus_port=args.port, + ) + pool.start_metrics_thread() + pool.start_result_thread() + pool.start_spawn_children_thread() + + while pool.ready_count < args.concurrency: + time.sleep(0.1) + print( + f"{args.concurrency} children ready. " + f"Scrape: curl -s localhost:{args.port}/metrics | grep occupancy" + ) + + stop_feeding = threading.Event() + + def feeder() -> None: + i = 0 + while not stop_feeding.is_set(): + # push_task blocks when the child queue is full, which naturally + # keeps the slots saturated. + pool.push_task(make_task(i, args.task_seconds)) + i += 1 + + threading.Thread(target=feeder, name="feeder", daemon=True).start() + + def report(phase: str, seconds: float) -> None: + deadline = time.monotonic() + seconds + while time.monotonic() < deadline: + busy = max(0, min(pool._busy_counter.value, args.concurrency)) + occ = busy / args.concurrency if args.concurrency else 0.0 + print(f"[{phase}] busy={busy}/{args.concurrency} occupancy={occ:.2f}") + time.sleep(1) + + report("feeding", args.feed_seconds) + print("--- stopping feeder, watching drain ---") + stop_feeding.set() + report("draining", args.drain_seconds) + + pool.shutdown() + + +if __name__ == "__main__": + main() diff --git a/clients/python/src/examples/cli.py b/clients/python/src/examples/cli.py index 1689ef60..79d38a0b 100644 --- a/clients/python/src/examples/cli.py +++ b/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 occupancy on this port for Prometheus 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( diff --git a/clients/python/src/taskbroker_client/worker/worker.py b/clients/python/src/taskbroker_client/worker/worker.py index 3c34cac4..bef38eef 100644 --- a/clients/python/src/taskbroker_client/worker/worker.py +++ b/clients/python/src/taskbroker_client/worker/worker.py @@ -875,6 +875,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", @@ -887,18 +899,6 @@ def metrics_thread() -> None: float(self._processed_tasks.qsize()), tags=tags, ) - - 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 - ) except Exception as e: logger.debug( "taskworker.worker.queue_gauges.error", diff --git a/uv.lock b/uv.lock index c02f2fb1..4ec4e7a6 100644 --- a/uv.lock +++ b/uv.lock @@ -485,6 +485,14 @@ wheels = [ { url = "https://pypi.devinfra.sentry.io/wheels/pre_commit-4.2.0-py2.py3-none-any.whl", hash = "sha256:a009ca7205f1eb497d10b845e52c838a98b6cdd2102a6c8e4540e94ee75c58bd" }, ] +[[package]] +name = "prometheus-client" +version = "0.25.0" +source = { registry = "https://pypi.devinfra.sentry.io/simple" } +wheels = [ + { url = "https://pypi.devinfra.sentry.io/wheels/prometheus_client-0.25.0-py3-none-any.whl", hash = "sha256:d5aec89e349a6ec230805d0df882f3807f74fd6c1a2fa86864e3c2279059fed1" }, +] + [[package]] name = "protobuf" version = "5.29.6" @@ -775,6 +783,7 @@ dependencies = [ { name = "grpcio", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, { name = "grpcio-health-checking", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, { name = "msgpack", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, + { name = "prometheus-client", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, { name = "protobuf", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, { name = "redis", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, { name = "sentry-arroyo", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, @@ -820,6 +829,7 @@ requires-dist = [ { name = "grpcio", specifier = ">=1.67.1" }, { name = "grpcio-health-checking", specifier = ">=1.67.1" }, { name = "msgpack", specifier = ">=1.0.0" }, + { name = "prometheus-client", specifier = ">=0.20" }, { name = "protobuf", specifier = ">=5.28.3" }, { name = "redis", specifier = ">=3.4.1" }, { name = "redis-py-cluster", marker = "extra == 'cluster'", specifier = ">=2.1.0" }, From 50e72f59828e3801f8c09281402dc6e95b46df4b Mon Sep 17 00:00:00 2001 From: Enoch Tang Date: Mon, 29 Jun 2026 16:04:23 -0400 Subject: [PATCH 3/4] remove extra file --- .../scripts/occupancy_prometheus_demo.py | 107 ------------------ 1 file changed, 107 deletions(-) delete mode 100644 clients/python/scripts/occupancy_prometheus_demo.py diff --git a/clients/python/scripts/occupancy_prometheus_demo.py b/clients/python/scripts/occupancy_prometheus_demo.py deleted file mode 100644 index 6f0ed804..00000000 --- a/clients/python/scripts/occupancy_prometheus_demo.py +++ /dev/null @@ -1,107 +0,0 @@ -""" -Local demo: watch worker occupancy on the Prometheus endpoint. - -Spins up a real TaskWorkerProcessingPool (no broker needed), feeds it slow -`examples.timed` tasks to saturate the child slots, and exposes occupancy on a -Prometheus /metrics endpoint. Use it to confirm the scrape endpoint works and -that occupancy tracks real load. - - python scripts/occupancy_prometheus_demo.py --port 9100 --concurrency 4 - -Then in another terminal: - - watch -n1 'curl -s localhost:9100/metrics | grep taskworker_worker_occupancy' - -You should see occupancy climb toward 1.0 while tasks are fed, then fall to 0 -once the feeder stops. -""" - -from __future__ import annotations - -import argparse -import multiprocessing as mp -import threading -import time - -import msgpack -from sentry_protos.taskbroker.v1.taskbroker_pb2 import TaskActivation - -from taskbroker_client.types import InflightTaskActivation -from taskbroker_client.worker.worker import TaskWorkerProcessingPool - - -def make_task(task_id: int, sleep_seconds: float) -> InflightTaskActivation: - return InflightTaskActivation( - host="localhost:0", - receive_timestamp=time.monotonic(), - activation=TaskActivation( - id=str(task_id), - taskname="examples.timed", - namespace="examples", - parameters_bytes=msgpack.packb( - {"args": [sleep_seconds], "kwargs": {}}, use_bin_type=True - ), - processing_deadline_duration=60, - ), - ) - - -def main() -> None: - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--port", type=int, default=9100, help="Prometheus port") - parser.add_argument("--concurrency", type=int, default=4, help="child processes") - parser.add_argument("--task-seconds", type=float, default=2.0, help="task duration") - parser.add_argument("--feed-seconds", type=float, default=20.0, help="how long to feed load") - parser.add_argument("--drain-seconds", type=float, default=15.0, help="observe drain after") - args = parser.parse_args() - - pool = TaskWorkerProcessingPool( - app_module="examples.app:app", - mp_context=mp.get_context("fork"), - # No broker in this demo: results are simply dropped. - send_result_fn=lambda results, is_draining: None, - concurrency=args.concurrency, - processing_pool_name="local-demo", - prometheus_port=args.port, - ) - pool.start_metrics_thread() - pool.start_result_thread() - pool.start_spawn_children_thread() - - while pool.ready_count < args.concurrency: - time.sleep(0.1) - print( - f"{args.concurrency} children ready. " - f"Scrape: curl -s localhost:{args.port}/metrics | grep occupancy" - ) - - stop_feeding = threading.Event() - - def feeder() -> None: - i = 0 - while not stop_feeding.is_set(): - # push_task blocks when the child queue is full, which naturally - # keeps the slots saturated. - pool.push_task(make_task(i, args.task_seconds)) - i += 1 - - threading.Thread(target=feeder, name="feeder", daemon=True).start() - - def report(phase: str, seconds: float) -> None: - deadline = time.monotonic() + seconds - while time.monotonic() < deadline: - busy = max(0, min(pool._busy_counter.value, args.concurrency)) - occ = busy / args.concurrency if args.concurrency else 0.0 - print(f"[{phase}] busy={busy}/{args.concurrency} occupancy={occ:.2f}") - time.sleep(1) - - report("feeding", args.feed_seconds) - print("--- stopping feeder, watching drain ---") - stop_feeding.set() - report("draining", args.drain_seconds) - - pool.shutdown() - - -if __name__ == "__main__": - main() From 2afd074dc4355a4ce4a8ba1a434830fb394186e0 Mon Sep 17 00:00:00 2001 From: Enoch Tang Date: Tue, 30 Jun 2026 11:15:45 -0400 Subject: [PATCH 4/4] generalize comment --- clients/python/src/examples/cli.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/clients/python/src/examples/cli.py b/clients/python/src/examples/cli.py index 79d38a0b..5e9d4e46 100644 --- a/clients/python/src/examples/cli.py +++ b/clients/python/src/examples/cli.py @@ -93,7 +93,7 @@ def scheduler() -> None: ) @click.option( "--prometheus-port", - help="Expose occupancy on this port for Prometheus scraping. Unset = disabled.", + help="Expose prometheus metrics on this port for scraping. Unset = disabled.", default=None, type=int, )