Re-enable Tests for OSS PR Builds (#1598)

Summary:
- Upgrade the FBGEMM_GPU release workflows to use the same build scripts framework for running the builds
- Update the build scripts framework to properly support running PyTests
- Disable certain tests from running in the **CPU-only build** as they are known to be failing at the moment (`jagged_tensor_ops_test`, `uvm_test`)
- Update both the nightly and release jobs to actually run tests prior to publishing to PyPI
- Fix bug with PyPI publishing for the FBGEMM_GPU-CPU nightly job
- Update the nightly build cron schedules to follow PST
- Remove some jobs from FBGEMM_GPU CI as they are now redundant with the FBGEMM_GPU nightly / release builds
- Update the FBGEMM_GPU-CPU CI job (Ubuntu) to use the build scripts framework for building and running tests

Pull Request resolved: #1598

Reviewed By: brad-mengchi

Differential Revision: D43448188

Pulled By: q10

fbshipit-source-id: 81ec5c0ee2d9c7d1246649a31f669b0ed585d23d

Author: q10

.github/scripts/setup_env.bash

@@ -93,6 +93,66 @@ test_env_var () {
   fi
 }
 
+install_system_packages () {
+  if [ $# -le 0 ]; then
+    echo "Usage: ${FUNCNAME[0]} PACKAGE_NAME ... "
+    echo "Example(s):"
+    echo "    ${FUNCNAME[0]} miopen-hip miopen-hip-dev"
+    return 1
+  fi
+
+  if which sudo; then
+    update_cmd=("sudo")
+    install_cmd=("sudo")
+  else
+    update_cmd=()
+    install_cmd=()
+  fi
+
+  if which apt-get; then
+    update_cmd+=(apt update -y)
+    install_cmd+=(apt install -y "$@")
+  elif which yum; then
+    update_cmd+=(yum update -y)
+    install_cmd+=(yum install -y "$@")
+  else
+    echo "[INSTALL] Could not find a system package installer to install packages!"
+    return 1
+  fi
+
+  echo "[INSTALL] Updating system repositories ..."
+  # shellcheck disable=SC2068
+  print_exec ${update_cmd[@]}
+
+  # shellcheck disable=SC2145
+  echo "[INSTALL] Installing system package(s): $@ ..."
+  # shellcheck disable=SC2068
+  print_exec ${install_cmd[@]}
+}
+
+run_python_test () {
+  env_name="$1"
+  python_test_file="$2"
+  if [ "$python_test_file" == "" ]; then
+    echo "Usage: ${FUNCNAME[0]} ENV_NAME PYTHON_TEST_FILE"
+    echo "Example(s):"
+    echo "    ${FUNCNAME[0]} build_env quantize_ops_test.py"
+    return 1
+  else
+    echo "################################################################################"
+    echo "# [$(date --utc +%FT%T.%3NZ)] Run Python Test Suite:"
+    echo "#   ${python_test_file}"
+    echo "################################################################################"
+  fi
+
+  if conda run -n "${env_name}" python -m pytest -v -s -W ignore::pytest.PytestCollectionWarning "${python_test_file}"; then
+    echo "[TEST] Python test suite PASSED: ${python_test_file}"
+  else
+    echo "[TEST] Python test suite FAILED: ${python_test_file}"
+    return 1
+  fi
+}
+
 print_system_info () {
   echo "################################################################################"
   echo "# Print System Info"
@@ -114,12 +174,7 @@ print_system_info () {
   print_exec cat /etc/os-release
 
   echo "[INFO] Check GPU info"
-  if which apt-get; then
-    print_exec sudo apt-get install -y lshw
-  else
-    print_exec sudo yum install -y lshw
-  fi
-
+  install_system_packages lshw
   print_exec sudo lshw -C display
 }
 
@@ -412,10 +467,6 @@ install_cuda () {
   # Ensure that the libraries are properly installed
   test_filepath "${env_name}" libnvToolsExt.so || return 1
 
-  # LIBNVTOOLSEXT
-  # CUDA_TOOLKIT_ROOT_DIR
-  # print_exec conda env config vars set -n "${env_name}" CUDNN_INCLUDE_DIR="${install_path}/include" CUDNN_LIBRARY="${install_path}/lib"
-
   # Print nvcc version
   print_exec conda run -n "${env_name}" nvcc --version
   echo "[INSTALL] Successfully installed CUDA ${cuda_version}"
@@ -440,9 +491,8 @@ install_cxx_compiler () {
   fi
 
   if [ "$use_yum" != "" ]; then
-    echo "[INSTALL] Installing C/C++ compilers through Yum ..."
-    print_exec sudo yum update -y
-    print_exec sudo yum install -y gcc gcc-c++
+    echo "[INSTALL] Installing C/C++ compilers through yum ..."
+    install_system_packages gcc gcc-c++
   else
     # Install gxx_linux-64 from main instead of cxx-compiler from conda-forge, as
     # the latter breaks builds:
@@ -770,6 +820,95 @@ build_fbgemm_gpu_install () {
   echo "[BUILD] FBGEMM-GPU build + install completed"
 }
 
+install_fbgemm_gpu_package () {
+  env_name="$1"
+  package_name="$2"
+  if [ "$package_name" == "" ]; then
+    echo "Usage: ${FUNCNAME[0]} ENV_NAME WHEEL_NAME"
+    echo "Example(s):"
+    echo "    ${FUNCNAME[0]} build_env fbgemm_gpu.whl     # Install the package (wheel)"
+    return 1
+  else
+    echo "################################################################################"
+    echo "# Install FBGEMM-GPU Package (Wheel)"
+    echo "#"
+    echo "# [TIMESTAMP] $(date --utc +%FT%T.%3NZ)"
+    echo "################################################################################"
+    echo ""
+  fi
+
+  echo "[BUILD] Installing FBGEMM-GPU wheel: ${package_name} ..."
+  conda run -n "${env_name}" python -m pip install "${package_name}"
+
+  echo "[BUILD] Checking imports ..."
+  test_python_import "${env_name}" fbgemm_gpu || return 1
+  test_python_import "${env_name}" fbgemm_gpu.split_embedding_codegen_lookup_invokers || return 1
+
+  echo "[BUILD] Wheel installation completed ..."
+}
+
+
+################################################################################
+# Publish Functions
+################################################################################
+
+run_fbgemm_gpu_tests () {
+  env_name="$1"
+  cpu_only="$2"
+  if [ "$env_name" == "" ]; then
+    echo "Usage: ${FUNCNAME[0]} ENV_NAME [CPU_ONLY]"
+    echo "Example(s):"
+    echo "    ${FUNCNAME[0]} build_env    # Run all tests"
+    echo "    ${FUNCNAME[0]} build_env 1  # Skip tests known to be broken in CPU-only mode"
+    return 1
+  else
+    echo "################################################################################"
+    echo "# Run FBGEMM-GPU Tests"
+    echo "#"
+    echo "# [TIMESTAMP] $(date --utc +%FT%T.%3NZ)"
+    echo "################################################################################"
+    echo ""
+  fi
+
+  # These are either non-tests or currently-broken tests in both FBGEMM_GPU and FBGEMM_GPU-CPU
+  files_to_skip=(
+    split_table_batched_embeddings_test.py
+    test_utils.py
+    ssd_split_table_batched_embeddings_test.py
+  )
+
+  if [ "$cpu_only" != "" ]; then
+    # These are tests that are currently broken in FBGEMM_GPU-CPU
+    unstable_tests=(
+      jagged_tensor_ops_test.py
+      uvm_test.py
+    )
+  else
+    unstable_tests=()
+  fi
+
+  echo "[TEST] Installing pytest ..."
+  print_exec conda install -n "${env_name}" -y pytest
+
+  echo "[BUILD] Checking imports ..."
+  test_python_import "${env_name}" fbgemm_gpu || return 1
+
+  # NOTE: These tests running on single CPU core with a less powerful testing
+  # GPU in GHA can take up to 5 hours.
+  for test_file in *.py; do
+    if echo "${files_to_skip[@]}" | grep "${test_file}"; then
+      echo "[TEST] Skipping test file known to be broken: ${test_file}"
+    elif echo "${unstable_tests[@]}" | grep "${test_file}"; then
+      echo "[TEST] Skipping test file: ${test_file}"
+    elif run_python_test "${env_name}" "${test_file}"; then
+      echo ""
+    else
+      return 1
+    fi
+  done
+}
+
+
 ################################################################################
 # Publish Functions
 ################################################################################

.github/workflows/fbgemm_gpu_ci.yml

@@ -14,81 +14,6 @@ on:
       - main
 
 jobs:
-  build_nvidia_gpu:
-    if: ${{ false }}  # Disable the job for now
-    runs-on: ${{ matrix.os }}
-    defaults:
-      run:
-        shell: bash
-    env:
-      PRELUDE: .github/scripts/setup_env.bash
-      BUILD_ENV: build_binary
-    strategy:
-      # Don't fast-fail all the other builds if one of the them fails
-      fail-fast: false
-      matrix:
-        os: [ ubuntu-20.04 ]
-        python-version: [ "3.8" ]
-        # As of version 2.0, PyTorch has dropped support for CUDA 11.6
-        cuda-version: [ 11.7.1 ]
-
-    steps:
-    - name: Checkout the Repository
-      uses: actions/checkout@v3
-      with:
-        submodules: true
-
-    - name: Display System Info
-      run: . $PRELUDE; print_system_info
-
-    - name: Setup Miniconda
-      run: |
-        . $PRELUDE; setup_miniconda $HOME/miniconda
-        echo "${HOME}/miniconda/bin" >> $GITHUB_PATH
-        echo "CONDA=${HOME}/miniconda" >> $GITHUB_PATH
-
-    - name: Create Conda Environment
-      run: |
-        . $PRELUDE
-        create_conda_environment $BUILD_ENV ${{ matrix.python-version }}
-
-        # This hack is needed to get builds running properly on Ubuntu 20.04
-        echo "[SETUP] Creating symlink \$CONDA_PREFIX/lib64 -> \$CONDA_PREFIX/lib ..."
-        conda_prefix=$(conda run -n "${env_name}" printenv CONDA_PREFIX)
-        ln -s "${conda_prefix}/lib" "${conda_prefix}/lib64"
-
-    # - name: Install C/C++ Compilers
-    #   run: . $PRELUDE; install_cxx_compiler $BUILD_ENV
-
-    - 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
-      run:  . $PRELUDE; install_pytorch_pip $BUILD_ENV nightly cuda "${{ matrix.cuda-version }}"
-
-    - name: Install cuDNN
-      run: . $PRELUDE; install_cudnn $BUILD_ENV "$(pwd)/build_only/cudnn" "${{ matrix.cuda-version }}"
-
-    - name: Prepare FBGEMM Build
-      run: . $PRELUDE; cd fbgemm_gpu; prepare_fbgemm_gpu_build $BUILD_ENV
-
-    - name: Build and Install FBGEMM_GPU
-      run: . $PRELUDE; cd fbgemm_gpu; build_fbgemm_gpu_install $BUILD_ENV
-
-    - name: Test FBGEMM_GPU installation
-      shell: bash
-      run: |
-        . $PRELUDE;
-        cd fbgemm_gpu/test
-        print_exec conda run -n $BUILD_ENV python input_combine_test.py -v
-        print_exec conda run -n $BUILD_ENV python quantize_ops_test.py -v
-        print_exec conda run -n $BUILD_ENV python sparse_ops_test.py -v
-        conda run -n $BUILD_ENV python -c "import fbgemm_gpu"
-        conda run -n $BUILD_ENV python -c "import fbgemm_gpu.split_embedding_codegen_lookup_invokers"
-
   build_amd_gpu:
     if: ${{ false }}  # Disable the job for now
     runs-on: ${{ matrix.os }}
@@ -200,71 +125,48 @@ jobs:
         "
         docker run $DOCKER_OPTIONS $DOCKER_IMAGE $JENKINS_REPO_DIR_DOCKER/.jenkins/rocm/build_and_test.sh $JENKINS_REPO_DIR_DOCKER
 
-  test_nvidia_gpu:
-    if: ${{ false }}  # Disable the job for now
-    uses: pytorch/test-infra/.github/workflows/linux_job.yml@main
-    with:
-      job-name: cuda 11.7, A10
-      runner: linux.g5.4xlarge.nvidia.gpu # A10
-      repository: pytorch/fbgemm
-      gpu-arch-type: cuda
-      gpu-arch-version: 11.7
-      timeout: 150
-      script: |
-        set -x
-        # Checkout FBGEMM_GPU
-        git submodule update --init
-
-        # Build FBGEMM_GPU with pytorch-nightly
-        CUDA_VERSION="11.7.1"
-        PYTHON_VERSION="3.10"
-        bash .github/scripts/build_wheel.bash -v -p "$PYTHON_VERSION" -o fbgemm_gpu_test -P pytorch-nightly -c "$CUDA_VERSION" -m /opt/conda
-
-        # Test FBGEMM_GPU using a generated wheel file
-        WHEEL_PATH="$(ls fbgemm_gpu/dist/*.whl)"
-        bash .github/scripts/test_wheel.bash -v -p "$PYTHON_VERSION" -P pytorch-nightly -c "$CUDA_VERSION" -w "$WHEEL_PATH" -m /opt/conda
-
-  build_cpu_only:
+  build_and_test_cpu:
     runs-on: ${{ matrix.os }}
+    defaults:
+      run:
+        shell: bash
+    env:
+      PRELUDE: .github/scripts/setup_env.bash
+      BUILD_ENV: build_binary
     strategy:
       matrix:
-        os: [ubuntu-latest]
+        os: [ ubuntu-20.04, ubuntu-latest ]
+        python-version: [ "3.8", "3.9", "3.10" ]
 
     steps:
-    - uses: actions/checkout@v3
+    - name: Checkout the Repository
+      uses: actions/checkout@v3
+      with:
+        submodules: true
 
-    - name: Install dependencies
-      shell: bash
-      run: |
-        sudo apt-get update
-        sudo apt-get -y install git pip python3-dev
-        sudo pip install cmake scikit-build ninja jinja2 numpy hypothesis --no-input
-        # Install PyTorch (nightly) as required by fbgemm_gpu
-        sudo pip install --pre torch -f https://download.pytorch.org/whl/nightly/cpu/torch_nightly.html
+    - name: Display System Info
+      run: . $PRELUDE; print_system_info
 
-    - name: Checkout submodules
-      shell: bash
+    - name: Setup Miniconda
       run: |
-        cd fbgemm_gpu
-        git submodule sync
-        git submodule update --init --recursive
+        . $PRELUDE; setup_miniconda $HOME/miniconda
+        echo "${HOME}/miniconda/bin" >> $GITHUB_PATH
+        echo "CONDA=${HOME}/miniconda" >> $GITHUB_PATH
 
-    - name: Build fbgemm_gpu
-      shell: bash
-      run: |
-        cd fbgemm_gpu
-        # to avoid "Permission denied" error in '/usr/local/lib/python3.8/dist-packages/' folder
-        sudo python setup.py install --cpu_only
+    - name: Create Conda Environment
+      run: . $PRELUDE; create_conda_environment $BUILD_ENV ${{ matrix.python-version }}
 
-    - name: Test fbgemm_gpu cpu-only installation
-      shell: bash
-      run: |
-        cd fbgemm_gpu
-        cd test
-        python batched_unary_embeddings_test.py -v
-        python input_combine_test.py -v
-        python layout_transform_ops_test.py -v
-        python merge_pooled_embeddings_test.py -v
-        python permute_pooled_embedding_modules_test.py -v
-        python quantize_ops_test.py -v
-        python sparse_ops_test.py -v
+    - name: Install Build Tools
+      run: . $PRELUDE; install_build_tools $BUILD_ENV
+
+    - name: Install PyTorch
+      run:  . $PRELUDE; install_pytorch_pip $BUILD_ENV nightly cpu
+
+    - name: Prepare FBGEMM Build
+      run: . $PRELUDE; cd fbgemm_gpu; prepare_fbgemm_gpu_build $BUILD_ENV
+
+    - name: Build and Install FBGEMM_GPU (CPU version)
+      run: . $PRELUDE; cd fbgemm_gpu; build_fbgemm_gpu_install $BUILD_ENV cpuonly
+
+    - name: Test with PyTest
+      run: . $PRELUDE; cd fbgemm_gpu/test; run_fbgemm_gpu_tests $BUILD_ENV cpuonly