Add pyyaml installation to ci prerequisites

.github/scripts/aiter_prebuild_upload.sh

@@ -12,13 +12,9 @@ set -euo pipefail
 
 # Derive ROCm version and aiter commit -> cache key
 ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)"
-if [ -n "${ROCM_PATH:-}" ]; then
-    true # Use provided ROCM_PATH
-elif [ -d "/opt/rocm/core" ]; then
-    ROCM_PATH="/opt/rocm/core"
-else
-    ROCM_PATH="/opt/rocm"
-fi
+DIR="${ROOT_DIR}/ci"
+. "${ROOT_DIR}/ci/_utils.sh"
+ROCM_PATH="$(resolve_rocm_path)"
 export ROCM_PATH
 ROCM_VER=`head -n1 "${ROCM_PATH}/.info/version" | cut -d. -f1`
 

build_tools/rocm_utils.cmake

@@ -7,25 +7,24 @@ include_guard(GLOBAL)
 if(NOT "$ENV{ROCM_PATH}" STREQUAL "")
     set(ROCM_PATH "$ENV{ROCM_PATH}")
 else()
+    set(_ROCM_SDK_ROOT "")
     find_program(ROCM_SDK_CLI rocm-sdk)
     if(ROCM_SDK_CLI)
         execute_process(
             COMMAND ${ROCM_SDK_CLI} path --root
             OUTPUT_VARIABLE _ROCM_SDK_ROOT
             OUTPUT_STRIP_TRAILING_WHITESPACE
             ERROR_QUIET
-            RESULT_VARIABLE _ROCM_SDK_RESULT
         )
-        if(_ROCM_SDK_RESULT EQUAL 0 AND EXISTS "${_ROCM_SDK_ROOT}/bin/hipcc")
-            set(ROCM_PATH "${_ROCM_SDK_ROOT}")
-        endif()
     endif()
-    if(NOT DEFINED ROCM_PATH)
-        if(EXISTS "/opt/rocm/core")
-            set(ROCM_PATH "/opt/rocm/core")
-        else()
-            set(ROCM_PATH "/opt/rocm")
-        endif()
+    if(NOT _ROCM_SDK_ROOT STREQUAL "" AND EXISTS "${_ROCM_SDK_ROOT}/bin/hipcc")
+        set(ROCM_PATH "${_ROCM_SDK_ROOT}")
+    elseif(EXISTS "/opt/rocm/core")
+        set(ROCM_PATH "/opt/rocm/core")
+    elseif(EXISTS "/opt/rocm")
+        set(ROCM_PATH "/opt/rocm")
+    else()
+        message(FATAL_ERROR "Could not find ROCm installation")
     endif()
 endif()
 

ci/_utils.sh

@@ -2,6 +2,32 @@
 #
 # See LICENSE for license information.
 
+# ROCM_PATH resolution
+resolve_rocm_path() {
+    if [ -n "${ROCM_PATH:-}" ]; then
+        echo "$ROCM_PATH"
+        return 0
+    fi
+    if command -v rocm-sdk >/dev/null 2>&1; then
+        local _root
+        _root="$(rocm-sdk path --root)"
+        if [ -n "$_root" ] && [ -f "${_root}/bin/hipcc" ]; then
+            echo "$_root"
+            return 0
+        fi
+    fi
+    if [ -d "/opt/rocm/core" ]; then
+        echo /opt/rocm/core
+        return 0
+    fi
+    if [ -d "/opt/rocm" ]; then
+        echo /opt/rocm
+        return 0
+    fi
+    echo "Could not find ROCm installation" >&2
+    exit 1
+}
+
 REALPATH=realpath
 realpath $DIR >/dev/null 2>/dev/null
 test $? -ne 0 && REALPATH=echo
