From 894e1a070067a7d9f22b032b53cdd3a3528dce53 Mon Sep 17 00:00:00 2001
From: Adam Florizone <72236049+adamflorizone@users.noreply.github.com>
Date: Wed, 3 Jan 2024 18:43:02 -0500
Subject: [PATCH] Docker: added build args for non AVX2 CPU (#5154)

---
 docker/docker-compose.yml | 17 +++++++++++++++++
 docker/nvidia/Dockerfile  |  3 ++-
 2 files changed, 19 insertions(+), 1 deletion(-)

diff --git a/docker/docker-compose.yml b/docker/docker-compose.yml
index 2aa6608..adaecc5 100644
--- a/docker/docker-compose.yml
+++ b/docker/docker-compose.yml
@@ -4,6 +4,23 @@ services:
     build:
       context: .
       args:
+        # Requirements file to use: 
+        # | GPU | CPU | requirements file to use |
+        # |--------|---------|---------|
+        # | NVIDIA | has AVX2 | `requirements.txt` |
+        # | NVIDIA | no AVX2 | `requirements_noavx2.txt` |
+        # | AMD | has AVX2 | `requirements_amd.txt` |
+        # | AMD | no AVX2 | `requirements_amd_noavx2.txt` |
+        # | CPU only | has AVX2 | `requirements_cpu_only.txt` |
+        # | CPU only | no AVX2 | `requirements_cpu_only_noavx2.txt` |
+        # | Apple | Intel | `requirements_apple_intel.txt` |
+        # | Apple | Apple Silicon | `requirements_apple_silicon.txt` |
+        # Default: requirements.txt`
+        # BUILD_REQUIREMENTS: requirements.txt
+        
+        # Extension requirements to build: 
+        # BUILD_EXTENSIONS: 
+        
         # specify which cuda version your card supports: https://developer.nvidia.com/cuda-gpus
         TORCH_CUDA_ARCH_LIST: ${TORCH_CUDA_ARCH_LIST:-7.5} 
         BUILD_EXTENSIONS: ${BUILD_EXTENSIONS:-}
diff --git a/docker/nvidia/Dockerfile b/docker/nvidia/Dockerfile
index 8da5564..6429d5e 100644
--- a/docker/nvidia/Dockerfile
+++ b/docker/nvidia/Dockerfile
@@ -3,6 +3,7 @@ FROM nvidia/cuda:12.1.1-devel-ubuntu22.04 as builder
 WORKDIR /builder
 ARG TORCH_CUDA_ARCH_LIST="${TORCH_CUDA_ARCH_LIST:-3.5;5.0;6.0;6.1;7.0;7.5;8.0;8.6+PTX}"
 ARG BUILD_EXTENSIONS="${BUILD_EXTENSIONS:-}"
+ARG BUILD_REQUIREMENTS="${BUILD_REQUIREMENTS:-requirements.txt}"
 ARG APP_UID="${APP_UID:-6972}"
 ARG APP_GID="${APP_GID:-6972}"
 # create / update build env
@@ -18,7 +19,7 @@ RUN --mount=type=cache,target=/root/.cache/pip,rw \
 USER app:app_grp
 # build wheels for runtime
 WORKDIR /home/app/build
-COPY --chown=app:app_grp requirements.txt /home/app/build
+COPY --chown=app:app_grp "$BUILD_REQUIREMENTS" /home/app/build/requirements.txt
 COPY --chown=app:app_grp extensions /home/app/build/extensions
 RUN --mount=type=cache,target=/root/.cache/pip,rw \
     # build all requirements files as wheel dists