guest-diag: vsock pin stream (guest pagemap -> host sema_map_add)

The guest-agent exec path is too slow under driver-init load (a 15k-page
window scan took 17-35 s to complete), so GR window pins landed after the
channels' entries had been consumed and after the client's cuCtxCreate
timeout. The guest now scans /proc/<pid>/pagemap locally and streams only
delta runs over vsock (no agent on the hot path); the host listener pins
them as they arrive. The scanner writes a pidfile so restarts do not kill
the new instance.

1 files changed, 134 insertions, 0 deletions

diff --git a/tools/guest-diag/guest_pin_stream.py b/tools/guest-diag/guest_pin_stream.py
new file mode 100644
index 0000000..169d048
--- /dev/null
+++ b/tools/guest-diag/guest_pin_stream.py
@@ -0,0 +1,134 @@
+#!/usr/bin/env python3
+# guest_pin_stream.py - stream the CUDA client's push-buffer VA->PA pairs to
+# the host over vsock, continuously and without the guest agent on the hot
+# path. Runs in the guest as root (pagemap reads).
+#
+# Protocol (text lines to the host):
+#   #pid <n>            the client pid just changed (host resets its state)
+#   0x<va>,0x<pa>,<n>   n contiguous pages (va+i*0x1000 -> pa+i*0x1000)
+#
+# The host side is tools/guest-diag/pin_gr_stream.py.
+#
+# Usage: sudo python3 guest_pin_stream.py [host-cid] [port] [lo-hex] [hi-hex] [pa-max-hex]
+import os
+import socket
+import sys
+import time
+
+CID = int(sys.argv[1]) if len(sys.argv) > 1 else 2
+PORT = int(sys.argv[2]) if len(sys.argv) > 2 else 45678
+LO = int(sys.argv[3], 16) if len(sys.argv) > 3 else 0x200400000
+HI = int(sys.argv[4], 16) if len(sys.argv) > 4 else 0x206000000
+PAMAX = int(sys.argv[5], 16) if len(sys.argv) > 5 else 0x180000000
+INTERVAL = float(os.environ.get("STREAM_INTERVAL", "0.25"))
+BATCH = 4096
+
+
+def cdt_pid():
+    for d in os.listdir('/proc'):
+        if not d.isdigit():
+            continue
+        try:
+            with open('/proc/%s/comm' % d) as f:
+                if f.read().strip() == 'cdt':
+                    return int(d)
+        except OSError:
+            pass
+    return None
+
+
+def scan(pid):
+    """Present pages of the VA range as {va: pa}."""
+    fd = os.open('/proc/%d/pagemap' % pid, os.O_RDONLY)
+    out = {}
+    try:
+        va = LO & ~0xFFF
+        while va < HI:
+            n = min(BATCH, (HI - va) // 0x1000)
+            if n <= 0:
+                break
+            os.lseek(fd, (va // 4096) * 8, 0)
+            data = os.read(fd, n * 8)
+            got = len(data) // 8
+            for i in range(got):
+                e = int.from_bytes(data[i * 8:i * 8 + 8], 'little')
+                if not ((e >> 63) & 1):
+                    continue
+                pfn = e & ((1 << 55) - 1)
+                pa = pfn << 12
+                if not pa or (PAMAX and pa >= PAMAX):
+                    continue
+                out[va + i * 0x1000] = pa
+            va += got * 0x1000
+    finally:
+        os.close(fd)
+    return out
+
+
+def main():
+    with open('/tmp/guest_pin_stream.pid', 'w') as f:
+        f.write('%d\n' % os.getpid())
+    sock = None
+    seen = {}
+    last_pid = None
+    while True:
+        pid = cdt_pid()
+        if pid is None:
+            last_pid = None
+            seen = {}
+            time.sleep(INTERVAL)
+            continue
+        if pid != last_pid:
+            last_pid = pid
+            seen = {}
+            if sock is not None:
+                try:
+                    sock.sendall(("#pid %d\n" % pid).encode())
+                except OSError:
+                    sock.close()
+                    sock = None
+        try:
+            cur = scan(pid)
+        except OSError:
+            time.sleep(INTERVAL)
+            continue
+        new = [(va, pa) for va, pa in cur.items() if seen.get(va) != pa]
+        seen = cur
+        if not new:
+            time.sleep(INTERVAL)
+            continue
+        new.sort()
+        lines = []
+        va0, pa0, n = new[0][0], new[0][1], 1
+        for va, pa in new[1:]:
+            if va == va0 + n * 0x1000 and pa == pa0 + n * 0x1000:
+                n += 1
+                continue
+            lines.append('0x%x,0x%x,%d\n' % (va0, pa0, n))
+            va0, pa0, n = va, pa, 1
+        lines.append('0x%x,0x%x,%d\n' % (va0, pa0, n))
+        payload = ''.join(lines).encode()
+        if sock is None:
+            try:
+                sock = socket.socket(socket.AF_VSOCK, socket.SOCK_STREAM)
+                sock.settimeout(5)
+                sock.connect((CID, PORT))
+                sock.settimeout(None)
+                if last_pid is not None:
+                    sock.sendall(("#pid %d\n" % last_pid).encode())
+            except OSError:
+                if sock:
+                    sock.close()
+                sock = None
+                time.sleep(INTERVAL)
+                continue
+        try:
+            sock.sendall(payload)
+        except OSError:
+            sock.close()
+            sock = None
+        time.sleep(INTERVAL)
+
+
+if __name__ == '__main__':
+    main()