Enhancements to PIP install scripts, pt 2 (pytorch#2220)

Summary: - Clean up PIP intstall scripts Pull Request resolved: pytorch#2220 Reviewed By: spcyppt Differential Revision: D52223334 Pulled By: q10 fbshipit-source-id: 2c3021bfb570cd71061e320f2aa784eadf890184
q10 · Dec 16, 2023 · a535f22 · a535f22
1 parent 51ee221
commit a535f22
Show file tree

Hide file tree

Showing 9 changed files with 161 additions and 114 deletions.
diff --git a/.github/scripts/fbgemm_gpu_install.bash b/.github/scripts/fbgemm_gpu_install.bash
@@ -68,10 +68,9 @@ install_fbgemm_gpu_wheel () {
 install_fbgemm_gpu_pip () {
   local env_name="$1"
   local fbgemm_gpu_channel_version="$2"
-  local fbgemm_gpu_variant_type="$3"
-  local fbgemm_gpu_variant_version="$4"
-  if [ "$fbgemm_gpu_variant_type" == "" ]; then
-    echo "Usage: ${FUNCNAME[0]} ENV_NAME FBGEMM_GPU_VERSION FBGEMM_GPU_VARIANT_TYPE [FBGEMM_GPU_VARIANT_VERSION]"
+  local fbgemm_gpu_variant_type_version="$3"
+  if [ "$fbgemm_gpu_variant_type_version" == "" ]; then
+    echo "Usage: ${FUNCNAME[0]} ENV_NAME FBGEMM_GPU_CHANNEL[/VERSION] FBGEMM_GPU_VARIANT_TYPE[/VARIANT_VERSION]"
     echo "Example(s):"
     echo "    ${FUNCNAME[0]} build_env 0.5.0 cpu                  # Install the CPU variant, specific version from release channel"
     echo "    ${FUNCNAME[0]} build_env release cuda 12.1.1        # Install the CUDA variant, latest version from release channel"
@@ -89,7 +88,7 @@ install_fbgemm_gpu_pip () {
 
   # Install the package from PyTorch PIP (not PyPI)
   # The package's canonical name is 'fbgemm-gpu' (hyphen, not underscore)
-  install_from_pytorch_pip "${env_name}" fbgemm_gpu "${fbgemm_gpu_channel_version}" "${fbgemm_gpu_variant_type}" "${fbgemm_gpu_variant_version}" || return 1
+  install_from_pytorch_pip "${env_name}" fbgemm_gpu "${fbgemm_gpu_channel_version}" "${fbgemm_gpu_variant_type_version}" || return 1
 
   # Run post-installation checks
   __fbgemm_gpu_post_install_checks "${env_name}" || return 1

diff --git a/.github/scripts/fbgemm_gpu_test.bash b/.github/scripts/fbgemm_gpu_test.bash
@@ -182,7 +182,7 @@ test_setup_conda_environment () {
   if [ "$pytorch_installer" == "conda" ]; then
     install_pytorch_conda     "${env_name}" "${pytorch_version}" "${pytorch_variant_type}" "${pytorch_variant_version}" || return 1
   else
-    install_pytorch_pip       "${env_name}" "${pytorch_version}" "${pytorch_variant_type}" "${pytorch_variant_version}" || return 1
+    install_pytorch_pip       "${env_name}" "${pytorch_version}" "${pytorch_variant_type}"/"${pytorch_variant_version}" || return 1
   fi
 }
 

diff --git a/.github/scripts/utils_pip.bash b/.github/scripts/utils_pip.bash
@@ -13,37 +13,15 @@
 # PyTorch PIP Install Functions
 ################################################################################
 
-__extract_pip_arguments () {
-  export env_name="$1"
-  export package_name_raw="$2"
-  export package_channel_version="$3"
-  export package_variant_type="$4"
-  export package_variant_version="$5"
-  if [ "$package_variant_type" == "" ]; then
-    echo "Usage: ${FUNCNAME[0]} ENV_NAME PACKAGE_NAME PACKAGE_CHANNEL_VERSION PACKAGE_VARIANT_TYPE [PACKAGE_VARIANT_VERSION]"
-    echo "Example(s):"
-    echo "    ${FUNCNAME[0]} build_env torch 1.11.0 cpu                       # Install the CPU variant, specific version from release channel"
-    echo "    ${FUNCNAME[0]} build_env torch release cpu                      # Install the CPU variant, latest version from release channel"
-    echo "    ${FUNCNAME[0]} build_env fbgemm_gpu test/0.6.0rc0 cuda 12.1.0   # Install the CUDA 12.1 variant, specific version from test channel"
-    echo "    ${FUNCNAME[0]} build_env fbgemm_gpu nightly rocm 5.3            # Install the ROCM 5.3 variant, latest version from nightly channel"
-    return 1
-  else
-    echo "################################################################################"
-    echo "# Extract PIP Arguments (PyTorch PIP)"
-    echo "#"
-    echo "# [$(date --utc +%FT%T.%3NZ)] + ${FUNCNAME[0]} ${*}"
-    echo "################################################################################"
-    echo ""
-  fi
-
-  test_network_connection || return 1
+__export_package_channel_info () {
+  local package_channel_version="$1"
 
   # Extract the package channel and version from the tuple-string
   if [ "$package_channel_version" == "nightly" ] || [ "$package_channel_version" == "test" ] || [ "$package_channel_version" == "release" ]; then
     export package_channel="$package_channel_version"
     export package_version=""
   else
-  # shellcheck disable=SC2207
+    # shellcheck disable=SC2207
     local package_channel_version_arr=($(echo "${package_channel_version}" | tr '/' '\n'))
     if [ ${#package_channel_version_arr[@]} -lt 2 ]; then
       export package_channel="release"
@@ -53,34 +31,73 @@ __extract_pip_arguments () {
       export package_version="${package_channel_version_arr[1]}"
     fi
   fi
-  echo "[INSTALL] Extracted package (channel, version): (${package_channel}, ${package_version})"
+  if [ "$package_channel" != "nightly" ] && [ "$package_channel" != "test" ] && [ "$package_channel" != "release" ]; then
+    echo "[INSTALL] Invalid PyTorch PIP package channel: ${package_channel}"
+    return 1
+  fi
+  echo "[INSTALL] Extracted package (channel, version): (${package_channel}, ${package_version:-LATEST})"
+}
 
-  # Replace underscores with hyphens to materialize the canonical name of the package
-  # shellcheck disable=SC2155
-  export package_name=$(echo "${package_name_raw}" | tr '_' '-')
+__export_package_variant_info () {
+  local package_variant_type_version="$1"
+
+  local FALLBACK_VERSION_CUDA="12.1.1"
+  local FALLBACK_VERSION_ROCM="5.7.0"
+
+  if [ "$package_variant_type_version" == "cuda" ]; then
+    # If "cuda", default to latest CUDA
+    local variant_type="cu"
+    local variant_version="$FALLBACK_VERSION_CUDA"
+
+  elif [ "$package_variant_type_version" == "rocm" ]; then
+    # If "rocm", default to latest ROCm
+    local variant_type="rocm"
+    local variant_version="$FALLBACK_VERSION_ROCM"
+
+  elif [ "$package_variant_type_version" == "cpu" ]; then
+    # If "cpu", default to latest cpu
+    local variant_type="cpu"
+    local variant_version=""
 
-  # Set the package variant
-  if [ "$package_variant_type" == "cuda" ]; then
-    # Extract the CUDA version or default to 12.1.0
-    local cuda_version="${package_variant_version:-12.1.0}"
-    # shellcheck disable=SC2206
-    local cuda_version_arr=(${cuda_version//./ })
-    # Convert, i.e. cuda 12.1.0 => cu121
-    export package_variant="cu${cuda_version_arr[0]}${cuda_version_arr[1]}"
-  elif [ "$package_variant_type" == "rocm" ]; then
-    # Extract the ROCM version or default to 5.7.0
-    local rocm_version="${package_variant_version:-5.7.0}"
-    # shellcheck disable=SC2206
-    local rocm_version_arr=(${rocm_version//./ })
-    # Convert, i.e. rocm 5.5.1 => rocm5.5
-    export package_variant="rocm${rocm_version_arr[0]}.${rocm_version_arr[1]}"
   else
-    echo "[INSTALL] Invalid package variant type $package_variant_type, defaulting to cpu"
-    export package_variant_type="cpu"
-    export package_variant="cpu"
+    # Split along '/', e.g. cuda/12.1.0
+    # shellcheck disable=SC2207
+    local package_variant_type_version_arr=($(echo "${package_variant_type_version}" | tr '/' '\n'))
+    local variant_type="${package_variant_type_version_arr[0]}"
+    local variant_version="${package_variant_type_version_arr[1]}"
+
+    if [ "$variant_type" == "cuda" ]; then
+      # Extract the CUDA version or set to default
+      local cuda_version="${variant_version:-${FALLBACK_VERSION_CUDA}}"
+      # shellcheck disable=SC2206
+      local cuda_version_arr=(${cuda_version//./ })
+      # Convert, i.e. cuda 12.1.0 => cu121
+      local variant_type="cu"
+      local variant_version="${cuda_version_arr[0]}${cuda_version_arr[1]}"
+
+    elif [ "$variant_type" == "rocm" ]; then
+      # Extract the ROCM version or set to default
+      local rocm_version="${variant_version:-${FALLBACK_VERSION_ROCM}}"
+      # shellcheck disable=SC2206
+      local rocm_version_arr=(${rocm_version//./ })
+      # Convert, i.e. rocm 5.5.1 => rocm5.5
+      local variant_type="rocm"
+      local variant_version="${rocm_version_arr[0]}.${rocm_version_arr[1]}"
+
+    else
+      echo "[INSTALL] Package variant type '$variant_type' is neither CUDA nor ROCm variant, falling back to cpu"
+      local variant_type="cpu"
+      local variant_version=""
+    fi
   fi
+
+  # Export the extracted information
+  export package_variant_type="${variant_type}"
+  export package_variant="${variant_type}${variant_version}"
   echo "[INSTALL] Extracted package variant: ${package_variant}"
+}
 
+__export_pip_arguments () {
   # Extract the PIP channel
   if [ "$package_channel" == "release" ]; then
     export pip_channel="https://download.pytorch.org/whl/${package_variant}/"
@@ -104,19 +121,52 @@ __extract_pip_arguments () {
   echo "[INSTALL] Extracted the full PIP package: ${pip_package}"
 }
 
+__prepare_pip_arguments () {
+  local package_name_raw="$1"
+  local package_channel_version="$2"
+  local package_variant_type_version="$3"
+  if [ "$package_variant_type_version" == "" ]; then
+    echo "Usage: ${FUNCNAME[0]} PACKAGE_NAME PACKAGE_CHANNEL[/VERSION] PACKAGE_VARIANT_TYPE[/VARIANT_VERSION]"
+    return 1
+  else
+    echo "################################################################################"
+    echo "# Prepare PIP Arguments (PyTorch PIP)"
+    echo "#"
+    echo "# [$(date --utc +%FT%T.%3NZ)] + ${FUNCNAME[0]} ${*}"
+    echo "################################################################################"
+    echo ""
+  fi
+
+  # Replace underscores with hyphens to materialize the canonical name of the
+  # package, and export variable to environment
+  # shellcheck disable=SC2155
+  export package_name=$(echo "${package_name_raw}" | tr '_' '-')
+
+  # Extract the package channel and package version from the tuple-string, and
+  # export variables to environment
+  __export_package_channel_info "$package_channel_version"
+
+  # Extract the package variant type and variant version from the tuple-string,
+  # and export variables to environment
+  __export_package_variant_info "${package_variant_type_version}"
+
+  # With all package_* variables exported, extract the arguments for PIP, and
+  # export variabels to environment
+  __export_pip_arguments
+}
+
 install_from_pytorch_pip () {
   local env_name="$1"
   local package_name_raw="$2"
   local package_channel_version="$3"
-  local package_variant_type="$4"
-  local package_variant_version="$5"
-  if [ "$package_variant_type" == "" ]; then
-    echo "Usage: ${FUNCNAME[0]} ENV_NAME PACKAGE_NAME PACKAGE_CHANNEL_VERSION PACKAGE_VARIANT_TYPE [PACKAGE_VARIANT_VERSION]"
+  local package_variant_type_version="$4"
+  if [ "$package_variant_type_version" == "" ]; then
+    echo "Usage: ${FUNCNAME[0]} ENV_NAME PACKAGE_NAME PACKAGE_CHANNEL[/VERSION] PACKAGE_VARIANT_TYPE[/VARIANT_VERSION]"
     echo "Example(s):"
     echo "    ${FUNCNAME[0]} build_env torch 1.11.0 cpu                       # Install the CPU variant, specific version from release channel"
     echo "    ${FUNCNAME[0]} build_env torch release cpu                      # Install the CPU variant, latest version from release channel"
-    echo "    ${FUNCNAME[0]} build_env fbgemm_gpu test/0.6.0rc0 cuda 12.1.0   # Install the CUDA 12.1 variant, specific version from test channel"
-    echo "    ${FUNCNAME[0]} build_env fbgemm_gpu nightly rocm 5.3            # Install the ROCM 5.3 variant, latest version from nightly channel"
+    echo "    ${FUNCNAME[0]} build_env fbgemm_gpu test/0.6.0rc0 cuda/12.1.0   # Install the CUDA 12.1 variant, specific version from test channel"
+    echo "    ${FUNCNAME[0]} build_env fbgemm_gpu nightly rocm/5.3            # Install the ROCM 5.3 variant, latest version from nightly channel"
     return 1
   else
     echo "################################################################################"
@@ -129,12 +179,12 @@ install_from_pytorch_pip () {
 
   test_network_connection || return 1
 
-  __extract_pip_arguments "$env_name" "$package_name_raw" "$package_channel_version" "$package_variant_type" "$package_variant_version"
+  __prepare_pip_arguments "$package_name_raw" "$package_channel_version" "$package_variant_type_version"
 
   # shellcheck disable=SC2155
   local env_prefix=$(env_name_or_prefix "${env_name}")
 
-  echo "[INSTALL] Attempting to install [${package_name}, ${package_version}+${package_variant}] from PyTorch PIP using channel ${pip_channel} ..."
+  echo "[INSTALL] Attempting to install [${package_name}, ${package_version:-LATEST}] from PyTorch PIP using channel ${pip_channel} ..."
   # shellcheck disable=SC2086
   (exec_with_retries 3 conda run ${env_prefix} pip install ${pip_package} --extra-index-url ${pip_channel}) || return 1
 
@@ -144,9 +194,9 @@ install_from_pytorch_pip () {
     # This test usually applies to the nightly builds
     # shellcheck disable=SC2086
     if conda run ${env_prefix} pip list "${package_name}" | grep "${package_name}" | grep "${package_variant}"; then
-      echo "[CHECK] The installed package [${package_name}, ${package_version}] is the correct variant (${package_variant})"
+      echo "[CHECK] The installed package [${package_name}, ${package_channel}/${package_version:-LATEST}] is the correct variant (${package_variant})"
     else
-      echo "[CHECK] The installed package [${package_name}, ${package_version}] appears to be an incorrect variant as it is missing references to ${package_variant}!"
+      echo "[CHECK] The installed package [${package_name}, ${package_channel}/${package_version:-LATEST}] appears to be an incorrect variant as it is missing references to ${package_variant}!"
       echo "[CHECK] This can happen if the variant of the package (e.g. GPU, nightly) for the MAJOR.MINOR version of CUDA or ROCm presently installed on the system is not available."
       return 1
     fi
@@ -161,15 +211,14 @@ download_from_pytorch_pip () {
   local env_name="$1"
   local package_name_raw="$2"
   local package_channel_version="$3"
-  local package_variant_type="$4"
-  local package_variant_version="$5"
-  if [ "$package_variant_type" == "" ]; then
-    echo "Usage: ${FUNCNAME[0]} ENV_NAME PACKAGE_NAME PACKAGE_CHANNEL_VERSION PACKAGE_VARIANT_TYPE [PACKAGE_VARIANT_VERSION]"
+  local package_variant_type_version="$4"
+  if [ "$package_variant_type_version" == "" ]; then
+    echo "Usage: ${FUNCNAME[0]} ENV_NAME PACKAGE_NAME PACKAGE_CHANNEL[/VERSION] PACKAGE_VARIANT_TYPE[/VARIANT_VERSION]"
     echo "Example(s):"
     echo "    ${FUNCNAME[0]} build_env torch 1.11.0 cpu                       # Download the CPU variant, specific version from release channel"
     echo "    ${FUNCNAME[0]} build_env torch release cpu                      # Download the CPU variant, latest version from release channel"
-    echo "    ${FUNCNAME[0]} build_env fbgemm_gpu test/0.6.0rc0 cuda 12.1.0   # Download the CUDA 12.1 variant, specific version from test channel"
-    echo "    ${FUNCNAME[0]} build_env fbgemm_gpu nightly rocm 5.3            # Download the ROCM 5.3 variant, latest version from nightly channel"
+    echo "    ${FUNCNAME[0]} build_env fbgemm_gpu test/0.6.0rc0 cuda/12.1.0   # Download the CUDA 12.1 variant, specific version from test channel"
+    echo "    ${FUNCNAME[0]} build_env fbgemm_gpu nightly rocm/5.3            # Download the ROCM 5.3 variant, latest version from nightly channel"
     return 1
   else
     echo "################################################################################"
@@ -182,7 +231,7 @@ download_from_pytorch_pip () {
 
   test_network_connection || return 1
 
-  __extract_pip_arguments "$env_name" "$package_name_raw" "$package_channel_version" "$package_variant_type" "$package_variant_version"
+  __prepare_pip_arguments "$package_name_raw" "$package_channel_version" "$package_variant_type_version"
 
   # shellcheck disable=SC2155
   local env_prefix=$(env_name_or_prefix "${env_name}")
@@ -191,7 +240,7 @@ download_from_pytorch_pip () {
   # shellcheck disable=SC2035
   rm -rf *.whl || return 1
 
-  echo "[DOWNLOAD] Attempting to download wheel [${package_name}, ${package_version}+${package_variant}] from PyTorch PIP using channel ${pip_channel} ..."
+  echo "[DOWNLOAD] Attempting to download wheel [${package_name}, ${package_version:-LATEST}] from PyTorch PIP using channel ${pip_channel} ..."
   # shellcheck disable=SC2086
   (exec_with_retries 3 conda run ${env_prefix} pip download ${pip_package} --extra-index-url ${pip_channel}) || return 1
 

diff --git a/.github/scripts/utils_pytorch.bash b/.github/scripts/utils_pytorch.bash
@@ -105,16 +105,15 @@ install_pytorch_conda () {
 
 install_pytorch_pip () {
   local env_name="$1"
-  local pytorch_version="$2"
-  local pytorch_variant_type="$3"
-  local pytorch_variant_version="$4"
-  if [ "$pytorch_variant_type" == "" ]; then
-    echo "Usage: ${FUNCNAME[0]} ENV_NAME PYTORCH_VERSION PYTORCH_VARIANT_TYPE [PYTORCH_VARIANT_VERSION]"
+  local pytorch_channel_version="$2"
+  local pytorch_variant_type_version="$3"
+  if [ "$pytorch_variant_type_version" == "" ]; then
+    echo "Usage: ${FUNCNAME[0]} ENV_NAME PYTORCH_CHANNEL[/VERSION] PYTORCH_VARIANT_TYPE[/VARIANT_VERSION]"
     echo "Example(s):"
-    echo "    ${FUNCNAME[0]} build_env 1.11.0 cpu         # Install the CPU variant for a specific version"
+    echo "    ${FUNCNAME[0]} build_env test/2.1.0rc0 cpu  # Install the CPU variant for a specific version"
     echo "    ${FUNCNAME[0]} build_env release cpu        # Install the CPU variant, latest release version"
-    echo "    ${FUNCNAME[0]} build_env test cuda 12.1.0   # Install the CUDA 12.1 variant, latest test version"
-    echo "    ${FUNCNAME[0]} build_env nightly rocm 5.3   # Install the ROCM 5.3 variant, latest nightly version"
+    echo "    ${FUNCNAME[0]} build_env test cuda/12.1.0   # Install the CUDA 12.1 variant, latest test version"
+    echo "    ${FUNCNAME[0]} build_env nightly rocm/5.3   # Install the ROCM 5.3 variant, latest nightly version"
     return 1
   else
     echo "################################################################################"
@@ -129,7 +128,7 @@ install_pytorch_pip () {
   local env_prefix=$(env_name_or_prefix "${env_name}")
 
   # Install the package from PyTorch PIP (not PyPI)
-  install_from_pytorch_pip "${env_name}" torch "${pytorch_version}" "${pytorch_variant_type}" "${pytorch_variant_version}" || return 1
+  install_from_pytorch_pip "${env_name}" torch "${pytorch_channel_version}" "${pytorch_variant_type_version}" || return 1
 
   # Check that PyTorch is importable
   (test_python_import_package "${env_name}" torch.distributed) || return 1

diff --git a/.github/workflows/fbgemm_gpu_ci_cuda.yml b/.github/workflows/fbgemm_gpu_ci_cuda.yml
@@ -99,7 +99,7 @@ jobs:
 
     # Install via PIP to avoid defaulting to the CPU variant if the GPU variant of the day is not ready
     - name: Install PyTorch Nightly
-      run: . $PRELUDE; install_pytorch_pip $BUILD_ENV nightly cuda ${{ matrix.cuda-version }}
+      run: . $PRELUDE; install_pytorch_pip $BUILD_ENV nightly cuda/${{ matrix.cuda-version }}
 
     - name: Collect PyTorch Environment Info
       if: ${{ success() || failure() }}
@@ -178,7 +178,7 @@ jobs:
 
     # Install via PIP to avoid defaulting to the CPU variant if the GPU variant of the day is not ready
     - name: Install PyTorch Nightly
-      run: . $PRELUDE; install_pytorch_pip $BUILD_ENV nightly cuda ${{ matrix.cuda-version }}
+      run: . $PRELUDE; install_pytorch_pip $BUILD_ENV nightly cuda/${{ matrix.cuda-version }}
 
     - name: Collect PyTorch Environment Info
       if: ${{ success() || failure() }}

diff --git a/.github/workflows/fbgemm_gpu_ci_rocm.yml b/.github/workflows/fbgemm_gpu_ci_rocm.yml
@@ -103,7 +103,7 @@ jobs:
       run: . $PRELUDE; install_build_tools $BUILD_ENV
 
     - name: Install PyTorch-ROCm Nightly
-      run:  . $PRELUDE; install_pytorch_pip $BUILD_ENV nightly rocm ${{ matrix.rocm-version }}
+      run:  . $PRELUDE; install_pytorch_pip $BUILD_ENV nightly rocm/${{ matrix.rocm-version }}
 
     - name: Collect PyTorch Environment Info
       if: ${{ success() || failure() }}
@@ -180,7 +180,7 @@ jobs:
       run: . $PRELUDE; install_build_tools $BUILD_ENV
 
     - name: Install PyTorch-ROCm Nightly
-      run:  . $PRELUDE; install_pytorch_pip $BUILD_ENV nightly rocm ${{ matrix.rocm-version }}
+      run:  . $PRELUDE; install_pytorch_pip $BUILD_ENV nightly rocm/${{ matrix.rocm-version }}
 
     - name: Collect PyTorch Environment Info
       if: ${{ success() || failure() }}