Fix PIP install workflows (pytorch#2026)

Summary: - Fix package name from fbgemm_gpu to fbgemm-gpu to make post-install checks pass - Add numpy as an install dependency for the fbgemm-gpu package - Rename workflow files for consistency Pull Request resolved: pytorch#2026 Reviewed By: spcyppt Differential Revision: D49408816 Pulled By: q10 fbshipit-source-id: aa7f90625c48a0cf0732c788050a189c28418479
q10 · Sep 19, 2023 · e41a2c5 · e41a2c5
1 parent aae88d6
commit e41a2c5
Show file tree

Hide file tree

Showing 11 changed files with 50 additions and 26 deletions.
diff --git a/.github/scripts/fbgemm_gpu_install.bash b/.github/scripts/fbgemm_gpu_install.bash
@@ -16,6 +16,10 @@
 ################################################################################
 
 __fbgemm_gpu_post_install_checks () {
+  local env_name="$1"
+  # shellcheck disable=SC2155
+  local env_prefix=$(env_name_or_prefix "${env_name}")
+
   echo "[INSTALL] Checking imports and symbols ..."
   (test_python_import_package "${env_name}" fbgemm_gpu) || return 1
   (test_python_import_package "${env_name}" fbgemm_gpu.split_embedding_codegen_lookup_invokers) || return 1
@@ -56,7 +60,7 @@ install_fbgemm_gpu_wheel () {
   # shellcheck disable=SC2086
   (exec_with_retries conda run ${env_prefix} python -m pip install "${wheel_path}") || return 1
 
-  __fbgemm_gpu_post_install_checks || return 1
+  __fbgemm_gpu_post_install_checks "${env_name}" || return 1
 
   echo "[INSTALL] FBGEMM-GPU installation through wheel completed ..."
 }
@@ -80,14 +84,12 @@ install_fbgemm_gpu_pip () {
     echo ""
   fi
 
-  # shellcheck disable=SC2155
-  local env_prefix=$(env_name_or_prefix "${env_name}")
-
   # 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_version}" "${fbgemm_gpu_variant_type}" "${fbgemm_gpu_variant_version}" || return 1
 
   # Run post-installation checks
-  __fbgemm_gpu_post_install_checks || return 1
+  __fbgemm_gpu_post_install_checks "${env_name}" || return 1
 
   echo "[INSTALL] Successfully installed FBGEMM-GPU through PyTorch PIP"
 }
diff --git a/.github/scripts/utils_base.bash b/.github/scripts/utils_base.bash
@@ -83,7 +83,7 @@ env_name_or_prefix () {
 }
 
 test_network_connection () {
-  wget --timeout 1 pypi.org -O /dev/null
+  wget -q --timeout 1 pypi.org -O /dev/null
   local exit_status=$?
 
   # https://man7.org/linux/man-pages/man1/wget.1.html

diff --git a/.github/scripts/utils_pip.bash b/.github/scripts/utils_pip.bash
@@ -15,7 +15,7 @@
 
 install_from_pytorch_pip () {
   local env_name="$1"
-  local package_name="$2"
+  local package_name_raw="$2"
   local package_version="$3"
   local package_variant_type="$4"
   local package_variant_version="$5"
@@ -29,7 +29,7 @@ install_from_pytorch_pip () {
     return 1
   else
     echo "################################################################################"
-    echo "# Install ${package_name} (PyTorch PIP)"
+    echo "# Install ${package_name_raw} (PyTorch PIP)"
     echo "#"
     echo "# [TIMESTAMP] $(date --utc +%FT%T.%3NZ)"
     echo "################################################################################"
@@ -38,6 +38,10 @@ install_from_pytorch_pip () {
 
   test_network_connection || return 1
 
+  # Replace underscores with hyphens to materialize the canonical name of the package
+  # shellcheck disable=SC2155
+  local package_name=$(echo "${package_name_raw}" | tr '_' '-')
+
   # Set the package variant
   if [ "$package_variant_type" == "cuda" ]; then
     # Extract the CUDA version or default to 11.8.0

diff --git a/.../workflows/build_wheels_aarch64_linux.yml → .../workflows/build_wheels_linux_aarch64.yml b/.../workflows/build_wheels_aarch64_linux.yml → .../workflows/build_wheels_linux_aarch64.yml
diff --git a/.github/workflows/build-wheels-linux.yml → .github/workflows/build_wheels_linux_x86.yml b/.github/workflows/build-wheels-linux.yml → .github/workflows/build_wheels_linux_x86.yml
@@ -1,4 +1,4 @@
-name: Build Linux Wheels
+name: Build x86 Linux Wheels
 
 on:
   pull_request:

diff --git a/.github/workflows/fbgemm_gpu_cpu_nightly.yml b/.github/workflows/fbgemm_gpu_cpu_nightly.yml
@@ -6,7 +6,7 @@
 name: FBGEMM_GPU-CPU Nightly Build
 
 on:
-  # PR Trigger (enabled only for debugging)
+  # PR Trigger (enabled for regression checks and debugging)
   #
   pull_request:
     branches:

diff --git a/.github/workflows/fbgemm_gpu_cpu_release.yml b/.github/workflows/fbgemm_gpu_cpu_release.yml
@@ -6,7 +6,7 @@
 name: FBGEMM_GPU-CPU Release Build
 
 on:
-  # PR Trigger (enabled only for debugging)
+  # PR Trigger (enabled for regression checks and debugging)
   #
   pull_request:
     branches:

diff --git a/.github/workflows/fbgemm_gpu_cuda_nightly.yml b/.github/workflows/fbgemm_gpu_cuda_nightly.yml
@@ -6,7 +6,7 @@
 name: FBGEMM_GPU-CUDA Nightly Build
 
 on:
-  # PR Trigger (enabled only for debugging)
+  # PR Trigger (enabled for regression checks and debugging)
   #
   pull_request:
     branches:

diff --git a/.github/workflows/fbgemm_gpu_cuda_release.yml b/.github/workflows/fbgemm_gpu_cuda_release.yml
@@ -6,7 +6,7 @@
 name: FBGEMM_GPU-CUDA Release Build
 
 on:
-  # PR Trigger (enabled only for debugging)
+  # PR Trigger (enabled for regression checks and debugging)
   #
   pull_request:
     branches:

diff --git a/.github/workflows/fbgemm_gpu_pip.yml b/.github/workflows/fbgemm_gpu_pip.yml
@@ -6,12 +6,22 @@
 name: FBGEMM_GPU PIP Install + Test
 
 on:
+  # PR Trigger (enabled for regression checks and debugging)
+  #
+  pull_request:
+    branches:
+      - main
+
   # Manual Trigger
   #
   workflow_dispatch:
     inputs:
+      pytorch_version:
+        description: PyTorch Version (e.g. '2.1.0', 'nightly', 'test')
+        type: string
+        required: true
       fbgemm_gpu_version:
-        description: FBGEMM-GPU Version (e.g. '0.5.0rc1')
+        description: FBGEMM-GPU Version (e.g. '0.5.0rc1', 'nightly', 'test')
         type: string
         required: true
       fbgemm_gpu_variant_type:
@@ -20,15 +30,11 @@ on:
         required: true
         options: [ "cpu", "cuda", "rocm" ]
         default: "cpu"
-      fbgemm_gpu_variant_version:
-        description: FBGEMM-GPU Variant Version (e.g. 'CUDA 12.1.1' --> 12.1.1)
-        type: string
-        required: false
 
 
 jobs:
   test_pypi_install_cpu:
-    if: ${{ github.event_name == 'workflow_dispatch' && github.event.inputs.fbgemm_gpu_variant_type == 'cpu' }}
+    if: ${{ github.event_name == 'pull_request' || (github.event_name == 'workflow_dispatch' && github.event.inputs.fbgemm_gpu_variant_type == 'cpu') }}
     runs-on: ${{ matrix.host-machine.instance }}
     container:
       image: amazonlinux:2023
@@ -67,19 +73,22 @@ jobs:
     - name: Create Conda Environment
       run: . $PRELUDE; create_conda_environment $BUILD_ENV ${{ matrix.python-version }}
 
+    - name: Install Build Tools
+      run: . $PRELUDE; install_build_tools $BUILD_ENV
+
     - name: Install PyTorch-CPU
-      run: . $PRELUDE; install_pytorch_pip $BUILD_ENV test cpu
+      run: . $PRELUDE; install_pytorch_pip $BUILD_ENV ${{ github.event.inputs.pytorch_version || 'nightly' }} cpu
 
     - name: Install FBGEMM_GPU-CPU
-      run: . $PRELUDE; cd fbgemm_gpu; install_fbgemm_gpu_pip $BUILD_ENV ${{ github.event.inputs.fbgemm_gpu_version }} cpu
+      run: . $PRELUDE; install_fbgemm_gpu_pip $BUILD_ENV ${{ github.event.inputs.fbgemm_gpu_version || 'nightly' }} cpu
 
     - name: Test with PyTest
       timeout-minutes: 10
       run: . $PRELUDE; cd fbgemm_gpu/test; run_fbgemm_gpu_tests $BUILD_ENV cpu
 
 
   test_pypi_install_cuda:
-    if: ${{ github.event_name == 'workflow_dispatch' && github.event.inputs.fbgemm_gpu_variant_type == 'cuda' }}
+    if: ${{ github.event_name == 'pull_request' || (github.event_name == 'workflow_dispatch' && github.event.inputs.fbgemm_gpu_variant_type == 'cuda') }}
     runs-on: ${{ matrix.host-machine.instance }}
     defaults:
       run:
@@ -118,14 +127,17 @@ jobs:
     - name: Create Conda Environment
       run: . $PRELUDE; create_conda_environment $BUILD_ENV ${{ matrix.python-version }}
 
+    - name: Install Build Tools
+      run: . $PRELUDE; install_build_tools $BUILD_ENV
+
     - name: Install CUDA
       run: . $PRELUDE; install_cuda $BUILD_ENV ${{ matrix.cuda-version }}
 
     - name: Install PyTorch-CUDA
-      run: . $PRELUDE; install_pytorch_pip $BUILD_ENV nightly cuda ${{ matrix.cuda-version }}
+      run: . $PRELUDE; install_pytorch_pip $BUILD_ENV ${{ github.event.inputs.pytorch_version || 'nightly' }} cuda ${{ matrix.cuda-version }}
 
     - name: Install FBGEMM_GPU-CUDA
-      run: . $PRELUDE; cd fbgemm_gpu; install_fbgemm_gpu_pip $BUILD_ENV ${{ github.event.inputs.fbgemm_gpu_version }} cuda ${{ github.event.inputs.fbgemm_gpu_variant_version }}
+      run: . $PRELUDE; install_fbgemm_gpu_pip $BUILD_ENV ${{ github.event.inputs.fbgemm_gpu_version || 'nightly' }} cuda ${{ matrix.cuda-version }}
 
     - name: Test with PyTest
       timeout-minutes: 10
@@ -184,10 +196,10 @@ jobs:
       run: . $PRELUDE; install_build_tools $BUILD_ENV
 
     - name: Install PyTorch-ROCm
-      run:  . $PRELUDE; install_pytorch_pip $BUILD_ENV nightly rocm ${{ matrix.rocm-version }}
+      run:  . $PRELUDE; install_pytorch_pip $BUILD_ENV ${{ github.event.inputs.pytorch_version || 'nightly' }} rocm ${{ matrix.rocm-version }}
 
     - name: Install FBGEMM_GPU-ROCm
-      run: . $PRELUDE; cd fbgemm_gpu; install_fbgemm_gpu_pip $BUILD_ENV ${{ github.event.inputs.fbgemm_gpu_version }} rocm ${{ github.event.inputs.fbgemm_gpu_variant_version }}
+      run: . $PRELUDE; install_fbgemm_gpu_pip $BUILD_ENV ${{ github.event.inputs.fbgemm_gpu_version || 'nightly' }} rocm ${{ matrix.rocm-version }}
 
     - name: Test FBGEMM_GPU-ROCm
       timeout-minutes: 15

diff --git a/fbgemm_gpu/setup.py b/fbgemm_gpu/setup.py
@@ -318,6 +318,12 @@ def main(argv: List[str]) -> None:
             "CUDA",
         ],
         packages=["fbgemm_gpu"],
+        install_requires=[
+            # Only specify numpy, as specifying torch will auto-install the
+            # release version of torch, which is not what we want for the
+            # nightly and test packages
+            "numpy",
+        ],
         cmake_args=cmake_environment_variables(args),
         cmdclass={
             "install": FbgemmGpuInstaller,