revert: restore execution.py to origin/main (remove local revert of alive_workers)

Author: hsuhanooi

lib/zephyr/src/zephyr/execution.py

@@ -380,11 +380,6 @@ def __init__(self):
 
         # Lock for accessing coordinator state from background thread
         self._lock = threading.Lock()
-        # Signaled on every state change that _wait_for_stage cares about
-        # (shard completed, task requeued, fatal error, worker registered).
-        # Allows _wait_for_stage to wake immediately instead of spinning up
-        # to its max 1-second poll interval.
-        self._stage_event = threading.Event()
 
         actor_ctx = current_actor()
         self._name = f"{actor_ctx.group_name}"
@@ -438,14 +433,13 @@ def register_worker(self, worker_id: str, worker_handle: ActorHandle) -> None:
                 # the worker as unhealthy via heartbeat and re-registration. If we do not requeue we may silently
                 # lose tasks.
                 self._maybe_requeue_worker_task(worker_id)
-            else:
-                self._worker_handles[worker_id] = worker_handle
-                self._worker_states[worker_id] = WorkerState.READY
-                self._last_seen[worker_id] = time.monotonic()
+                return
+
+            self._worker_handles[worker_id] = worker_handle
+            self._worker_states[worker_id] = WorkerState.READY
+            self._last_seen[worker_id] = time.monotonic()
 
-                logger.info("Worker %s registered, total: %d", worker_id, len(self._worker_handles))
-        # Wake _wait_for_stage so it re-evaluates the alive-worker count.
-        self._stage_event.set()
+            logger.info("Worker %s registered, total: %d", worker_id, len(self._worker_handles))
 
     def _coordinator_loop(self) -> None:
         """Background loop for heartbeat checking and worker job monitoring."""
@@ -682,9 +676,6 @@ def report_result(
             # Zero the in-flight counters but keep the generation watermark
             # so late heartbeats from this task are rejected.
             self._worker_counters[worker_id] = CounterSnapshot.empty(counter_snapshot.generation)
-        # Wake _wait_for_stage immediately so it advances to the next stage
-        # without waiting out the polling backoff interval.
-        self._stage_event.set()
 
     def report_error(self, worker_id: str, shard_idx: int, error_info: str) -> None:
         """Worker reports a task failure. Re-queues up to MAX_SHARD_FAILURES."""
@@ -693,8 +684,6 @@ def report_error(self, worker_id: str, shard_idx: int, error_info: str) -> None:
             self._assert_in_flight_consistent(worker_id, shard_idx)
             aborted = self._record_shard_failure(worker_id, ShardFailureKind.TASK, error_info)
             self._worker_states[worker_id] = WorkerState.DEAD if aborted else WorkerState.READY
-        # Wake _wait_for_stage so it re-evaluates the alive-worker count and fatal error.
-        self._stage_event.set()
 
     def heartbeat(self, worker_id: str, counter_snapshot: CounterSnapshot | None = None) -> None:
         self._last_seen[worker_id] = time.monotonic()
@@ -760,8 +749,6 @@ def abort(self, reason: str) -> None:
             if self._fatal_error is None:
                 logger.error("Coordinator aborted: %s", reason)
                 self._fatal_error = reason
-        # Wake _wait_for_stage so it raises ZephyrWorkerError immediately.
-        self._stage_event.set()
 
     def _start_stage(self, stage_name: str, tasks: list[ShardTask], is_last_stage: bool = False) -> None:
         """Load a new stage's tasks into the queue."""
@@ -781,9 +768,6 @@ def _start_stage(self, stage_name: str, tasks: list[ShardTask], is_last_stage: b
             # Only reset in-flight worker snapshots; completed snapshots
             # accumulate across stages for full pipeline visibility.
             self._worker_counters = {}
-        # Clear after releasing the lock so _wait_for_stage can't miss
-        # a signal set between lock release and the clear call.
-        self._stage_event.clear()
 
     def _wait_for_stage(self) -> None:
         """Block until current stage completes or error occurs."""
@@ -792,6 +776,7 @@ def _wait_for_stage(self) -> None:
         start_time = time.monotonic()
         all_dead_since: float | None = None
         no_workers_timeout = self._no_workers_timeout
+        stage_done = threading.Event()
 
         while True:
             with self._lock:
@@ -835,13 +820,7 @@ def _wait_for_stage(self) -> None:
                 last_log_completed = completed
                 backoff.reset()
 
-            # Wait for a state-change signal (shard completed, error, worker registered)
-            # with a timeout as backstop for the alive-worker check and log lines.
-            # _stage_event is set by report_result/report_error/abort/register_worker,
-            # so the stage transition is detected within microseconds of the last
-            # shard completing rather than after up to the full backoff interval.
-            self._stage_event.wait(timeout=backoff.next_interval())
-            self._stage_event.clear()
+            stage_done.wait(timeout=backoff.next_interval())
 
     def _collect_results(self) -> dict[int, TaskResult]:
         """Return results for the completed stage."""