1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
|
#!/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; "
"if [ -f /tmp/guest_pin_stream.pid ]; then "
"p=$(cat /tmp/guest_pin_stream.pid 2>/dev/null); "
"[ -n \"$p\" ] && grep -q guest_pin_stream "
"/proc/$p/cmdline 2>/dev/null && sudo kill $p 2>/dev/null; "
"fi; "
"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()
|
