tools: kill-cascade purge + stall watchdogs + monotonic client identity (tools-fix c8fb5377)

Dead-pid purge via in-band #nopid (zombie-aware probe), monotonic
(pid,starttime) client identity so trace-ring replays can never reset an
active client (fixes the m178 mid-run reset regression avenue), pre-capture
and post-capture stall watchdogs that re-assert pins and escalate loudly,
heartbeat threads for both tools, supervisor STALE 75s/GRACE 120s.

1 files changed, 201 insertions, 13 deletions

diff --git a/tools/guest-diag/pin_gr_stream.py b/tools/guest-diag/pin_gr_stream.py
index 0072b20..2f78f2a 100644
--- a/tools/guest-diag/pin_gr_stream.py
+++ b/tools/guest-diag/pin_gr_stream.py
@@ -32,8 +32,20 @@ LO = os.environ.get("PINWIN_LO", "0x200000000")
 HI = os.environ.get("PINWIN_HI", "0x206000000")
 PAMAX = os.environ.get("PINWIN_PA_MAX", "0x180000000")
 SEMSTREAM = os.environ.get("SEM_STREAM_FILE", "/tmp/opencode/sem_stream.txt")
-# Liveness file checked by pin_supervisor.sh (mtime must stay < 30 s).
+# Liveness file checked by pin_supervisor.sh (mtime must stay < 75 s; the
+# supervisor grants a 120 s startup grace on top).
 HEARTBEAT_FILE = os.environ.get("PIN_GR_HEARTBEAT", "/tmp/opencode/pin_gr_heartbeat")
+# Inter-tool re-assert request (written by semfix_uvmt.py's stall watchdog):
+# when its mtime moves, re-assert every pair this tool has ever pinned.
+REASSERT_REQ = os.environ.get("PIN_REASSERT_FILE",
+                              "/tmp/opencode/pin_reassert_request")
+# Stall watchdog: a known client with no pin/data activity for this long gets
+# an automatic recovery (full re-assert + demand of the fixed counter/pool
+# VAs). Pre-capture (no pin yet for this client) uses the tighter limit.
+STALL_PRE_S = float(os.environ.get("PIN_STALL_PRE", "60"))
+STALL_POST_S = float(os.environ.get("PIN_STALL_POST", "120"))
+STALL_RETRY_S = float(os.environ.get("PIN_STALL_RETRY", "30"))
+STALL_ESCALATE_S = float(os.environ.get("PIN_STALL_ESCALATE", "180"))
 
 
 def log(msg):
@@ -48,6 +60,16 @@ def _heartbeat():
         pass
 
 
+def _hb_loop():
+    """Refresh the heartbeat on a fixed cadence, independent of data-path
+    activity: recv can be legally idle for minutes between client phases and
+    the supervisor used to read that silence as a hung tool (r198 restarts at
+    34 s staleness while injection bursts were running)."""
+    while True:
+        _heartbeat()
+        time.sleep(5)
+
+
 def _mirror(line):
     """Append guest control lines (#pid/#sem) to the shared host file so the
     single-QMP watcher can consume them without guest-agent round trips."""
@@ -67,6 +89,30 @@ def vexec(cmd, timeout=40):
         return ""
 
 
