[iris] Add user-defined counters (MapReduce-style per-job stats) (#4085)

Add map<string, int64> counters to TaskStatus, JobStatus, and
WorkerTaskStatus in cluster.proto. Task code calls
iris.counters.increment(name, value) which writes to a JSON file in the
workdir; the worker monitor loop reads it each poll cycle and forwards
values through the heartbeat to the controller, which stores them per
task in a new counters_json DB column. get_job_status() sums counters
across all tasks. Includes migration 0013 and three unit tests.

Fixes #3995

---------

Co-authored-by: Rafal Wojdyla <ravwojdyla@gmail.com>
Co-authored-by: claude[bot] <41898282+claude[bot]@users.noreply.github.com>
Co-authored-by: Rafal Wojdyla <ravwojdyla@users.noreply.github.com>
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>

Author: 4 people

lib/iris/src/iris/cluster/controller/transitions.py

@@ -1047,7 +1047,6 @@ def _apply_single_heartbeat(
                     "UPDATE tasks SET resource_usage_proto = ? WHERE task_id = ?",
                     (usage_payload, update.task_id.to_wire()),
                 )
-
             terminal_ms: int | None = None
             started_ms: int | None = None
             task_state = prior_state

lib/iris/src/iris/cluster/worker/task_attempt.py

@@ -444,7 +444,6 @@ def to_proto(self) -> cluster_pb2.TaskStatus:
             proto.build_metrics.build_started.CopyFrom(self.build_started.to_proto())
         if self.build_finished is not None:
             proto.build_metrics.build_finished.CopyFrom(self.build_finished.to_proto())
-
         return proto
 
     def _check_cancelled(self) -> None:

lib/zephyr/src/zephyr/__init__.py

@@ -5,6 +5,7 @@
 
 import logging
 
+from zephyr import counters
 from zephyr.dataset import Dataset, ShardInfo
 from zephyr.execution import WorkerContext, ZephyrContext, zephyr_worker_ctx
 from zephyr.expr import Expr, col, lit
@@ -25,6 +26,7 @@
     "atomic_rename",
     "col",
     "compute_plan",
+    "counters",
     "lit",
     "load_file",
     "load_jsonl",

lib/zephyr/src/zephyr/counters.py

@@ -0,0 +1,49 @@
+# Copyright The Marin Authors
+# SPDX-License-Identifier: Apache-2.0
+
+"""User-defined counters for Zephyr tasks.
+
+Task code can increment named counters during execution; counters are
+aggregated across all tasks and exposed via the coordinator's ``get_counters()``
+actor method.
+
+Usage::
+
+    from zephyr import counters
+
+    counters.increment("documents_processed", 100)
+    counters.increment("validation_errors")
+
+Counter values are accumulated in-memory on each worker and sent to the
+coordinator via the heartbeat loop.
+
+Outside of a Zephyr worker context, all calls are silent no-ops.
+"""
+
+import logging
+
+from zephyr.execution import _worker_ctx_var
+
+logger = logging.getLogger(__name__)
+
+
+def increment(name: str, value: int = 1) -> None:
+    """Increment a named counter by ``value`` (default 1).
+
+    O(1) in-memory update. Thread-safe. No-op outside a Zephyr worker.
+    """
+    worker = _worker_ctx_var.get()
+    if worker is None:
+        return
+    worker.increment_counter(name, value)
+
+
+def get_counters() -> dict[str, int]:
+    """Return a snapshot of the current task's counters.
+
+    Returns an empty dict outside a Zephyr worker context.
+    """
+    worker = _worker_ctx_var.get()
+    if worker is None:
+        return {}
+    return worker.get_counter_snapshot()

lib/zephyr/src/zephyr/execution.py

@@ -293,6 +293,8 @@ class ZephyrWorkerError(RuntimeError):
 
 class WorkerContext(Protocol):
     def get_shared(self, name: str) -> Any: ...
+    def increment_counter(self, name: str, value: int = 1) -> None: ...
+    def get_counter_snapshot(self) -> dict[str, int]: ...
 
 
 _worker_ctx_var: ContextVar[ZephyrWorker | None] = ContextVar("zephyr_worker_ctx", default=None)
@@ -351,6 +353,9 @@ def __init__(self):
         self._chunk_prefix: str = ""
         self._execution_id: str = ""
         self._no_workers_timeout: float = 60.0
+        # User-defined counters: in-flight per-worker snapshots and global accumulator.
+        self._worker_counters: dict[str, dict[str, int]] = {}
+        self._global_counters: dict[str, int] = {}
 
         # Worker management state (workers self-register via register_worker)
         self._worker_handles: dict[str, ActorHandle] = {}
@@ -487,6 +492,9 @@ def _maybe_requeue_worker_task(self, worker_id: str) -> None:
             self._task_attempts[task.shard_idx] += 1
             self._task_queue.append(task)
             self._retries += 1
+        # Discard in-flight counter snapshot so it doesn't double-count when the
+        # shard is retried on another worker.
+        self._worker_counters.pop(worker_id, None)
 
     def _check_worker_heartbeats(self, timeout: float = 120.0) -> None:
         """Internal heartbeat check (called with lock held)."""
@@ -565,20 +573,27 @@ def report_result(self, worker_id: str, shard_idx: int, attempt: int, result: Ta
             self._completed_shards += 1
             self._in_flight.pop(worker_id, None)
             self._worker_states[worker_id] = WorkerState.READY
+            # Accumulate final counters for this task into the global total
+            for name, value in self._worker_counters.pop(worker_id, {}).items():
+                self._global_counters[name] = self._global_counters.get(name, 0) + value
 
     def report_error(self, worker_id: str, shard_idx: int, error_info: str) -> None:
         """Worker reports a task failure. All errors are fatal."""
         with self._lock:
             self._last_seen[worker_id] = time.monotonic()
             self._assert_in_flight_consistent(worker_id, shard_idx)
             self._in_flight.pop(worker_id, None)
+            self._worker_counters.pop(worker_id, None)
             self._fatal_error = error_info
             self._worker_states[worker_id] = WorkerState.DEAD
 
-    def heartbeat(self, worker_id: str) -> None:
-        # No lock needed: _last_seen is only read by _check_worker_heartbeats
+    def heartbeat(self, worker_id: str, counters: dict[str, int] | None = None) -> None:
+        # No lock needed for _last_seen: only read by _check_worker_heartbeats
         # (which holds the lock), and monotonic float writes are atomic on CPython.
         self._last_seen[worker_id] = time.monotonic()
+        if counters:
+            with self._lock:
+                self._worker_counters[worker_id] = counters
 
     def get_status(self) -> JobStatus:
         with self._lock:
@@ -600,6 +615,15 @@ def get_status(self) -> JobStatus:
                 },
             )
 
