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):