fix(python): Make worker occupancy time-weighted and warm-aware

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

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)