# --------------------------------------------------------------------------------------
# Part of the interTwin Project: https://www.intertwin.eu/
#
# Created by: Matteo Bunino
#
# Credit:
# - Matteo Bunino <matteo.bunino@cern.ch> - CERN
# --------------------------------------------------------------------------------------

# Singularity definition file for ROCm with PyTorch, starting from ROCm base image provided
# by LUMI and AMD.

Bootstrap: docker
From: registry.egi.eu/dev.intertwin.eu/itwinai-dev:lumi-pytorch-rocm-6.1.3-python-3.12-pytorch-v2.4.1

# # To build from local image
# Bootstrap: localimage
# From: /path/to/base.sif

%files
    pyproject.toml /app/pyproject.toml
    src /app/src

%post
    # ------------------------------------------------------------------
    # 1. system build dependencies (compiler, cmake, make, etc.)
    #    + mpich-devel compiled with GCC (avoids NVIDIA HPC‑SDK runtime)
    # ------------------------------------------------------------------
    zypper --non-interactive refresh && \
    zypper --non-interactive install --no-recommends make cmake gcc11 gcc11-c++ mpich-devel && \
    zypper clean --all && rm -rf /var/cache/zypp/*

    # ------------------------------------------------------------------
    # 2. activate the pre-installed conda environment provided by LUMI
    # ------------------------------------------------------------------
    $WITH_CONDA

    # ------------------------------------------------------------------
    # 3. Build Horovod from source for ROCm **with patches**
    #    https://github.com/ROCm/ROCm/issues/3987#issuecomment-2479684637
    #    Extra safeguards so CUDA on the build host cannot leak in.
    # ------------------------------------------------------------------

    #‑‑‑ Safety: make ROCm toolchain come first and stop CUDA detection
    export PATH=/opt/rocm/bin:$PATH   # put HIP/nvcc-stub ahead of host CUDA
    unset CUDA_HOME CUDACXX CUDAHOSTCXX

    # Horovod build environment variables
    export HOROVOD_GPU=ROCM
    export HOROVOD_ROCM_HOME=/opt/rocm
    export HOROVOD_CMAKE_FLAGS="-DHOROVOD_GPU_ROCM=ON -DHOROVOD_GPU_CUDA=OFF"
    export HOROVOD_WITH_PYTORCH=1
    export HOROVOD_WITHOUT_TENSORFLOW=1
    export HOROVOD_WITHOUT_MXNET=1
    export HOROVOD_MPI_THREADS_DISABLE=1
    export HOROVOD_CPU_OPERATIONS=MPI
    export HOROVOD_GPU_ALLREDUCE=NCCL
    export HOROVOD_GPU_BROADCAST=NCCL
    export HOROVOD_NCCL_LINK=SHARED

    export ROCM_PATH=/opt/rocm
    export CC=/usr/bin/gcc-11
    export CXX=/usr/bin/g++-11
    export CXXFLAGS="-std=c++17"

    # clone + patch + build wheel (j16 keeps image build time reasonable)
    git clone --recursive https://github.com/horovod/horovod.git \
        && cd horovod \
        && ln -s $ROCM_PATH/lib/cmake/hip/FindHIP* cmake/Modules/ \
        && sed -i 's/rccl\.h/rccl\/rccl\.h/' horovod/common/ops/nccl_operations.h \
        && MAKEFLAGS=-j16 python setup.py bdist_wheel \
        && pip install dist/*.whl \
        && cd .. && rm -rf horovod

    # quick check (non‑fatal)
    horovodrun --check-build || true

    # ------------------------------------------------------------------
    # 4. Install the project and the rest of the Python stack
    #    --no-binary=horovod prevents pip replacing our ROCm wheel.
    # ------------------------------------------------------------------
    cd /app

    pip install --no-cache-dir --no-binary=horovod \
        .[torch,amd] \
        --extra-index-url https://download.pytorch.org/whl/rocm6.1 \
        "prov4ml[amd]@git+https://github.com/matbun/ProvML@new-main" \
        pytest pytest-xdist psutil wheel

%labels
    Author Matteo Bunino
    Project interTwin - itwinai
    Version 0.3.1

%test
    # Activate conda env
    $WITH_CONDA

    echo "Running post-build test..."
    itwinai sanity-check --torch \
        --optional-deps deepspeed \
        --optional-deps horovod \
        --optional-deps prov4ml \
        --optional-deps ray
