From 892c1b2228b3891eedb18c617ab4e8a4b364dbb3 Mon Sep 17 00:00:00 2001
From: Eric Work <work.eric@gmail.com>
Date: Sat, 12 Sep 2026 01:38:41 -0700
Subject: [PATCH] build: Support x86_64 when using `docker compose`

You can use a command such as the following,

```docker compose build --build-arg ARCH=$(uname -m)```

to build the CI container for x86_64 (or aarch64).

- The default CPU architecture remains `aarch64` so CI doesn't break.
- Prefix `mpirun` with `HWLOC_COMPONENTS=-gl` to avoid a hang when graphics
  are present on the GPU.

Signed-off-by: Eric Work <work.eric@gmail.com>
---
 Dockerfile         | 6 ++++--
 docker-compose.yml | 2 ++
 2 files changed, 6 insertions(+), 2 deletions(-)

diff --git a/Dockerfile b/Dockerfile
index eabbfd4af0..9de139a0f2 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -3,6 +3,8 @@
 
 FROM nvidia/cuda:13.3.0-devel-ubuntu24.04@sha256:ef2203909e80b8b976cfc672f7e2ae2b00bc0e25c404ee86d89e10a3802f1c52
 
+ARG ARCH=aarch64
+
 ENV DEBIAN_FRONTEND=noninteractive
 ENV TORCH_CUDA_ARCH_LIST=10.0
 ENV OMPI_ALLOW_RUN_AS_ROOT=1
@@ -81,7 +83,7 @@ RUN python3.12 -m venv "$VIRTUAL_ENV" \
 ENV TRT_LIB_DIR=/opt/venv/lib/python3.12/site-packages/tensorrt_libs
 ENV NCCL_LIB_DIR=/opt/venv/lib/python3.12/site-packages/nvidia/nccl/lib
 ENV TVM_FFI_LIB_DIR=/opt/venv/lib/python3.12/site-packages/tvm_ffi/lib
-ENV TRT_INC_DIR=/usr/include/aarch64-linux-gnu
+ENV TRT_INC_DIR=/usr/include/$ARCH-linux-gnu
 ENV LD_LIBRARY_PATH=$TRT_LIB_DIR:$NCCL_LIB_DIR:$TVM_FFI_LIB_DIR:/usr/local/cuda/lib64
 ENV LD_PRELOAD=/usr/local/cuda/lib64/libcublas.so.13
 
@@ -92,7 +94,7 @@ RUN python3.12 -c \
     && test -f "$TRT_LIB_DIR/libnvonnxparser.so.11" \
     && test -f "$NCCL_LIB_DIR/libnccl.so.2" \
     && test -f "$TVM_FFI_LIB_DIR/libtvm_ffi.so" \
-    && mpirun --tag-output -np 1 true
+    && HWLOC_COMPONENTS=-gl mpirun --tag-output -np 1 true
 
 WORKDIR /workspace/tensorrt-model-connect
 CMD ["bash"]
diff --git a/docker-compose.yml b/docker-compose.yml
index 5dcf01e6d5..f7ccc8538c 100644
--- a/docker-compose.yml
+++ b/docker-compose.yml
@@ -6,6 +6,8 @@ services:
     build:
       context: .
       dockerfile: Dockerfile
+      args:
+        - ARCH=${ARCH:-aarch64}
     container_name: trtmc-dev
     working_dir: /workspace/tensorrt-model-connect
     tty: true