+    def get_counters(self) -> dict[str, int]:
+        """Return global counter totals: completed-task values plus current in-flight snapshots."""
+        with self._lock:
+            totals = dict(self._global_counters)
+            for ctrs in self._worker_counters.values():
+                for name, value in ctrs.items():
+                    totals[name] = totals.get(name, 0) + value
+            return totals
+
     def get_fatal_error(self) -> str | None:
         with self._lock:
             return self._fatal_error
@@ -628,6 +652,8 @@ def _start_stage(self, stage_name: str, tasks: list[ShardTask], is_last_stage: b
             self._task_attempts = {task.shard_idx: 0 for task in tasks}
             self._fatal_error = None
             self._is_last_stage = is_last_stage
+            self._worker_counters = {}
+            self._global_counters = {}
 
     def _wait_for_stage(self) -> None:
         """Block until current stage completes or error occurs."""
@@ -876,6 +902,9 @@ def __init__(self, coordinator_handle: ActorHandle):
         self._shutdown_event = threading.Event()
         self._chunk_prefix: str = ""
         self._execution_id: str = ""
+        self._counters: dict[str, int] = {}
+        self._counters_lock = threading.Lock()
+        self._last_reported_counters: dict[str, int] = {}
 
         # Build descriptive worker ID from actor context
         actor_ctx = current_actor()
@@ -911,6 +940,28 @@ def get_shared(self, name: str) -> Any:
             )
         return self._shared_data_cache[name]
 
