gw#601: generic worker-activity progress — mint/warmup report liveness, typed activity_failed

proto/worker_scheduler.proto

@@ -41,14 +41,15 @@ service WorkerScheduler {
 
 message WorkerMessage {
   oneof msg {
-    Hello         hello          = 1;
-    StateDelta    state_delta    = 2;
-    JobAccepted   job_accepted   = 3;
-    JobResult     job_result     = 4;
-    JobProgress   job_progress   = 5;
-    ModelEvent    model_event    = 6;
-    FnUnavailable fn_unavailable = 7;
-    FnDegraded    fn_degraded    = 8;
+    Hello          hello           = 1;
+    StateDelta     state_delta     = 2;
+    JobAccepted    job_accepted    = 3;
+    JobResult      job_result      = 4;
+    JobProgress    job_progress    = 5;
+    ModelEvent     model_event     = 6;
+    FnUnavailable  fn_unavailable  = 7;
+    FnDegraded     fn_degraded     = 8;
+    ActivityUpdate activity_update = 9;
   }
 }
 
@@ -575,6 +576,40 @@ enum ModelState {
   MODEL_STATE_HOST_CAPACITY_PROGRESS = 8;
 }
 
+// ---------------------------------------------------------------------------
+// Worker activity (th#929 / gw#601)
+// ---------------------------------------------------------------------------
+
+// Worker -> orchestrator: generic progress for the worker's current internal
+// job (self-mint compile, warmup, ...). Liveness contract: while an activity
+// RUNS the worker keeps seq advancing — phase transitions, step advances, or
+// watchdog heartbeats — and the hub enforces ONE generic stall rule on
+// silence; there are no per-activity wall-clock budgets. A silent death is a
+// bug by contract: terminal FAILED carries the exception. Model downloads
+// keep their richer ModelEvent byte shape; the hub folds them into the same
+// last-progress store. Additive proto3 oneof field: old hubs ignore it.
+message ActivityUpdate {
+  string kind  = 1; // "self_mint_compile" | "warmup" | ...
+  string phase = 2; // e.g. "load", "inductor_compile", "warmup_forward", "seal_publish"
+  int64 step        = 3; // k in "phase k/K"; 0 = not stepwise
+  int64 total_steps = 4; // K; 0 = unknown
+  // Monotonic per worker process across all activities. The hub records
+  // last-progress on every increase; heartbeats bump it without changing
+  // phase/step.
+  uint64 seq = 5;
+  ActivityState state = 6;
+  string error  = 7; // set when state = FAILED (the typed activity_failed terminal)
+  string detail = 8; // human-readable elaboration
+  int64 updated_at_unix_ms = 9; // worker clock, informational; hub stamps receipt
+}
+
+enum ActivityState {
+  ACTIVITY_STATE_UNSPECIFIED = 0;
+  ACTIVITY_STATE_RUNNING     = 1;
+  ACTIVITY_STATE_COMPLETED   = 2;
+  ACTIVITY_STATE_FAILED      = 3;
+}
+
 // ---------------------------------------------------------------------------
 // Availability / drain
 // ---------------------------------------------------------------------------

src/gen_worker/activity.py

@@ -0,0 +1,246 @@
+"""Generic worker-activity progress (gw#601 / th#929).
+
+The worker reports whatever internal job it is doing — self-mint compile,
+warmup — as ActivityUpdate envelopes on the existing worker->hub stream.
+Liveness contract: while an activity runs, seq keeps advancing (phase
+transitions, step advances, or watchdog heartbeats around long silent calls);
+the hub enforces ONE generic stall rule on silence. A silent death is a bug
+by contract: terminal FAILED carries the exception.
+
+Kind/phase strings are wire-shared with tensorhub
+(internal/orchestrator/grpc/worker_activity.go) — keep them identical.
+
+Without a bound transport sink (cozy-local, tests) reports land on the
+logger, which IS the local progress UI.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+import threading
+import time
+from types import TracebackType
+from typing import Awaitable, Callable, Optional
+
+from .pb import worker_scheduler_pb2 as pb
+
+logger = logging.getLogger(__name__)
+
+KIND_SELF_MINT_COMPILE = "self_mint_compile"
+KIND_WARMUP = "warmup"
+
+PHASE_LOAD = "load"
+PHASE_TRACE_GRAPH = "trace_graph"
+PHASE_INDUCTOR_COMPILE = "inductor_compile"
+PHASE_WARMUP_FORWARD = "warmup_forward"
+PHASE_SEAL_PUBLISH = "seal_publish"
+
+# Default watchdog cadence; the hub's stall rule (~10 min) tolerates many
+# missed beats.
+HEARTBEAT_INTERVAL_S = 60.0
+# Minimum evidence advance (process CPU seconds) per interval for a heartbeat:
+# a hung (blocked/deadlocked) call stops accruing CPU and the beat stops with
+# it, which is exactly the silence the hub enforces on.
+_EVIDENCE_EPS = 0.05
+
+_lock = threading.Lock()
+_seq = 0
+_sink: Optional[Callable[[pb.ActivityUpdate], None]] = None
+_current: Optional["Activity"] = None
+
+
+def bind_sink(
+    emit: Callable[["pb.WorkerMessage"], Awaitable[None]],
+    loop: asyncio.AbstractEventLoop,
+) -> None:
+    """Route reports onto the worker->hub stream: emit is the async
+    WorkerMessage sender, loop the transport loop. Thread-safe emission."""
+    def sink(update: pb.ActivityUpdate) -> None:
+        async def _ship() -> None:
+            await emit(pb.WorkerMessage(activity_update=update))
+        try:
+            running = asyncio.get_running_loop()
+        except RuntimeError:
+            running = None
+        if running is loop:
+            loop.create_task(_ship())
+        elif not loop.is_closed():
+            asyncio.run_coroutine_threadsafe(_ship(), loop)
+    global _sink
+    with _lock:
+        _sink = sink
+
+
+def _next_seq() -> int:
+    global _seq
+    with _lock:
+        _seq += 1
+        return _seq
+
+
+def _emit(update: pb.ActivityUpdate) -> None:
+    with _lock:
+        sink = _sink
+    try:
+        if sink is not None:
+            sink(update)
+        else:
+            state = pb.ActivityState.Name(update.state)
+            logger.info(
+                "[activity] %s %s %s/%s %s %s", update.kind, update.phase,
+                update.step, update.total_steps, state, update.error or update.detail,
+            )
+    except Exception:  # reporting must never break the work it reports on
+        logger.debug("activity report dropped", exc_info=True)
+
+
+class Activity:
+    """One running activity. Use begin() / the context manager `running()`."""
+
+    def __init__(self, kind: str) -> None:
+        self.kind = kind
+        self._phase = ""
+        self._step = 0
+        self._total = 0
+        self._done = False
+
+    def _report(self, state: "pb.ActivityState", error: str = "", detail: str = "") -> None:
+        _emit(pb.ActivityUpdate(
+            kind=self.kind, phase=self._phase, step=self._step,
+            total_steps=self._total, seq=_next_seq(), state=state,
+            error=error, detail=detail,
+            updated_at_unix_ms=int(time.time() * 1000),
+        ))
+
+    def phase(self, phase: str, step: int = 0, total: int = 0) -> None:
+        self._phase, self._step, self._total = phase, step, total
+        self._report(pb.ActivityState.ACTIVITY_STATE_RUNNING)
+
+    def heartbeat(self) -> None:
+        """Re-report the current phase with a fresh seq (liveness proof)."""
+        self._report(pb.ActivityState.ACTIVITY_STATE_RUNNING)
+
+    def completed(self) -> None:
+        if not self._done:
+            self._done = True
+            self._report(pb.ActivityState.ACTIVITY_STATE_COMPLETED)
+        _end(self)
+
+    def failed(self, exc: BaseException) -> None:
+        """The typed activity_failed terminal — a silent death is a bug."""
+        if not self._done:
+            self._done = True
+            self._report(
+                pb.ActivityState.ACTIVITY_STATE_FAILED,
+                error=f"{type(exc).__name__}: {exc}"[:2000],
+            )
+        _end(self)
+
+
+def begin(kind: str, phase: str = "") -> Activity:
+    global _current
+    act = Activity(kind)
+    with _lock:
+        _current = act
+    act.phase(phase) if phase else act.heartbeat()
+    return act
+
+
+def current_phase(phase: str, step: int = 0, total: int = 0) -> None:
+    """Report a phase on the current activity; no-op when none is running.
+    Setups serialize under the executor load lock, so one current is enough."""
+    with _lock:
+        act = _current
+    if act is not None and not act._done:
+        act.phase(phase, step, total)
+
+
+def _end(act: Activity) -> None:
+    global _current
+    with _lock:
+        if _current is act:
+            _current = None
+
+
+class running:
+    """Context manager: begin() on enter; COMPLETED on clean exit, FAILED
+    (carrying the exception) on raise."""
+
+    def __init__(self, kind: str, phase: str = "") -> None:
+        self._kind, self._phase = kind, phase
+        self.activity: Optional[Activity] = None
+
+    def __enter__(self) -> Activity:
+        self.activity = begin(self._kind, self._phase)
+        return self.activity
+
+    def __exit__(
+        self,
+        exc_type: Optional[type[BaseException]],
+        exc: Optional[BaseException],
+        tb: Optional[TracebackType],
+    ) -> None:
+        assert self.activity is not None
+        if exc is not None:
+            self.activity.failed(exc)
+        else:
+            self.activity.completed()
+        _end(self.activity)
+
+
+def _process_cpu_evidence() -> float:
+    return time.process_time()
+
+
+class watchdog:
+    """Bracket for a long call that may stay wire-silent (inductor compile,
+    large fuse): a background thread samples an evidence counter every
+    interval and heartbeats the activity ONLY while evidence advances. A hung
+    call stops accruing evidence, the beat stops within one interval, and the
+    hub's stall rule takes it from there.
+
+    Default evidence is process CPU seconds; pass a monotonic callable
+    (e.g. compile-wall-seconds) for calls with better signals."""
+
+    def __init__(
+        self,
+        act: Activity,
+        *,
+        interval_s: float = HEARTBEAT_INTERVAL_S,
+        evidence: Optional[Callable[[], float]] = None,
+    ) -> None:
+        self._act = act
+        self._interval = interval_s
+        self._evidence = evidence or _process_cpu_evidence
+        self._stop = threading.Event()
+        self._thread = threading.Thread(
+            target=self._run, name="activity-watchdog", daemon=True,
+        )
+
+    def _run(self) -> None:
+        try:
+            last = self._evidence()
+        except Exception:
+            last = 0.0
+        while not self._stop.wait(self._interval):
+            try:
+                now = self._evidence()
+            except Exception:
+                continue
+            if now - last >= _EVIDENCE_EPS:
+                last = now
+                self._act.heartbeat()
+
+    def __enter__(self) -> "watchdog":
+        self._thread.start()
+        return self
+
+    def __exit__(
+        self,
+        exc_type: Optional[type[BaseException]],
+        exc: Optional[BaseException],
+        tb: Optional[TracebackType],
+    ) -> None:
+        self._stop.set()
+        self._thread.join(timeout=5)

src/gen_worker/executor.py

@@ -29,6 +29,7 @@
 
 import msgspec
 
+from . import activity as activity_mod
 from .api.binding import ModelRef, wire_ref
 from .api.errors import (
     ArtifactTransferError,
@@ -2979,6 +2980,10 @@ async def ensure_setup(
         if spec.cls is None:
             return None  # function-shaped endpoint: no instance, no setup
         self.store.bind_loop()
+        try:
+            activity_mod.bind_sink(self._send, asyncio.get_running_loop())
+        except RuntimeError:
+            pass
         rec = self._class_record(spec)
         async with rec.lock:
             if rec.ready and not rec.stale:
@@ -3061,9 +3066,20 @@ async def ensure_setup(
                     spec.name,
                 )
                 await self._rollback_failed_setup(rec)
+            # gw#601: setup+warmup is one reportable activity. The watchdog
+            # heartbeats through long wire-silent calls (inductor etc.) while
+            # they provably burn CPU; a hang stops the beat within one
+            # interval and the hub's generic stall rule owns termination.
+            act = activity_mod.begin(
+                activity_mod.KIND_SELF_MINT_COMPILE if spec.compile is not None
+                else activity_mod.KIND_WARMUP,
+                activity_mod.PHASE_LOAD,
+            )
             try:
-                instance = await self._setup_locked(spec, rec, snapshots)
+                with activity_mod.watchdog(act):
+                    instance = await self._setup_locked(spec, rec, snapshots)
             except BaseException as exc:
+                act.failed(exc)
                 # Setup is a transaction: endpoint construction, tenant
                 # setup/warmup, residency registration, and compile-target
                 # publication either all reach READY or all ownership is
@@ -3096,6 +3112,7 @@ async def ensure_setup(
                     self.unavailable.pop(s.name, None)
             rec.instance = instance
             rec.ready = True
+            act.completed()
             self._clear_recovered_compile_failures(rec)
             self._on_state_change()
             return instance
@@ -3179,7 +3196,9 @@ async def _run_synthesized_warmup(
             logger.info("boot warmup skipped for %s: %s", skip.spec.name, skip.reason)
         objects = tuple({id(obj): obj for obj in proof_objects}.values())
         evidence = _WarmupEvidence()
-        for wj in jobs:
+        for wj_index, wj in enumerate(jobs, start=1):
+            activity_mod.current_phase(
+                activity_mod.PHASE_WARMUP_FORWARD, wj_index, len(jobs))
             before = {
                 id(obj): (
                     compile_cache.execution_count(obj),
@@ -3463,6 +3482,9 @@ async def run_warmup() -> Tuple[int, Dict[int, set[str]]]:
                 )
                 return evidence.count, evidence.functions_by_object
 
+            activity_mod.current_phase(
+                activity_mod.PHASE_INDUCTOR_COMPILE if inj.pending_self_mints
+                else activity_mod.PHASE_WARMUP_FORWARD)
             compile_seconds_before = (
                 compile_cache.compile_wall_seconds() if proves_inductor else 0.0)
             if inj.active_compile_artifacts:
@@ -3517,6 +3539,8 @@ async def run_warmup() -> Tuple[int, Dict[int, set[str]]]:
                         # advertising/publishing this boot's capture.
                         from . import fleet_cells as fleet_cells_mod
 
+                        activity_mod.current_phase(
+                            activity_mod.PHASE_SEAL_PUBLISH)
                         finalized = fleet_cells_mod.finalize_self_mint(
                             pipe, pending_mint)
                         inj.pending_self_mints.pop(id(pipe), None)