Skip to content

Make Docker builds architecture-agnostic (x86_64 / aarch64) and modernize Dockerfile syntax #112

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
azrael417 merged 5 commits into from
Jul 16, 2026
Merged

Make Docker builds architecture-agnostic (x86_64 / aarch64) and modernize Dockerfile syntax #112

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
a498b1d
Make Dockerfiles architecture agnostic (x86_64 / aarch64)
azrael417 Jul 15, 2026
bd14bc5
updating ENV signature in all dockerfiles to modern standard
azrael417 Jul 15, 2026
80bd9be
increasing test tolerance for a test failing on arm. this is because …
azrael417 Jul 15, 2026
f553ede
cleaning up first LD_LIBRARY_PATH exports
azrael417 Jul 15, 2026
54eaa0e
Use buildkit exported TARGETARCH variable to automatically detect arc…
romerojosh Jul 15, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
26 changes: 17 additions & 9 deletions docker/Dockerfile
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,12 +1,20 @@
FROM nvcr.io/nvidia/nvhpc:26.5-devel-cuda13.2-ubuntu24.04
FROM nvcr.io/nvidia/nvhpc:26.5-devel-cuda13.2-ubuntu24.04 AS base

FROM base AS arch-amd64
ENV ARCH=x86_64

FROM base AS arch-arm64
ENV ARCH=aarch64

FROM arch-${TARGETARCH} AS final

SHELL ["/bin/bash", "-c"]

ENV CUDA_HOME /opt/nvidia/hpc_sdk/Linux_x86_64/26.5/cuda
ENV LD_LIBRARY_PATH /opt/nvidia/hpc_sdk/Linux_x86_64/26.5/cuda/13.2/extras/CUPTI/lib64:$LD_LIBRARY_PATH
ENV CUDA_HOME=/opt/nvidia/hpc_sdk/Linux_${ARCH}/26.5/cuda
ENV LD_LIBRARY_PATH=/opt/nvidia/hpc_sdk/Linux_${ARCH}/26.5/cuda/13.2/extras/CUPTI/lib64${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}

# Install System Dependencies
ENV DEBIAN_FRONTEND noninteractive
ENV DEBIAN_FRONTEND=noninteractive
RUN apt update -y && \
apt install -y wget cmake && \
apt install -y python3 python-is-python3 python3-pip python3-pybind11 && \
Expand All @@ -18,10 +26,10 @@ RUN pip3 install --break-system-packages typing_extensions && \
pip3 install --break-system-packages --no-deps nvidia-cudnn-cu13==9.20.0.48 nvidia-cusparselt-cu13==0.8.1 nvidia-nccl-cu13==2.29.7 nvidia-cufile==1.18.1.6

# Remove conflicting NCCL version from NVHPC SDK, add libnccl.so symlink to pip installed NCCL
RUN rm -rf /opt/nvidia/hpc_sdk/Linux_x86_64/26.5/comm_libs/nccl/ && \
RUN rm -rf /opt/nvidia/hpc_sdk/Linux_${ARCH}/26.5/comm_libs/nccl/ && \
cd /usr/local/lib/python3.12/dist-packages/nvidia/nccl/lib && \
ln -s libnccl.so.2 libnccl.so
ENV LD_LIBRARY_PATH /usr/local/lib/python3.12/dist-packages/nvidia/nccl/lib:$LD_LIBRARY_PATH
ENV LD_LIBRARY_PATH=/usr/local/lib/python3.12/dist-packages/nvidia/nccl/lib:$LD_LIBRARY_PATH

# Install yaml-cpp
RUN git clone https://github.com/jbeder/yaml-cpp.git --branch 0.8.0 && \
Expand All @@ -32,7 +40,7 @@ RUN git clone https://github.com/jbeder/yaml-cpp.git --branch 0.8.0 && \
-DBUILD_SHARED_LIBS=OFF \
-DCMAKE_POSITION_INDEPENDENT_CODE=ON .. && \
make -j$(nproc) && make install
ENV LD_LIBRARY_PATH /opt/yaml-cpp/lib:${LD_LIBRARY_PATH}
ENV LD_LIBRARY_PATH=/opt/yaml-cpp/lib:${LD_LIBRARY_PATH}

