install_pytorch.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_from_src() {
  # check if running in docker container by checking for /.dockerenv
  if [ ! -f /.dockerenv ]; then
    echo "Looks like we're not running inside a Docker container. Can't continue"
    exit 1
  fi
  
  export PYTORCH_ROCM_ARCH=$(get_gpu_architecture)

  # pytorch
  cd /build
  LATEST_TAG=$(get_latest_release_tag pytorch/pytorch)
  git clone --recursive -b "$LATEST_TAG"  https://github.com/pytorch/pytorch.git
  cd pytorch
  pip install -r requirements.txt --break-system-packages
  pip install mkl-static mkl-include --break-system-packages
  export CMAKE_PREFIX_PATH="/opt/rocm:${CMAKE_PREFIX_PATH}"
  python tools/amd_build/build_amd.py
  HIPCC_COMPILE_FLAGS_APPEND="-Wno-unused-command-line-argument -Wno-unused-result" \
    HIP_CLANG_FLAGS="-Wno-unused-command-line-argument -Wno-unused-result" \
    python3 -m pip wheel --no-build-isolation -v -e .
  cp torch*.whl /wheels

  # we need to install torch to compile vision/audio
  WHEEL_FILE=$(ls /wheels/torch-*.whl)
  pip -v install $WHEEL_FILE --break-system-packages

  LATEST_TAG=$(get_latest_release_tag pytorch/vision)
  cd /build
  git clone --recursive -b "$LATEST_TAG"  https://github.com/pytorch/vision.git
  cd vision
  pip wheel -e . -v --no-build-isolation --no-deps
  mv torchvision*.whl /wheels/

  cd /build
  LATEST_TAG=$(get_latest_release_tag pytorch/audio)
  git clone --recursive -b "$LATEST_TAG"  https://github.com/pytorch/audio.git
  cd audio
  pip wheel -e . -v --no-build-isolation --no-deps
  mv torchaudio*.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}

case "${1:-}" in
    "wheels")
        echo "Installing PyTorch from wheels..."
        install_from_wheels
        ;;
    "source")
        echo "Preparing to compile PyTorch from source..."

        # 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 build PyTorch from source
        # Note script looks back on itself calling 'build'
        echo "Starting Docker container to build PyTorch..."
        docker run --rm \
            -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_pytorch.sh build
        ;;
    "build")
        echo "Building PyTorch from source inside Docker container..."
        install_from_src
        ;;
    *)
        echo "Usage: $0 {wheels|source|build}"
        echo "  wheels - Install PyTorch from pre-built wheels"
        echo "  source - Compile PyTorch from source code using Docker"
        echo "  build  - Internal: Build PyTorch inside Docker container"
        exit 1
        ;;
esac