[ROCm] Parameterize the triton build dir (#125420)

pruthvistony · pytorchmergebot · commit a1a22a22d56d · 2024-05-08T04:46:52.000Z
- Removes hard coding and helps in internal builds Pull Request resolved: #125420 Approved by: https://github.com/malfet
diff --git a/.github/scripts/amd/patch_triton_wheel.sh b/.github/scripts/amd/patch_triton_wheel.sh
@@ -1,7 +1,11 @@
 #!/bin/bash
 set -x
 
-WHEELHOUSE_DIR=/artifacts
+if [ -z "$1" ]; then
+    echo "Need wheel location argument" && exit 1
+fi
+
+WHEELHOUSE_DIR=$1
 PATCHELF_BIN=patchelf
 ROCM_LIB=backends/amd/lib
 ROCM_LD=backends/amd/llvm/bin
diff --git a/.github/scripts/build_triton_wheel.py b/.github/scripts/build_triton_wheel.py
@@ -157,10 +157,10 @@ def build_triton(
 
         if build_rocm:
             check_call(
-                [f"{SCRIPT_DIR}/amd/patch_triton_wheel.sh"],
+                [f"{SCRIPT_DIR}/amd/patch_triton_wheel.sh", Path.cwd()],
                 cwd=triton_basedir,
-                shell=True,
             )
+
         return Path.cwd() / whl_path.name
 
 
