Make torchtext & torchaudio load libomp

http://b/255757999

Author: djherbis

Dockerfile.tmpl

@@ -90,7 +90,7 @@ RUN conda config --add channels nvidia && \
 # b/232247930: uninstall pyarrow to avoid double installation with the GPU specific version.
 RUN pip uninstall -y pyarrow && \
     conda install -c conda-forge mamba && \
-    mamba install cudf cuml cudatoolkit && \
+    mamba install -y cudf cuml cudatoolkit && \
     /tmp/clean-layer.sh
 {{ end }}
 
@@ -108,6 +108,13 @@ RUN conda install implicit && \
 COPY --from=torch_whl /tmp/whl/*.whl /tmp/torch/
 RUN conda install -c pytorch magma-cuda${CUDA_MAJOR_VERSION}${CUDA_MINOR_VERSION} && \
     pip install /tmp/torch/*.whl && \
+    # b/255757999 openmp (libomp.so) is an dependency of libtorchtext and libtorchaudio but
+    # the built from source versions don't seem to properly link it in. This forces the dep
+    # which makes sure that libomp is loaded when these libraries are loaded.
+    conda install -y openmp && \
+    pip install patchelf && \
+    patchelf --add-needed libomp.so.5 /opt/conda/lib/python3.7/site-packages/torchtext/lib/libtorchtext.so && \
+    patchelf --add-needed libomp.so.5 /opt/conda/lib/python3.7/site-packages/torchaudio/lib/libtorchaudio.so && \
     rm -rf /tmp/torch && \
     /tmp/clean-layer.sh
 {{ else }}