diff --git a/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/__init__.py b/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/__init__.py
index ab406e7e0d87..de851516d3bf 100644
--- a/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/__init__.py
+++ b/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/__init__.py
@@ -1,5 +1,4 @@
 import argparse
-import atexit
 import json
 import logging
 import os
@@ -293,8 +292,6 @@ def execute(argv: list[str]) -> None:
     if not WITH_NIX_2_18:
         logger.warning("you're using Nix <2.18, some features will not work correctly")
 
-    atexit.register(cleanup_ssh)
-
     common_flags = vars(args_groups["common_flags"])
     common_build_flags = common_flags | vars(args_groups["common_build_flags"])
     build_flags = common_build_flags | vars(args_groups["classic_build_flags"])
diff --git a/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/process.py b/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/process.py
index 337fae18b2ff..f7c83f714937 100644
--- a/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/process.py
+++ b/pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/process.py
@@ -1,3 +1,4 @@
+import atexit
 import logging
 import os
 import shlex
@@ -76,6 +77,9 @@ def cleanup_ssh() -> None:
         )
 
 
+atexit.register(cleanup_ssh)
+
+
 def run_wrapper(
     args: Sequence[str | bytes | os.PathLike[str] | os.PathLike[bytes]],
     *,