diff --git a/clients/python/pyproject.toml b/clients/python/pyproject.toml index e1ffe9cb..45d1b6be 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/examples/cli.py b/clients/python/src/examples/cli.py index 1689ef60..5e9d4e46 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 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( diff --git a/clients/python/src/taskbroker_client/worker/worker.py b/clients/python/src/taskbroker_client/worker/worker.py index 5141d0b4..794c7d8c 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, 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() diff --git a/clients/python/src/taskbroker_client/worker/workerchild.py b/clients/python/src/taskbroker_client/worker/workerchild.py index 41141189..1483b492 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], @@ -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 diff --git a/clients/python/tests/worker/test_worker.py b/clients/python/tests/worker/test_worker.py index 2c8d9149..d06f1d02 100644 --- a/clients/python/tests/worker/test_worker.py +++ b/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", diff --git a/uv.lock b/uv.lock index f779d2eb..91145946 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" },