Skip to content

Commit 92a1093

Browse files
nipung90facebook-github-bot
nipung90
authored and
facebook-github-bot
committed
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
1 parent c70539a commit 92a1093

File tree

2 files changed

+122
-34
lines changed

2 files changed

+122
-34
lines changed

‎.github/scripts/test_torchrec.bash

100644100755
Lines changed: 44 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -8,20 +8,15 @@
88
# Exit on failure
99
set -e
1010

11-
# shellcheck source=/dev/null
12-
. "$(dirname "$(realpath -s "$0")")/setup_env.bash"
13-
1411
verbose=0
15-
env_name=test_binary
1612
torchrec_package_name=""
1713
python_version=""
1814
cuda_version="x"
19-
fbgemm_wheel_path="x"
2015
miniconda_prefix="${HOME}/miniconda"
2116

2217
usage () {
2318
# shellcheck disable=SC2086
24-
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]"
19+
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]"
2520
echo "-v : verbose"
2621
echo "-h : help"
2722
echo "PACKAGE_NAME : output package name of TorchRec (e.g., torchrec_nightly)"
@@ -30,14 +25,13 @@ usage () {
3025
echo "PYTHON_VERSION : Python version (e.g., 3.10)"
3126
echo "PYTORCH_CHANNEL_NAME: PyTorch's channel name (e.g., pytorch-nightly, pytorch-test (=pre-release), pytorch (=stable release))"
3227
echo "CUDA_VERSION : PyTorch's CUDA version (e.g., 12.4)"
33-
echo "FBGEMM_WHEEL_PATH : path to FBGEMM_GPU's wheel file"
3428
echo "MINICONDA_PREFIX : path to install Miniconda (default: \$HOME/miniconda)"
3529
echo "Example: Python 3.10 + PyTorch nightly (CUDA 12.4), install miniconda at \$HOME/miniconda, using dist/fbgemm_gpu_nightly.whl"
3630
# shellcheck disable=SC2086
3731
echo " bash $(basename ${BASH_SOURCE[0]}) -v -o torchrec_nightly -p 3.10 -P pytorch-nightly -c 11.7 -w dist/fbgemm_gpu_nightly.whl"
3832
}
3933

40-
while getopts vho:p:P:c:m:w: flag
34+
while getopts vho:p:P:c:m:b: flag
4135
do
4236
case "$flag" in
4337
v) verbose="1";;
@@ -46,18 +40,20 @@ do
4640
P) pytorch_channel_name="${OPTARG}";;
4741
c) cuda_version="${OPTARG}";;
4842
m) miniconda_prefix="${OPTARG}";;
49-
w) fbgemm_wheel_path="${OPTARG}";;
43+
b) build_env="${OPTARG}";;
5044
h) usage
5145
exit 0;;
5246
*) usage
5347
exit 1;;
5448
esac
5549
done
5650

57-
if [ "$torchrec_package_name" == "" ] || [ "$python_version" == "" ] || [ "$cuda_version" == "x" ] || [ "$miniconda_prefix" == "" ] || [ "$pytorch_channel_name" == "" ] || [ "$fbgemm_wheel_path" == "" ]; then
51+
if [ "$torchrec_package_name" == "" ] || [ "$python_version" == "" ] || [ "$cuda_version" == "x" ] || [ "$miniconda_prefix" == "" ] || [ "$pytorch_channel_name" == "" ] || [ "$build_env" == "" ]; then
5852
usage
5953
exit 1
6054
fi
55+
56+
env_name=$build_env
6157
python_tag="${python_version//\./}"
6258

6359
if [ "$verbose" == "1" ]; then
@@ -74,44 +70,59 @@ if [ ! -d "torchrec" ]; then
7470
exit 1
7571
fi
7672

77-
################################################################################
78-
echo "## 1. Set up Miniconda"
79-
################################################################################
80-
81-
setup_miniconda "$miniconda_prefix"
73+
# Install PyTorch
74+
conda run -n "$env_name" pip install torch --index-url https://download.pytorch.org/whl/nightly/cpu
75+
conda run -n "$env_name" python -c "import torch"
8276

