功能：完善通用云管工作台、审计与诊断能力

tools/pve-shared-gpu-hook.py

@@ -0,0 +1,124 @@
+#!/usr/bin/python3
+"""Host-specific PVE hook: re-enumerate the idle GTX 965M before guest start.
+
+Install on starhe as local:snippets/vdi-shared-gpu-reset.py for VMs 100-102.
+An incomplete preparation stays reserved until an administrator inspects it.
+"""
+import json
+import os
+from pathlib import Path
+import socket
+import subprocess
+import sys
+import time
+
+MEMBERS = (100, 101, 102)
+PCI = Path('/sys/bus/pci/devices')
+GPU = PCI / '0000:01:00.0'
+BRIDGE = PCI / '0000:00:01.0'
+STATE_DIR = Path('/run/vdi-shared-gpu')
+BOOT_ID = Path('/proc/sys/kernel/random/boot_id')
+
+
+def stopped():
+    for vmid in MEMBERS:
+        status = subprocess.check_output(
+            ['qm', 'status', str(vmid)], text=True, timeout=10).strip()
+        if status != 'status: stopped':
+            raise RuntimeError(f'VM {vmid} is not stopped; refusing GPU reset')
+
+
+def unused_gpu():
+    parent = BRIDGE.resolve(strict=True)
+    children = sorted(p.name for p in parent.glob('0000:*')
+                      if (p / 'vendor').is_file())
+    if children != [GPU.name]:
+        raise RuntimeError(f'unexpected PCI bridge devices: {children}')
+    if ((GPU / 'vendor').read_text().strip() != '0x10de'
+            or (GPU / 'device').read_text().strip() != '0x13d9'
+            or (GPU / 'driver').resolve(strict=True).name != 'vfio-pci'):
+        raise RuntimeError('unexpected GPU identity or driver')
+    group = (GPU / 'iommu_group').resolve(strict=True)
+    devices = sorted(p.name for p in (group / 'devices').iterdir())
+    if devices != [BRIDGE.name, GPU.name]:
+        raise RuntimeError(f'unexpected IOMMU group: {devices}')
+    result = subprocess.run(['fuser', '/dev/vfio/' + group.name],
+                            capture_output=True, timeout=10)
+    if result.returncode != 1:
+        raise RuntimeError('GPU group is in use or usage could not be checked')
+
+
+def reenumerate():
+    (GPU / 'power/control').write_text('on')
+    (GPU / 'remove').write_text('1')
+    time.sleep(2)
+    (BRIDGE / 'rescan').write_text('1')
+    for _ in range(20):
+        if GPU.exists() and (GPU / 'driver').exists():
+            break
+        time.sleep(.5)
+    unused_gpu()
+    (GPU / 'power/control').write_text('on')
+    if (GPU / 'power_state').read_text().strip() != 'D0':
+        raise RuntimeError('GPU did not return to D0')
+
+
+def save(state):
+    temporary = STATE_DIR / 'reservation.tmp'
+    temporary.write_text(json.dumps(state) + '\n')
+    temporary.replace(STATE_DIR / 'reservation.json')
+
+
+def handle(vmid, phase):
+    path = STATE_DIR / 'reservation.json'
+    boot = BOOT_ID.read_text().strip()
+    state = json.loads(path.read_text()) if path.exists() else None
+    if state and state.get('boot') != boot:
+        state = None
+    if phase == 'pre-start':
+        stopped()
+        unused_gpu()
+        if state and state.get('phase') != 'running':
+            raise RuntimeError('previous GPU preparation is unconfirmed; inspect '
+                               '/run/vdi-shared-gpu/reservation.json before retry')
+        state = dict(vmid=vmid, phase='preparing', boot=boot)
+        save(state)
+        try:
+            reenumerate()
+        except Exception:
+            save(dict(state, phase='failed'))
+            raise
+        save(dict(state, phase='prepared'))
+        print(f'VDI GPU: exclusive re-enumeration complete for VM {vmid}', flush=True)
+    elif phase == 'post-start':
+        if not state or state.get('vmid') != vmid or state.get('phase') != 'prepared':
+            raise RuntimeError('GPU start reservation does not match this VM')
+        save(dict(state, phase='running'))
+    elif phase == 'post-stop':
+        if state and state.get('vmid') == vmid and state.get('phase') in ('prepared', 'running'):
+            stopped()
+            unused_gpu()
+            path.unlink()
+    elif phase != 'pre-stop':
+        raise RuntimeError('unsupported hook phase')
+
+
+def main():
+    import fcntl
+
+    if os.geteuid() != 0 or socket.gethostname().split('.')[0] != 'starhe':
+        raise RuntimeError('this hook is restricted to root on starhe')
+    if len(sys.argv) != 3 or int(sys.argv[1]) not in MEMBERS:
+        raise RuntimeError('unexpected VM or hook arguments')
+    STATE_DIR.mkdir(mode=0o700, exist_ok=True)
+    with (STATE_DIR / 'lock').open('a') as lock:
+        fcntl.flock(lock, fcntl.LOCK_EX)
+        handle(int(sys.argv[1]), sys.argv[2])
+
+
+if __name__ == '__main__':
+    try:
+        main()
+    except Exception as error:
+        print(f'VDI GPU preparation blocked: {error}', file=sys.stderr, flush=True)
+        sys.exit(1)