# Install additional Python dependencies
RUN pip3 install --break-system-packages wandb ruamel-yaml matplotlib pygame moviepy
Expand All @@ -52,5 +60,5 @@ RUN cd /torchfort && mkdir build && cd build && \
.. && \
make -j$(nproc) install && \
cd / && rm -rf torchfort
ENV LD_LIBRARY_PATH /opt/torchfort/lib:${LD_LIBRARY_PATH}
ENV LD_LIBRARY_PATH /usr/local/lib/python3.12/dist-packages/torch/lib:${LD_LIBRARY_PATH}
ENV LD_LIBRARY_PATH=/opt/torchfort/lib:${LD_LIBRARY_PATH}
ENV LD_LIBRARY_PATH=/usr/local/lib/python3.12/dist-packages/torch/lib:${LD_LIBRARY_PATH}
36 changes: 22 additions & 14 deletions docker/Dockerfile_gnu
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,11 +1,19 @@
FROM nvcr.io/nvidia/cuda:13.2.0-devel-ubuntu24.04
FROM nvcr.io/nvidia/cuda:13.2.0-devel-ubuntu24.04 AS base

FROM base AS arch-amd64
ENV ARCH=x86_64

FROM base AS arch-arm64
ENV ARCH=aarch64

FROM arch-${TARGETARCH} AS final

SHELL ["/bin/bash", "-c"]

ENV CUDA_HOME /usr/local/cuda
ENV CUDA_HOME=/usr/local/cuda

# Install System Dependencies
ENV DEBIAN_FRONTEND noninteractive
ENV DEBIAN_FRONTEND=noninteractive
RUN apt update -y && \
apt install -y wget cmake && \
apt install -y python3 python-is-python3 python3-pip python3-pybind11 && \
Expand All @@ -14,13 +22,13 @@ RUN apt update -y && \

# Download HPCX
RUN cd /opt && \
wget https://content.mellanox.com/hpc/hpc-x/v2.24.1_cuda13/hpcx-v2.24.1-gcc-doca_ofed-ubuntu24.04-cuda13-x86_64.tbz && \
tar xjf hpcx-v2.24.1-gcc-doca_ofed-ubuntu24.04-cuda13-x86_64.tbz && \
mv hpcx-v2.24.1-gcc-doca_ofed-ubuntu24.04-cuda13-x86_64 hpcx && \
cd /opt && rm hpcx-v2.24.1-gcc-doca_ofed-ubuntu24.04-cuda13-x86_64.tbz
wget https://content.mellanox.com/hpc/hpc-x/v2.24.1_cuda13/hpcx-v2.24.1-gcc-doca_ofed-ubuntu24.04-cuda13-${ARCH}.tbz && \
tar xjf hpcx-v2.24.1-gcc-doca_ofed-ubuntu24.04-cuda13-${ARCH}.tbz && \
mv hpcx-v2.24.1-gcc-doca_ofed-ubuntu24.04-cuda13-${ARCH} hpcx && \
cd /opt && rm hpcx-v2.24.1-gcc-doca_ofed-ubuntu24.04-cuda13-${ARCH}.tbz

ENV PATH /opt/hpcx/ompi/bin:$PATH
ENV LD_LIBRARY_PATH /opt/hpcx/ompi/lib:$LD_LIBRARY_PATH
ENV PATH=/opt/hpcx/ompi/bin:$PATH
ENV LD_LIBRARY_PATH=/opt/hpcx/ompi/lib${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}

RUN echo "source /opt/hpcx/hpcx-init.sh; hpcx_load" >> /root/.bashrc

Expand All @@ -32,9 +40,9 @@ RUN pip3 install --break-system-packages typing_extensions && \
# Add libnccl.so symlink to pip installed NCCL (CUDA devel image ships no NCCL)
RUN cd /usr/local/lib/python3.12/dist-packages/nvidia/nccl/lib && \
ln -s libnccl.so.2 libnccl.so
ENV LD_LIBRARY_PATH /usr/local/lib/python3.12/dist-packages/nvidia/nccl/lib:$LD_LIBRARY_PATH
ENV LD_LIBRARY_PATH=/usr/local/lib/python3.12/dist-packages/nvidia/nccl/lib:$LD_LIBRARY_PATH
# NVSHMEM is a torch runtime dependency (libnvshmem_host.so.3) not present in the CUDA image
ENV LD_LIBRARY_PATH /usr/local/lib/python3.12/dist-packages/nvidia/nvshmem/lib:$LD_LIBRARY_PATH
ENV LD_LIBRARY_PATH=/usr/local/lib/python3.12/dist-packages/nvidia/nvshmem/lib:$LD_LIBRARY_PATH

