1 files changed, 152 insertions, 0 deletions
diff --git a/tools/guest-diag/pin_gr_stream.py b/tools/guest-diag/pin_gr_stream.py
new file mode 100644
index 0000000..22c202b
--- /dev/null
+++ b/tools/guest-diag/pin_gr_stream.py
@@ -0,0 +1,152 @@
+#!/usr/bin/env python3
+# pin_gr_stream.py - pin the CUDA client's GR push-buffer pages from a vsock
+# stream produced by guest_pin_stream.py (guest truth via /proc/<pid>/pagemap).
+#
+# The guest scanner runs in the guest and streams only deltas, so the host
+# learns a page the moment it is mapped - no guest-agent round trips on the
+# critical path (the agent is only used once to install/start the scanner).
+#
+# Usage: tools/guest-diag/pin_gr_stream.py
+# Env: PINWIN_LO/PINWIN_HI/PINWIN_PA_MAX as in pin_gr_window.py,
+#      PINWIN_PORT vsock port (default 45678)
+import os
+import re
+import socket
+import subprocess
+import sys
+import time
+
+HERE = os.path.dirname(os.path.abspath(__file__))
+VMEXEC = os.path.normpath(os.path.join(HERE, "..", "..", "scripts", "vgpu_vm.py"))
+PARAM = "/sys/module/open_iov_vfio/parameters"
+GUEST_SCRIPT = os.path.join(HERE, "guest_pin_stream.py")
+
+PORT = int(os.environ.get("PINWIN_PORT", "45678"))
+LO = os.environ.get("PINWIN_LO", "0x200400000")
+HI = os.environ.get("PINWIN_HI", "0x206000000")
+PAMAX = os.environ.get("PINWIN_PA_MAX", "0x180000000")
+
+
+def log(msg):
+    print("[%s] %s" % (time.strftime("%H:%M:%S"), msg), flush=True)
+
+
+def vexec(cmd, timeout=40):
+    try:
+        r = subprocess.run([VMEXEC, "exec", "--timeout", str(timeout - 5), cmd],
+                           capture_output=True, text=True, timeout=timeout)
+        return r.stdout
+    except Exception:
+        return ""
+
+
+def install_and_start():
+    with open(GUEST_SCRIPT) as f:
+        b64 = subprocess.run(["base64", "-w0"], stdin=f,
+                             capture_output=True, text=True).stdout
+    cmd = ("echo '%s' | base64 -d > /tmp/guest_pin_stream.py; "
+           "[ -f /tmp/guest_pin_stream.pid ] && sudo kill "
+           "$(cat /tmp/guest_pin_stream.pid) 2>/dev/null; "
+           "sudo rm -f /tmp/guest_pin_stream.pid; "
+           "sudo setsid nohup python3 /tmp/guest_pin_stream.py %d %d %s %s %s "
+           ">/tmp/guest_pin_stream.log 2>&1 & sleep 2; "
+           "sudo cat /tmp/guest_pin_stream.pid 2>/dev/null"
+           % (b64, 2, PORT, LO, HI, PAMAX))
+    out = vexec(cmd, 60)
+    return any(ch.isdigit() for ch in out)
+
+
+def inject(va, pa, ln):
+    cmd = "printf '0x%x,0x%x,0x%x,1' > %s/sema_map_add" % (va, pa, ln, PARAM)
+    for _ in range(60):
+        r = subprocess.run(["sudo", "-n", "sh", "-c", cmd],
+                           capture_output=True, text=True)
+        if r.returncode == 0:
+            return True
+        time.sleep(0.05)
+    return False
+
+
+def main():
+    if not install_and_start():
+        log("WARN: could not start guest scanner (will retry via vsock wait)")
+    srv = socket.socket(socket.AF_VSOCK, socket.SOCK_STREAM)
+    srv.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+    srv.bind((socket.VMADDR_CID_ANY, PORT))
+    srv.listen(4)
+    log("listening on vsock port %d (window %s..%s)" % (PORT, LO, HI))
+
+    seen = {}
+    pin_total = 0
+    buf = b""
+    while True:
+        try:
+            conn, _ = srv.accept()
+        except OSError as e:
+            log("accept failed: %s" % e)
+            time.sleep(1)
+            continue
+        conn.settimeout(5)
+        log("guest connected")
+        buf = b""
+        runs = []
+        while True:
+            try:
+                chunk = conn.recv(1 << 20)
+            except socket.timeout:
+                if runs:
+                    _flush(runs, seen)
+                    runs = []
+                continue
+            except OSError:
+                break
+            if not chunk:
+                break
+            buf += chunk
+            while b"\n" in buf:
+                line, buf = buf.split(b"\n", 1)
+                text = line.decode(errors="replace").strip()
+                if not text:
+                    continue
+                if text.startswith("#pid"):
+                    m = re.search(r"(\d+)", text)
+                    seen = {}
+                    log("guest reports cdt pid %s" % (m.group(1) if m else "?"))
+                    continue
+                m = re.match(r"^(0x[0-9a-f]+),(0x[0-9a-f]+)(?:,(\d+))?$", text)
+                if m:
+                    runs.append((int(m.group(1), 16), int(m.group(2), 16),
+                                 int(m.group(3)) if m.group(3) else 1))
+            if runs:
+                pin_total += _flush(runs, seen)
+                runs = []
+        log("guest disconnected (pins this session: %d)" % pin_total)
+        try:
+            conn.close()
+        except Exception:
+            pass
+
+
+def _flush(runs, seen):
+    """Inject run lines, return number of pages pinned."""
+    pages = 0
+    for va, pa, n in runs:
+        if seen.get(va) == pa and n == 1:
+            continue
+        for i in range(n):
+            seen[va + i * 0x1000] = pa + i * 0x1000
+        if not inject(va, pa, n * 0x1000):
+            for i in range(n):
+                if inject(va + i * 0x1000, pa + i * 0x1000, 0x1000):
+                    pages += 1
+            continue
+        pages += n
+    if pages:
+        log("pinned %d page(s) in %d run(s)" % (pages, len(runs)))
+        for va, pa, n in runs[:3]:
+            log("  va=0x%x -> pa=0x%x (%d)" % (va, pa, n))
+    return pages
+
+
+if __name__ == "__main__":
+    main()