Skip to content

Commit c7bdd3b

Browse files
q10q10
committed
[fbgemm_gpu] Update Nova jobs
- Update Nova jobs
1 parent 4587ad0 commit c7bdd3b

7 files changed

+156
-8
lines changed

‎.github/scripts/nova_postscript.bash

+12-2
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,11 @@ echo "[NOVA] Current working directory: $(pwd)"
2525
# Record time for each step
2626
start_time=$(date +%s)
2727

28+
echo "################################################################################"
29+
echo "Environment Variables:"
30+
printenv
31+
echo "################################################################################"
32+
2833
# Collect PyTorch environment information
2934
collect_pytorch_env_info "${BUILD_ENV_NAME}"
3035
end_time=$(date +%s)
@@ -42,8 +47,13 @@ echo "[NOVA] Time taken to install wheel: ${runtime} seconds"
4247
# Test with PyTest
4348
echo "[NOVA] Current working directory: $(pwd)"
4449
if [[ $CU_VERSION = cu* ]]; then
45-
echo "[NOVA] Testing the CUDA variant of FBGEMM_GPU ..."
46-
export fbgemm_variant="cuda"
50+
if [[ ${BUILD_TARGET} == "genai" ]]; then
51+
echo "[NOVA] Testing the GenAI variant of FBGEMM_GPU ..."
52+
export fbgemm_variant="genai"
53+
else
54+
echo "[NOVA] Testing the CUDA variant of FBGEMM_GPU ..."
55+
export fbgemm_variant="cuda"
56+
fi
4757

4858
elif [[ $CU_VERSION = rocm* ]]; then
4959
echo "[NOVA] Testing the ROCm variant of FBGEMM_GPU ..."

‎.github/scripts/nova_prescript.bash

+12-2
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,11 @@ BUILD_ENV_NAME=${CONDA_ENV}
2121
# Record time for each step
2222
start_time=$(date +%s)
2323

24+
echo "################################################################################"
25+
echo "Environment Variables:"
26+
printenv
27+
echo "################################################################################"
28+
2429
# Display System Info
2530
print_system_info
2631
end_time=$(date +%s)
@@ -94,8 +99,13 @@ if [[ $CU_VERSION = cu* ]]; then
9499
start_time=${end_time}
95100
echo "[NOVA] Time taken to find NVML_LIB_PATH: ${runtime} seconds"
96101

97-
echo "[NOVA] Building the CUDA variant of FBGEMM_GPU ..."
98-
export fbgemm_variant="cuda"
102+
if [[ ${BUILD_TARGET} == "genai" ]]; then
103+
echo "[NOVA] Building the GenAI variant of FBGEMM_GPU ..."
104+
export fbgemm_variant="genai"
105+
else
106+
echo "[NOVA] Building the CUDA variant of FBGEMM_GPU ..."
107+
export fbgemm_variant="cuda"
108+
fi
99109

100110
elif [[ $CU_VERSION = rocm* ]]; then
101111
echo "[NOVA] Building the ROCm variant of FBGEMM_GPU ..."

‎.github/workflows/build_wheels_genai_linux_aarch64.yml

+66
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,66 @@
1+
name: Build FBGEMM GenAI Aarch64 Linux Wheels
2+
3+
on:
4+
pull_request:
5+
push:
6+
branches:
7+
- nightly
8+
- main
9+
# Release candidate branch look like: v1.11.0-release
10+
- v[0-9]+.[0-9]+.[0-9]+-release+
11+
tags:
12+
# Release candidate tag look like: v1.11.0-rc1
13+
- v[0-9]+.[0-9]+.[0-9]+-rc[0-9]+
14+
- v[0-9]+.[0-9]+.[0-9]+
15+
workflow_dispatch:
16+
17+
concurrency:
18+
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
19+
cancel-in-progress: true
20+
21+
permissions:
22+
id-token: write
23+
contents: read
24+
25+
jobs:
26+
generate-matrix:
27+
if: ${{ github.repository_owner == 'pytorch' }}
28+
uses: pytorch/test-infra/.github/workflows/generate_binary_build_matrix.yml@main
29+
with:
30+
package-type: wheel
31+
os: linux-aarch64
32+
test-infra-repository: pytorch/test-infra
33+
test-infra-ref: main
34+
with-cuda: disable
35+
36+
build:
37+
if: ${{ github.repository_owner == 'pytorch' }}
38+
needs: generate-matrix
39+
strategy:
40+
fail-fast: false
41+
matrix:
42+
include:
43+
- repository: pytorch/FBGEMM
44+
smoke-test-script: ""
45+
pre-script: ../.github/scripts/nova_prescript.bash
46+
post-script: ../.github/scripts/nova_postscript.bash
47+
env-var-script: .github/scripts/nova_dir.bash
48+
package-name: fbgemm_gpu
49+
name: ${{ matrix.repository }}
50+
uses: pytorch/test-infra/.github/workflows/build_wheels_linux.yml@main
51+
with:
52+
repository: ${{ matrix.repository }}
53+
ref: ""
54+
test-infra-repository: pytorch/test-infra
55+
test-infra-ref: main
56+
build-matrix: ${{ needs.generate-matrix.outputs.matrix }}
57+
pre-script: ${{ matrix.pre-script }}
58+
post-script: ${{ matrix.post-script }}
59+
package-name: ${{ matrix.package-name }}
60+
build-target: genai
61+
env-var-script: ${{ matrix.env-var-script }}
62+
smoke-test-script: ${{ matrix.smoke-test-script }}
63+
trigger-event: ${{ github.event_name }}
64+
architecture: aarch64
65+
setup-miniconda: false
66+
timeout: 210

