Error while building docker image: ERROR: Could not find a version that satisfies the requirement jaxlib==0.1.69+cuda111

[MurphyDavid]

Commands to replicate on Ubuntu 20.04:

git clone https://github.com/deepmind/alphafold.git
cd alphafold/
docker build -f docker/Dockerfile -t alphafold .

The error:

WARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv
Looking in links: https://storage.googleapis.com/jax-releases/jax_releases.html
Requirement already satisfied: jax==0.2.14 in /opt/conda/lib/python3.7/site-packages (0.2.14)
ERROR: Could not find a version that satisfies the requirement jaxlib==0.1.69+cuda111 (from versions: 0.1.23, 0.1.29, 0.1.32, 0.1.40, 0.1.41, 0.1.42, 0.1.43, 0.1.44, 0.1.46, 0.1.50, 0.1.51, 0.1.52, 0.1.55, 0.1.56, 0.1.57, 0.1.58, 0.1.59, 0.1.60, 0.1.61, 0.1.62, 0.1.63, 0.1.64, 0.1.65, 0.1.66, 0.1.67, 0.1.68, 0.1.69, 0.1.70, 0.1.71, 0.1.72, 0.1.73, 0.1.74, 0.1.75, 0.1.76, 0.3.0, 0.3.2, 0.3.5, 0.3.7, 0.3.8, 0.3.10)
ERROR: No matching distribution found for jaxlib==0.1.69+cuda111
The command '/bin/bash -c pip3 install --upgrade pip     && pip3 install -r /app/alphafold/requirements.txt     && pip3 install --upgrade       jax==0.2.14       jaxlib==0.1.69+cuda$(cut -f1,2 -d. <<< ${CUDA} | sed 's/\.//g')       -f https://storage.googleapis.com/jax-releases/jax_releases.html' returned a non-zero code: 1

[Hoda121]

Changing this line in DockerFile: && pip3 install --upgrade jax jaxlib==0.1.69+cuda${CUDA/./} -f
to this: && pip3 install --upgrade jax==0.2.14 jaxlib>=0.1.69+cuda${CUDA/./} -f
solves this error

[brianloyal]

I think this may have something to do with it as well: jax-ml/jax#11142. Seems like the jaxlib+cuda source changed earlier this week

[xtalyama]

you can find the jaxlib==0.1.69+cuda111 version
by changing this line in Dockerfile: -f https://storage.googleapis.com/jax-releases/jax_releases.html
to: -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

[Phage-structure-geek]

you can find the jaxlib==0.1.69+cuda111 version by changing this line in Dockerfile: -f https://storage.googleapis.com/jax-releases/jax_releases.html to: -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

I confirm that this fix works on Ubuntu 20.04.4 LTS. Docker image compiles without errors and runs fine. Thank you, @xtalyama!

[BlankSama]

I tried changing the URL to https://storage.googleapis.com/jax-releases/jax_cuda_releases.html in Dockerfile but it is showing the error :

@Phage-structure-geek @Augustin-Zidek @xtalyama any help is appreciated

Edit: the problem is solved and alpha fold 2.2 is running fine with GPU

[kaiwanxiao1996]

Dear Alphafold team, i Changed this line in DockerFile: && pip3 install --upgrade jax jaxlib==0.1.69+cuda${CUDA/./} -f
to this: && pip3 install --upgrade jax==0.2.14 jaxlib>=0.1.69+cuda${CUDA/./} -f. Although this step ran successfully, a new problem showed when running alphafold2. "run alphafold2 error"AttributeError: module 'jaxlib.pocketfft' has no attribute 'pocketfft'" ".

[swanss]

@kaiwanxiao1996 Hi, I had the exact same problem. When I checked the conda environment I noticed that the jaxlib version was much newer than 0.1.69.

This worked for me:
pip3 install --upgrade jax==0.2.14 jaxlib==0.1.69+cuda111 -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

[alirezaomidi]

Simply downgrading the jaxlib to 0.1.69 worked for me.