Call torchrec cpu tests from fbgemm test gha (pytorch#4424)

Summary:

This diff kicks off the torchrec cpu unittests along with fbgemm's cpu tests when there is a pull request on fbgemm

Differential Revision: D77598015

Author: nipung90

.github/scripts/test_torchrec.bash

@@ -16,12 +16,11 @@ env_name=test_binary
 torchrec_package_name=""
 python_version=""
 cuda_version="x"
-fbgemm_wheel_path="x"
 miniconda_prefix="${HOME}/miniconda"
 
 usage () {
   # shellcheck disable=SC2086
-  echo "Usage: bash $(basename ${BASH_SOURCE[0]}) -o PACKAGE_NAME -p PYTHON_VERSION -P PYTORCH_CHANNEL_NAME -c CUDA_VERSION -w FBGEMM_WHEEL_PATH [-m MINICONDA_PREFIX] [-v] [-h]"
+  echo "Usage: bash $(basename ${BASH_SOURCE[0]}) -o PACKAGE_NAME -p PYTHON_VERSION -P PYTORCH_CHANNEL_NAME -c CUDA_VERSION [-m MINICONDA_PREFIX] [-v] [-h]"
   echo "-v                  : verbose"
   echo "-h                  : help"
   echo "PACKAGE_NAME        : output package name of TorchRec (e.g., torchrec_nightly)"
@@ -30,7 +29,6 @@ usage () {
   echo "PYTHON_VERSION      : Python version (e.g., 3.10)"
   echo "PYTORCH_CHANNEL_NAME: PyTorch's channel name (e.g., pytorch-nightly, pytorch-test (=pre-release), pytorch (=stable release))"
   echo "CUDA_VERSION        : PyTorch's CUDA version (e.g., 12.4)"
-  echo "FBGEMM_WHEEL_PATH   : path to FBGEMM_GPU's wheel file"
   echo "MINICONDA_PREFIX    : path to install Miniconda (default: \$HOME/miniconda)"
   echo "Example: Python 3.10 + PyTorch nightly (CUDA 12.4), install miniconda at \$HOME/miniconda, using dist/fbgemm_gpu_nightly.whl"
   # shellcheck disable=SC2086
@@ -46,15 +44,14 @@ do
         P) pytorch_channel_name="${OPTARG}";;
         c) cuda_version="${OPTARG}";;
         m) miniconda_prefix="${OPTARG}";;
-        w) fbgemm_wheel_path="${OPTARG}";;
         h) usage
            exit 0;;
         *) usage
            exit 1;;
     esac
 done
 
-if [ "$torchrec_package_name" == "" ] || [ "$python_version" == "" ] || [ "$cuda_version" == "x" ] || [ "$miniconda_prefix" == "" ] || [ "$pytorch_channel_name" == "" ] || [ "$fbgemm_wheel_path" == "" ]; then
+if [ "$torchrec_package_name" == "" ] || [ "$python_version" == "" ] || [ "$cuda_version" == "x" ] || [ "$miniconda_prefix" == "" ] || [ "$pytorch_channel_name" == "" ]; then
   usage
   exit 1
 fi
@@ -93,11 +90,10 @@ fi
 # shellcheck disable=SC2086
 test_setup_conda_environment "$env_name" gcc "$python_version" pip "$pytorch_channel_name" $pytorch_variant
 
-# Comment out FBGEMM_GPU since we will install it from "$fbgemm_wheel_path"
+# Comment out FBGEMM_GPU since we should pre-install it from the downloaded wheel file
 sed -i 's/fbgemm-gpu/#fbgemm-gpu/g' requirements.txt
 conda run -n "$env_name" python -m pip install -r requirements.txt
-# Install FBGEMM_GPU from a local wheel file.
-conda run -n "$env_name" python -m pip install "$fbgemm_wheel_path"
+# Check FBGemm can be imported.
 conda run -n "$env_name" python -c "import fbgemm_gpu"
 
 ################################################################################

.github/workflows/fbgemm_gpu_ci_cpu.yml

@@ -200,3 +200,80 @@ jobs:
       env:
         PYPI_TOKEN: ${{ secrets.PYPI_TOKEN }}
       run: . $PRELUDE; publish_to_pypi $BUILD_ENV "$PYPI_TOKEN" *.whl
+
+
+  # Run torchrec CPU tests
+  torchrec_cpu_tests:
+    if: ${{ github.repository_owner == 'pytorch' }}
+    runs-on: ${{ matrix.host-machine.instance }}
+    container:
+      image: amazonlinux:2023
+      options: --user root
+    defaults:
+      run:
+        shell: bash
+    env:
+      PRELUDE: .github/scripts/setup_env.bash
+      BUILD_ENV: build_binary
+      BUILD_TARGET: ${{ matrix.build-target }}
+      BUILD_VARIANT: cpu
+    strategy:
+      fail-fast: false
+      matrix:
+        host-machine: [
+          { arch: x86, instance: "linux.4xlarge", timeout: 20 },
+          { arch: arm, instance: "linux.arm64.2xlarge", timeout: 30 },
+        ]
+        build-target: [ "default" ]
+        python-version: [ "3.9", "3.10", "3.11", "3.12", "3.13" ]
+        compiler: [ "gcc", "clang" ]
+    needs: build_artifact
+
+    steps:
+    - name: Setup Build Container
+      run: yum update -y; yum install -y binutils findutils git pciutils sudo wget which
+
+    - name: Checkout the Repository
+      uses: actions/checkout@v4
+
+    - name: Download Wheel Artifact from GHA
+      uses: actions/download-artifact@v4
+      with:
+        name: fbgemm_${{ matrix.build-target }}_${{ matrix.host-machine.arch }}_${{ matrix.compiler }}_py${{ matrix.python-version }}_cpu.whl
+
+    - name: Display System Info
+      run: . $PRELUDE; print_system_info; print_ec2_info
+
+    - name: Display GPU Info
+      run: . $PRELUDE; print_gpu_info
+
+    - name: Setup Miniconda
+      run: . $PRELUDE; setup_miniconda $HOME/miniconda
+
+    - name: Create Conda Environment
+      run: . $PRELUDE; create_conda_environment $BUILD_ENV ${{ matrix.python-version }}
+
+    - name: Install C/C++ Compilers for Updated LIBGCC
+      run: . $PRELUDE; install_cxx_compiler $BUILD_ENV ${{ matrix.compiler }}
+
+    - name: Install PyTorch-CPU Nightly
+      run: . $PRELUDE; install_pytorch_pip $BUILD_ENV ${{ (github.event_name == 'workflow_dispatch' && github.event.inputs.pytorch_channel_version) || 'nightly' }} cpu
+
+    - name: Collect PyTorch Environment Info
+      if: ${{ success() || failure() }}
+      run: if . $PRELUDE && which conda; then collect_pytorch_env_info $BUILD_ENV; fi
+
+    - name: Prepare FBGEMM_GPU Build
+      run: . $PRELUDE; cd fbgemm_gpu; prepare_fbgemm_gpu_build $BUILD_ENV
+
+    - name: Install FBGEMM_GPU Wheel
+      run: . $PRELUDE; install_fbgemm_gpu_wheel $BUILD_ENV *.whl
+
+    - name: Clone torchrec
+      uses: actions/checkout@v4
+      with:
+        repository: pytorch/torchrec
+        path: torchrec
+
+    - name: Run Pytorchrec CPU tests
+      run: . $PRELUDE; cd torchrec/torchrec; ../../.github/scripts/test_torchrec.bash -p ${{ matrix.python-version }} -P pytorchtest -c ${{ matrix.compiler }}