fix distrib tests

tests/common-test-functionality.sh

@@ -14,6 +14,7 @@ run_tests() {
     local trap_deselected_exit_code=1
     local use_last_failed=0
     local use_coverage=0
+    local world_size=0
     # Always clean up pytest.ini
     trap 'rm -f pytest.ini' RETURN
     # Parse arguments
@@ -56,6 +57,11 @@ run_tests() {
             shift
             shift
             ;;
+            --world_size)
+            world_size="$2"
+            shift
+            shift
+            ;;
             *)
             echo "Error: Unknown argument $key"
             exit 1
@@ -82,6 +88,10 @@ run_tests() {
     if [ "${use_coverage}" -eq "1" ]; then
         pytest_args="--cov ignite --cov-append --cov-report term-missing --cov-report xml ${pytest_args}"
     fi
+    if [ ! "${world_size}" -eq "0" ]; then
+        export WORLD_SIZE="${world_size}"
+        pytest_args="--dist=each --tx ${WORLD_SIZE}*popen//python=python ${pytest_args}"
+    fi
 
     # Run the command
     if [ "$trap_deselected_exit_code" -eq "1" ]; then

tests/run_cpu_tests.sh

@@ -22,7 +22,8 @@ fi
 
 # Run 2 processes with --dist=each
 run_tests \
-    --core_args "--dist=each --tx 2*popen//python=python -m distributed -vvv tests" \
+    --core_args "-m distributed -vvv tests" \
+    --world_size 2 \
     --cache_dir ".cpu-distrib" \
     --skip_distrib_tests 0 \
     --use_coverage 1 \

tests/run_gpu_tests.sh

@@ -37,7 +37,8 @@ run_tests \
 
 if [ ${ngpus} -gt 1 ]; then
     run_tests \
-        --core_args "--tx ${ngpus}*popen//python=python --dist=each -vvv -m distributed tests" \
+        --core_args "-vvv -m distributed tests" \
+        --world_size "${ngpus}" \
         --cache_dir ".gpu-distrib-multi" \
         --skip_distrib_tests 0 \
         --use_coverage 1 \