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 }}
@@ -57,22 +63,41 @@ runs:
           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++) {
-              if (z != "") {
-                  ze = xpu_list[i+1];
-              } else {
-                  ze = i;
+        if [ "${{ inputs.ut_name }}" == "xpu_distributed" ];then
+          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);
               }
-              printf(" --tx popen//env:ZE_AFFINITY_MASK=%d//env:OMP_NUM_THREADS=%d//python=\"numactl -l -C %d-%d python\"",
-                      ze, cx, i*cx, (i+1)*cx-1);
+            }else {
+              printf(" -n 1 ");
             }
-          }else {
-            printf(" -n 1 ");
-          }
-        }')"
+          }')"
+        else
+          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++) {
+                if (z != "") {
+                    ze = xpu_list[i+1];
+                } else {
+                    ze = i;
+                }
+                printf(" --tx popen//env:ZE_AFFINITY_MASK=%d//env:OMP_NUM_THREADS=%d//python=\"numactl -l -C %d-%d python\"",
+                        ze, cx, i*cx, (i+1)*cx-1);
+              }
+            }else {
+              printf(" -n 1 ");
+            }
+          }')"
+        fi
         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}

.github/workflows/_linux_ut.yml

@@ -39,7 +39,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 }}
@@ -54,6 +54,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:
@@ -105,7 +107,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):