add back transformer engine install

Signed-off-by: Peter St. John <[email protected]>

Dockerfile

@@ -16,12 +16,6 @@ RUN rustup set profile minimal && \
 
 FROM ${BASE_IMAGE} AS bionemo2-base
 
-# Install NeMo dependencies.
-WORKDIR /build
-
-ARG MAX_JOBS=4
-ENV MAX_JOBS=${MAX_JOBS}
-
 # Install core apt packages.
 RUN --mount=type=cache,id=apt-cache,target=/var/cache/apt,sharing=locked \
   --mount=type=cache,id=apt-lib,target=/var/lib/apt,sharing=locked \
@@ -34,13 +28,17 @@ apt-get install -qyy \
   git \
   curl \
   pre-commit \
-  sudo
+  sudo \
+  gnupg
 apt-get upgrade -qyy \
   rsync
 rm -rf /tmp/* /var/tmp/*
 EOF
 
-RUN apt-get install -y gnupg
+ARG TE_COMMIT=2215fa5c7557b66034068816020f9f611019e457
+RUN NVTE_FRAMEWORK=pytorch NVTE_WITH_USERBUFFERS=1 MPI_HOME=/usr/local/mpi \
+  pip --disable-pip-version-check --no-cache-dir install \
+  git+https://github.com/NVIDIA/TransformerEngine.git@${TE_COMMIT}
 
 # Check the nemo dependency for causal conv1d and make sure this checkout
 # tag matches. If not, update the tag in the following line.
@@ -57,9 +55,7 @@ RUN pip install nemo_run@git+https://github.com/NVIDIA/NeMo-Run.git@${NEMU_RUN_T
 
 RUN mkdir -p /workspace/bionemo2/
 
-# Delete the temporary /build directory.
 WORKDIR /workspace
-RUN rm -rf /build
 
 # Addressing Security Scan Vulnerabilities
 RUN rm -rf /opt/pytorch/pytorch/third_party/onnx
@@ -79,7 +75,7 @@ ENV UV_LINK_MODE=copy \
 # Install the bionemo-geometric requirements ahead of copying over the rest of the repo, so that we can cache their
 # installation. These involve building some torch extensions, so they can take a while to install.
 RUN --mount=type=bind,source=./sub-packages/bionemo-geometric/requirements.txt,target=/requirements-pyg.txt \
-    --mount=type=cache,target=/root/.cache \
+  --mount=type=cache,target=/root/.cache \
   uv pip install --no-build-isolation -r /requirements-pyg.txt
 
 COPY --from=rust-env /usr/local/cargo /usr/local/cargo
@@ -99,9 +95,9 @@ COPY ./sub-packages /workspace/bionemo2/sub-packages
 # Includes a hack to install tensorstore 0.1.45, which doesn't distribute a pypi wheel for python 3.12, and the metadata
 # in the source distribution doesn't match the expected pypi version.
 RUN --mount=type=bind,source=./.git,target=./.git \
-    --mount=type=bind,source=./requirements-test.txt,target=/requirements-test.txt \
-    --mount=type=bind,source=./requirements-cve.txt,target=/requirements-cve.txt \
-    --mount=type=cache,target=/root/.cache <<EOF
+  --mount=type=bind,source=./requirements-test.txt,target=/requirements-test.txt \
+  --mount=type=bind,source=./requirements-cve.txt,target=/requirements-cve.txt \
+  --mount=type=cache,target=/root/.cache <<EOF
 set -eo pipefail
 
 uv pip install maturin --no-build-isolation
@@ -124,8 +120,8 @@ EOF
 FROM ${BASE_IMAGE} AS dev
 
 RUN --mount=type=cache,id=apt-cache,target=/var/cache/apt,sharing=locked \
-    --mount=type=cache,id=apt-lib,target=/var/lib/apt,sharing=locked \
-    <<EOF
+  --mount=type=cache,id=apt-lib,target=/var/lib/apt,sharing=locked \
+  <<EOF
 set -eo pipefail
 apt-get update -qy
 apt-get install -qyy \
@@ -168,7 +164,7 @@ ENV PATH="/usr/local/cargo/bin:/usr/local/rustup/bin:${PATH}"
 ENV RUSTUP_HOME="/usr/local/rustup"
 
 RUN --mount=type=bind,source=./requirements-dev.txt,target=/workspace/bionemo2/requirements-dev.txt \
-    --mount=type=cache,target=/root/.cache <<EOF
+  --mount=type=cache,target=/root/.cache <<EOF
   set -eo pipefail
   uv pip install -r /workspace/bionemo2/requirements-dev.txt
   rm -rf /tmp/*