Fixes CI build script (#1350)

* Fix the parsing of /etc/os-release

The old code parses OS_DISTRO as 'PRETTY_Ubuntu' on Ubuntu and thus
never links to libtinfo correctly.

* Configurable CMAKE_PREFIX_PATH in CI script.

Author: xinyazhang

.ci/pytorch/build.sh

@@ -81,7 +81,20 @@ if ! which conda; then
     export USE_MKLDNN=0
   fi
 else
-  export CMAKE_PREFIX_PATH=/opt/conda
+  # CMAKE_PREFIX_PATH precedences
+  # 1. $CONDA_PREFIX, if defined. This follows the pytorch official build instructions.
+  # 2. /opt/conda/envs/py_${ANACONDA_PYTHON_VERSION}, if ANACONDA_PYTHON_VERSION defined.
+  #    This is for CI, which defines ANACONDA_PYTHON_VERSION but not CONDA_PREFIX.
+  # 3. $(conda info --base). The fallback value of pytorch official build
+  #    instructions actually refers to this.
+  #    Commonly this is /opt/conda/
+  if [[ -v CONDA_PREFIX ]]; then
+    export CMAKE_PREFIX_PATH=${CONDA_PREFIX}
+  elif [[ -v ANACONDA_PYTHON_VERSION ]]; then
+    export CMAKE_PREFIX_PATH="/opt/conda/envs/py_${ANACONDA_PYTHON_VERSION}"
+  else
+    export CMAKE_PREFIX_PATH="$(conda info --base)"
+  fi
 
   # Workaround required for MKL library linkage
   # https://github.com/pytorch/pytorch/issues/119557

cmake/Dependencies.cmake

@@ -1244,9 +1244,19 @@ if(USE_ROCM)
   # Currently only active for Ubuntu 20.04 and greater versions.
   if(UNIX AND EXISTS "/etc/os-release")
     file(STRINGS /etc/os-release OS_RELEASE)
-    string(REGEX REPLACE ".*NAME=\"([A-Za-z]+).*" "\\1" OS_DISTRO ${OS_RELEASE})
-    string(REGEX REPLACE ".*VERSION_ID=\"([0-9\.]+).*" "\\1" OS_VERSION ${OS_RELEASE})
-    if(OS_DISTRO STREQUAL "Ubuntu" AND OS_VERSION VERSION_GREATER_EQUAL "20.04")
+    set(DISTRO_NAME "")
+    set(DISTRO_VERSION "")
+    foreach(line ${OS_RELEASE})
+      string(REGEX MATCH "^NAME=" DISTRO_NAME_MATCH ${line})
+      if (NOT DISTRO_NAME_MATCH STREQUAL "")
+        string(REGEX REPLACE "^NAME=\"(.*)\"" "\\1" DISTRO_NAME ${line})
+      endif()
+      string(REGEX MATCH "^VERSION_ID=" DISTRO_VERSION_MATCH ${line})
+      if (NOT DISTRO_VERSION_MATCH STREQUAL "")
+        string(REGEX REPLACE "^VERSION_ID=\"(.*)\"" "\\1" DISTRO_VERSION ${line})
+      endif()
+    endforeach()
+    if(DISTRO_NAME STREQUAL "Ubuntu" AND DISTRO_VERSION VERSION_GREATER_EQUAL "20.04")
       find_library(LIBTINFO_LOC tinfo NO_CMAKE_PATH NO_CMAKE_ENVIRONMENT_PATH)
       if(LIBTINFO_LOC)
         get_filename_component(LIBTINFO_LOC_PARENT ${LIBTINFO_LOC} DIRECTORY)