fix bugs when torchrun fail main process exit but remain some zombie processes keep using gpu

Author: klhhhhh

fastvideo/tests/distributed/test_ring_attention.py

@@ -10,7 +10,7 @@
 # If the current environment has known NCCL P2P / IPC issues (e.g., some Docker
 # containers), the test can be launched with NCCL socket fallback:
 #
-#   CUDA_VISIBLE_DEVICES=0,1 \
+#   CUDA_VISIBLE_DEVICES=0,1,2,3 \
 #   NCCL_CUMEM_ENABLE=0 \
 #   NCCL_CUMEM_HOST_ENABLE=0 \
 #   NCCL_P2P_DISABLE=1 \
@@ -706,24 +706,36 @@ def _run_torchrun(
     )
 
     try:
-        # Do not capture stdout/stderr: per-rank progress logs must remain
-        # visible so a communication hang can be located immediately.
-        returncode = proc.wait(timeout=120)
-    except subprocess.TimeoutExpired as exc:
-        os.killpg(os.getpgid(proc.pid), signal.SIGKILL)
+        try:
+            # Do not capture stdout/stderr: per-rank progress logs must
+            # remain visible so a communication hang can be located
+            # immediately.
+            returncode = proc.wait(timeout=120)
+        except subprocess.TimeoutExpired as exc:
+            raise RuntimeError(
+                "The multi-GPU Ring Attention worker timed out after 120 "
+                "seconds and its process group was killed. Use the last "
+                "printed per-rank stage to determine whether the hang "
+                "occurred during initialization, Ring P2P communication, "
+                "all_gather, or the final barrier."
+            ) from exc
+
+        if returncode != 0:
+            raise RuntimeError(
+                f"Ring Attention worker exited with code {returncode}."
+            )
+    finally:
+        # torchrun exiting -- whether cleanly, with a nonzero code, or via
+        # the timeout above -- does not guarantee every rank worker in its
+        # process group has exited. A rank stuck in a NCCL collective after
+        # a sibling crashes can outlive the launcher, reparent to init, and
+        # keep pegging the GPUs for every test that runs afterward. Always
+        # sweep the whole session's process group before returning control.
+        try:
+            os.killpg(os.getpgid(proc.pid), signal.SIGKILL)
+        except ProcessLookupError:
+            pass
         proc.wait()
-        raise RuntimeError(
-            "The multi-GPU Ring Attention worker timed out after 120 seconds "
-            "and its process group was killed. Use the last printed per-rank "
-            "stage to determine whether the hang occurred during "
-            "initialization, Ring P2P communication, all_gather, or the "
-            "final barrier."
-        ) from exc
-
-    if returncode != 0:
-        raise RuntimeError(
-            f"Ring Attention worker exited with code {returncode}."
-        )
 
 
 @pytest.mark.parametrize(
@@ -744,7 +756,11 @@ def test_torchrun_debug_defaults_are_opt_in(
     captured_env: dict[str, str] = {}
 
     class _CompletedProcess:
-        pid = 1
+        # A pid that cannot correspond to a real process: _run_torchrun's
+        # cleanup now unconditionally calls os.getpgid()/os.killpg() on this
+        # pid in a finally block, and this test must not let that touch a
+        # real process group (e.g. pid 1's).
+        pid = 999999999
 
         @staticmethod
         def wait(timeout: int | None = None) -> int: