From 8f935f12200ae3280649cbfc12fa84e78a36d62e Mon Sep 17 00:00:00 2001 From: Enoch Tang Date: Tue, 4 Aug 2026 13:48:37 -0400 Subject: [PATCH 1/5] improve occupancy --- .../src/taskbroker_client/worker/worker.py | 17 ++++--- clients/python/tests/worker/test_worker.py | 51 +++++++++++++++++++ 2 files changed, 61 insertions(+), 7 deletions(-) diff --git a/clients/python/src/taskbroker_client/worker/worker.py b/clients/python/src/taskbroker_client/worker/worker.py index 18e497d9..0de3a1a5 100644 --- a/clients/python/src/taskbroker_client/worker/worker.py +++ b/clients/python/src/taskbroker_client/worker/worker.py @@ -842,13 +842,16 @@ def _emit_periodic_metrics(self) -> None: bounded_busy = max(0, min(busy, self._concurrency)) occupancy = bounded_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) + if state_counts["running"] > 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 + ) # Emit number of children in each state for state, count in state_counts.items(): diff --git a/clients/python/tests/worker/test_worker.py b/clients/python/tests/worker/test_worker.py index 9e18b1cb..74a427b4 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,56 @@ 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: bool) -> TrackedChild: + return TrackedChild( + process=mock.Mock(), + state=state, # type: ignore[arg-type] + release=mock.Mock(), + busy=busy, + ) + + +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", busy=False) + pool._children[uuid4()] = _make_tracked_child("pending", busy=False) + + 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") + + +def test_emit_periodic_metrics_reports_occupancy_once_warm() -> None: + pool = _make_result_thread_pool(_SendResultCapture(), concurrency=4) + pool._metrics = mock.Mock() + + with pool._children_lock: + pool._children[uuid4()] = _make_tracked_child("running", busy=True) + pool._children[uuid4()] = _make_tracked_child("running", busy=True) + pool._children[uuid4()] = _make_tracked_child("running", busy=False) + # A still-warming child does not block reporting for the warm ones. + pool._children[uuid4()] = _make_tracked_child("pending", busy=False) + + pool._emit_periodic_metrics() + + occupancy_calls = _gauge_calls(pool._metrics, "taskworker.worker.occupancy") + assert len(occupancy_calls) == 1 + # 2 busy children over the configured concurrency of 4. + assert occupancy_calls[0].args[1] == pytest.approx(0.5) + + def test_spawn_children_counts_pending_children_toward_concurrency() -> None: fake_context = _FakeContext() pool = _make_fake_context_pool(fake_context, concurrency=2) From 856f8f50b25db57fe4879b3352ed9e226d5fdebe Mon Sep 17 00:00:00 2001 From: Enoch Tang Date: Tue, 4 Aug 2026 15:10:51 -0400 Subject: [PATCH 2/5] make occupancy time weighted --- .../src/taskbroker_client/worker/worker.py | 50 ++++++-- clients/python/tests/worker/test_worker.py | 117 ++++++++++++++++-- 2 files changed, 144 insertions(+), 23 deletions(-) diff --git a/clients/python/src/taskbroker_client/worker/worker.py b/clients/python/src/taskbroker_client/worker/worker.py index 0de3a1a5..7881698d 100644 --- a/clients/python/src/taskbroker_client/worker/worker.py +++ b/clients/python/src/taskbroker_client/worker/worker.py @@ -156,7 +156,13 @@ class TrackedChild: process: BaseProcess state: ChildState release: Event - busy: bool = False + # Time-weighted busy tracking. `busy_since` is the monotonic timestamp of the + # currently-open busy segment (None while idle); `busy_accumulated` is the busy + # seconds banked since the last occupancy flush. Together they let us report the + # fraction of the interval a child spent executing, rather than a single + # instantaneous busy/idle sample. + busy_since: float | None = None + busy_accumulated: float = 0.0 class PushTaskWorker: @@ -787,6 +793,8 @@ def __init__( self._children: Dict[UUID, TrackedChild] = {} self._exiting_children: Deque[UUID] = deque() self._children_lock = threading.Lock() + # Start of the interval currently being accumulated for occupancy. + 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,7 +834,7 @@ 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 + now = time.monotonic() with self._children_lock: state_counts: dict[ChildState, int] = { "pending": 0, @@ -834,15 +842,25 @@ def _emit_periodic_metrics(self) -> None: "exiting": 0, } + busy_time = 0.0 for child in self._children.values(): state_counts[child.state] += 1 - busy = sum(1 for child in self._children.values() if child.busy) + if child.busy_since is not None: + child.busy_accumulated += now - child.busy_since + child.busy_since = now + busy_time += child.busy_accumulated + child.busy_accumulated = 0.0 + 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 - if state_counts["running"] > 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 = max(0.0, min(occupancy, 1.0)) self._metrics.gauge( "taskworker.worker.occupancy", occupancy, @@ -853,6 +871,12 @@ def _emit_periodic_metrics(self) -> None: occupancy ) + self._metrics.gauge( + "taskworker.worker.concurrency", + float(self._concurrency), + tags=tags, + ) + # Emit number of children in each state for state, count in state_counts.items(): self._metrics.gauge( @@ -1020,13 +1044,19 @@ 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. + # Guard against duplicate "busy" so we don't lose the + # original start time. elif message.event == "busy": - child.busy = True + if child.busy_since is None: + child.busy_since = 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. Guard against a stray "idle". elif message.event == "idle": - child.busy = False + if child.busy_since is not None: + child.busy_accumulated += time.monotonic() - child.busy_since + child.busy_since = None 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 74a427b4..30cde78e 100644 --- a/clients/python/tests/worker/test_worker.py +++ b/clients/python/tests/worker/test_worker.py @@ -871,12 +871,18 @@ 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: bool) -> TrackedChild: +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=busy, + busy_since=busy_since, + busy_accumulated=busy_accumulated, ) @@ -890,8 +896,8 @@ def test_emit_periodic_metrics_skips_occupancy_during_warmup() -> None: # A freshly started pod has only pending children; none are consuming yet. with pool._children_lock: - pool._children[uuid4()] = _make_tracked_child("pending", busy=False) - pool._children[uuid4()] = _make_tracked_child("pending", busy=False) + pool._children[uuid4()] = _make_tracked_child("pending") + pool._children[uuid4()] = _make_tracked_child("pending") pool._emit_periodic_metrics() @@ -900,25 +906,110 @@ def test_emit_periodic_metrics_skips_occupancy_during_warmup() -> None: 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_reports_occupancy_once_warm() -> None: +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=True) - pool._children[uuid4()] = _make_tracked_child("running", busy=True) - pool._children[uuid4()] = _make_tracked_child("running", busy=False) - # A still-warming child does not block reporting for the warm ones. - pool._children[uuid4()] = _make_tracked_child("pending", busy=False) + pool._children[uuid4()] = _make_tracked_child("running", busy_accumulated=1.0) + pool._children[uuid4()] = _make_tracked_child("exiting", busy_accumulated=1.0) - pool._emit_periodic_metrics() + 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 children over the configured concurrency of 4. - assert occupancy_calls[0].args[1] == pytest.approx(0.5) + 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: From 9e1e77729074844ff6f03cae08241ffbd55d5b6d Mon Sep 17 00:00:00 2001 From: Enoch Tang Date: Tue, 4 Aug 2026 16:27:53 -0400 Subject: [PATCH 3/5] comment --- .../python/src/taskbroker_client/worker/worker.py | 12 ++++-------- 1 file changed, 4 insertions(+), 8 deletions(-) diff --git a/clients/python/src/taskbroker_client/worker/worker.py b/clients/python/src/taskbroker_client/worker/worker.py index 7881698d..1ce22d5d 100644 --- a/clients/python/src/taskbroker_client/worker/worker.py +++ b/clients/python/src/taskbroker_client/worker/worker.py @@ -156,13 +156,9 @@ class TrackedChild: process: BaseProcess state: ChildState release: Event - # Time-weighted busy tracking. `busy_since` is the monotonic timestamp of the - # currently-open busy segment (None while idle); `busy_accumulated` is the busy - # seconds banked since the last occupancy flush. Together they let us report the - # fraction of the interval a child spent executing, rather than a single - # instantaneous busy/idle sample. - busy_since: float | None = None - busy_accumulated: float = 0.0 + # 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 class PushTaskWorker: @@ -793,7 +789,6 @@ def __init__( self._children: Dict[UUID, TrackedChild] = {} self._exiting_children: Deque[UUID] = deque() self._children_lock = threading.Lock() - # Start of the interval currently being accumulated for occupancy. self._last_occupancy_flush_at = time.monotonic() self._shutdown_event = self._mp_context.Event() self._prometheus_port = prometheus_port @@ -834,6 +829,7 @@ def _emit_periodic_metrics(self) -> None: extra={"error": e, "processing_pool": self._processing_pool_name}, ) + # Calculate time-weighted occupancy. now = time.monotonic() with self._children_lock: state_counts: dict[ChildState, int] = { From fa80d6e0df6b30f9cd03ccb1081986cdf8feca97 Mon Sep 17 00:00:00 2001 From: Enoch Tang Date: Wed, 5 Aug 2026 12:11:22 -0400 Subject: [PATCH 4/5] pull busy logic to TrackedChild --- .../src/taskbroker_client/worker/worker.py | 54 +++++++++++++------ 1 file changed, 39 insertions(+), 15 deletions(-) diff --git a/clients/python/src/taskbroker_client/worker/worker.py b/clients/python/src/taskbroker_client/worker/worker.py index 1ce22d5d..46188993 100644 --- a/clients/python/src/taskbroker_client/worker/worker.py +++ b/clients/python/src/taskbroker_client/worker/worker.py @@ -160,6 +160,40 @@ class TrackedChild: 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: _mp_context: ForkContext | SpawnContext | ForkServerContext @@ -841,12 +875,7 @@ def _emit_periodic_metrics(self) -> None: busy_time = 0.0 for child in self._children.values(): state_counts[child.state] += 1 - - if child.busy_since is not None: - child.busy_accumulated += now - child.busy_since - child.busy_since = now - busy_time += child.busy_accumulated - child.busy_accumulated = 0.0 + busy_time += child.drain_busy(now) exiting_children = len(self._exiting_children) @@ -856,7 +885,7 @@ def _emit_periodic_metrics(self) -> None: running_count = state_counts["running"] if running_count > 0 and elapsed > 0: occupancy = busy_time / (elapsed * running_count) - occupancy = max(0.0, min(occupancy, 1.0)) + occupancy = min(occupancy, 1.0) self._metrics.gauge( "taskworker.worker.occupancy", occupancy, @@ -1041,18 +1070,13 @@ def spawn_children_thread() -> None: self._exiting_children.append(message.child_id) # This child started executing a task: open a busy segment. - # Guard against duplicate "busy" so we don't lose the - # original start time. elif message.event == "busy": - if child.busy_since is None: - child.busy_since = time.monotonic() + child.mark_busy(time.monotonic()) # This child finished a task: close the open busy segment - # and bank the elapsed time. Guard against a stray "idle". + # and bank the elapsed time. elif message.event == "idle": - if child.busy_since is not None: - child.busy_accumulated += time.monotonic() - child.busy_since - child.busy_since = None + child.mark_idle(time.monotonic()) while True: # Compute how many children are still running From c865a94f29a56d6826932add88dd4c472fb0d39b Mon Sep 17 00:00:00 2001 From: Enoch Tang Date: Wed, 5 Aug 2026 12:28:42 -0400 Subject: [PATCH 5/5] read clock after acquiring lock --- clients/python/src/taskbroker_client/worker/worker.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/clients/python/src/taskbroker_client/worker/worker.py b/clients/python/src/taskbroker_client/worker/worker.py index 46188993..30f7c298 100644 --- a/clients/python/src/taskbroker_client/worker/worker.py +++ b/clients/python/src/taskbroker_client/worker/worker.py @@ -863,9 +863,8 @@ def _emit_periodic_metrics(self) -> None: extra={"error": e, "processing_pool": self._processing_pool_name}, ) - # Calculate time-weighted occupancy. - now = time.monotonic() with self._children_lock: + now = time.monotonic() state_counts: dict[ChildState, int] = { "pending": 0, "running": 0,