‎.github/workflows/build_wheels_genai_linux_x86.yml

+56
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,56 @@
1+
name: Build FBGEMM GenAI x86 Linux Wheels
2+
3+
on:
4+
pull_request:
5+
push:
6+
branches:
7+
- nightly
8+
- main
9+
# Release candidate branch look like: v1.11.0-rc1
10+
- v[0-9]+.[0-9]+.[0-9]+-release+
11+
tags:
12+
# Release candidate tag look like: v1.11.0-rc1
13+
- v[0-9]+.[0-9]+.[0-9]+-rc[0-9]+
14+
- v[0-9]+.[0-9]+.[0-9]+
15+
workflow_dispatch:
16+
17+
concurrency:
18+
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
19+
cancel-in-progress: true
20+
21+
permissions:
22+
id-token: write
23+
contents: read
24+
25+
jobs:
26+
generate-matrix:
27+
if: ${{ github.repository_owner == 'pytorch' }}
28+
uses: pytorch/test-infra/.github/workflows/generate_binary_build_matrix.yml@main
29+
with:
30+
package-type: wheel
31+
os: linux
32+
test-infra-repository: pytorch/test-infra
33+
test-infra-ref: main
34+
with-cuda: enable
35+
with-rocm: disable
36+
with-cpu: disable
37+
38+
build:
39+
if: ${{ github.repository_owner == 'pytorch' }}
40+
needs: generate-matrix
41+
name: pytorch/FBGEMM
42+
uses: pytorch/test-infra/.github/workflows/build_wheels_linux.yml@main
43+
with:
44+
repository: pytorch/FBGEMM
45+
ref: ""
46+
pre-script: ../.github/scripts/nova_prescript.bash
47+
post-script: ../.github/scripts/nova_postscript.bash
48+
smoke-test-script: ""
49+
env-var-script: .github/scripts/nova_dir.bash
50+
package-name: fbgemm_gpu
51+
build-target: genai
52+
test-infra-repository: pytorch/test-infra
53+
test-infra-ref: main
54+
build-matrix: ${{ needs.generate-matrix.outputs.matrix }}
55+
trigger-event: ${{ github.event_name }}
56+
timeout: 240

‎.github/workflows/build_wheels_linux_aarch64.yml

+2-1
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
name: Build Aarch64 Linux Wheels
1+
name: Build FBGEMM_GPU Aarch64 Linux Wheels
22

33
on:
44
pull_request:
@@ -32,6 +32,7 @@ jobs:
3232
test-infra-repository: pytorch/test-infra
3333
test-infra-ref: main
3434
with-cuda: disable
35+
3536
build:
3637
if: ${{ github.repository_owner == 'pytorch' }}
3738
needs: generate-matrix

‎.github/workflows/build_wheels_linux_x86.yml

+2-1
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
name: Build x86 Linux Wheels
1+
name: Build FBGEMM_GPU x86 Linux Wheels
22

33
on:
44
pull_request:
@@ -34,6 +34,7 @@ jobs:
3434
with-cuda: enable
3535
with-rocm: enable
3636
with-cpu: enable
37+
3738
build:
3839
if: ${{ github.repository_owner == 'pytorch' }}
3940
needs: generate-matrix

‎fbgemm_gpu/setup.py

+6-2
Original file line numberDiff line numberDiff line change
@@ -116,8 +116,12 @@ def package_name(self) -> str:
116116
sys.exit(0)
117117

118118
elif self.nova_flag() == 0:
119-
# The package name is the same for all build variants in Nova
120-
pass
119+
# In Nova, we are publishing genai packages separately from the main
120+
# fbgemm_gpu package, so if the package variant is genai, we need to
121+
# update the package name accordingly. Otherwise, the package name
122+
# is the same for all other build variants in Nova
123+
if self.args.package_variant == "genai":
124+
pkg_name = "fbgemm_gpu_genai"
121125

122126
else:
123127
# If running outside of Nova workflow context, append the channel

0 commit comments

Comments
 (0)