add distributed UT xdist

Author: zxd1997066

.github/actions/get-runner/action.yml

@@ -1,5 +1,11 @@
 name: Get Runner Infos
 
+inputs:
+  ut_name:
+    required: true
+    type: string
+    description: Which ut to launch
+
 outputs:
   runner_id:
     value: ${{ steps.runner.outputs.runner_id }}
@@ -21,6 +27,7 @@ runs:
   steps:
     - name: Get runner
       shell: bash -xe {0}
+      if: ${{ ! contains(inputs.ut_name, 'distributed') }}
       id: runner
       run: |
         # get test runner
@@ -76,6 +83,64 @@ runs:
         echo "xpu_num=${xpu_num}" |tee -a ${GITHUB_OUTPUT}
         echo "cpus_per_xpu=${cpus_per_xpu}" |tee -a ${GITHUB_OUTPUT}
         echo "pytest_extra_args=${pytest_extra_args}" |tee -a ${GITHUB_OUTPUT}
+    - name: Get distributed runner
+      shell: bash -xe {0}
+      if: ${{ contains(inputs.ut_name, 'distributed') }}
+      id: runner-dist
+      run: |
+        # get test runner
+        echo "runner_id=$(echo ${RUNNER_NAME} |sed 's/\-[0-9]$//')" |tee -a ${GITHUB_OUTPUT}
+        echo "user_id=$(id -u)" |tee -a ${GITHUB_OUTPUT}
+        echo "render_id=$(getent group render |cut -d: -f3)" |tee -a ${GITHUB_OUTPUT}
+        echo "hostname=$(hostname)" |tee -a ${GITHUB_OUTPUT}
+        # show host info
+        lscpu
+        lshw -C display
+        free -h
+        df -h
+        cat /etc/os-release
+        uname -a
+        # clinfo hang and reboot system to recover
+        timeout 120 clinfo --list || sudo reboot
+        scaling_governor=$(cat /sys/devices/system/cpu/cpu*/cpufreq/scaling_governor |sort |uniq)
+        if [ $(sudo -n true > /dev/null 2>&1 && echo $? || echo $?) -eq 0 ];then
+            if [ "${scaling_governor}" != "performance" ];then
+                # set frequency governor to performance mode
+                sudo apt-get update
+                sudo apt-get install -y linux-tools-common linux-tools-$(uname -r) linux-cloud-tools-$(uname -r)
+                sudo cpupower set -b 0
+                sudo cpupower frequency-set -g performance
+            fi
+            # clean cache
+            sync; sudo sh -c "echo 3 > /proc/sys/vm/drop_caches" || true
+        else
+            echo "[INFO] You do NOT have ROOT permission to set system config."
+            echo "       The frequency governor is ${scaling_governor}."
+        fi
+        cpu_num="$(lscpu |grep -E 'Core\(s\) per socket:|Socket\(s\):' |awk 'BEGIN{sum=1}{sum*=$NF}END{printf sum}')"
+        xpu_num="$(clinfo --list |awk 'BEGIN{gpu=0;}{
+          if(gpu==1 && $0~/Platform/){gpu=0}; if(gpu==1){print $0}; if($0~/Platform.*Graphics/){gpu=1}
+        }' |wc -l)"
+        cpus_per_xpu="$(echo |awk -v c="${cpu_num}" -v x="${xpu_num}" '{printf c/x}')"
+        pytest_extra_args="$(echo |awk -v x="${xpu_num}" -v z="${ZE_AFFINITY_MASK}" -v cx="${cpus_per_xpu}" '{
+          if (x > 0) {
+            split(z, xpu_list, ",");
+            for (i=0;i<x;i=i+4) {
+              if (z != "") {
+                  ze = xpu_list[i+1];
+              } else {
+                  ze = i;
+              }
+              printf(" --tx popen//env:ZE_AFFINITY_MASK=%d,%d,%d,%d//env:OMP_NUM_THREADS=%d//python=\"numactl -l -C %d-%d python\"",
+                      ze,ze+1,ze+2,ze+3,4*cx,i*cx,(i+4)*cx-1);
+            }
+          }else {
+            printf(" -n 1 ");
+          }
+        }')"
+        echo "xpu_num=${xpu_num}" |tee -a ${GITHUB_OUTPUT}
+        echo "cpus_per_xpu=${cpus_per_xpu}" |tee -a ${GITHUB_OUTPUT}
+        echo "pytest_extra_args=${pytest_extra_args}" |tee -a ${GITHUB_OUTPUT}
     - name: Cleanup host
       shell: bash -xe {0}
       run: |

.github/workflows/_linux_ut.yml

@@ -38,7 +38,7 @@ env:
 jobs:
   runner:
     runs-on: ${{ inputs.runner }}
-    name: get-runner
+    name: get-runner 
     outputs:
       runner_id: ${{ steps.runner-info.outputs.runner_id }}
       user_id: ${{ steps.runner-info.outputs.user_id }}
@@ -53,6 +53,8 @@ jobs:
         uses: actions/checkout@v4
       - name: Get runner
         id: runner-info
+        with:
+          ut_name: ${{ inputs.ut }}
         uses: ./.github/actions/get-runner
 
   test-in-container:
@@ -104,7 +106,7 @@ jobs:
     runs-on: ${{ needs.runner.outputs.runner_id }}
     env:
       AGENT_TOOLSDIRECTORY: /tmp/xpu-tool
-      PYTEST_ADDOPTS: -v --timeout 3600 --timeout_method=thread -n 1
+      PYTEST_ADDOPTS: -v --timeout 3600 --timeout_method=thread --dist worksteal ${{ needs.runner.outputs.pytest_extra_args }}
     steps:
       - name: Checkout torch-xpu-ops
         uses: actions/checkout@v4

test/xpu/run_distributed.py

@@ -1,4 +1,3 @@
-import os
 import subprocess
 import sys
 
@@ -9,42 +8,6 @@
 res2 = 0
 fail_test = []
 
-# Get the xelink group card affinity
-ret = os.system("xpu-smi topology -m 2>&1|tee topology.log")
-if ret == 0:
-    gpu_dict = {}
-    with open("topology.log") as file:
-        lines = file.readlines()
-        for line in lines:
-            if "CPU Affinity" in line:
-                continue
-            line = line.strip()
-            if line.startswith("GPU "):
-                items = line.split(" ")
-                items = [x for x in items if x]
-                gpu_id = items[1]
-                i = gpu_id.split("/")[0]
-                affinity = ""
-                for j, item in enumerate(items):
-                    if "SYS" not in item and ("XL" in item or "S" in item):
-                        if len(affinity) == 0:
-                            affinity = str(j - 2)
-                        else:
-                            affinity = affinity + "," + str(j - 2)
-                gpu_dict[i] = affinity
-
-    max_affinity = ""
-    for key, value in gpu_dict.items():
-        if len(value) > len(max_affinity):
-            max_affinity = value
-
-    os.environ["ZE_AFFINITY_MASK"] = str(max_affinity)
-    print(str("ZE_AFFINITY_MASK=" + os.environ.get("ZE_AFFINITY_MASK")))
-
-else:
-    print("xpu-smi topology failed")
-    sys.exit(255)
-
 
 # run python test
 def run(test_command):