---
title: "AlphaFold Docker 설치 시 JAX와 CUDA 호환성 문제 해결하기"
description: "AlphaFold Docker 설치 중 JAX와 CUDA 호환성 문제를 해결하기 위한 과정이 설명된다. JAX 라이브러리가 CUDA를 인식하지 못하는 문제를 분석하고, nvidia-modprobe 설치, 시스템 재설치, CUDA 버전 변경 등의 방법을 시도한 후, 최종적으로 "
date: "2024-10-24"
last_modified: "2026-05-15T08:36:00.000Z"
type: "Post"
tags:
  - "alphafold"
  - "docker"
  - "python"
  - "troubleshooting"
  - "gpu"
  - "ubuntu"
  - "conda"
  - "install"
categories:
  - "🤖 Computer Science"
canonical_url: "https://blog.pieroot.xyz/alphafold-jax-cuda"
markdown_url: "https://blog.pieroot.xyz/alphafold-jax-cuda.md"
---

# AlphaFold Docker 설치 시 JAX와 CUDA 호환성 문제 해결하기

AlphaFold Docker 설치 중 JAX와 CUDA 호환성 문제를 해결하기 위한 과정이 설명된다. JAX 라이브러리가 CUDA를 인식하지 못하는 문제를 분석하고, nvidia-modprobe 설치, 시스템 재설치, CUDA 버전 변경 등의 방법을 시도한 후, 최종적으로 

### 서론: AlphaFold Docker 설치의 고난

AlphaFold를 사용하려고 Docker 이미지를 빌드하다 보면 의존성 문제로 골머리를 앓게 됩니다. 특히 conda를 통해 CUDA 드라이버와 툴킷을 설치할 때, 설치에 실패하거나 의존성에 맞지 않는 드라이버가 자동으로 설치되는 경우가 빈번합니다.

