From 503ad91705861eee9f599c2f928a254b0f22c66f Mon Sep 17 00:00:00 2001
From: Yu-Hang 'Maxin' Tang <Tang.Maxin@gmail.com>
Date: Tue, 18 Jul 2023 22:35:59 +0100
Subject: [PATCH] use RUN heredoc

---
 .github/container/Dockerfile.jax | 14 ++++++++------
 1 file changed, 8 insertions(+), 6 deletions(-)

diff --git a/.github/container/Dockerfile.jax b/.github/container/Dockerfile.jax
index c06896660..b0d2972dd 100644
--- a/.github/container/Dockerfile.jax
+++ b/.github/container/Dockerfile.jax
@@ -61,14 +61,16 @@ COPY --from=jax-builder ${SRC_PATH_JAX}-no-git ${SRC_PATH_JAX}
 COPY --from=jax-builder ${SRC_PATH_XLA}-no-git ${SRC_PATH_XLA}
 COPY --from=jax-builder ${SRC_PATH_TE}-no-git ${SRC_PATH_TE}
 
+RUN <<EOF
 # Transformer Engine installation dependencies
-RUN pip install --no-cache-dir pybind11 ninja packaging && rm -rf ~/.cache/pip/
+pip install --no-cache-dir pybind11 ninja packaging
 # Install JAX + Transformer Engine
-RUN NVTE_FRAMEWORK=jax pip --disable-pip-version-check install -e \
-        ${SRC_PATH_JAX}/dist/*.whl \
-        ${SRC_PATH_JAX} \
-        ${SRC_PATH_TE} \
-    && rm -rf ~/.cache/pip/
+NVTE_FRAMEWORK=jax pip --disable-pip-version-check --no-cache-dir install -e \
+    ${SRC_PATH_JAX}/dist/*.whl \
+    ${SRC_PATH_JAX} \
+    ${SRC_PATH_TE}
+rm -rf ~/.cache/pip/
+EOF
 
 # Install software stack in JAX ecosystem
 # Made this optional since tensorstore cannot build on Ubuntu 20.04 + ARM