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
|
#!/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 re
import socket
import subprocess
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
REGION = 0x200000
SPAN = int(os.environ.get("STREAM_SPAN", "0x10000"), 16)
TRACE = ("/sys/kernel/tracing/trace"
if os.path.exists("/sys/kernel/tracing/trace")
else "/sys/kernel/debug/tracing/trace")
SEM_EVERY = int(os.environ.get("STREAM_SEM_EVERY", "4"))
def new_sems(seen):
"""New tracking-sem kptrs (uvmt lines) in the trace tail.
Read in the guest so the host does not need a guest-agent round trip for
every semaphore discovery; the trace is only tailed while a client runs.
"""
try:
out = subprocess.run(["tail", "-c", "262144", TRACE],
capture_output=True, timeout=10).stdout
except Exception:
return []
found = []
for m in re.finditer(rb"sem=0x([0-9a-fA-F]+)", out):
v = int(m.group(1), 16)
if v not in seen:
seen.add(v)
found.append(v)
return found
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 in the first SPAN bytes of every REGION, as {va: pa}.
The pushes sit at the start of each per-channel allocation (2/4 MB
aligned); scanning only the allocation heads keeps the per-loop page
table work small, which is what makes the stream fast while the client
is mapping and the guest is loaded.
"""
fd = os.open('/proc/%d/pagemap' % pid, os.O_RDONLY)
out = {}
try:
base = LO & ~(REGION - 1)
while base < HI:
va = max(base, LO) & ~0xFFF
end = min(base + SPAN, HI)
while va < end:
n = min(BATCH, (end - 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
base += REGION
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 = {}
seen_sems = set()
last_pid = None
last_log = 0.0
loop = 0
while True:
pid = cdt_pid()
loop += 1
if pid is None:
last_pid = None
seen = {}
seen_sems = set()
time.sleep(INTERVAL)
continue
if pid != last_pid:
last_pid = pid
seen = {}
seen_sems = set()
if sock is not None:
try:
sock.sendall(("#pid %d\n" % pid).encode())
except OSError:
sock.close()
sock = None
t0 = time.time()
try:
cur = scan(pid)
except OSError:
time.sleep(INTERVAL)
continue
dt = time.time() - t0
new = [(va, pa) for va, pa in cur.items() if seen.get(va) != pa]
seen = cur
lines = []
if loop % SEM_EVERY == 0:
for sem in new_sems(seen_sems):
lines.append('#sem 0x%x\n' % sem)
now = time.time()
if now - last_log > 1.0:
last_log = now
try:
with open('/tmp/guest_pin_stream.log', 'a') as f:
f.write('%s pid=%s present=%d new=%d sem=%d scan=%.3fs\n'
% (time.strftime('%T'), pid, len(cur), len(new),
len(seen_sems), dt))
except OSError:
pass
if not new and not lines:
time.sleep(INTERVAL)
continue
if new:
new.sort()
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()
|
