tools: arm UVM kprobes deterministically (modprobe+wait+verify) and retry pushbuffer pins

Root cause of the cudaMalloc/UVM-init stall on our bring-up path:
guest_kprobes.sh ran while nvidia_uvm was not yet loaded (it loads lazily
with the first CUDA client), so every probe add failed silently (symbols
not in kallsyms) and the host never got the #sem -> pushbuffer-pin
pipeline. The CE push page (0x120021db8) then never translated and the
guest spun forever at DEVICE.

- guest_kprobes.sh: modprobe nvidia_uvm, wait (bounded 30 s) for the
  uvmt symbol, install, and print KP_ARMED/KP_FAIL for bring-up checks.
- semfix_watch.pin_pushbuffer: retry transient/partial gva2gpa misses
  in-call (bounded passes with a short sleep) instead of aborting on the
  first all-miss block; report unresolved page counts.

Validated: 2/2 cold boots (vmkeep6/7) arm the probes and run
victory_mail green in 15-20 s (MAIL + VADD bad=0 checksum=12204516 +
STATUS OK, oracle-matched); pre-fix boots with unarmed probes stalled
4/4.

2 files changed, 61 insertions, 30 deletions

diff --git a/tools/guest-diag/guest_kprobes.sh b/tools/guest-diag/guest_kprobes.sh
index 75b2ba3..cf34833 100644
--- a/tools/guest-diag/guest_kprobes.sh
+++ b/tools/guest-diag/guest_kprobes.sh
@@ -22,6 +22,19 @@ if [ ! -d "$TR/events" ]; then
         mount -t debugfs debugfs "$TR" 2>/dev/null || true
 fi
 
+# The UVM symbols only exist after nvidia_uvm is loaded. It loads lazily with
+# the first CUDA client, but the probes must be armed BEFORE that client's
+# first CE push - otherwise no #sem reaches the host, the push page never
+# translates and the client stalls in cudaMalloc forever. modprobe it and
+# wait (bounded) for the tracking-symbol to appear.
+modprobe nvidia_uvm 2>/dev/null || true
+i=0
+while [ $i -lt 30 ]; do
+    grep -qw uvm_gpu_tracking_semaphore_update_completed_value /proc/kallsyms 2>/dev/null && break
+    i=$((i+1))
+    sleep 1
+done
+
 # An enabled kprobe event cannot be deleted (EBUSY); disable the group first.
 echo 0 > "$TR/events/kprobes/enable" 2>/dev/null || true
 
@@ -81,3 +94,11 @@ echo 65536 > "$TR/buffer_size_kb" 2>/dev/null || true
 echo > "$TR/trace" 2>/dev/null || true
 echo "--- installed probes ---"
 grep -E 'uvmt|uvch|uvbp|uvpbf|uvtw|uvsv' "$TR/kprobe_events" 2>/dev/null || echo "(none)"
+
+# Arming check: the host recovery pipeline (#sem -> pushbuffer pin) depends on
+# these probes; make success/failure explicit for the bring-up scripts.
+if grep -q "uvmt" "$TR/kprobe_events" 2>/dev/null; then
+    echo "KP_ARMED"
+else
+    echo "KP_FAIL uvmt missing"
+fi
diff --git a/tools/guest-diag/semfix_watch.py b/tools/guest-diag/semfix_watch.py
index 5740b42..0a2a6fb 100644
--- a/tools/guest-diag/semfix_watch.py
+++ b/tools/guest-diag/semfix_watch.py
@@ -424,48 +424,58 @@ def resolve_sem_pool(q, sem):
     return [(gv, gpa) for gv in gpu_vas]
 
 
-def pin_pushbuffer(q, gpu_vas, cpu, gpa, size=0):
+def pin_pushbuffer(q, gpu_vas, cpu, gpa, size=0, retries=8):
     """Pin the pushbuffer pages (the UVM allocation is 16 x 1 MB chunks =
     16 MB = 4096 pages; the page layout is only piecewise contiguous, so
     group ascending-contiguous runs into one bulk injection each).
 
     gva2gpa of a kernel vmalloc VA intermittently fails while the sampled
     vCPU executes in user mode (page-table isolation hides the kernel half),
-    so a pass can miss a run of pages. Do not stop at the first hole: scan
-    the whole range (bailing only when essentially everything misses) and
-    return the pinned count, so the caller can retry the remainder later.
+    so a pass can miss a run of pages. Retry the missing pages in-call
+    (bounded, with a short sleep) before reporting; the caller also retries
+    incomplete pins on later loops.
     """
     pages = min(max(size >> 12, 1) if size else PUSHBUF_MAX_PAGES,
                 PUSHBUF_MAX_PAGES)
     pinned = 0
     for gv in gpu_vas:
-        miss = 0
-        i = 0
-        while i < pages:
-            g = q.gva2gpa(cpu + i * 0x1000)
-            if not g:
-                miss += 1
-                if miss >= pages:
-                    print("semfix: pushbuf va=0x%x: no page resolved (%d "
-                          "tried), wrong cpu base" % (gv, pages), flush=True)
+        pending = list(range(pages))
+        attempts = 0
+        while pending and attempts < retries:
+            attempts += 1
+            rest = []
+            i = 0
+            while i < len(pending):
+                page = pending[i]
+                g = q.gva2gpa(cpu + page * 0x1000)
+                if not g:
+                    rest.append(page)
+                    i += 1
+                    continue
+                n = 1
+                last = g
+                while i + n < len(pending):
+                    npg = pending[i + n]
+                    if npg != page + n:
+                        break
+                    gn = q.gva2gpa(cpu + npg * 0x1000)
+                    if gn != last + 0x1000:
+                        break
+                    last = gn
+                    n += 1
+                if not pin_bulk(gv + page * 0x1000, g, n * 0x1000):
+                    print("semfix: pushbuf va=0x%x: bulk pin failed at page %d"
+                          % (gv, page), flush=True)
                     return pinned
-                i += 1
-                continue
-            miss = 0
-            j = i + 1
-            last = g
-            while j < pages:
-                gj = q.gva2gpa(cpu + j * 0x1000)
-                if gj != last + 0x1000:
-                    break
-                last = gj
-                j += 1
-            n = j - i
-            if not pin_bulk(gv + i * 0x1000, g, n * 0x1000):
-                print("semfix: pushbuf va=0x%x: bulk pin failed at page %d" % (gv, i), flush=True)
-                return pinned
-            pinned += n
-            i = j
+                pinned += n
+                i += n
+            if rest:
+                time.sleep(0.1)
+            pending = rest
+        if pending:
+            print("semfix: pushbuf va=0x%x: %d/%d pages unresolved after %d "
+                  "attempts (wrong cpu base or PTI)"
+                  % (gv, len(pending), pages, attempts), flush=True)
         print("semfix: pushbuf va=0x%x cpu=0x%x gpa=0x%x size=0x%x: pinned %d page(s)"
               % (gv, cpu, gpa, size, pinned), flush=True)
     return pinned