Customize NCCL for base container

[DwarKapex]

Add ability to build a base container with a custom NCCL version.
2 main scenarios are considered:

• Default build (JAX_NCCL_VERSION is not defined): JAX_NCCL_VERSION will be assigned to the value of an existing NCCL_VERSION env var of the base image. If NCCL_VERSION is not defined in the base image, then JAX_NCCL_VERSION will be assigned to the version of libnccl-dev library if it exists, or left empty.
• Custom build with user defined JAX_NCCL_VERSION: libnccl2 and libnccl-dev will be installed with the provided version of NCCL if NCCL not in the system or updated if NCCL comes with the base image.

Installation of NCCL happens if and only if either JAX_NCCL_VERSION or NCCL_VERSION is defined.

[olupton]

Where would a default be specified? I think the situation today is that we want the default nightly/CI behaviour to be "12.6.1 base image with NCCL 2.23 installed on top", so it should be possible to express that.

I know that for the base image, there is special treatment of latest to enable this.

[yhtang]

So the "special treatment" is to embed a default in the Dockerfile, and then conditionally override it via an Actions conditional expression:

JAX-Toolbox/.github/workflows/_build_base.yaml

Line 135 in 61d8446

${{ inputs.BASE_IMAGE != 'latest' && format('BASE_IMAGE={0}', inputs.BASE_IMAGE) || '' }}