Install Miniconda flavor that is suitable for the current CPU arch.

[martin-g]

This is a small first step for supporting Linux ARM64. Related-to: https://github.com/deepmind/alphafold/issues/528 and https://github.com/deepmind/alphafold/issues/530

The Docker image builds fail later with:

...
 Downloading jax-0.2.14.tar.gz (669 kB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 669.2/669.2 kB 32.0 MB/s eta 0:00:00
  Preparing metadata (setup.py): started
  Preparing metadata (setup.py): finished with status 'done'
ERROR: Could not find a version that satisfies the requirement jaxlib==0.1.69+cuda111 (from versions: none)
ERROR: No matching distribution found for jaxlib==0.1.69+cuda111
...

I tried to upgrade jax from 0.2.14 to 0.3.16 (current latest) in requirements.txt but it is not picked up by the Docker build. I will work on this in a separate PR because it seems it will be a bigger change.

[martin-g]

The jaxlib issue was due to https://github.com/deepmind/alphafold/blob/22fc0eea17f3f2941e0bee23d2ca0d72a0a89e1e/docker/Dockerfile#L69-L70

but there are no arm64 wheels at https://storage.googleapis.com/jax-releases/jax_cuda_releases.html I will see how this could be improved with the jaxlib project!

Update: https://github.com/conda-forge/jaxlib-feedstock/issues/125

[martin-g]

Any feedback on the proposed changes ?

After this improvement the x86_64 users won't notice any change in the behavior! But it will help to Linux aarch64 and s390x users, because Miniconda provides binaries for those. For any other CPU architectures, e.g. ppcle, the error message will be more clear - it will fail to download Miniconda with 404. The current behavior is that it will download the x86_64 binary and fail much later with error like standard_init_linux.go:211: exec user process caused “exec format error” which does not tell what exactly went wrong.