install_packages.sh

#!/bin/bash

get_latest_release_tag() {
    if ! command -v jq &> /dev/null; then
        print_warning "jq is not installed. Cannot dynamically fetch latest release tag. Falling back to hardcoded tag."
        echo "v1.22.2" # Fallback to a known stable version
        return
    fi
    curl -s https://api.github.com/repos/${1}/releases/latest | jq -r '.tag_name'
}

get_gpu_architecture() {
  rocm_agent_enumerator | head -1
}


install_basic() {

  # check running under docker
  if [ ! -f /.dockerenv ]; then
    echo "Looks like we're not running inside a Docker container. Can't continue."
    exit 1
  fi  
  
  
  # ONNX
  cd /build
  LATEST_TAG=$(get_latest_release_tag onnx/onnx)
  git clone --recursive -b "$LATEST_TAG" https://github.com/onnx/onnx.git
  cd onnx
  git submodule update --init --recursive
  CMAKE_ARGS="-DONNX_USE_PROTOBUF_SHARED_LIBS=ON" \
    pip install -e . -v --break-system-packages --force-reinstall --no-cache-dir
  
  # onnxruntime
  cd /build
  LATEST_TAG=$(get_latest_release_tag microsoft/onnxruntime)
  git clone --recursive -b "$LATEST_TAG" https://github.com/microsoft/onnxruntime.git
  cd onnxruntime
  echo "==== Pausing here to tell you this will take a while ===="
  sleep 10
  # note we need to specify 'gfx90a' as an additional architecture to have enough dependencies 
  # generated to permit composable kernel compilation
  pip install packaging --break-system-packages
  ./build.sh --config Release \
  --build_wheel \
  --parallel=$(nproc --ignore=1) \
  --use_rocm \
  --rocm_home /opt/rocm \
  --use_migraphx \
  --migraphx_home /opt/rocm \
  --allow_running_as_root \
  --compile_no_warning_as_error \
  --build_shared_lib \
  --skip_tests \
  --cmake_extra_defines "onnx_SOURCE_DIR=/build/onnx"  \
  --cmake_extra_defines "CMAKE_HIP_ARCHITECTURES=gfx90a;$DETECTED_GPU" \
  --cmake_extra_defines "CMAKE_CXX_FLAGS='-Wno-array-bounds -Wno-shorten-64-to-32 -Wno-pass-failed -Wno-unused-parameter -Wno-inconsistent-missing-override'"
  mv build/Linux/Release/dist/*.whl /wheels

  # triton
  cd /build
  pip install ninja cmake pybind11 --break-system-packages
  git clone https://github.com/ROCm/triton.git
  cd triton/python
  TRITON_CODEGEN_BACKENDS=amd \
    TRITON_BUILD_TUTORIALS=off \
    GPU_ARCHS=$DETECTED_GPU \
    MAX_JOBS=$MAX_JOBS \
    python3 setup.py bdist_wheel
  mv dist/triton*.whl /wheels/

}

install_extra() {
  # check running under docker
  if [ ! -f /.dockerenv ]; then
    echo "Looks like we're not running inside a Docker container. Can't continue."
    exit 1
  fi

  # Check for torch wheel and install it if present
  if ls /wheels/torch-*.whl 1> /dev/null 2>&1; then
    echo "Found torch wheel, installing..."
    pip install /wheels/torch-*.whl --break-system-packages --force-reinstall
  else
    echo "Error: torch wheel not found in /wheels/"
    echo "Please run the torch install script first before building additional packages"
    exit 1
  fi

  # megatron
  cd /build
  LATEST_TAG=$(get_latest_release_tag NVIDIA/Megatron-LM)
  git clone --recursive -b "$LATEST_TAG" https://github.com/NVIDIA/Megatron-LM.git
  cd Megatron-LM
  pip wheel --no-deps -e .
  mv megatron_core*.whl /wheels/

  #flash attention with triton backend
  cd /build
  LATEST_TAG=$(get_latest_release_tag ROCm/flash-attention)
  git clone  --recursive -b "$LATEST_TAG" https://github.com/ROCm/flash-attention.git
  cd flash-attention
  FLASH_ATTENTION_TRITON_AMD_ENABLE="TRUE" GPU_ARCHS=$DETECTED_GPU python3 setup.py bdist_wheel 
  mv dist/flash_attn*.whl /wheels/

  # bitsandbytes
  cd /build
  git clone -b multi-backend-refactor https://github.com/bitsandbytes-foundation/bitsandbytes.git
  cd bitsandbytes
  cmake -DCOMPUTE_BACKEND=hip -S .
  make
  pip -v  wheel --no-deps -e .
  cp bitsandbytes*.whl /wheels/

  # torch-migraph
  cd /build
  pip install tabulate --break-system-packages
  git clone https://github.com/ROCmSoftwarePlatform/torch_migraphx.git
  cd torch_migraphx/py
  rm requirements.py
  pip wheel . --no-deps --no-build-isolation
  mv torch_migraphx*.whl /wheels/

  # aiter and vllm 
  cd /build
  LATEST_TAG=$(get_latest_release_tag ROCm/aiter)
  git clone --recursive -b "$LATEST_TAG" https://github.com/ROCm/aiter.git
  cd aiter
  python3 setup.py bdist_wheel
  cp dist/aiter*.whl /wheels/
  pip install dist/aiter*.whl --break-system-packages

  pip install setuptools-scm --break-system-packages
  cd /build
  LATEST_TAG=$(get_latest_release_tag vllm-project/vllm)
  git clone --recursive -b "$LATEST_TAG" https://github.com/vllm-project/vllm.git
  cd vllm
  # fix license file
  sed -i 's/license = "Apache-2.0"/license = {file = "LICENSE"}/' pyproject.toml
  sed -i '/license-files/d' pyproject.toml

  CXX_FLAGS="-Wno-vla-cxx-extension -Wno-undefined-internal" \
  VLLM_TARGET_DEVICE=rocm \
  PYTORCH_ROCM_ARCH=$DETECTED_GPU GPU_ARCHS=$DETECTED_GPU \
  python3 setup.py bdist_wheel
  cp dist/vllm*.whl /wheels/


  # rocm-smi
  cd /build
  LATEST_TAG=$(get_latest_release_tag ROCm/amdsmi)
  git clone --recursive -b "$LATEST_TAG" https://github.com/ROCm/amdsmi.git
  cd amdsmi
  mkdir -p build
  cd build
  cmake ..
  make
  cd py-interface/python_package
  pip wheel -e .
  cp amdsmi*.whl /wheels/
}

# Main execution logic

# set rocm docker default version - check if second parameter is provided, otherwise use default
DOCKER_ROCM_VERSION=${2:-6.4.3}

DETECTED_GPU=$(get_gpu_architecture)
echo "Building for $DETECTED_GPU"

MAX_JOBS=$(nproc --ignore=1)

case "${1:-}" in
    "basic")
        echo "Preparing to install basic packages..."

        # Check if docker image exists
        if ! docker image inspect build-env >/dev/null 2>&1; then
            echo "Docker image 'build-env' not found. Building it now..."
            docker build --build-arg DOCKER_ROCM_VERSION=${DOCKER_ROCM_VERSION} -f Dockerfile.build -t build-env .
            if [ $? -ne 0 ]; then
                echo "Failed to build Docker image"
                exit 1
            fi
        else
            echo "Docker image 'build-env' already exists"
        fi

        # Create wheels directory if it doesn't exist
        mkdir -p ./wheels

        # Run the container to install basic packages
        echo "Starting Docker container to build basic packages..."
        docker run --rm -t \
            -v "$(pwd)/wheels:/wheels" \
            -v "$(pwd):/workspace" \
            -w /workspace \
            --device=/dev/kfd \
            --device=/dev/dri \
            --group-add video \
            --group-add render \
            --security-opt seccomp=unconfined \
            build-env \
            ./install_packages.sh build-basic
        ;;
    "additional")
        echo "Preparing to install additional packages..."

        # Check if docker image exists
        if ! docker image inspect build-env >/dev/null 2>&1; then
            echo "Docker image 'build-env' not found. Building it now..."
            docker build --build-arg DOCKER_ROCM_VERSION=${DOCKER_ROCM_VERSION} -f Dockerfile.build -t build-env .
            if [ $? -ne 0 ]; then
                echo "Failed to build Docker image"
                exit 1
            fi
        else
            echo "Docker image 'build-env' already exists"
        fi

        # Create wheels directory if it doesn't exist
        mkdir -p ./wheels

        # Run the container to install additional packages
        echo "Starting Docker container to build additional packages..."
        docker run --rm -t \
            -v "$(pwd)/wheels:/wheels" \
            -v "$(pwd):/workspace" \
            -w /workspace \
            --device=/dev/kfd \
            --device=/dev/dri \
            --group-add video \
            --group-add render \
            --security-opt seccomp=unconfined \
            build-env \
            ./install_packages.sh build-additional
        ;;
    "build-basic")
        echo "Building basic packages inside Docker container..."
        install_basic
        ;;
    "build-additional")
        echo "Building additional packages inside Docker container..."
        install_extra
        ;;
    *)
        echo "Usage: $0 {basic|additional|build-basic|build-additional}"
        echo "  basic            - Install ONNX, ONNX Runtime, torch-migraphx, and triton"
        echo "  additional       - Install Megatron, flash-attention, bitsandbytes, aiter, vllm, and rocm-smi"
        echo "  build-basic      - Internal: Build basic packages inside Docker container"
        echo "  build-additional - Internal: Build additional packages inside Docker container"
        exit 1
        ;;
esac