Skip to content

Commit

Permalink
[fbgemm_gpu] Enable building under python -m build
Browse files Browse the repository at this point in the history
- Update build scripts to enable building fbgemm_gpu under `python -m build`
instead of `python setup.py bdist_wheel`
  • Loading branch information
q10 committed Apr 11, 2024
1 parent 818a26f commit 071ae38
Show file tree
Hide file tree
Showing 5 changed files with 43 additions and 57 deletions.
75 changes: 23 additions & 52 deletions .github/scripts/fbgemm_gpu_build.bash
Original file line number Diff line number Diff line change
Expand Up @@ -64,7 +64,7 @@ __configure_fbgemm_gpu_build_clang () {
local conda_prefix=$(conda run ${env_prefix} printenv CONDA_PREFIX)
# shellcheck disable=SC2206
build_args+=(
--cxxprefix ${conda_prefix}
--cxxprefix=${conda_prefix}
)
}

Expand Down Expand Up @@ -258,6 +258,11 @@ __configure_fbgemm_gpu_build () {
__configure_fbgemm_gpu_build_clang
fi

# Set verbosity
build_args+=(
--verbose
)

# shellcheck disable=SC2145
echo "[BUILD] FBGEMM_GPU build arguments have been set: ${build_args[@]}"
}
Expand Down Expand Up @@ -443,15 +448,26 @@ build_fbgemm_gpu_package () {
echo "################################################################################"
echo ""

# Distribute Python extensions as wheels on Linux
# Set packaging options
build_args+=(
--package_channel="${fbgemm_release_channel}"
--python-tag="${python_tag}"
--plat-name="${python_plat_name}"
)

# Prepend build options correctly for `python -m build`
# https://build.pypa.io/en/stable/index.html
# https://gregoryszorc.com/blog/2023/10/30/my-user-experience-porting-off-setup.py/
for i in "${!build_args[@]}"; do
build_args[i]="--config-setting=--build-option=${build_args[i]}"
done

# Build the wheel. Invoke using `python -m build`
# https://blog.ganssle.io/articles/2021/10/setup-py-deprecated.html
echo "[BUILD] Building FBGEMM-GPU wheel (VARIANT=${fbgemm_variant}) ..."
# shellcheck disable=SC2086
print_exec conda run --no-capture-output ${env_prefix} \
python setup.py "${run_multicore}" bdist_wheel \
--package_channel="${fbgemm_release_channel}" \
--python-tag="${python_tag}" \
--plat-name="${python_plat_name}" \
--verbose \
python -m build --wheel --no-isolation \
"${build_args[@]}"

# Run checks on the built libraries
Expand Down Expand Up @@ -503,7 +519,6 @@ build_fbgemm_gpu_install () {
# shellcheck disable=SC2086
print_exec conda run --no-capture-output ${env_prefix} \
python setup.py "${run_multicore}" install \
--verbose \
"${build_args[@]}"

# Run checks on the built libraries
Expand All @@ -519,47 +534,3 @@ build_fbgemm_gpu_install () {

echo "[BUILD] FBGEMM-GPU build + install completed"
}

build_fbgemm_gpu_develop () {
env_name="$1"
fbgemm_variant="$2"
fbgemm_variant_targets="$3"
if [ "$fbgemm_variant" == "" ]; then
echo "Usage: ${FUNCNAME[0]} ENV_NAME VARIANT [TARGETS]"
echo "Example(s):"
echo " ${FUNCNAME[0]} build_env cpu # CPU-only variant"
echo " ${FUNCNAME[0]} build_env cuda # CUDA variant for default target(s)"
echo " ${FUNCNAME[0]} build_env cuda '7.0;8.0' # CUDA variant for custom target(s)"
echo " ${FUNCNAME[0]} build_env rocm # ROCm variant for default target(s)"
echo " ${FUNCNAME[0]} build_env rocm 'gfx906;gfx908;gfx90a' # ROCm variant for custom target(s)"
return 1
fi

# shellcheck disable=SC2155
local env_prefix=$(env_name_or_prefix "${env_name}")

# Set up and configure the build
__build_fbgemm_gpu_common_pre_steps || return 1
__configure_fbgemm_gpu_build "${fbgemm_variant}" "${fbgemm_variant_targets}" || return 1

echo "################################################################################"
echo "# Build + Install FBGEMM-GPU Package (Develop)"
echo "#"
echo "# [$(date --utc +%FT%T.%3NZ)] + ${FUNCNAME[0]} ${*}"
echo "################################################################################"
echo ""

# Parallelism may need to be limited to prevent the build from being
# canceled for going over ulimits
echo "[BUILD] Building (develop) FBGEMM-GPU (VARIANT=${fbgemm_variant}) ..."
# shellcheck disable=SC2086
print_exec conda run --no-capture-output ${env_prefix} \
python setup.py "${run_multicore}" build develop \
--verbose \
"${build_args[@]}"

# Run checks on the built libraries
(run_fbgemm_gpu_postbuild_checks "${fbgemm_variant}") || return 1

echo "[BUILD] FBGEMM-GPU build + develop completed"
}
1 change: 1 addition & 0 deletions .github/scripts/utils_build.bash
Original file line number Diff line number Diff line change
Expand Up @@ -246,6 +246,7 @@ install_build_tools () {
# shellcheck disable=SC2086
(exec_with_retries 3 conda install ${env_prefix} -c conda-forge -y \
bazel \
build \
click \
cmake \
hypothesis \
Expand Down
5 changes: 4 additions & 1 deletion fbgemm_gpu/CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -223,7 +223,6 @@ endif()
file(GLOB_RECURSE asmjit_sources
"${CMAKE_CURRENT_SOURCE_DIR}/../third_party/asmjit/src/asmjit/*/*.cpp")


################################################################################
# Optimizer Group Definitions
################################################################################
Expand Down Expand Up @@ -803,6 +802,10 @@ if(NVML_LIB_PATH)
target_link_libraries(fbgemm_gpu_py ${NVML_LIB_PATH})
endif()

# Silence warnings in asmjit
target_compile_options(fbgemm_gpu_py PRIVATE -Wno-deprecated-anon-enum-enum-conversion)
target_compile_options(fbgemm_gpu_py PRIVATE -Wno-deprecated-declarations)


################################################################################
# FBGEMM_GPU Install
Expand Down
1 change: 1 addition & 0 deletions fbgemm_gpu/requirements.txt
Original file line number Diff line number Diff line change
Expand Up @@ -17,5 +17,6 @@ mpmath==1.3.0
ninja
numpy
scikit-build
setuptools
setuptools_git_versioning
tabulate
18 changes: 14 additions & 4 deletions fbgemm_gpu/setup.py
Original file line number Diff line number Diff line change
Expand Up @@ -118,7 +118,17 @@ def package_name(self) -> str:
def variant_version(self) -> str:
pkg_vver: str = ""

if self.nova_flag() is None:
if "egg_info" in self.other_args:
# If build is invoked through `python -m build` instead of
# `python setup.py`, this script is invoked twice, once as
# `setup.py egg_info`, and once as `setup.py bdist_wheel`.
# Ignore determining the variant_version for the first case.
print(
"[SETUP.PY] Script was invoked as `setup.py egg_info`, ignoring variant_version"
)
return pkg_vver

elif self.nova_flag() is None:
# If not running in a Nova workflow, then use the
# `fbgemm_gpu-<variant>` naming convention for the package, since
# PyPI does not accept version+xx in the naming convention.
Expand All @@ -134,7 +144,7 @@ def variant_version(self) -> str:
pkg_vver = f"+cu{cuda_version[0]}{cuda_version[1]}"
else:
sys.exit(
"[SETUP.PY] Installed PyTorch variant is not CUDA; cannot determine the CUDA version!"
"[SETUP.PY] The installed PyTorch variant is not CUDA; cannot determine the CUDA version!"
)

elif self.args.package_variant == "rocm":
Expand All @@ -143,7 +153,7 @@ def variant_version(self) -> str:
pkg_vver = f"+rocm{rocm_version[0]}.{rocm_version[1]}"
else:
sys.exit(
"[SETUP.PY] Installed PyTorch variant is not ROCm; cannot determine the ROCm version!"
"[SETUP.PY] The installed PyTorch variant is not ROCm; cannot determine the ROCm version!"
)

else:
Expand Down Expand Up @@ -485,5 +495,5 @@ def main(argv: List[str]) -> None:


if __name__ == "__main__":
print(f"[SETUP.PY] {sys.argv}")
print(f"[SETUP.PY] ARGV: {sys.argv}")
main(sys.argv[1:])

0 comments on commit 071ae38

Please sign in to comment.