Fix k2 and torchaudio installation (Docker, macOS) (NVIDIA#6094)

* torchaudio and k2: add flags, make optional Signed-off-by: Vladimir Bataev <[email protected]> * Fix syntax Signed-off-by: Vladimir Bataev <[email protected]> * Fix return code Signed-off-by: Vladimir Bataev <[email protected]> * Fix k2 installation Signed-off-by: Vladimir Bataev <[email protected]> * Fix torchaudio + k2 installation in docker Signed-off-by: Vladimir Bataev <[email protected]> * Fix flags Signed-off-by: Vladimir Bataev <[email protected]> * Fix torchaudio installation Signed-off-by: Vladimir Bataev <[email protected]> * Fix branch ckeck Signed-off-by: Vladimir Bataev <[email protected]> * Fix latest torchaudio installation Signed-off-by: Vladimir Bataev <[email protected]> --------- Signed-off-by: Vladimir Bataev <[email protected]>
titu1994 · Mar 24, 2023 · 015aad5 · 015aad5
1 parent 302b584
commit 015aad5
Show file tree

Hide file tree

Showing 3 changed files with 64 additions and 24 deletions.
diff --git a/Dockerfile b/Dockerfile
@@ -21,15 +21,23 @@ ARG BASE_IMAGE=nvcr.io/nvidia/pytorch:23.01-py3
 # image (by specifying build target as `nemo-deps`)
 FROM ${BASE_IMAGE} as nemo-deps
 
+# dependency flags; should be declared after FROM
+# torchaudio: not required by default
+ARG REQUIRE_TORCHAUDIO=false
+# k2: not required by default
+ARG REQUIRE_K2=false
+
 # Ensure apt-get won't prompt for selecting options
 ENV DEBIAN_FRONTEND=noninteractive
+# libavdevice-dev rerquired for latest torchaudio
 RUN apt-get update && \
     apt-get upgrade -y && \
     apt-get install -y \
     libsndfile1 sox \
     libfreetype6 \
     swig \
-    ffmpeg && \
+    ffmpeg \
+    libavdevice-dev && \
     rm -rf /var/lib/apt/lists/*
 
 WORKDIR /tmp/
@@ -47,7 +55,14 @@ RUN pip3 uninstall -y sacrebleu torchtext
 # build torchaudio
 WORKDIR /tmp/torchaudio_build
 COPY scripts/installers /tmp/torchaudio_build/scripts/installers/
-RUN /bin/bash /tmp/torchaudio_build/scripts/installers/install_torchaudio_latest.sh
+RUN INSTALL_MSG=$(/bin/bash /tmp/torchaudio_build/scripts/installers/install_torchaudio_latest.sh); INSTALL_CODE=$?; \
+    echo ${INSTALL_MSG}; \
+    if [ ${INSTALL_CODE} -ne 0 ]; then \
+      echo "torchaudio installation failed";  \
+      if [ "${REQUIRE_TORCHAUDIO}" = true ]; then \
+        exit ${INSTALL_CODE};  \
+      else echo "Skipping failed torchaudio installation"; fi \
+    else echo "torchaudio installed successfully"; fi
 
 # install nemo dependencies
 WORKDIR /tmp/nemo
@@ -56,7 +71,14 @@ RUN for f in $(ls requirements*.txt); do pip3 install --disable-pip-version-chec
 
 # install k2, skip if installation fails
 COPY scripts /tmp/nemo/scripts/
-RUN /bin/bash /tmp/nemo/scripts/speech_recognition/k2/setup.sh || exit 0
+RUN INSTALL_MSG=$(/bin/bash /tmp/nemo/scripts/speech_recognition/k2/setup.sh); INSTALL_CODE=$?; \
+    echo ${INSTALL_MSG}; \
+    if [ ${INSTALL_CODE} -ne 0 ]; then \
+      echo "k2 installation failed";  \
+      if [ "${REQUIRE_K2}" = true ]; then \
+        exit ${INSTALL_CODE};  \
+      else echo "Skipping failed k2 installation"; fi \
+    else echo "k2 installed successfully"; fi
 
 # copy nemo source into a scratch image
 FROM scratch as nemo-src

diff --git a/scripts/installers/install_torchaudio_latest.sh b/scripts/installers/install_torchaudio_latest.sh
@@ -22,30 +22,46 @@ TORCHAUDIO_REPO=https://github.com/pytorch/audio
 LATEST_RELEASE=$(git -c 'versionsort.suffix=-' \
     ls-remote --exit-code --refs --sort='version:refname' --heads ${TORCHAUDIO_REPO} 'release/*.*' \
     | tail --lines=1 \
-    | cut --delimiter='/' --fields=3,4)
-# expected TORCHAUDIO_BUILD_VERSION=*.**.*
-TORCHAUDIO_BUILD_VERSION=${LATEST_RELEASE:8:1}${PYTORCH_VERSION:1:5}
+    | cut -d '/' -f 3,4)
+TORCHAUDIO_LATEST_MAJOR_VERSION=$(python3 -c "major_version = (\"${LATEST_RELEASE}\".split('/')[-1]).split('.')[0]; print(major_version)")
+TORCHAUDIO_LATEST_MINOR_VERSION=$(python3 -c "minor_version = \"${LATEST_RELEASE}\".rsplit('.')[-1]; print(minor_version)")
 
-TORCH_MAJOR_VERSION=$(python3 -c "major_version = \"${PYTORCH_VERSION}\".split('.')[0]; print(major_version)")
-TORCH_MINOR_VERSION=$(python3 -c "minor_version = \"${PYTORCH_VERSION}\".split('.')[1]; print(minor_version)")
-TORCHAUDIO_MINOR_VERSION=$(python3 -c "minor_version = \"${LATEST_RELEASE}\".rsplit('.')[-1]; print(minor_version)")
+# avoid checking PYTORCH_VERSION variable, not available everywhere
+TORCH_FULL_VERSION=$(python3 -c "import torch; print(torch.__version__)")
+TORCH_MAIN_VERSION=$(python3 -c "import torch, re; print(re.search(r'(\d+\.?)+', torch.__version__).group(0))")
+TORCH_MAJOR_VERSION=$(python3 -c "major_version = \"${TORCH_MAIN_VERSION}\".split('.')[0]; print(major_version)")
+TORCH_MINOR_VERSION=$(python3 -c "minor_version = \"${TORCH_MAIN_VERSION}\".split('.')[1]; print(minor_version)")
+TORCH_FIX_VERSION=$(python3 -c "minor_version = \"${TORCH_MAIN_VERSION}\".split('.')[2]; print(minor_version)")
 
-if [[ $TORCH_MAJOR_VERSION -ne 1 ]]; then
-    echo "WARNING: Pytorch major version different from 1 not supported"
-fi
 
-echo "Latest torchaudio release: ${LATEST_RELEASE:8:4}"
-echo "Pytorch version: ${PYTORCH_VERSION:0:6}"
-echo "Torchaudio build version: ${TORCHAUDIO_BUILD_VERSION}"
+echo "Latest torchaudio release: ${TORCHAUDIO_LATEST_MAJOR_VERSION}.${TORCHAUDIO_LATEST_MINOR_VERSION}"
+echo "Pytorch version: ${TORCH_MAIN_VERSION:0:6}"
 
-if [[ "$TORCH_MINOR_VERSION" -lt "$TORCHAUDIO_MINOR_VERSION" ]]; then
-    # for old containers, we need to install matching torchaudio version
+if [[ $TORCH_MAJOR_VERSION -eq 1 ]]; then
+  if [[ $TORCH_MINOR_VERSION -le 13 ]]; then
     INSTALL_BRANCH="release/0.${TORCH_MINOR_VERSION}"
-else
-    # for new containers use latest release
-    INSTALL_BRANCH=${LATEST_RELEASE}
+  else
+    # fix for PyTorch 1.14 (no official release)
+    INSTALL_BRANCH="release/2.0"
+  fi
+  TORCHAUDIO_MAJOR_VERSION=0
+else  # version 2 expected
+  TORCHAUDIO_MAJOR_VERSION=${TORCH_MAJOR_VERSION}
+  INSTALL_BRANCH="release/${TORCH_MAJOR_VERSION}.${TORCH_MINOR_VERSION}"
 fi
 
+
+# check if install branch exists
+if [[ $(git ls-remote --heads ${TORCHAUDIO_REPO} ${INSTALL_BRANCH} | wc -l) -eq 0 ]]
+then
+  echo "Branch ${INSTALL_BRANCH} does not exist in torchaudio repo. Using latest release."
+  INSTALL_BRANCH=${LATEST_RELEASE}
+fi
+
+# expected TORCHAUDIO_BUILD_VERSION=*.**.*
+TORCHAUDIO_BUILD_VERSION="${TORCHAUDIO_MAJOR_VERSION}.${TORCH_MINOR_VERSION}.${TORCH_FIX_VERSION}"
+
+echo "Torchaudio build version: ${TORCHAUDIO_BUILD_VERSION}"
 echo "Installing torchaudio from branch: ${INSTALL_BRANCH}"
 
 # we need parameterized to run torchaudio tests
@@ -56,7 +72,7 @@ pip install parameterized
 git clone --depth 1 --branch ${INSTALL_BRANCH} https://github.com/pytorch/audio.git && \
 cd audio && \
 git submodule update --init --recursive && \
-BUILD_SOX=1 BUILD_VERSION=${TORCHAUDIO_BUILD_VERSION} python setup.py install && \
+USE_FFMPEG=1 BUILD_SOX=1 BUILD_VERSION=${TORCHAUDIO_BUILD_VERSION} python setup.py install && \
 cd .. && \
 pytest -rs audio/test/torchaudio_unittest/transforms/torchscript_consistency_cpu_test.py -k 'test_MFCC' || \
 { echo "ERROR: Failed to install torchaudio!"; exit 1; };

diff --git a/scripts/speech_recognition/k2/setup.sh b/scripts/speech_recognition/k2/setup.sh
@@ -18,7 +18,9 @@ K2_REPO=https://github.com/k2-fsa/k2
 LATEST_RELEASE=$(git -c 'versionsort.suffix=-' \
     ls-remote --exit-code --refs --sort='version:refname' --tags ${K2_REPO} '*.*' \
     | tail --lines=1 \
-    | cut --delimiter='/' --fields=3)
+    | cut -d '/' -f 3)
+# "cut --delimiter '/' --fields 3" doesn't work on macOS, use "-d ... -f ..." instead
 
-K2_MAKE_ARGS="-j" pip install -v git+${K2_REPO}@${LATEST_RELEASE}#egg=k2 || (echo "k2 could not be installed!"; exit 1)
-python3 -m k2.version > /dev/null || (echo "k2 installed with errors! Please check installation manually."; exit 1) && echo "k2 installed successfully!"
+K2_MAKE_ARGS="-j" pip install -v "git+${K2_REPO}@${LATEST_RELEASE}#egg=k2" || { echo "k2 could not be installed!"; exit 1; }
+python3 -m k2.version > /dev/null || { echo "k2 installed with errors! Please check installation manually."; exit 1; }
+echo "k2 installed successfully!"