Update the dockerfile base image to cuda-dl-base

[Steboss]

Update the base docker image, so we can use cuda-dl-base from nvcri.

[Steboss]

So I run a check in the cuda-dl-base image, and I can see that:

• for nccl we need to create symlink for include and lib directories, so they're mapped in opt/nvidia/nccl
• same for cudNN
• we can safely remove install-ofed.sh
• and amazon efa support

For the symlink, we'd just need this part of the install-nccl.sh script (and the counterpart in install-cudnn.sh script):

arch=$(uname -m)-linux-gnu
for nccl_file in $(dpkg -L libnccl2 libnccl-dev | sort -u); do
  # Real files and symlinks are linked into $prefix
  if [[ -f "${nccl_file}" || -h "${nccl_file}" ]]; then
    # Replace /usr with $prefix and remove arch-specific lib directories
    nosysprefix="${nccl_file#"/usr/"}"
    noarchlib="${nosysprefix/#"lib/${arch}"/lib}"
    link_name="${prefix}/${noarchlib}"
    link_dir=$(dirname "${link_name}")
    mkdir -p "${link_dir}"
    ln -s "${nccl_file}" "${link_name}"
  else
    echo "Skipping ${nccl_file}"
  fi
done

@DwarKapex does it sound right to you?

[Steboss]

@DwarKapex

• Updated the based Dockerfile to have cuda-dl-base image
• to avoid having conflicts and re-install of nccl and cudnn i've modified install-cudnn.sh and install-nccl.sh
• for both script I added a part where we're doing the symlink step to that resources are present in /opt/nvidia/{package}

[olupton]

The nsys-jax test failures are because the 24.12 cuda-dl-base includes Nsight Systems 2024.7, whereas we currently install 2024.6 because of some issues with 2024.7 (#1176 is the - pending - attempt to move to 2024.7). A possible workaround would be to use the 24.11 cuda-dl-base for the moment.

[olupton]

LGTM, only one minor nit left.
@yhtang and/or @chaserileyroberts to review the GCP networking relevant parts.

[yhtang]

What is the symlink-xyz scripts modeled after? How do other DLFW containers accomodate the dl core container?

[Steboss]

@yhtang
The symlink-xyz scripts are meant to create a symlink for nvcc and cudnn. These packages are already install in cuda-dl-base image, but they're not linked to /opt/nvidia/ folder, as we were doing before.
This is all a jax/xla thing, that's why we might not need this in the other DLFW

[Steboss]

@olupton @yhtang
i think it may be wise to add an additional step in the CI, that check for the very latest cuda-dl-base image, so we can avoid updating this manually, and we'll have an automatic system that does it for us.

[yhtang]

@olupton @yhtang
i think it may be wise to add an additional step in the CI, that check for the very latest cuda-dl-base image, so we can avoid updating this manually, and we'll have an automatic system that does it for us.

Bumping the CUDA base image is usually not a light job, as it may break things. Hence why we always update it via a PR. The DL base image is only updated once a month so IMHO we can live with it.