+def _scanner_alive():
+    """True when the in-guest scanner process still exists.
+
+    install_and_start() pkill's and restarts the scanner, which resets every
+    in-guest latch; probing first avoids killing a healthy scanner that is
+    merely silent because no client exists yet (r198 pre-client churn)."""
+    out = vexec("sudo sh -c 'kill -0 $(cat /tmp/guest_pin_stream.pid) "
+                "2>/dev/null && echo ALIVE'")
+    return "ALIVE" in (out or "")
+
+
+def _demand_send(vas):
+    """Append VA hints to the guest scanner's demand file (best-effort).
+
+    Same channel pin_gr_stream._demand_loop uses for host-dmesg stalls; the
+    stall watchdog uses it to ask the scanner to re-resolve the fixed
+    pre-capture VAs (counter page / UVM pool) by content/pagemap."""
+    lines = "".join("0x%x\n" % va for va in vas)
+    try:
+        vexec("printf '%%s' '%s' >> /tmp/guest_want_vas.txt" % lines, 30)
+    except Exception:
+        pass
+
+
 def _start_cmd():
     with open(GUEST_SCRIPT) as f:
         b64 = subprocess.run(["base64", "-w0"], stdin=f,
@@ -187,11 +233,19 @@ def _demand_loop():
 
 
 _DEMAND_STARTED = False
+_HB_STARTED = False
+_REASSERT_REQ_SEEN = 0.0
 
 
 def main():
-    global _DEMAND_STARTED
+    global _DEMAND_STARTED, _HB_STARTED
     _heartbeat()
+    if not _HB_STARTED:
+        _HB_STARTED = True
+        try:
+            threading.Thread(target=_hb_loop, daemon=True).start()
+        except Exception as e:
+            log("WARN: heartbeat thread failed to start: %s" % e)
     if not install_and_start():
         log("WARN: could not start guest scanner (will retry via vsock wait)")
     if not _DEMAND_STARTED:
@@ -211,6 +265,25 @@ def main():
     pin_total = 0
     last_conn = time.time()
     last_reinstall = 0.0
+    # Stall-watchdog state for the connected client. client_pid is the
+    # scanner-reported pid (None until a client appears); last_activity marks
+    # the last pin/data/control event; pins_since_client counts pages pinned
+    # for the current client (0 = pre-capture).
+    client_pid = None
+    client_start = None
+    last_activity = time.time()
+    pins_since_client = 0
+    stall_since = None
+    recover_next = 0.0
+    escalated = False
+    no_conn_warn_next = 0.0
+
+    def _activity():
+        nonlocal last_activity, stall_since, escalated
+        last_activity = time.time()
+        stall_since = None
+        escalated = False
+
     while True:
         try:
             conn, _ = srv.accept()
@@ -218,9 +291,20 @@ def main():
             now = time.time()
             # No guest scanner connection for 30 s: it died or never started
             # (a boot without it streams no pins and every client then hangs).
-            # Re-run install_and_start() at most once per 15 s.
-            if now - last_conn > 30.0 and now - last_reinstall >= 15.0:
+            # Probe the scanner first instead of reinstalling blindly:
+            # install_and_start() pkill's a live scanner and resets all
+            # in-guest latches, and pre-client silence is normal (r198 churn).
+            if now - last_conn > 30.0 and now - last_reinstall >= 60.0:
                 last_reinstall = now
+                if _scanner_alive():
+                    log("no guest connection for %.0fs but scanner alive; "
+                        "not reinstalling" % (now - last_conn))
+                    if now - last_conn > 180.0 and now >= no_conn_warn_next:
+                        no_conn_warn_next = now + 60.0
+                        log("STALL WATCHDOG: no guest scanner connection for "
+                            "%.0fs (scanner process alive) - pins/demand may "
+                            "be stranded" % (now - last_conn))
+                    continue
                 log("reinstall attempt: no guest connection for %.0fs"
                     % (now - last_conn))
                 try:
@@ -244,8 +328,43 @@ def main():
                     chunk = conn.recv(1 << 20)
                 except socket.timeout:
                     if runs:
-                        _flush(runs, seen)
+                        pages = _flush(runs, seen)
+                        pin_total += pages
+                        pins_since_client += pages
+                        if pages:
+                            _activity()
                         runs = []
+                    _maybe_reassert_request(seen)
+                    now = time.time()
+                    # Stall watchdog: a known client with no pin/data
+                    # activity gets a bounded, idempotent recovery (re-assert
+                    # everything + demand the fixed counter/pool VAs), then a
+                    # loud escalation if the stall persists. This is the
+                    # "first cudaMalloc / first GR capture must not stall
+                    # silently" guard. recover_next rate-limits attempts (the
+                    # loop wakes every 5 s on the recv timeout).
+                    if client_pid is not None:
+                        limit = STALL_PRE_S if pins_since_client == 0 \
+                            else STALL_POST_S
+                        idle = now - last_activity
+                        if idle >= limit and now >= recover_next:
+                            recover_next = now + STALL_RETRY_S
+                            if stall_since is None:
+                                stall_since = now
+                            log("STALL WATCHDOG: client pid %d - no pin "
+                                "activity for %.0fs (pre-capture=%s) - "
+                                "re-asserting all known pairs + demanding "
+                                "counter/pool VAs"
+                                % (client_pid, idle, pins_since_client == 0))
+                            if seen:
+                                _flush([], seen, full=True)
+                            _demand_send([0x20400f000, 0x121000000])
+                        if not escalated and stall_since is not None \
+                                and now - stall_since >= STALL_ESCALATE_S:
+                            escalated = True
+                            log("STALL WATCHDOG: recovery FAILED for pid %d "
+                                "after %.0fs - stall persists, needs manual "
+                                "attention" % (client_pid, now - stall_since))
                     continue
                 except OSError:
                     break
@@ -259,21 +378,45 @@ def main():
                     if not text:
                         continue
                     if text.startswith("#pid"):
-                        m = re.search(r"(\d+)", text)
+                        m = re.match(r"#pid\s+(\d+)(?:\s+(\d+))?", text)
                         seen = {}
-                        log("guest reports cdt pid %s"
-                            % (m.group(1) if m else "?"))
+                        client_pid = int(m.group(1)) if m else None
+                        client_start = (int(m.group(2)) if m and m.group(2)
+                                        else None)
+                        pins_since_client = 0
+                        _activity()
+                        log("guest reports cdt pid %s%s"
+                            % (client_pid if client_pid is not None else "?",
+                               (" start=%s" % client_start
+                                if client_start is not None else "")))
+                        _mirror(text)
+                        continue
+                    if text.startswith("#nopid"):
+                        # In-band client death: the guest scanner tells us the
+                        # client is gone (with its identity). Drop the per-pair
+                        # dedupe map so the next client's pairs are injected
+                        # unconditionally; semfix_uvmt consumes the mirrored
+                        # line and purges its per-client state too.
+                        client_pid = None
+                        client_start = None
+                        seen = {}
+                        pins_since_client = 0
+                        _activity()
+                        log("guest reports client gone (%s)" % text)
                         _mirror(text)
                         continue
                     if text.startswith("#sem"):
                         _mirror(text)
+                        _activity()
                         continue
                     if text.startswith("#pb"):
                         _mirror(text)
+                        _activity()
                         continue
                     if text.startswith("#put"):
                         log(text)
                         _mirror(text)
+                        _activity()
                         continue
                     if text.startswith("#scrub"):
                         m = re.search(r"0x([0-9a-fA-F]+),0x([0-9a-fA-F]+)",
@@ -290,7 +433,11 @@ def main():
                                      int(m.group(2), 16),
                                      int(m.group(3)) if m.group(3) else 1))
                 if runs:
