diff --git a/.github/scripts/nova_postscript.bash b/.github/scripts/nova_postscript.bash index 52c3def34e..40448b2f55 100644 --- a/.github/scripts/nova_postscript.bash +++ b/.github/scripts/nova_postscript.bash @@ -25,6 +25,11 @@ echo "[NOVA] Current working directory: $(pwd)" # Record time for each step start_time=$(date +%s) +echo "################################################################################" +echo "Environment Variables:" +printenv +echo "################################################################################" + # Collect PyTorch environment information collect_pytorch_env_info "${BUILD_ENV_NAME}" end_time=$(date +%s) @@ -42,8 +47,13 @@ echo "[NOVA] Time taken to install wheel: ${runtime} seconds" # Test with PyTest echo "[NOVA] Current working directory: $(pwd)" if [[ $CU_VERSION = cu* ]]; then - echo "[NOVA] Testing the CUDA variant of FBGEMM_GPU ..." - export fbgemm_variant="cuda" + if [[ ${BUILD_TARGET} == "genai" ]]; then + echo "[NOVA] Testing the GenAI variant of FBGEMM_GPU ..." + export fbgemm_variant="genai" + else + echo "[NOVA] Testing the CUDA variant of FBGEMM_GPU ..." + export fbgemm_variant="cuda" + fi elif [[ $CU_VERSION = rocm* ]]; then echo "[NOVA] Testing the ROCm variant of FBGEMM_GPU ..." diff --git a/.github/scripts/nova_prescript.bash b/.github/scripts/nova_prescript.bash index d9a38d6a79..66e08b3ef9 100644 --- a/.github/scripts/nova_prescript.bash +++ b/.github/scripts/nova_prescript.bash @@ -21,6 +21,11 @@ BUILD_ENV_NAME=${CONDA_ENV} # Record time for each step start_time=$(date +%s) +echo "################################################################################" +echo "Environment Variables:" +printenv +echo "################################################################################" + # Display System Info print_system_info end_time=$(date +%s) @@ -94,8 +99,13 @@ if [[ $CU_VERSION = cu* ]]; then start_time=${end_time} echo "[NOVA] Time taken to find NVML_LIB_PATH: ${runtime} seconds" - echo "[NOVA] Building the CUDA variant of FBGEMM_GPU ..." - export fbgemm_variant="cuda" + if [[ ${BUILD_TARGET} == "genai" ]]; then + echo "[NOVA] Building the GenAI variant of FBGEMM_GPU ..." + export fbgemm_variant="genai" + else + echo "[NOVA] Building the CUDA variant of FBGEMM_GPU ..." + export fbgemm_variant="cuda" + fi elif [[ $CU_VERSION = rocm* ]]; then echo "[NOVA] Building the ROCm variant of FBGEMM_GPU ..." diff --git a/.github/workflows/build_wheels_genai_linux_aarch64.yml b/.github/workflows/build_wheels_genai_linux_aarch64.yml new file mode 100644 index 0000000000..a8ef63a95e --- /dev/null +++ b/.github/workflows/build_wheels_genai_linux_aarch64.yml @@ -0,0 +1,64 @@ +name: Build FBGEMM GenAI Aarch64 Linux Wheels + +on: + pull_request: + push: + branches: + - nightly + - main + tags: + # Release candidate tag look like: v1.11.0-rc1 + - v[0-9]+.[0-9]+.[0-9]+-rc[0-9]+ + - v[0-9]+.[0-9]+.[0-9]+ + workflow_dispatch: + +concurrency: + group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }} + cancel-in-progress: true + +permissions: + id-token: write + contents: read + +jobs: + generate-matrix: + if: ${{ github.repository_owner == 'pytorch' }} + uses: pytorch/test-infra/.github/workflows/generate_binary_build_matrix.yml@main + with: + package-type: wheel + os: linux-aarch64 + test-infra-repository: pytorch/test-infra + test-infra-ref: main + with-cuda: disable + + build: + if: ${{ github.repository_owner == 'pytorch' }} + needs: generate-matrix + strategy: + fail-fast: false + matrix: + include: + - repository: pytorch/FBGEMM + smoke-test-script: "" + pre-script: ../.github/scripts/nova_prescript.bash + post-script: ../.github/scripts/nova_postscript.bash + env-var-script: .github/scripts/nova_dir.bash + package-name: fbgemm_gpu + name: ${{ matrix.repository }} + uses: pytorch/test-infra/.github/workflows/build_wheels_linux.yml@main + with: + repository: ${{ matrix.repository }} + ref: "" + test-infra-repository: pytorch/test-infra + test-infra-ref: main + build-matrix: ${{ needs.generate-matrix.outputs.matrix }} + pre-script: ${{ matrix.pre-script }} + post-script: ${{ matrix.post-script }} + package-name: ${{ matrix.package-name }} + build-target: genai + env-var-script: ${{ matrix.env-var-script }} + smoke-test-script: ${{ matrix.smoke-test-script }} + trigger-event: ${{ github.event_name }} + architecture: aarch64 + setup-miniconda: false + timeout: 210 diff --git a/.github/workflows/build_wheels_genai_linux_x86.yml b/.github/workflows/build_wheels_genai_linux_x86.yml new file mode 100644 index 0000000000..2224ed3447 --- /dev/null +++ b/.github/workflows/build_wheels_genai_linux_x86.yml @@ -0,0 +1,54 @@ +name: Build FBGEMM GenAI x86 Linux Wheels + +on: + pull_request: + push: + branches: + - nightly + - main + tags: + # Release candidate tag look like: v1.11.0-rc1 + - v[0-9]+.[0-9]+.[0-9]+-rc[0-9]+ + - v[0-9]+.[0-9]+.[0-9]+ + workflow_dispatch: + +concurrency: + group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }} + cancel-in-progress: true + +permissions: + id-token: write + contents: read + +jobs: + generate-matrix: + if: ${{ github.repository_owner == 'pytorch' }} + uses: pytorch/test-infra/.github/workflows/generate_binary_build_matrix.yml@main + with: + package-type: wheel + os: linux + test-infra-repository: pytorch/test-infra + test-infra-ref: main + with-cuda: enable + with-rocm: disable + with-cpu: disable + + build: + if: ${{ github.repository_owner == 'pytorch' }} + needs: generate-matrix + name: pytorch/FBGEMM + uses: pytorch/test-infra/.github/workflows/build_wheels_linux.yml@main + with: + repository: pytorch/FBGEMM + ref: "" + pre-script: ../.github/scripts/nova_prescript.bash + post-script: ../.github/scripts/nova_postscript.bash + smoke-test-script: "" + env-var-script: .github/scripts/nova_dir.bash + package-name: fbgemm_gpu + build-target: genai + test-infra-repository: pytorch/test-infra + test-infra-ref: main + build-matrix: ${{ needs.generate-matrix.outputs.matrix }} + trigger-event: ${{ github.event_name }} + timeout: 240 diff --git a/.github/workflows/build_wheels_linux_aarch64.yml b/.github/workflows/build_wheels_linux_aarch64.yml index 187b91acaa..89edc67552 100644 --- a/.github/workflows/build_wheels_linux_aarch64.yml +++ b/.github/workflows/build_wheels_linux_aarch64.yml @@ -1,4 +1,4 @@ -name: Build Aarch64 Linux Wheels +name: Build FBGEMM_GPU Aarch64 Linux Wheels on: pull_request: @@ -6,8 +6,6 @@ on: branches: - nightly - main - # Release candidate branch look like: v1.11.0-release - - v[0-9]+.[0-9]+.[0-9]+-release+ tags: # Release candidate tag look like: v1.11.0-rc1 - v[0-9]+.[0-9]+.[0-9]+-rc[0-9]+ @@ -32,6 +30,7 @@ jobs: test-infra-repository: pytorch/test-infra test-infra-ref: main with-cuda: disable + build: if: ${{ github.repository_owner == 'pytorch' }} needs: generate-matrix diff --git a/.github/workflows/build_wheels_linux_x86.yml b/.github/workflows/build_wheels_linux_x86.yml index c303a034e5..2df211ae50 100644 --- a/.github/workflows/build_wheels_linux_x86.yml +++ b/.github/workflows/build_wheels_linux_x86.yml @@ -1,4 +1,4 @@ -name: Build x86 Linux Wheels +name: Build FBGEMM_GPU x86 Linux Wheels on: pull_request: @@ -6,8 +6,6 @@ on: branches: - nightly - main - # Release candidate branch look like: v1.11.0-rc1 - - v[0-9]+.[0-9]+.[0-9]+-release+ tags: # Release candidate tag look like: v1.11.0-rc1 - v[0-9]+.[0-9]+.[0-9]+-rc[0-9]+ @@ -34,6 +32,7 @@ jobs: with-cuda: enable with-rocm: enable with-cpu: enable + build: if: ${{ github.repository_owner == 'pytorch' }} needs: generate-matrix diff --git a/fbgemm_gpu/setup.py b/fbgemm_gpu/setup.py index b45a8696ef..a44b8458c7 100644 --- a/fbgemm_gpu/setup.py +++ b/fbgemm_gpu/setup.py @@ -116,8 +116,12 @@ def package_name(self) -> str: sys.exit(0) elif self.nova_flag() == 0: - # The package name is the same for all build variants in Nova - pass + # In Nova, we are publishing genai packages separately from the main + # fbgemm_gpu package, so if the package variant is genai, we need to + # update the package name accordingly. Otherwise, the package name + # is the same for all other build variants in Nova + if self.args.package_variant == "genai": + pkg_name = "fbgemm_gpu_genai" else: # If running outside of Nova workflow context, append the channel