+    def increment_counter(self, name: str, value: int = 1) -> None:
+        with self._counters_lock:
+            self._counters[name] = self._counters.get(name, 0) + value
+
+    def get_counter_snapshot(self) -> dict[str, int]:
+        with self._counters_lock:
+            return dict(self._counters)
+
+    def _reset_counters(self) -> None:
+        """Clear counters for a new task."""
+        with self._counters_lock:
+            self._counters.clear()
+
+    def _counters_changed(self) -> bool:
+        """Return True if counters have changed since the last heartbeat report."""
+        with self._counters_lock:
+            current = dict(self._counters)
+        if current == self._last_reported_counters:
+            return False
+        self._last_reported_counters = current
+        return True
+
     def _run_polling(self, coordinator: ActorHandle) -> None:
         """Main polling loop. Runs in a background thread started by __init__."""
         logger.info("[%s] Starting polling loop", self._worker_id)
@@ -941,8 +992,13 @@ def _heartbeat_loop(
         while not self._shutdown_event.is_set():
             try:
                 # Block on result to avoid congesting the coordinator RPC pipe
-                # with fire-and-forget heartbeats.
-                coordinator.heartbeat.remote(self._worker_id).result()
+                # with fire-and-forget heartbeats. Only send counter snapshot
+                # when values have changed.
+                snapshot = self.get_counter_snapshot() if self._counters_changed() else None
+                coordinator.heartbeat.remote(
+                    self._worker_id,
+                    snapshot,
+                ).result()
                 heartbeat_count += 1
                 consecutive_failures = 0
                 if heartbeat_count % 10 == 1:
@@ -1049,6 +1105,9 @@ def _execute_shard(self, task: ShardTask, config: dict) -> TaskResult:
         self._chunk_prefix = config["chunk_prefix"]
         self._execution_id = config["execution_id"]
 
+        # Reset counters for the new task
+        self._reset_counters()
+
         _worker_ctx_var.set(self)
 
         logger.info(

lib/zephyr/tests/test_counters.py

@@ -0,0 +1,53 @@
+# Copyright The Marin Authors
+# SPDX-License-Identifier: Apache-2.0
+
+"""Tests for Zephyr user-defined counters: worker API and heartbeat plumbing."""
+
+import threading
+
+from zephyr import counters
+from zephyr.execution import _worker_ctx_var
+
+
+class FakeWorker:
+    """Minimal WorkerContext implementation for testing counters."""
+
+    def __init__(self):
+        self._counters: dict[str, int] = {}
+        self._counters_lock = threading.Lock()
+
+    def get_shared(self, name: str):
+        raise NotImplementedError
+
+    def increment_counter(self, name: str, value: int = 1) -> None:
+        with self._counters_lock:
+            self._counters[name] = self._counters.get(name, 0) + value
+
+    def get_counter_snapshot(self) -> dict[str, int]:
+        with self._counters_lock:
+            return dict(self._counters)
+
+
+def test_counters_increment_and_snapshot():
+    """increment() accumulates in-memory; get_counter_snapshot() returns current values."""
+    worker = FakeWorker()
+    token = _worker_ctx_var.set(worker)
+    try:
+        counters.increment("docs", 10)
+        counters.increment("docs", 5)
+        counters.increment("errors", 1)
+
+        snapshot = counters.get_counters()
+        assert snapshot == {"docs": 15, "errors": 1}
+    finally:
+        _worker_ctx_var.reset(token)
+
+
+def test_counters_noop_outside_worker():
+    """increment() is a no-op when not inside a Zephyr worker context."""
+    token = _worker_ctx_var.set(None)
+    try:
+        counters.increment("anything", 999)  # should not raise
+        assert counters.get_counters() == {}
+    finally:
+        _worker_ctx_var.reset(token)