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
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
|
#!/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"))
# 0x200000000 covers the GR GPFIFO rings/UVM channel mappings below the
# old 0x200400000 start; without them the host learns bogus vidmem pairs
# for the rings and never decodes GR pushes.
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")
def log(msg):
print("[%s] %s" % (time.strftime("%H:%M:%S"), msg), flush=True)
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."""
try:
with open(SEMSTREAM, "a") as f:
f.write(line + "\n")
except OSError:
pass
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 "?"))
_mirror(text)
continue
if text.startswith("#sem"):
_mirror(text)
continue
if text.startswith("#pb"):
_mirror(text)
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 (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.
"""
todo = []
for va, pa, n in runs:
if seen.get(va) == pa and n == 1:
continue
todo.append((va, pa, n))
if not todo:
return 0
def mark(va, pa, n):
for i in range(n):
seen[va + i * 0x1000] = pa + i * 0x1000
script = []
for va, pa, n in todo:
script.append(
"printf '0x%x,0x%x,0x%x,1' > %s/sema_map_add 2>/dev/null || "
"echo 'FAIL 0x%x 0x%x 0x%x'"
% (va, pa, n * 0x1000, PARAM, va, pa, n * 0x1000))
failed = set()
try:
r = subprocess.run(["sudo", "-n", "sh", "-c", "\n".join(script)],
capture_output=True, text=True, timeout=60)
for line in r.stdout.splitlines():
m = re.match(r"FAIL (0x[0-9a-f]+) (0x[0-9a-f]+) (0x[0-9a-f]+)",
line.strip())
if m:
failed.add((int(m.group(1), 16), int(m.group(2), 16),
int(m.group(3), 16) // 0x1000))
except Exception as e:
log("batch inject failed: %s" % e)
failed = set(todo)
pages = 0
for va, pa, n in todo:
if (va, pa, n) in failed:
continue
mark(va, pa, n)
pages += n
for va, pa, n in failed:
for i in range(n):
if inject(va + i * 0x1000, pa + i * 0x1000, 0x1000):
mark(va + i * 0x1000, pa + i * 0x1000, 1)
pages += 1
if pages:
log("pinned %d page(s) in %d run(s)" % (pages, len(todo)))
for va, pa, n in todo[:3]:
log(" va=0x%x -> pa=0x%x (%d)" % (va, pa, n))
return pages
if __name__ == "__main__":
main()
|