83-
################################################################################
84-
echo "## 2. Create Conda environment"
85-
################################################################################
77+
# Import torch.distributed
78+
conda run -n "$env_name" python -c "import torch.distributed"
8679

87-
if [ "${cuda_version}" != "" ]; then
88-
pytorch_variant="cuda ${cuda_version}"
89-
else
90-
pytorch_variant="cpu"
91-
fi
80+
# Import fbgemm_gpu
9281

93-
# shellcheck disable=SC2086
94-
test_setup_conda_environment "$env_name" gcc "$python_version" pip "$pytorch_channel_name" $pytorch_variant
82+
echo "[INSTALL] Installing FBGEMM-GPU wheel: ${wheel_path} ..."
83+
# cd ../fbgemm_gpu
84+
# conda run -n "$env_name" python setup.py clean
85+
# conda run -n "$env_name" python setup.py bdist_wheel --python-tag="py${python_tag}"
86+
# conda run -n "$env_name" pip install fbgemm-gpu ../*.whl
87+
conda run -n "$env_name" python -c "import fbgemm_gpu"
88+
# cd ../torchrec
9589

96-
# Comment out FBGEMM_GPU since we will install it from "$fbgemm_wheel_path"
90+
################################################################################
91+
echo "## 1. Install TorchRec Requirements"
92+
################################################################################
93+
# Comment out FBGEMM_GPU since we should pre-install it from the downloaded wheel file
9794
sed -i 's/fbgemm-gpu/#fbgemm-gpu/g' requirements.txt
9895
conda run -n "$env_name" python -m pip install -r requirements.txt
99-
# Install FBGEMM_GPU from a local wheel file.
100-
conda run -n "$env_name" python -m pip install "$fbgemm_wheel_path"
101-
conda run -n "$env_name" python -c "import fbgemm_gpu"
96+
10297

10398
################################################################################
104-
echo "## 3. Build TorchRec"
99+
echo "## 2. Build TorchRec"
105100
################################################################################
106101

107102
rm -rf dist
108-
conda run -n "$env_name" python setup.py bdist_wheel --package_name "${torchrec_package_name}" --python-tag="py${python_tag}"
103+
conda run -n "$env_name" python setup.py bdist_wheel --python-tag="py${python_tag}"
109104

110105
################################################################################
111-
echo "## 4. Import TorchRec"
106+
echo "## 3. Import TorchRec"
112107
################################################################################
113108

114-
conda run -n "$env_name" python -m pip install dist/"${torchrec_package_name}"*.whl
115109
conda run -n "$env_name" python -c "import torchrec"
116110

117111
echo "Test succeeded"
112+
113+
################################################################################
114+
echo "## 4. Run TorchRec tests"
115+
################################################################################
116+
117+
conda install -n "$env_name" -y pytest
118+
# Read the list of tests to skip from a file, ignoring empty lines and comments
119+
skip_expression=$(awk '!/^($|#)/ {printf " and not %s", $0}' ./.github/scripts/tests_to_skip.txt)
120+
# Check if skip_expression is effectively empty
121+
if [ -z "$skip_expression" ]; then
122+
skip_expression=""
123+
else
124+
skip_expression=${skip_expression:5} # Remove the leading " and "
125+
fi
126+
conda run -n "$env_name" \
127+
python -m pytest torchrec -v -s -W ignore::pytest.PytestCollectionWarning --continue-on-collection-errors \
128+
--ignore-glob=**/test_utils/ -k "$skip_expression"

‎.github/workflows/fbgemm_gpu_ci_cpu.yml

Lines changed: 78 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -114,7 +114,7 @@ jobs:
114114
run: . $PRELUDE; cd fbgemm_gpu; prepare_fbgemm_gpu_build $BUILD_ENV
115115

116116
- name: Build FBGEMM_GPU Wheel
117-
run: . $PRELUDE; cd fbgemm_gpu; build_fbgemm_gpu_package $BUILD_ENV nightly ${{ matrix.build-target }}/cpu
117+
run: . $PRELUDE; cd fbgemm_gpu; build_fbgemm_gpu_package $BUILD_ENV nightly ${{ matrix.build-target }}/cuda
118118

