Unable to Initialize the Cuda in Alphafold

[ravisaroch]

Error Description
Unable to initialize backend 'cuda': jaxlib/cuda/versions_helpers.cc:98: operation cuInit(0) failed: Unknown CUDA error 303; cuGetErrorName failed. This probably means that JAX was unable to load the CUDA libraries.

Command
python3 docker/run_docker.py --fasta_paths=/Data/Ind_NEG/negative_9.fasta --max_template_date=2024-08-01 --data_dir=/Data/AFDB_AUG2024 --output_dir=/Data/OUT_Ind_NEG

Jax and jaxlib vrsion
0.4.24

nvidia-smi
Fri Mar 28 12:37:33 2025

+-----------------------------------------------------------------------------------------+
| NVIDIA-SMI 550.120                Driver Version: 550.120        CUDA Version: 12.4     |
|-----------------------------------------+------------------------+----------------------+
| GPU  Name                 Persistence-M | Bus-Id          Disp.A | Volatile Uncorr. ECC |
| Fan  Temp   Perf          Pwr:Usage/Cap |           Memory-Usage | GPU-Util  Compute M. |
|                                         |                        |               MIG M. |
|=========================================+========================+======================|
|   0  NVIDIA RTX A6000               Off |   00000000:21:00.0  On |                  Off |
| 30%   47C    P8             22W /  300W |     536MiB /  49140MiB |      2%      Default |
|                                         |                        |                  N/A |
+-----------------------------------------+------------------------+----------------------+
                                                                                         
+-----------------------------------------------------------------------------------------+
| Processes:                                                                              |
|  GPU   GI   CI        PID   Type   Process name                              GPU Memory |
|        ID   ID                                                               Usage      |
|=========================================================================================|
|    0   N/A  N/A      4977      G   /usr/lib/xorg/Xorg                            245MiB |

+-----------------------------------------------------------------------------------------+

Alphafold build command
docker run --rm --gpus all nvidia/cuda:12.4-.0base nvidia-smi

docker build -f docker/Dockerfile -t alphafold .

output Alphafold build command
[+] Building 2.2s (18/18) FINISHED docker:default

 => [internal] load build definition from Dockerfile                                                                                                                                                                                                                          
 => => transferring dockerfile: 3.60kB                                                                                                                                                                                                                                          
 => [internal] load metadata for docker.io/nvidia/cuda:12.2.2-cudnn8-runtime-ubuntu20.04                                                                                                                                                                                        
 => [auth] nvidia/cuda:pull token for registry-1.docker.io                                                                                                                                                                                                                      
 => [internal] load .dockerignore                                                                                                                                                                                                                                                 => transferring context: 72B                                                                                                                                                                                                                                                
 => [ 1/12] FROM _docker.io/nvidia/cuda:12.2.2-cudnn8-runtime-ubuntu20.04@sha256:8d29e0f835bf0c038a105444fa6f07fcbc0081ca11f6a9fa72fd67ac82a1de9f_                                                                                                                                
 => [internal] load build context                                                                                                                                                                                                                                               
 => transferring context:                                                                                                                                                                                                                                         
 => CACHED [ 2/12] RUN apt-get update --quiet     && apt-get install --no-install-recommends --yes --quiet         build-essential         cmake         cuda-command-line-tools-$(cut -f1,2 -d- <<< 12-2-2)         git         hmmer         kalign         tzdata         w  0.0s
 => CACHED [ 3/12] RUN git clone --branch v3.3.0 --single-branch https://github.com/soedinglab/hh-suite.git /tmp/hh-suite     && mkdir /tmp/hh-suite/build     && pushd /tmp/hh-suite/build     && cmake -DCMAKE_INSTALL_PREFIX=/opt/hhsuite ..     && make --jobs 4 && make i 
 => CACHED [ 4/12] RUN wget -q -P /tmp   https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh     && bash /tmp/Miniconda3-latest-Linux-x86_64.sh -b -p /opt/conda     && rm /tmp/Miniconda3-latest-Linux-x86_64.sh                                             
 => CACHED [ 5/12] RUN conda install --quiet --yes conda==24.11.1 pip python=3.11     && conda install --quiet --yes --channel nvidia cuda=12.2.2     && conda install --quiet --yes --channel conda-forge openmm=8.0.0 pdbfixer     && conda clean --all --force-pkgs-dirs -
 => CACHED [ 6/12] COPY . /app/alphafold                                                                                                                                                                                                                                        
 => CACHED [ 7/12] RUN wget -q -P /app/alphafold/alphafold/common/   https://git.scicore.unibas.ch/schwede/openstructure/-/raw/7102c63615b64735c4941278d92b554ec94415f8/modules/mol/alg/src/stereo_chemical_props.txt   
 => CACHED [ 8/12] RUN pip3 install --upgrade pip --no-cache-dir     && pip3 install -r /app/alphafold/requirements.txt --no-cache-dir     && pip3 install --upgrade --no-cache-dir       jax==0.4.26       jaxlib==0.4.26+cuda12.cudnn89       -f https://storage.googleapis.  0.0s
 => CACHED [ 9/12] RUN chmod u+s /sbin/ldconfig.real                                                                                                                                                                                                                            
 => CACHED [10/12] RUN ln -sf /usr/lib/x86_64-linux-gnu/libffi.so.7 /opt/conda/lib/libffi.so.7                                                                                                                                                                                  
 => CACHED [11/12] WORKDIR /app/alphafold                                                                                                                                                                                                                                       
 => CACHED [12/12] RUN echo $'#!/bin/bash\nldconfig\npython /app/alphafold/run_alphafold.py 
 = > /app/run_alphafold.sh   && chmod +x /app/run_alphafold.sh 
 => exporting to image                                                                                                                                                                                                                                                          
 => exporting layers                                                                                                                                                                                                                                                         
 => writing image sha256:4c2c004aa30be040fcc7682e9dfe45e556f6f61f3c28cfee1e624d8dae91052e                                                                                                                                                                                    
 => naming to docker.io/library/alphafold