Add a per-worker logging thread

Author: wmoss

lib/zephyr/src/zephyr/execution.py

@@ -61,6 +61,9 @@
 # unbounded. `_check_worker_group` backstops if workers fully exhaust Iris retries.
 MAX_SHARD_FAILURES = 3
 
+ZEPHYR_STAGE_ITEM_COUNT_KEY = "zephyr/stage/{stage_name}/item_count"
+ZEPHYR_STAGE_BYTES_PROCESSED_KEY = "zephyr/stage/{stage_name}/bytes_processed"
+
 
 class ShardFailureKind(enum.StrEnum):
     """TASK failures count toward MAX_SHARD_FAILURES; INFRA failures (preemption) do not."""
@@ -496,8 +499,8 @@ def _log_status(self) -> None:
         dead = sum(1 for s in states if s in {WorkerState.FAILED, WorkerState.DEAD})
 
         totals = self.get_counters()
-        items = totals.get(f"zephyr/stage/{self._stage_name}/item_count", 0)
-        bytes_processed = totals.get(f"zephyr/stage/{self._stage_name}/bytes_processed", 0)
+        items = totals.get(ZEPHYR_STAGE_ITEM_COUNT_KEY.format(stage_name=self._stage_name), 0)
+        bytes_processed = totals.get(ZEPHYR_STAGE_BYTES_PROCESSED_KEY.format(stage_name=self._stage_name), 0)
         elapsed = time.monotonic() - (
             self._stage_monotonic_start if self._stage_monotonic_start is not None else float("inf")
         )

lib/zephyr/src/zephyr/subprocess_worker.py

@@ -19,6 +19,7 @@
 import re
 import sys
 import threading
+import time
 import traceback
 from contextlib import suppress
 from typing import Any, TypeVar
