diff options
| -rw-r--r-- | tools/guest-diag/guest_pin_stream.py | 48 |
1 files changed, 44 insertions, 4 deletions
diff --git a/tools/guest-diag/guest_pin_stream.py b/tools/guest-diag/guest_pin_stream.py index 031271c..29ef114 100644 --- a/tools/guest-diag/guest_pin_stream.py +++ b/tools/guest-diag/guest_pin_stream.py @@ -632,18 +632,42 @@ def new_sems(seen): return found +def nvidia_fds(pid): + """CUDA device fds this pid holds: 'uvm', 'gpu', or None. + + /dev/nvidia-uvm marks a compute client (the CUDA runtime opens it at + context init; nvidia-smi and persistenced do not), so a process holding + it is a real client regardless of its binary name. + """ + kind = None + try: + fddir = '/proc/%s/fd' % pid + for fd in os.listdir(fddir): + try: + tgt = os.readlink(os.path.join(fddir, fd)) + except OSError: + continue + if tgt.startswith('/dev/nvidia-uvm'): + return 'uvm' + if tgt.startswith('/dev/nvidia'): + kind = 'gpu' + except OSError: + pass + return kind + + def cdt_pid(): prefixes = tuple(p for p in os.environ.get("CLIENT_PREFIXES", "cdt,va,victory").split(",") if p) best = None + uvm = None + gpu = None for d in os.listdir('/proc'): if not d.isdigit(): continue try: with open('/proc/%s/comm' % d) as f: name = f.read().strip() - if not any(name.startswith(p) for p in prefixes): - continue # Skip zombies: a defunct previous client keeps its comm and # would otherwise shadow the live client (pagemap reads then fail # and no pins are ever streamed). @@ -652,10 +676,26 @@ def cdt_pid(): rp = st.rfind(')') if rp >= 0 and st[rp + 2:rp + 3] == 'Z': continue - best = int(d) + if any(name.startswith(p) for p in prefixes): + best = int(d) + continue + # Arbitrary CUDA binaries (any comm) are clients too: match by + # the CUDA device they hold open so streaming never depends on + # the binary's name. Prefer a UVM holder (compute client) over a + # plain /dev/nvidia0 holder (nvidia-smi, persistenced); the last + # match wins, which is the newest client when several overlap. + kind = nvidia_fds(d) + if kind == 'uvm': + uvm = int(d) + elif kind == 'gpu': + gpu = int(d) except OSError: pass - return best + if best is not None: + return best + if uvm is not None: + return uvm + return gpu def scan_window(fd, lo, hi, span, out): |