---
- name: "Install NCCL Dependencies"
  ansible.builtin.yum:
    name:
      - git
      - make
    state: present


- name: Clone Nvidia NCCL git repository
  ansible.builtin.git:
    repo: https://github.com/{{ nccl_repo }}/nccl.git
    dest: /opt/nccl
    version: '{{ nccl_version }}'

- name: Build NCCL
  ansible.builtin.shell: |
    make -j src.build CUDA_HOME=/usr/local/cuda NVCC_GENCODE='-gencode=arch=compute_70,code=sm_70 -gencode=arch=compute_75,code=sm_75 -gencode=arch=compute_80,code=sm_80'
  args:
    chdir: /opt/nccl

- name: Adding the NCCL profile file
  ansible.builtin.blockinfile:
    path: /etc/profile.d/nccl.sh
    create: yes
    block: |
      export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/opt/nccl/build/lib
      export NCCL_PROTO=simple

- name: Adding the NCCL preload profile file
  ansible.builtin.blockinfile:
    path: /etc/profile.d/nccl-preload.sh
    create: yes
    block: |
      export LD_PRELOAD=/opt/nccl/build/lib/libnccl.so:/opt/aws-ofi-nccl/lib/libnccl-net.so

- name: Clone Nvidia NCCL tests git repository
  when: install_nccl_tests
  block:
    - name: Clone Nvidia NCCL tests repository
      ansible.builtin.git:
        repo: https://github.com/NVIDIA/nccl-tests.git
        dest: /tmp/nccl-tests

    - name: Build NCCL tests
      ansible.builtin.shell: |
        make MPI=1 CUDA_HOME=/usr/local/cuda MPI_HOME=/opt/amazon/openmpi NCCL_HOME=/opt/nccl/build
      args:
        chdir: /tmp/nccl-tests