119119
- name: Upload Built Wheel as GHA Artifact
120120
uses: actions/upload-artifact@v4
@@ -200,3 +200,80 @@ jobs:
200200
env:
201201
PYPI_TOKEN: ${{ secrets.PYPI_TOKEN }}
202202
run: . $PRELUDE; publish_to_pypi $BUILD_ENV "$PYPI_TOKEN" *.whl
203+
204+
205+
# Run torchrec CPU tests
206+
torchrec_cpu_tests:
207+
if: ${{ github.repository_owner == 'pytorch' }}
208+
runs-on: ${{ matrix.host-machine.instance }}
209+
container:
210+
image: amazonlinux:2023
211+
options: --user root
212+
defaults:
213+
run:
214+
shell: bash
215+
env:
216+
PRELUDE: .github/scripts/setup_env.bash
217+
BUILD_ENV: build_binary
218+
BUILD_TARGET: ${{ matrix.build-target }}
219+
BUILD_VARIANT: cpu
220+
strategy:
221+
fail-fast: false
222+
matrix:
223+
host-machine: [
224+
# { arch: arm, instance: "linux.arm64.2xlarge", timeout: 30 },
225+
{ arch: x86, instance: "linux.4xlarge", timeout: 20 },
226+
]
227+
build-target: [ "default" ]
228+
python-version: [ "3.9", "3.10", "3.11", "3.12", "3.13" ]
229+
compiler: [ "gcc", "clang" ]
230+
needs: build_artifact
231+
232+
steps:
233+
- name: Setup Build Container
234+
run: yum update -y; yum install -y binutils findutils git pciutils sudo wget which
235+
236+
- name: Checkout the Repository
237+
uses: actions/checkout@v4
238+
239+
- name: Display System Info
240+
run: . $PRELUDE; print_system_info; print_ec2_info
241+
242+
- name: Display GPU Info
243+
run: . $PRELUDE; print_gpu_info
244+
245+
- name: Setup Miniconda
246+
run: . $PRELUDE; setup_miniconda $HOME/miniconda
247+
248+
- name: Create Conda Environment
249+
run: . $PRELUDE; create_conda_environment $BUILD_ENV ${{ matrix.python-version }}
250+
251+
- name: Install C/C++ Compilers for Updated LIBGCC
252+
run: . $PRELUDE; install_cxx_compiler $BUILD_ENV ${{ matrix.compiler }}
253+
254+
- name: Install PyTorch-CPU Nightly
255+
run: . $PRELUDE; install_pytorch_pip $BUILD_ENV ${{ (github.event_name == 'workflow_dispatch' && github.event.inputs.pytorch_channel_version) || 'nightly' }} cpu
256+
257+
- name: Collect PyTorch Environment Info
258+
if: ${{ success() || failure() }}
259+
run: if . $PRELUDE && which conda; then collect_pytorch_env_info $BUILD_ENV; fi
260+
261+
- name: Download Wheel Artifact from GHA
262+
uses: actions/download-artifact@v4
263+
with:
264+
name: fbgemm_${{ matrix.build-target }}_${{ matrix.host-machine.arch }}_${{ matrix.compiler }}_py${{ matrix.python-version }}_cpu.whl
265+
266+
- name: Prepare FBGEMM_GPU Build
267+
run: . $PRELUDE; cd fbgemm_gpu; prepare_fbgemm_gpu_build $BUILD_ENV
268+
269+
- name: Install FBGEMM_GPU Wheel
270+
run: . $PRELUDE; install_fbgemm_gpu_wheel $BUILD_ENV *.whl
271+
272+
- name: Clone torchrec
273+
uses: actions/checkout@v4
274+
with:
275+
repository: pytorch/torchrec
276+
path: torchrec
277+
278+
- name: Run torchrec CPU tests
279+
run: . $PRELUDE; cd torchrec; ../.github/scripts/test_torchrec.bash -o torchrec_nightly -p ${{ matrix.python-version }} -P pytorch-test -c ${{ matrix.compiler }} -b $BUILD_ENV -v 1

0 commit comments

Comments
 (0)