이 포스트에서는 AlphaFold Docker 이미지 빌드 중 발생한 다양한 문제들과 해결 과정을 공유합니다. **결론만 보고 싶다면** [아무튼 해결](https://app.notion.com/p/21b067c015d0808abadcffde92288e21#21b067c015d0805bb5b3cca24b23edf2) 섹션으로 바로 이동하세요.

> **참고 자료**
> 
> - [AlphaFold GitHub Repository](https://github.com/google-deepmind/alphafold)
> 
> - [JAX Installation Guide](https://jax.readthedocs.io/en/latest/installation.html)
> 
> - [NVIDIA JAX Release Notes](https://docs.nvidia.com/deeplearning/frameworks/jax-release-notes/)

---

### 문제 원인 분석

AlphaFold Docker 이미지를 빌드하던 중, **JAX 라이브러리가 CUDA를 인식하지 못하는 문제**가 발생했습니다.

![image](https://blog.pieroot.xyz/api/image-proxy?id=b9601a70-1b2d-4415-84cb-453630e0449f&kind=s3&pageId=21b067c0-15d0-808a-badc-ffde92288e21&source=block&blockId=21b067c0-15d0-8011-bb97-ee0f70800f4e)

아래는 AlphaFold GitHub에서 제공하는 `Dockerfile`의 JAX 설치 구문입니다:

```docker
# Install pip packages.
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.com/jax-releases/jax_cuda_releases.html
```

> **JAX와 jaxlib의 관계**
> 
> - **JAX**: NumPy 스타일의 Python 코드를 GPU에서 빠르게 실행할 수 있게 해주는 라이브러리
> 
> - **jaxlib**: JAX의 백엔드 라이브러리로, CUDA/cuDNN 지원을 위한 바이너리 포함
> 
> - JAX는 CPU 버전으로, jaxlib는 CUDA를 지원하기 위한 버전으로 각각 설치됩니다

과거에는 정상적으로 설치되던 라이브러리가 갑자기 CUDA 버전과 종속성이 맞지 않는 문제가 발생했습니다. 이는 JAX의 버전 업데이트와 CUDA/cuDNN 버전 간의 호환성 변화 때문입니다.

---

### 트러블슈팅 과정

#### 1. JAX Toolbox 사용 시도

보다 완벽한 해결을 위해 JAX에서 공식 지원하는 **JAX Toolbox**를 사용하려 했습니다.

> **실패 원인**
> 
> - 서버급 GPU만 지원 (RTX 시리즈 미지원)
> 
> - apt 패키지를 설치하지 못하는 문제 발생
> 
> - JAX Toolbox는 Rosetta라는 자체 이미지를 사용하여 커스터마이징이 어려움

#### 2. nvidia-modprobe 패키지 설치

NVIDIA 드라이버가 시스템에 제대로 잡히지 않을 경우를 대비하여 `nvidia-modprobe`를 설치했습니다.

```shell
sudo apt install nvidia-modprobe
```

이 모듈은 NVIDIA 드라이버를 안전하게 커널에 로드하는 역할을 합니다. Dockerfile에 다음 구문을 추가했습니다:

```docker
RUN apt-get update \
    && DEBIAN_FRONTEND=noninteractive apt-get install --no-install-recommends -y \
        build-essential \
        cmake \
        cuda-command-line-tools-$(cut -f1,2 -d- <<< ${CUDA//./-}) \
        nvidia-modprobe \
        git \
        hmmer \
        kalign \
        tzdata \
        wget \
    && rm -rf /var/lib/apt/lists/* \
    && apt-get autoremove -y \
    && apt-get clean
```

> 근본적인 문제를 해결하지는 못했지만, 없는 것보다는 있는 것이 나으므로 유지하기로 했습니다.

#### 3. 시스템 재설치

혹시 서버 설치 과정에서 문제가 발생했을 수 있다고 생각하여 시스템 초기화를 진행했습니다.

**결과**: 동일한 문제 발생. 시스템 문제가 아님을 확인했습니다.

#### 4. CUDA 버전 변경

AlphaFold의 기본 설정을 분석해 보았습니다:

```docker
ARG CUDA=12.2.2
FROM nvidia/cuda:${CUDA}-cudnn8-runtime-ubuntu20.04

jaxlib==0.4.26+cuda12.cudnn89
```

AlphaFold는 **CUDA 12.2.2**와 **cuDNN 8.9** 버전을 사용하며, Docker 이미지로 runtime 이미지를 사용하기 때문에 CUDA dev 드라이버를 포함하지 않습니다.

![image](https://blog.pieroot.xyz/api/image-proxy?id=ad87f594-5e22-4120-b029-99e3dfb9b13b&kind=s3&pageId=21b067c0-15d0-808a-badc-ffde92288e21&source=block&blockId=21b067c0-15d0-805d-81c8-c2126b9c1b24)

**NVIDIA 권장 사항 (JAX 0.4.26 기준)**

AlphaFold 설정과 NVIDIA 권장 사항이 다른 것을 확인했습니다. Docker 이미지를 **CUDA 12.4.1**, **Ubuntu 22.04** 기반으로 변경하기로 결정했습니다.

```docker
ARG CUDA=12.4.1
# FROM nvidia/cuda:${CUDA}-cudnn8-runtime-ubuntu20.04
FROM nvidia/cuda:12.4.1-cudnn-devel-ubuntu22.04
```

![image](https://blog.pieroot.xyz/api/image-proxy?id=842f3f90-a583-4433-aef6-6183e7a23977&kind=s3&pageId=21b067c0-15d0-808a-badc-ffde92288e21&source=block&blockId=21b067c0-15d0-80b2-9bb0-e691a0f6d6d4)

추가로 cuDNN 9.1.0.70을 포함하는 jaxlib 0.4.26 버전의 pip 패키지를 찾아야 했습니다.

![image](https://blog.pieroot.xyz/api/image-proxy?id=01ca9f77-b83b-4334-a782-4eb45fdfa18b&kind=s3&pageId=21b067c0-15d0-808a-badc-ffde92288e21&source=block&blockId=21b067c0-15d0-804a-bb66-f6f36c3bdc73)

> **문제 발견**
> 
> CUDA 12를 지원하면서 cuDNN 9.1을 지원하는 pip 패키지가 존재하지 않았습니다.

#### 5. JAX 설치 방법 재확인

JAX 공식 문서를 다시 확인했습니다.

![image](https://blog.pieroot.xyz/api/image-proxy?id=ad58c4e4-4d20-4ab5-b358-f05706899bff&kind=s3&pageId=21b067c0-15d0-808a-badc-ffde92288e21&source=block&blockId=21b067c0-15d0-8024-bf30-e1d77f0c7662)

> **JAX GPU 설치 권장 방법**
> 
> JAX에서는 CUDA를 사용하고 싶을 때 `jax[cuda]`를 사용하라고 권장합니다. 하지만 AlphaFold는 특정 CUDA/cuDNN 버전을 사용하기 위해 CPU 버전의 JAX를 설치하고 나중에 cuDNN을 지원하는 jaxlib를 별도로 설치하는 방식을 사용했습니다.

![image](https://blog.pieroot.xyz/api/image-proxy?id=5ad9f53a-8727-45f4-b966-b72e0f114d97&kind=s3&pageId=21b067c0-15d0-808a-badc-ffde92288e21&source=block&blockId=21b067c0-15d0-8097-998c-eb5436fb161d)

예전 버전을 사용하고 싶을 경우 jaxlib를 직접 설치하는 것을 권장하고 있습니다.

![image](https://blog.pieroot.xyz/api/image-proxy?id=24265e59-1742-4ecd-b672-8959ec14d93d&kind=s3&pageId=21b067c0-15d0-808a-badc-ffde92288e21&source=block&blockId=21b067c0-15d0-8081-b9c6-ff0b4c995f53)

**JAX 0.4.29 버전 이후**부터는 `jax[cuda12]`로 설치하는 것을 권장합니다.

![image](https://blog.pieroot.xyz/api/image-proxy?id=e87a9937-f8d4-4dfd-b3f9-335dd385b3fb&kind=s3&pageId=21b067c0-15d0-808a-badc-ffde92288e21&source=block&blockId=21b067c0-15d0-8040-b69d-f15e91a93e1f)

다행히 0.4.29 버전은 cuDNN 9.1을 지원하므로 이 버전을 사용하기로 결정했습니다. 추가로 pip install 시 버전을 명시하지 않으면 의존성 확인 후 최신 버전으로 설치되도록 시도해 보았습니다.

---

### 아무튼 해결

위의 모든 방법을 시도했지만 해결되지 않았습니다. **진짜 원인은 다른 곳에 있었습니다.**

![image](https://blog.pieroot.xyz/api/image-proxy?id=b692d2ad-c7e6-4fca-8f5c-0ffc8fedc318&kind=s3&pageId=21b067c0-15d0-808a-badc-ffde92288e21&source=block&blockId=21b067c0-15d0-808d-9f46-e5d8a398675c)

AlphaFold의 `run_`[`alphafold.py`](http://alphafold.py/) 파일에는 Docker를 실행하기 위한 구문이 있습니다. `device_requests` 변수가 호스트 PC의 GPU를 감지하고, [`containers.run`](http://containers.run/)을 실행하면 자동으로 마운트하는 시스템입니다.

> **이상한 상황**
> 
> - 호스트 PC에서 `nvidia-smi` 명령어 정상 동작 ✅
> 
> - Docker 실행 시 `--gpus all` 옵션 추가하면 컨테이너 내부에서 GPU 인식 ✅
> 
> - 컨테이너 내부에서 JAX 라이브러리가 GPU 정상 인식 ✅
> 
> - **그런데 AlphaFold만 안 됨** ❌

모든 드라이버와 설정이 정상인데 AlphaFold만 동작하지 않는 상태였습니다.

#### 진짜 해결책

보통은 `docker.types.DeviceRequest`가 GPU를 전부 마운트해 주어야 하지만, 간혹 마운트되지 못하는 문제가 있습니다.

**`count=-1`**** 옵션을 추가**해 주면 정상적으로 컨테이너 내부에서 GPU를 사용할 수 있습니다.

```python
import docker

client = docker.from_env()
device_requests = [docker.types.DeviceRequest(driver="nvidia", capabilities=[["gpu"]], count=-1)]
print(device_requests)

logs = client.containers.run(
    "nvidia/cuda:12.2.2-runtime-ubuntu20.04",
    "nvidia-smi",
    runtime="nvidia",
    device_requests=device_requests,
    remove=True,
)

print(logs.decode("utf-8"))
```

> **정상 출력 예시**
> 
> 위 명령어를 실행했을 때 아래와 같이 GPU 정보가 출력되면 성공입니다.

```shell
[{'Driver': 'nvidia', 'Count': -1, 'DeviceIDs': [], 'Capabilities': [['gpu']], 'Options': {}}]

==========
== CUDA ==
==========

CUDA Version 12.2.2

Container image Copyright (c) 2016-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.

+-----------------------------------------------------------------------------------------+
| NVIDIA-SMI 550.90.12              Driver Version: 550.90.12      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. |
|=========================================+========================+======================|
|   0  NVIDIA GeForce RTX 3060        Off |   00000000:3E:00.0  On |                  N/A |
|  0%   49C    P8             20W /  170W |       2MiB /  12288MiB |      0%      Default |
+-----------------------------------------+------------------------+----------------------+
```

---

### 완성된 Dockerfile

아래는 최종적으로 동작하는 Dockerfile 내용입니다.

> **실행 위치 주의**
> 
> 아래 Dockerfile은 AlphaFold 디렉토리가 아닌 **프로젝트 root( / ) 디렉토리**에서 실행해야 합니다.

```docker
# Copyright 2021 DeepMind Technologies Limited
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

ARG CUDA=12.2.2
FROM nvidia/cuda:${CUDA}-cudnn8-runtime-ubuntu20.04
# FROM directive resets ARGS, so we specify again (the value is retained if
# previously set).
ARG CUDA

# Use bash to support string substitution.
SHELL ["/bin/bash", "-o", "pipefail", "-c"]

RUN apt-get update \
  && DEBIAN_FRONTEND=noninteractive apt-get install --no-install-recommends -y \
  build-essential \
  cmake \
  cuda-command-line-tools-12-2 \
  git \
  hmmer \
  kalign \
  tzdata \
  wget \
  && rm -rf /var/lib/apt/lists/* \
  && apt-get autoremove -y \
  && apt-get clean

# Compile HHsuite from source.
RUN git clone --branch v3.3.0 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 -j 4 && make install \
  && ln -s /opt/hhsuite/bin/* /usr/bin \
  && popd \
  && rm -rf /tmp/hh-suite

# Install Miniconda package manager.
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

# Install conda packages.
ENV PATH="/opt/conda/bin:$PATH"
ENV LD_LIBRARY_PATH="/opt/conda/lib:$LD_LIBRARY_PATH"
RUN conda install -qy conda==24.1.2 pip python=3.11 \
  # && conda install -y -c nvidia cuda=${CUDA_VERSION} \
  && conda install -y nvidia/label/cuda-${CUDA_VERSION}::cuda \
  && conda install -y -c conda-forge openmm=8.0.0 pdbfixer ncurses \
  && conda clean --all --force-pkgs-dirs --yes

COPY ./bundle/alphafold /app/alphafold
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

# Install pip packages.
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[cuda12]==0.4.26 \
  jaxlib==0.4.26+cuda12.cudnn89 \
  -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

# Add SETUID bit to the ldconfig binary so that non-root users can run it.
RUN chmod u+s /sbin/ldconfig.real

# Currently needed to avoid undefined_symbol error.
RUN ln -sf /usr/lib/x86_64-linux-gnu/libffi.so.7 /opt/conda/lib/libffi.so.7

# We need to run `ldconfig` first to ensure GPUs are visible, due to some quirk
# with Debian. See https://github.com/NVIDIA/nvidia-docker/issues/1399 for
# details.
WORKDIR /app/alphafold
RUN echo $'#!/bin/bash\n\
  ldconfig\n\
  python /app/alphafold/run_alphafold.py "$@"' > /app/run_alphafold.sh \
  && chmod +x /app/run_alphafold.sh
ENTRYPOINT ["/app/run_alphafold.sh"]
```

---

### 핵심 정리

이번 트러블슈팅에서 배운 점을 정리하면 다음과 같습니다:

> **최종 해결책**
> 
> `docker.types.DeviceRequest`에 `count=-1` 옵션을 추가하면 모든 GPU가 정상적으로 마운트됩니다. 이 간단한 수정으로 며칠간의 삽질이 해결되었습니다.

---

### 참고 링크

- [AlphaFold GitHub](https://github.com/google-deepmind/alphafold)

- [JAX Installation Documentation](https://jax.readthedocs.io/en/latest/installation.html)

- [NVIDIA Docker Issues #1399](https://github.com/NVIDIA/nvidia-docker/issues/1399)