# Install yaml-cpp
RUN git clone https://github.com/jbeder/yaml-cpp.git --branch 0.8.0 && \
Expand All @@ -45,7 +53,7 @@ RUN git clone https://github.com/jbeder/yaml-cpp.git --branch 0.8.0 && \
-DBUILD_SHARED_LIBS=OFF \
-DCMAKE_POSITION_INDEPENDENT_CODE=ON .. && \
make -j$(nproc) && make install
ENV LD_LIBRARY_PATH /opt/yaml-cpp/lib:${LD_LIBRARY_PATH}
ENV LD_LIBRARY_PATH=/opt/yaml-cpp/lib:${LD_LIBRARY_PATH}

# Install additional Python dependencies
RUN pip3 install --break-system-packages wandb ruamel-yaml matplotlib pygame moviepy
Expand All @@ -64,5 +72,5 @@ RUN source /opt/hpcx/hpcx-init.sh && hpcx_load && \
.. && \
make -j$(nproc) install && \
cd / && rm -rf torchfort
ENV LD_LIBRARY_PATH /opt/torchfort/lib:${LD_LIBRARY_PATH}
ENV LD_LIBRARY_PATH /usr/local/lib/python3.12/dist-packages/torch/lib:${LD_LIBRARY_PATH}
ENV LD_LIBRARY_PATH=/opt/torchfort/lib:${LD_LIBRARY_PATH}
ENV LD_LIBRARY_PATH=/usr/local/lib/python3.12/dist-packages/torch/lib:${LD_LIBRARY_PATH}
8 changes: 4 additions & 4 deletions docker/Dockerfile_gnu_cpuonly
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@ FROM ubuntu:24.04
SHELL ["/bin/bash", "-c"]

# Install System Dependencies
ENV DEBIAN_FRONTEND noninteractive
ENV DEBIAN_FRONTEND=noninteractive
RUN apt update -y && \
apt install -y build-essential && \
apt install -y wget cmake && \
Expand All @@ -24,7 +24,7 @@ RUN git clone https://github.com/jbeder/yaml-cpp.git --branch 0.8.0 && \
-DBUILD_SHARED_LIBS=OFF \
-DCMAKE_POSITION_INDEPENDENT_CODE=ON .. && \
make -j$(nproc) && make install
ENV LD_LIBRARY_PATH /opt/yaml-cpp/lib:${LD_LIBRARY_PATH}
ENV LD_LIBRARY_PATH=/opt/yaml-cpp/lib

# Install additional Python dependencies
RUN pip3 install --break-system-packages wandb ruamel-yaml matplotlib pygame moviepy
Expand All @@ -42,5 +42,5 @@ RUN cd /torchfort && mkdir build && cd build && \
.. && \
make -j$(nproc) install && \
cd / && rm -rf torchfort
ENV LD_LIBRARY_PATH /opt/torchfort/lib:${LD_LIBRARY_PATH}
ENV LD_LIBRARY_PATH /usr/local/lib/python3.12/dist-packages/torch/lib:${LD_LIBRARY_PATH}
ENV LD_LIBRARY_PATH=/opt/torchfort/lib:${LD_LIBRARY_PATH}
ENV LD_LIBRARY_PATH=/usr/local/lib/python3.12/dist-packages/torch/lib:${LD_LIBRARY_PATH}
9 changes: 8 additions & 1 deletion tests/rl/test_replay_buffer.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -258,7 +258,14 @@ TEST_P(ReplayBuffer, NStepConsistency) {
}

EXPECT_FLOAT_EQ(sdiff, 0.);
EXPECT_FLOAT_EQ(rdiff, 0.);
// rdiff accumulates float32 rounding from the n-step discounted reward sum
// (r0 + r1*gamma + r2*gamma^2 + ...) over the whole batch. The buffer and this
// manual recomputation use mathematically equivalent but differently-rounded
// paths (std::pow + torch tensor ops vs. repeated scalar float multiply), which
// diverge by a few ULP per term and differ across architectures (e.g. FMA
// contraction on aarch64). Compare against a small absolute tolerance rather
// than exact float equality.
EXPECT_NEAR(rdiff, 0., 1e-4);
}

TEST_P(ReplayBuffer, SaveRestore) {
Expand Down
Loading