ARG BASE_DOCKER="rocm/pytorch-nightly:latest"
FROM $BASE_DOCKER
ARG CK_PYTORCH_BRANCH="develop"
RUN groupadd -g 109 render && \
    usermod -u 1001 jenkins && \
    groupmod -g 1001 jenkins && \
    pip install --upgrade pandas && \
    cd /tmp/pytorch && \
    rm -rf build && \
    cd /tmp/pytorch/third_party && \
    rm -rf composable_kernel && \
    git clone -b "$CK_PYTORCH_BRANCH" https://github.com/ROCm/composable_kernel.git && \
    cd /tmp/pytorch/third_party/aiter/3rdparty && \
    rm -rf composable_kernel && \
    git clone -b "$CK_PYTORCH_BRANCH" https://github.com/ROCm/composable_kernel.git && \
    cd /tmp/pytorch/third_party/fbgemm/external && \
    rm -rf composable_kernel && \
    git clone -b "$CK_PYTORCH_BRANCH" https://github.com/ROCm/composable_kernel.git && \
    cd /tmp/pytorch/third_party/flash-attention/csrc && \
    rm -rf composable_kernel && \
    git clone -b "$CK_PYTORCH_BRANCH" https://github.com/ROCm/composable_kernel.git && \
    mkdir -p /var/jenkins/workspace/pytorch && \
    cp -r /tmp/pytorch/* /var/jenkins/workspace/pytorch/ && \
    chown -R jenkins:jenkins /var/jenkins/workspace/pytorch && \
    chmod -R a+rwx /var/jenkins/workspace/pytorch && \
    sudo usermod -aG irc jenkins