fix pytorch build

Author: djherbis

Dockerfile.tmpl

@@ -10,7 +10,7 @@ ARG TORCHVISION_VERSION
 
 {{ if eq .Accelerator "gpu" }}
 FROM gcr.io/kaggle-images/python-lightgbm-whl:${GPU_BASE_IMAGE_NAME}-${BASE_IMAGE_TAG}-${LIGHTGBM_VERSION} AS lightgbm_whl
-#FROM gcr.io/kaggle-images/python-torch-whl:${GPU_BASE_IMAGE_NAME}-${BASE_IMAGE_TAG}-${TORCH_VERSION} AS torch_whl
+FROM gcr.io/kaggle-images/python-torch-whl:${GPU_BASE_IMAGE_NAME}-${BASE_IMAGE_TAG}-${TORCH_VERSION} AS torch_whl
 FROM ${BASE_IMAGE_REPO}/${GPU_BASE_IMAGE_NAME}:${BASE_IMAGE_TAG}
 {{ else }}
 FROM ${BASE_IMAGE_REPO}/${CPU_BASE_IMAGE_NAME}:${BASE_IMAGE_TAG}
@@ -111,15 +111,13 @@ RUN mamba install implicit && \
 {{ if eq .Accelerator "gpu" }}
 #COPY --from=torch_whl /tmp/whl/*.whl /tmp/torch/
 RUN mamba install -c pytorch magma-cuda${CUDA_MAJOR_VERSION}${CUDA_MINOR_VERSION} && \
-    #pip install /tmp/torch/*.whl && \
+    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.
-    mamba install pytorch==1.12.0 torchvision==0.13.0 torchaudio==0.12.0 torchtext cudatoolkit=11.3 -c pytorch && \
+    mamba install -y openmp && \
     #pip install patchelf && \
     #patchelf --add-needed libomp.so /opt/conda/lib/python3.7/site-packages/torchtext/lib/libtorchtext.so && \
     #patchelf --add-needed libomp.so /opt/conda/lib/python3.7/site-packages/torchaudio/lib/libtorchaudio.so && \
-    #rm -rf /tmp/torch && \
+    rm -rf /tmp/torch && \
     /tmp/clean-layer.sh
 {{ else }}
 RUN pip install \

Jenkinsfile

@@ -22,6 +22,27 @@ pipeline {
   stages {
     stage('Pre-build Packages from Source') {
       parallel {
+        stage('torch') {
+          options {
+            timeout(time: 180, unit: 'MINUTES')
+          }
+          steps {
+            sh '''#!/bin/bash
+              set -exo pipefail
+              source config.txt
+              cd packages/
+              ./build_package --base-image $BASE_IMAGE_REPO/$GPU_BASE_IMAGE_NAME:$BASE_IMAGE_TAG \
+                --package torch \
+                --version $TORCH_VERSION \
+                --build-arg TORCHAUDIO_VERSION=$TORCHAUDIO_VERSION \
+                --build-arg TORCHTEXT_VERSION=$TORCHTEXT_VERSION \
+                --build-arg TORCHVISION_VERSION=$TORCHVISION_VERSION \
+                --build-arg CUDA_MAJOR_VERSION=$CUDA_MAJOR_VERSION \
+                --build-arg CUDA_MINOR_VERSION=$CUDA_MINOR_VERSION \
+                --push
+            '''
+          }
+        }
         stage('lightgbm') {
           options {
             timeout(time: 10, unit: 'MINUTES')

config.txt

@@ -3,9 +3,9 @@ BASE_IMAGE_TAG=m103
 CPU_BASE_IMAGE_NAME=tf2-cpu.2-11
 GPU_BASE_IMAGE_NAME=tf2-gpu.2-11
 LIGHTGBM_VERSION=3.3.2
-TORCH_VERSION=1.12.0
-TORCHAUDIO_VERSION=0.12.0
-TORCHTEXT_VERSION=0.13.0
-TORCHVISION_VERSION=0.13.0
+TORCH_VERSION=1.13.0
+TORCHAUDIO_VERSION=0.13.0
+TORCHTEXT_VERSION=0.14.0
+TORCHVISION_VERSION=0.14.0
 CUDA_MAJOR_VERSION=11
 CUDA_MINOR_VERSION=3

packages/torch.Dockerfile

@@ -55,8 +55,10 @@ RUN sudo apt-get update && \
     cd audio && \
     git checkout tags/v$TORCHAUDIO_VERSION && \
     git submodule sync && \
-    git submodule update --init --recursive --jobs 1 && \
-    python setup.py bdist_wheel
+    git submodule update --init --recursive --jobs 1
+# https://github.com/pytorch/audio/issues/936#issuecomment-702990346
+RUN sed -i 's/set(envs/set(envs\n  "LIBS=-ltinfo"/' /usr/local/src/audio/third_party/sox/CMakeLists.txt 
+RUN cd /usr/local/src/audio && python setup.py bdist_wheel
 
 # Build torchtext
 # Instructions: https://github.com/pytorch/text#building-from-source