/proc/44797/task/44797/stack
[<0>] do_sys_poll+0x3ab/0x5b0
[<0>] __x64_sys_ppoll+0xbc/0x150
[<0>] do_syscall_64+0x33/0x40
[<0>] entry_SYSCALL_64_after_hwframe+0x44/0xa9
/proc/44797/task/44805/stack
[<0>] futex_wait_queue_me+0xb6/0x110
[<0>] futex_wait+0xe9/0x240
[<0>] do_futex+0x174/0xbf0
[<0>] __x64_sys_futex+0x146/0x1c0
[<0>] do_syscall_64+0x33/0x40
[<0>] entry_SYSCALL_64_after_hwframe+0x44/0xa9
/proc/44797/task/44810/stack
[<0>] kvm_vcpu_block+0x58/0x2f0 [kvm]
[<0>] kvm_arch_vcpu_ioctl_run+0x6c4/0x1720 [kvm]
[<0>] kvm_vcpu_ioctl+0x211/0x5a0 [kvm]
[<0>] __x64_sys_ioctl+0x83/0xb0
[<0>] do_syscall_64+0x33/0x40
[<0>] entry_SYSCALL_64_after_hwframe+0x44/0xa9
/proc/44797/stack
[<0>] do_sys_poll+0x3ab/0x5b0
[<0>] __x64_sys_ppoll+0xbc/0x150
[<0>] do_syscall_64+0x33/0x40
[<0>] entry_SYSCALL_64_after_hwframe+0x44/0xa9

diff --git a/nixos/lib/test-driver/test-driver.py b/nixos/lib/test-driver/test-driver.py
index 3ee8b3227c6..dd50696bfce 100755
--- a/nixos/lib/test-driver/test-driver.py
+++ b/nixos/lib/test-driver/test-driver.py
@@ -8,7 +8,6 @@ import queue
 import io
 import threading
 import argparse
-import atexit
 import base64
 import codecs
 import os
@@ -1128,11 +1127,13 @@ class Driver:
             for cmd in cmd(start_scripts)
         ]

-        @atexit.register
-        def clean_up() -> None:
-            with rootlog.nested("clean up"):
-                for machine in self.machines:
-                    machine.release()
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *_):
+        with rootlog.nested("clean up"):
+            for machine in self.machines:
+                machine.release()

     def subtest(self, name: str) -> Iterator[None]:
         """Group logs under a given test name"""
@@ -1307,14 +1308,13 @@ if __name__ == "__main__":
     if not args.keep_vm_state:
         rootlog.info("Machine state will be reset. To keep it, pass --keep-vm-state")

-    driver = Driver(
+    with Driver(
         args.start_scripts, args.vlans, args.testscript.read_text(), args.keep_vm_state
-    )
-
-    if args.interactive:
-        ptpython.repl.embed(driver.test_symbols(), {})
-    else:
-        tic = time.time()
-        driver.run_tests()
-        toc = time.time()
-        rootlog.info(f"test script finished in {(toc-tic):.2f}s")
+    ) as driver:
+        if args.interactive:
+            ptpython.repl.embed(driver.test_symbols(), {})
+        else:
+            tic = time.time()
+            driver.run_tests()
+            toc = time.time()
+            rootlog.info(f"test script finished in {(toc-tic):.2f}s")

    def __enter__(self) -> "Driver":
        return self

    def __exit__(self, *_: Any) -> None:
        with rootlog.nested("clean up"):
            for machine in self.machines:
                machine.release()