-                    pin_total += _flush(runs, seen)
+                    pages = _flush(runs, seen)
+                    pin_total += pages
+                    pins_since_client += pages
+                    if pages:
+                        _activity()
                     runs = []
         except Exception as e:
             # A transient error on one connection must not kill the process:
@@ -303,15 +450,52 @@ def main():
             pass
 
 
-def _flush(runs, seen):
+def _maybe_reassert_request(seen):
+    """Honour a full-re-assert request written by semfix_uvmt's watchdog."""
+    global _REASSERT_REQ_SEEN
+    try:
+        m = os.stat(REASSERT_REQ).st_mtime
+    except OSError:
+        return
+    if m > _REASSERT_REQ_SEEN:
+        _REASSERT_REQ_SEEN = m
+        if seen:
+            log("semfix requested a full pin re-assert - re-asserting %d "
+                "pair(s)" % len(seen))
+            _flush([], seen, full=True)
+
+
+def _coalesce_seen(seen):
+    """Coalesce the seen {va: pa} per-page map into (va, pa, npages) runs."""
+    items = sorted(seen.items())
+    runs = []
+    i = 0
+    while i < len(items):
+        va, pa = items[i]
+        n = 1
+        while (i + n < len(items)
+               and items[i + n][0] == va + n * 0x1000
+               and items[i + n][1] == pa + n * 0x1000):
+            n += 1
+        runs.append((va, pa, n))
+        i += n
+    return runs
+
+
+def _flush(runs, seen, full=False):
     """Inject run lines (one batched sudo pass), return number of pages pinned.
 
     The old implementation spawned one sudo process per run; with ~400 runs
     that cost ~4 s, which landed the GR pins after the client's context-init
     deadline. The batch writes all runs in a single shell, retrying only the
     ones the kernel queue rejected (ENOSPC) individually.
+
+    full=True re-asserts every pair ever pinned (coalesced), not just this
+    batch's runs: the kernel session flush wipes the pin table when a client's
+    channels close, and the stall watchdog / semfix re-assert request use this
+    to heal it without touching the guest. Writes are idempotent.
     """
-    todo = []
+    todo = {}
     for va, pa, n in runs:
         # The CE window (0x120000000..0x121000000) and the UVM semaphore
         # pool page range (0x121000000..0x122000000) are re-asserted by the
@@ -324,9 +508,13 @@ def _flush(runs, seen):
         # pool/tracker page (pending release starved, every client's first
         # seconds).
         ce_win = 0x120000000 <= va < 0x122000000
-        if not ce_win and seen.get(va) == pa and n == 1:
+        if not full and not ce_win and seen.get(va) == pa and n == 1:
             continue
-        todo.append((va, pa, n))
+        todo[(va, pa, n)] = True
+    if full:
+        for va, pa, n in _coalesce_seen(seen):
+            todo[(va, pa, n)] = True
+    todo = list(todo)
     if not todo:
         return 0