fix varlen ci test hang

Author: liangel-02

.github/workflows/integration_test_8gpu_rl.yaml

@@ -118,11 +118,13 @@ jobs:
         # (both TP=2) run as separate torchrun invocations so each backend gets a
         # fresh process group and vLLM engine; nproc matches each config's
         # tensor_parallel_degree (the test builds ParallelDims from it).
-        RL_TEST_DUMP_FOLDER="$RUNNER_TEMP/artifacts-to-be-uploaded" \
-          HF_ASSETS_PATH="$MODEL_PATH" torchrun --nproc-per-node=2 -m pytest \
+        PYTHONFAULTHANDLER=1 RL_TEST_WATCHDOG_SEC=1200 \
+          RL_TEST_DUMP_FOLDER="$RUNNER_TEMP/artifacts-to-be-uploaded" \
+          HF_ASSETS_PATH="$MODEL_PATH" timeout -s ABRT 1500 torchrun --nproc-per-node=2 -m pytest \
           torchtitan/experiments/rl/tests/test_bitwise_parity.py::TestBitwiseParityVarlen -v
-        RL_TEST_DUMP_FOLDER="$RUNNER_TEMP/artifacts-to-be-uploaded" \
-          HF_ASSETS_PATH="$MODEL_PATH" torchrun --nproc-per-node=2 -m pytest \
+        PYTHONFAULTHANDLER=1 RL_TEST_WATCHDOG_SEC=1200 \
+          RL_TEST_DUMP_FOLDER="$RUNNER_TEMP/artifacts-to-be-uploaded" \
+          HF_ASSETS_PATH="$MODEL_PATH" timeout -s ABRT 1500 torchrun --nproc-per-node=2 -m pytest \
           torchtitan/experiments/rl/tests/test_bitwise_parity.py::TestBitwiseParityFlex -v
 
         # Run E2E RL integration tests (up to 8 GPUs): includes the MoE TP=4 EP=4

torchtitan/experiments/rl/tests/test_bitwise_parity.py

@@ -30,6 +30,7 @@
         torchtitan/experiments/rl/tests/test_bitwise_parity.py::TestBitwiseParityFlex -v
 """
 
+import faulthandler
 import gc
 import logging
 import os
@@ -606,6 +607,10 @@ class BitwiseParityTestBase(unittest.TestCase):
 
     @classmethod
     def setUpClass(cls):
+        watchdog_sec = os.environ.get("RL_TEST_WATCHDOG_SEC")
+        if watchdog_sec:
+            faulthandler.dump_traceback_later(int(watchdog_sec), exit=True)
+
         world_size = (
             dist.get_world_size()
             if dist.is_initialized()
@@ -679,13 +684,19 @@ def setUpClass(cls):
 
     @classmethod
     def tearDownClass(cls):
+        if dist.is_initialized():
+            dist.barrier()
         if hasattr(cls, "engine"):
-            cls.engine.engine_core.shutdown()
+            renderer = getattr(cls.engine, "renderer", None)
+            if renderer is not None:
+                renderer.shutdown()
             del cls.engine
         if hasattr(cls, "model"):
             del cls.model
         gc.collect()
         torch.cuda.empty_cache()
+        if dist.is_initialized():
+            dist.destroy_process_group()
 
     def _assert_logprobs_equal(self, name, a, b, label_a="A", label_b="B"):
         """Assert two logprob sequences are bitwise identical."""