@@ -30,6 +31,8 @@
 from zephyr import counters
 from zephyr.execution import (
     CounterSnapshot,
+    ZEPHYR_STAGE_BYTES_PROCESSED_KEY,
+    ZEPHYR_STAGE_ITEM_COUNT_KEY,
     _shared_data_path,
     _worker_ctx_var,
     _write_stage_output,
@@ -57,8 +60,11 @@ def __init__(self, stage_name: str) -> None:
 
     def wrap(self, gen: Iterator[T]) -> Iterator[T]:
         for item in gen:
-            counters.increment(f"zephyr/stage/{self._stage_name}/item_count", 1)
-            counters.increment(f"zephyr/stage/{self._stage_name}/bytes_processed", sys.getsizeof(item))
+            counters.increment(ZEPHYR_STAGE_ITEM_COUNT_KEY.format(stage_name=self._stage_name), 1)
+            counters.increment(
+                ZEPHYR_STAGE_BYTES_PROCESSED_KEY.format(stage_name=self._stage_name),
+                sys.getsizeof(item),
+            )
             yield item
 
 
@@ -123,6 +129,45 @@ def _periodic_counter_writer(
             logger.warning("Failed to flush counter file to %s", counter_file, exc_info=True)
 
 
+def _periodic_status_logger(
+    stop_event: threading.Event,
+    ctx: _SubprocessWorkerContext,
+    stage_name: str,
+    execution_id: str,
+    shard_idx: int,
+    total_shards: int,
+    monotonic_start: float,
+    interval: float,
+) -> None:
+    """Log ``item_count`` / ``bytes_processed`` rates on a fixed interval (cf. coordinator ``_log_status``).
+
+    Runs in a dedicated daemon thread so logs are attributed to that thread name.
+    Reads ``ctx._counters`` the same way as the counter flusher (shallow copy).
+    """
+    item_key = ZEPHYR_STAGE_ITEM_COUNT_KEY.format(stage_name=stage_name)
+    byte_key = ZEPHYR_STAGE_BYTES_PROCESSED_KEY.format(stage_name=stage_name)
+    while not stop_event.wait(timeout=interval):
+        if sys.is_finalizing():
+            return
+        items = ctx._counters.get(item_key, 0)
+        bytes_processed = ctx._counters.get(byte_key, 0)
+        elapsed = time.monotonic() - monotonic_start
+        item_rate = items / elapsed
+        byte_rate = bytes_processed / elapsed
+        logger.info(
+            "[%s] [%s] [%s] shard %d/%d; items=%d (%.1f/s), bytes_processed=%.1fMiB (%.1fMiB/s)",
+            execution_id,
+            stage_name,
+            threading.current_thread().name,
+            shard_idx,
+            total_shards,
+            items,
+            item_rate,
+            bytes_processed / (1024 * 1024),
+            byte_rate / (1024 * 1024),
+        )
+
+
 def execute_shard(task_file: str, result_file: str) -> None:
     """Entry point for subprocess shard execution.
 
@@ -153,6 +198,7 @@ def execute_shard(task_file: str, result_file: str) -> None:
     counter_file = f"{result_file}.counters"
     stop_event = threading.Event()
     flusher: threading.Thread | None = None
+    status_logger: threading.Thread | None = None
     result_or_error: Any
     ctx: _SubprocessWorkerContext | None = None
     try:
@@ -162,6 +208,8 @@ def execute_shard(task_file: str, result_file: str) -> None:
         ctx = _SubprocessWorkerContext(chunk_prefix, execution_id)
         _worker_ctx_var.set(ctx)
 
+        shard_monotonic_start = time.monotonic()
+
         flusher = threading.Thread(
             target=_periodic_counter_writer,
             args=(stop_event, ctx, counter_file, SUBPROCESS_COUNTER_FLUSH_INTERVAL),
@@ -170,6 +218,23 @@ def execute_shard(task_file: str, result_file: str) -> None:
         )
         flusher.start()
 
+        status_logger = threading.Thread(
+            target=_periodic_status_logger,
+            args=(
+                stop_event,
+                ctx,
+                task.stage_name,
+                execution_id,
+                task.shard_idx,
+                task.total_shards,
+                shard_monotonic_start,
+                SUBPROCESS_COUNTER_FLUSH_INTERVAL,
+            ),
+            daemon=True,
+            name="zephyr-subprocess-status-logger",
+        )
+        status_logger.start()
+
         stage_ctx = StageContext(
             shard=task.shard,
             shard_idx=task.shard_idx,
@@ -206,6 +271,8 @@ def execute_shard(task_file: str, result_file: str) -> None:
         stop_event.set()
         if flusher is not None and flusher.is_alive():
             flusher.join(timeout=2.0)
+        if status_logger is not None and status_logger.is_alive():
+            status_logger.join(timeout=2.0)
 
     with open(result_file, "wb") as f:
         counters_out = dict(ctx._counters) if ctx is not None else {}

lib/zephyr/tests/subprocess_worker_test.py

@@ -6,7 +6,12 @@
 import cloudpickle
 
 import zephyr.subprocess_worker as sw
-from zephyr.execution import ListShard, ShardTask
+from zephyr.execution import (
+    ZEPHYR_STAGE_BYTES_PROCESSED_KEY,
+    ZEPHYR_STAGE_ITEM_COUNT_KEY,
+    ListShard,
+    ShardTask,
+)
 from zephyr.shuffle import MemChunk
 
 
@@ -33,5 +38,5 @@ def test_execute_shard_sets_stage_scoped_output_counters(tmp_path):
     with open(result_file, "rb") as f:
         _result_or_error, counters_out = cloudpickle.load(f)
 
-    assert counters_out[f"zephyr/stage/{stage_name}/item_count"] == 10
-    assert counters_out[f"zephyr/stage/{stage_name}/bytes_processed"] > 0
+    assert counters_out[ZEPHYR_STAGE_ITEM_COUNT_KEY.format(stage_name=stage_name)] == 10
+    assert counters_out[ZEPHYR_STAGE_BYTES_PROCESSED_KEY.format(stage_name=stage_name)] > 0