fix varlen ci test hang

Author: liangel-02

.github/workflows/integration_test_8gpu_rl.yaml

@@ -97,36 +97,15 @@ jobs:
         # but the CI docker image only has CUDA runtime, not the toolkit.
         uv pip install nvidia-cuda-nvcc
 
-        # RL Loss Guard runs FIRST, in a clean container, so a lingering port from
-        # the E2E suite can't collide with the generator's distributed init
-        # (EADDRINUSE). Deterministic loss regression check on the REAL Qwen3-0.6B
-        # (--hf-assets-path=$MODEL_PATH): real weights -> varied rewards -> non-zero
-        # advantage -> meaningful, full-precision losses (random init gives all
-        # zeros). The golden is arch-specific; regenerate on A10G if it changes.
-        LOSS_FILE="torchtitan/experiments/rl/tests/assets/losses/rl_grpo_cuda.txt"
-        # Override to trainer TP=4 + 1 generator TP=4: the config's TP=2 OOMs with
-        # batch-invariant on A10G; the golden is from TP=4.
-        python torchtitan/experiments/rl/scripts/loss_compare.py \
-          --hf-assets-path "$MODEL_PATH" \
-          --dump-folder "$RUNNER_TEMP/artifacts-to-be-uploaded/rl_loss_guard" \
-          --config rl_grpo_qwen3_0_6b_varlen_batch_invariant \
-          --import-result "$LOSS_FILE" \
-          --assert-equal \
-          --options='--trainer.parallelism.tensor_parallel_degree 4 --generator.parallelism.tensor_parallel_degree 4 --num_generators 1'
-
         # Bitwise trainer/generator parity tests (batch-invariant). Varlen and Flex
         # (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=540 \
+          RL_TEST_DUMP_FOLDER="$RUNNER_TEMP/artifacts-to-be-uploaded" \
+          HF_ASSETS_PATH="$MODEL_PATH" timeout -s ABRT 600 torchrun --nproc-per-node=2 -m pytest -s \
           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=540 \
+          RL_TEST_DUMP_FOLDER="$RUNNER_TEMP/artifacts-to-be-uploaded" \
+          HF_ASSETS_PATH="$MODEL_PATH" timeout -s ABRT 600 torchrun --nproc-per-node=2 -m pytest -s \
           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
-        # random-init test and the batch-invariant debug tests (both need 8 GPUs).
-        python -m torchtitan.experiments.rl.tests.integration_tests \
-          $RUNNER_TEMP/artifacts-to-be-uploaded --ngpu 8 \
-          --hf_assets_path "$MODEL_PATH"

torchtitan/experiments/rl/tests/test_bitwise_parity.py

@@ -30,10 +30,12 @@
         torchtitan/experiments/rl/tests/test_bitwise_parity.py::TestBitwiseParityFlex -v
 """
 
+import faulthandler
 import gc
 import logging
 import os
 import shutil
+import signal
 import tempfile
 import unittest
 
@@ -606,6 +608,16 @@ class BitwiseParityTestBase(unittest.TestCase):
 
     @classmethod
     def setUpClass(cls):
+        watchdog_sec = os.environ.get("RL_TEST_WATCHDOG_SEC")
+        if watchdog_sec:
+            faulthandler.enable()
+            for sig in (signal.SIGTERM, signal.SIGABRT):
+                try:
+                    faulthandler.register(sig, chain=True)
+                except (AttributeError, ValueError, OSError):
+                    pass
+            faulthandler.dump_traceback_later(int(watchdog_sec), exit=True)
+
         world_size = (
             dist.get_world_size()
             if dist.is_initialized()
@@ -679,13 +691,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."""