diff --git a/run_unit_tests.sh b/run_unit_tests.sh index 54ba85a9de0..82d32826ed2 100755 --- a/run_unit_tests.sh +++ b/run_unit_tests.sh @@ -19,6 +19,12 @@ if [[ "$HIP_ARCHITECTURES" == "gfx90a" ]]; then PYTEST_MARKERS="$PYTEST_MARKERS and not failing_on_rocm_mi250" fi +# Per-test-file timeout in seconds (30 minutes). +# Prevents a single hung torchrun from consuming the entire CI time budget. +# --kill-after=300: if SIGTERM is ignored, SIGKILL after 5 more minutes. +TEST_TIMEOUT=${TEST_TIMEOUT:-1800} +KILL_AFTER=${KILL_AFTER:-300} + # Find all test files recursively TEST_FILES=$(find tests/unit_tests -type f -name "test_*.py") @@ -28,19 +34,24 @@ for file in $TEST_FILES; do # Create unique filename by replacing slashes with underscores and removing tests/unit_tests/ prefix # E.g., tests/unit_tests/dist_checkpointing/test_optimizer.py -> dist_checkpointing_test_optimizer test_name=$(echo "$file" | sed 's|tests/unit_tests/||' | sed 's|/|_|g' | sed 's|\.py$||') - + csv_file="$OUT_DIR/test_report_${test_name}.csv" xml_file="$OUT_DIR/junit_report_${test_name}.xml" echo "Running test file: $file" - torchrun --standalone --nproc_per_node=$NUM_GPUS -m pytest \ - --showlocals --tb=long -v -s -m "$PYTEST_MARKERS" \ - --csv "$csv_file" \ - --junitxml "$xml_file" \ - $file + timeout --kill-after="$KILL_AFTER" "$TEST_TIMEOUT" \ + torchrun --standalone --nproc_per_node=$NUM_GPUS -m pytest \ + --showlocals --tb=long -v -s -m "$PYTEST_MARKERS" \ + --csv "$csv_file" \ + --junitxml "$xml_file" \ + $file rc=$? - if [[ $rc -ne 0 ]]; then + # exit code 124 = SIGTERM timeout; 137 = SIGKILL (--kill-after triggered) + if [[ $rc -eq 124 ]] || [[ $rc -eq 137 ]]; then + echo "TIMEOUT: $file exceeded ${TEST_TIMEOUT}s (kill-after=${KILL_AFTER}s) — marking as failed." + ANY_FAIL=1 + elif [[ $rc -ne 0 ]]; then echo "Test failed in $file." ANY_FAIL=1 fi