diff --git a/clients/python/src/taskbroker_client/worker/worker.py b/clients/python/src/taskbroker_client/worker/worker.py index 18e497d9..30f7c298 100644 --- a/clients/python/src/taskbroker_client/worker/worker.py +++ b/clients/python/src/taskbroker_client/worker/worker.py @@ -156,7 +156,43 @@ class TrackedChild: process: BaseProcess state: ChildState release: Event - busy: bool = False + # Time-weighted busy tracking + busy_since: float | None = None # monotonic timestamp of the currently-open busy segment + busy_accumulated: float = 0.0 # the busy seconds banked since the last occupancy flush + + def mark_busy(self, now: float) -> None: + """Open a busy segment when the child starts a task. + + `busy`/`idle` strictly alternate per child today, so a segment should + never already be open; the guard is defensive and keeps the original + start time if that invariant ever drifts. + """ + if self.busy_since is None: + self.busy_since = now + + def mark_idle(self, now: float) -> None: + """Close the open busy segment and bank its elapsed seconds. + + Guarded so an unexpected `idle` with no open segment is a no-op rather + than a crash. + """ + if self.busy_since is not None: + self.busy_accumulated += now - self.busy_since + self.busy_since = None + + def drain_busy(self, now: float) -> float: + """Return busy seconds since the last drain and reset the counter. + + Any segment still open is folded in up to `now` and left open (its + start advanced to `now`) so a task spanning multiple intervals keeps + contributing to each one. + """ + if self.busy_since is not None: + self.busy_accumulated += now - self.busy_since + self.busy_since = now + banked = self.busy_accumulated + self.busy_accumulated = 0.0 + return banked class PushTaskWorker: @@ -787,6 +823,7 @@ def __init__( self._children: Dict[UUID, TrackedChild] = {} self._exiting_children: Deque[UUID] = deque() self._children_lock = threading.Lock() + self._last_occupancy_flush_at = time.monotonic() self._shutdown_event = self._mp_context.Event() self._prometheus_port = prometheus_port self._prom: WorkerPrometheusMetrics | None = None @@ -826,29 +863,43 @@ def _emit_periodic_metrics(self) -> None: extra={"error": e, "processing_pool": self._processing_pool_name}, ) - # Count the number of children in each state and waiting for exit with self._children_lock: + now = time.monotonic() state_counts: dict[ChildState, int] = { "pending": 0, "running": 0, "exiting": 0, } + busy_time = 0.0 for child in self._children.values(): state_counts[child.state] += 1 + busy_time += child.drain_busy(now) - busy = sum(1 for child in self._children.values() if child.busy) exiting_children = len(self._exiting_children) - bounded_busy = max(0, min(busy, self._concurrency)) - occupancy = bounded_busy / self._concurrency if self._concurrency else 0.0 + elapsed = now - self._last_occupancy_flush_at + self._last_occupancy_flush_at = now + + running_count = state_counts["running"] + if running_count > 0 and elapsed > 0: + occupancy = busy_time / (elapsed * running_count) + occupancy = min(occupancy, 1.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 + ) + self._metrics.gauge( - "taskworker.worker.occupancy", - occupancy, + "taskworker.worker.concurrency", + float(self._concurrency), tags=tags, ) - if self._prom is not None: - self._prom.occupancy.labels(processing_pool=self._processing_pool_name).set(occupancy) # Emit number of children in each state for state, count in state_counts.items(): @@ -1017,13 +1068,14 @@ def spawn_children_thread() -> None: elif message.event == "exiting": self._exiting_children.append(message.child_id) - # This child is executing a task + # This child started executing a task: open a busy segment. elif message.event == "busy": - child.busy = True + child.mark_busy(time.monotonic()) - # This child isn't doing anything right now + # This child finished a task: close the open busy segment + # and bank the elapsed time. elif message.event == "idle": - child.busy = False + child.mark_idle(time.monotonic()) while True: # Compute how many children are still running diff --git a/clients/python/tests/worker/test_worker.py b/clients/python/tests/worker/test_worker.py index 9e18b1cb..30cde78e 100644 --- a/clients/python/tests/worker/test_worker.py +++ b/clients/python/tests/worker/test_worker.py @@ -47,6 +47,7 @@ PushTaskWorker, TaskWorker, TaskWorkerProcessingPool, + TrackedChild, WorkerServicer, ) from taskbroker_client.worker.workerchild import ChildMessage @@ -870,6 +871,147 @@ def test_start_does_not_serve_when_shutdown_during_warmup() -> None: fake_server.wait_for_termination.assert_not_called() +def _make_tracked_child( + state: str, + *, + busy_since: float | None = None, + busy_accumulated: float = 0.0, +) -> TrackedChild: + return TrackedChild( + process=mock.Mock(), + state=state, # type: ignore[arg-type] + release=mock.Mock(), + busy_since=busy_since, + busy_accumulated=busy_accumulated, + ) + + +def _gauge_calls(metrics: mock.Mock, name: str) -> list[Any]: + return [c for c in metrics.gauge.call_args_list if c.args[0] == name] + + +def test_emit_periodic_metrics_skips_occupancy_during_warmup() -> None: + pool = _make_result_thread_pool(_SendResultCapture(), concurrency=4) + pool._metrics = mock.Mock() + + # A freshly started pod has only pending children; none are consuming yet. + with pool._children_lock: + pool._children[uuid4()] = _make_tracked_child("pending") + pool._children[uuid4()] = _make_tracked_child("pending") + + pool._emit_periodic_metrics() + + # Occupancy must not be emitted while warming up, otherwise fresh pods + # publish misleading zeros that drag down the fleet-wide average. + assert _gauge_calls(pool._metrics, "taskworker.worker.occupancy") == [] + # Other gauges still fire so warmup stays observable. + assert _gauge_calls(pool._metrics, "taskworker.worker.children") + # Concurrency is static and emitted even before any child is warm. + concurrency_calls = _gauge_calls(pool._metrics, "taskworker.worker.concurrency") + assert len(concurrency_calls) == 1 + assert concurrency_calls[0].args[1] == pytest.approx(4.0) + + +def test_emit_periodic_metrics_time_weights_busy_over_the_interval() -> None: + # Worked example: interval [10.0, 11.0], 3 running children. + # A: 0.19s banked, idle at flush + # B: 0.30s banked + open segment since 10.70 -> +0.30 across the boundary + # C: 0.45s banked, idle at flush + # busy_time = 0.19 + 0.60 + 0.45 = 1.24 -> occupancy = 1.24 / (1.0 * 3) + pool = _make_result_thread_pool(_SendResultCapture(), concurrency=8) + pool._metrics = mock.Mock() + pool._last_occupancy_flush_at = 10.0 + + child_b = uuid4() + with pool._children_lock: + pool._children[uuid4()] = _make_tracked_child("running", busy_accumulated=0.19) + pool._children[child_b] = _make_tracked_child( + "running", busy_accumulated=0.30, busy_since=10.70 + ) + pool._children[uuid4()] = _make_tracked_child("running", busy_accumulated=0.45) + + with mock.patch("taskbroker_client.worker.worker.time.monotonic", return_value=11.0): + pool._emit_periodic_metrics() + + occupancy_calls = _gauge_calls(pool._metrics, "taskworker.worker.occupancy") + assert len(occupancy_calls) == 1 + assert occupancy_calls[0].args[1] == pytest.approx(1.24 / 3) + + # The open segment is carried into the next interval; banks are drained. + assert pool._children[child_b].busy_since == pytest.approx(11.0) + for child in pool._children.values(): + assert child.busy_accumulated == 0.0 + assert pool._last_occupancy_flush_at == pytest.approx(11.0) + + +def test_emit_periodic_metrics_divides_by_running_children() -> None: + # Two children busy for the whole 1s interval, one idle, one still warming. + pool = _make_result_thread_pool(_SendResultCapture(), concurrency=8) + pool._metrics = mock.Mock() + pool._last_occupancy_flush_at = 10.0 + + with pool._children_lock: + pool._children[uuid4()] = _make_tracked_child("running", busy_accumulated=1.0) + pool._children[uuid4()] = _make_tracked_child("running", busy_accumulated=1.0) + pool._children[uuid4()] = _make_tracked_child("running", busy_accumulated=0.0) + # Excluded from both numerator and denominator. + pool._children[uuid4()] = _make_tracked_child("pending") + + with mock.patch("taskbroker_client.worker.worker.time.monotonic", return_value=11.0): + pool._emit_periodic_metrics() + + occupancy_calls = _gauge_calls(pool._metrics, "taskworker.worker.occupancy") + assert len(occupancy_calls) == 1 + # 2 busy-child-seconds over 3 running slots for a 1s interval. + assert occupancy_calls[0].args[1] == pytest.approx(2 / 3) + + +def test_emit_periodic_metrics_clamps_occupancy_to_one() -> None: + # A draining child can still be mid-task, so busy-time can exceed the running + # capacity for the interval; occupancy must clamp to 1.0. + pool = _make_result_thread_pool(_SendResultCapture(), concurrency=4) + pool._metrics = mock.Mock() + pool._last_occupancy_flush_at = 10.0 + + with pool._children_lock: + pool._children[uuid4()] = _make_tracked_child("running", busy_accumulated=1.0) + pool._children[uuid4()] = _make_tracked_child("exiting", busy_accumulated=1.0) + + with mock.patch("taskbroker_client.worker.worker.time.monotonic", return_value=11.0): + pool._emit_periodic_metrics() + + occupancy_calls = _gauge_calls(pool._metrics, "taskworker.worker.occupancy") + assert len(occupancy_calls) == 1 + assert occupancy_calls[0].args[1] == pytest.approx(1.0) + + +def test_spawn_children_tracks_busy_and_idle_transitions() -> None: + fake_context = _FakeContext() + pool = _make_fake_context_pool(fake_context, concurrency=1) + + pool.start_spawn_children_thread() + try: + _wait_for(lambda: len(fake_context.processes) == 1) + messages = fake_context.queues[-1] + child_id = fake_context.processes[0].args[0] + + messages.put(ChildMessage(child_id, "running")) + _wait_for(lambda: pool.ready_count == 1) + + # "busy" opens a segment. + messages.put(ChildMessage(child_id, "busy")) + _wait_for(lambda: pool._children[child_id].busy_since is not None) + + # "idle" closes it and banks a positive amount of busy-time. + messages.put(ChildMessage(child_id, "idle")) + _wait_for( + lambda: pool._children[child_id].busy_since is None + and pool._children[child_id].busy_accumulated > 0 + ) + finally: + pool.shutdown() + + def test_spawn_children_counts_pending_children_toward_concurrency() -> None: fake_context = _FakeContext() pool = _make_fake_context_pool(fake_context, concurrency=2)