@@ -287,15 +313,9 @@ check_test_filter() {
 
 start_message() {
     echo "Started with TEST_LEVEL=$TEST_LEVEL sGPU='$TEST_SGPU' mGPU='$TEST_MGPU' at `date`"
-    if [ -n "$ROCM_PATH" ]; then
-        _rocm_path="$ROCM_PATH"
-    elif [ -d "/opt/rocm/core" ]; then
-        _rocm_path="/opt/rocm/core"
-    else
-        _rocm_path="/opt/rocm"
-    fi
-    _rocm_path=`$REALPATH "$_rocm_path"`
-    test -d "$_rocm_path" && echo "ROCm: $_rocm_path" || echo "ROCm path not found"
+    _rocm_path=$(resolve_rocm_path)
+    _rocm_path=`$REALPATH "$_rocm_path" 2>/dev/null || echo "$_rocm_path"`
+    echo "ROCM PATH: $_rocm_path"
     python3 --version
 }
 

ci/pytorch.sh

@@ -12,7 +12,7 @@ TEST_DIR=${TE_PATH}tests/pytorch
 #: ${TEST_WORKERS:=4}
 
 install_prerequisites() {
-    pip install 'numpy>=1.22.4' pandas safetensors pytest-timeout
+    pip install 'numpy>=1.22.4' pandas safetensors pyyaml pytest-timeout
     rc=$?
     if [ $rc -ne 0 ]; then
         script_error "Failed to install test prerequisites"

tests/cpp/CMakeLists.txt

@@ -10,12 +10,7 @@ option(USE_ROCM "Use ROCm" ON)
 set(USE_CUDA OFF)
 
 if (USE_ROCM)
-  #Check if ROCm exists
-  if( ((EXISTS "/opt/rocm/") OR (EXISTS $ENV{ROCM_PATH})) )
-    message("Using AMD Platform")
-  else()
-    message(FATAL_ERROR "Could not ROCm installation")
-  endif()
+  include("${CMAKE_CURRENT_SOURCE_DIR}/../../build_tools/rocm_utils.cmake")
 else()
   #Check if CUDA exists
   if (EXISTS "$ENV{CUDA_HOME}/bin/nvcc")
@@ -33,7 +28,6 @@ if(USE_ROCM)
   # Disable Asserts In Code (Can't use asserts on HIP stack.)
   add_definitions(-DNDEBUG)
   add_definitions(-DUSE_ROCM)
-  include("${CMAKE_CURRENT_SOURCE_DIR}/../../build_tools/rocm_utils.cmake")
 else()
   if(NOT DEFINED CMAKE_CUDA_ARCHITECTURES)
     if (CUDAToolkit_VERSION VERSION_GREATER_EQUAL 12.8)

transformer_engine/common/CMakeLists.txt

@@ -13,13 +13,7 @@ option(USE_HIPKITTENS_GEMM "Use HipKittens blockwise FP8 and MXFP8 GEMM kernels"
 set(USE_CUDA OFF)
 
 if (USE_ROCM)
-  #Check if ROCm exists
-  if( ((EXISTS "/opt/rocm/") OR (EXISTS $ENV{ROCM_PATH})) )
-    message("Using AMD Platform")
-  else()
-    message(FATAL_ERROR "Could not find ROCm installation")
-  endif()
-
+  include("${CMAKE_CURRENT_SOURCE_DIR}/../../build_tools/rocm_utils.cmake")
 else()
   #Check if CUDA exists
   if (EXISTS "$ENV{CUDA_HOME}/bin/nvcc")
@@ -136,7 +130,6 @@ set(CUTLASS_TOOLS_INCLUDE_DIR
 else()
   set(CMAKE_CXX_STANDARD 17)
   #define GPU ARCH before enabling HIP to avoid potential failures with arch detection
-  include("${CMAKE_CURRENT_SOURCE_DIR}/../../build_tools/rocm_utils.cmake")
   project(transformer_engine LANGUAGES HIP CXX)
 
   # If using ROCm, ninja must be installed to build TE.