From 1595e2abf1f0ec1c0c4094bba9bdbd2750c34980 Mon Sep 17 00:00:00 2001 From: Benson Ma Date: Wed, 27 Sep 2023 21:43:50 -0700 Subject: [PATCH] Fix package variant versioning (#2045) Summary: - Fix package variant versioning for ROCm Nova builds Pull Request resolved: https://github.com/pytorch/FBGEMM/pull/2045 Reviewed By: sryap Differential Revision: D49716498 Pulled By: q10 fbshipit-source-id: ff2e9f98c2acfe79fee2fb7af703e8c21587164d --- .github/scripts/fbgemm_gpu_build.bash | 9 +- fbgemm_gpu/docs/BuildInstructions.md | 9 +- fbgemm_gpu/setup.py | 177 ++++++++++++++++---------- 3 files changed, 126 insertions(+), 69 deletions(-) diff --git a/.github/scripts/fbgemm_gpu_build.bash b/.github/scripts/fbgemm_gpu_build.bash index b783651d8..269641cef 100644 --- a/.github/scripts/fbgemm_gpu_build.bash +++ b/.github/scripts/fbgemm_gpu_build.bash @@ -58,7 +58,9 @@ prepare_fbgemm_gpu_build () { __configure_fbgemm_gpu_build_cpu () { # Update the package name and build args depending on if CUDA is specified echo "[BUILD] Setting CPU-only build args ..." - build_args=(--cpu_only) + build_args=( + --package_variant=cpu + ) } __configure_fbgemm_gpu_build_rocm () { @@ -88,7 +90,9 @@ __configure_fbgemm_gpu_build_rocm () { print_exec conda env config vars set ${env_prefix} PYTORCH_ROCM_ARCH="${arch_list}" echo "[BUILD] Setting ROCm build args ..." - build_args=() + build_args=( + --package_variant=rocm + ) } __configure_fbgemm_gpu_build_cuda () { @@ -133,6 +137,7 @@ __configure_fbgemm_gpu_build_cuda () { # shellcheck disable=SC2155,SC2086 local nvml_lib_path=$(conda run --no-capture-output ${env_prefix} printenv NVML_LIB_PATH) build_args=( + --package_variant=cuda --nvml_lib_path="${nvml_lib_path}" -DTORCH_CUDA_ARCH_LIST="'${arch_list}'" ) diff --git a/fbgemm_gpu/docs/BuildInstructions.md b/fbgemm_gpu/docs/BuildInstructions.md index e67c5a60f..04090a1b7 100644 --- a/fbgemm_gpu/docs/BuildInstructions.md +++ b/fbgemm_gpu/docs/BuildInstructions.md @@ -354,6 +354,7 @@ unset TORCH_CUDA_ARCH_LIST # Build the wheel artifact only python setup.py bdist_wheel \ --package_name="${package_name}" \ + --package_variant=cuda \ --python-tag="${python_tag}" \ --plat-name="manylinux1_${ARCH}" \ --nvml_lib_path=${NVML_LIB_PATH} \ @@ -361,6 +362,7 @@ python setup.py bdist_wheel \ # Build and install the library into the Conda environment python setup.py install \ + --package_variant=cuda \ --nvml_lib_path=${NVML_LIB_PATH} \ -DTORCH_CUDA_ARCH_LIST="${cuda_arch_list}" ``` @@ -384,11 +386,12 @@ package_name=fbgemm_gpu_rocm # Build the wheel artifact only python setup.py bdist_wheel \ --package_name="${package_name}" \ + --package_variant=rocm \ --python-tag="${python_tag}" \ --plat-name="manylinux1_${ARCH}" # Build and install the library into the Conda environment -python setup.py install develop +python setup.py install develop --package_variant=rocm ``` ### CPU-Only Build @@ -406,12 +409,12 @@ package_name=fbgemm_gpu_cpu # Build the wheel artifact only python setup.py bdist_wheel \ --package_name="${package_name}" \ + --package_variant=cpu \ --python-tag="${python_tag}" \ --plat-name="manylinux1_${ARCH}" \ - --cpu_only # Build and install the library into the Conda environment -python setup.py install --cpu_only +python setup.py install --package_variant=cpu ``` ### Post-Build Checks (For Developers) diff --git a/fbgemm_gpu/setup.py b/fbgemm_gpu/setup.py index dfe98d384..161890a7f 100644 --- a/fbgemm_gpu/setup.py +++ b/fbgemm_gpu/setup.py @@ -24,47 +24,20 @@ from tabulate import tabulate -def generate_package_version(package_name: str, version_variant: str): - print("[SETUP.PY] Generating the package version ...") - - if "nightly" in package_name: - # Use date stamp for nightly versions - print("[SETUP.PY] Package is for NIGHTLY; using timestamp for the versioning") - today = date.today() - version = f"{today.year}.{today.month}.{today.day}" - - elif "test" in package_name: - # Use date stamp for nightly versions - print("[SETUP.PY] Package is for TEST: using random number for the versioning") - version = (f"0.0.{random.randint(0, 1000)}",) - - else: - # Use git tag / branch / commit info to generate a PEP-440-compliant version string - print("[SETUP.PY] Package is for RELEASE: using git info for the versioning") - print( - f"[SETUP.PY] TAG: {gitversion.get_tag()}, BRANCH: {gitversion.get_branch()}, SHA: {gitversion.get_sha()}" - ) - # Remove the local version identifier, if any (e.g. 0.4.0rc0.post0+git.6a63116c.dirty => 0.4.0rc0.post0) - # Then remove post0 (keep postN for N > 0) (e.g. 0.4.0rc0.post0 => 0.4.0rc0) - version = re.sub(".post0$", "", gitversion.version_from_git().split("+")[0]) - version = str(version) + version_variant - print(f"[SETUP.PY] Setting the package version: {version}") - return version - - def parse_args(argv: List[str]) -> argparse.Namespace: parser = argparse.ArgumentParser(description="fbgemm_gpu setup") parser.add_argument( - "--cpu_only", - dest="cpu_only", - action="store_true", - help="build for cpu-only (no GPU support)", + "--package_variant", + type=str, + choices=["cpu", "cuda", "rocm"], + default="cpu", + help="The FBGEMM_GPU variant to build.", ) parser.add_argument( "--package_name", type=str, default="fbgemm_gpu", - help="the name of this output wheel", + help="The candidate name of the output wheel.", ) parser.add_argument( "--nvml_lib_path", @@ -171,7 +144,7 @@ def _get_cxx11_abi(): os.environ["CMAKE_BUILD_PARALLEL_LEVEL"] = str(os.cpu_count() // 2) cmake_args = [f"-DCMAKE_PREFIX_PATH={torch_root}", _get_cxx11_abi()] - if args.cpu_only: + if args.package_variant == "cpu": cmake_args.append("-DFBGEMM_CPU_ONLY=ON") if args.nvml_lib_path: cmake_args.append(f"-DNVML_LIB_PATH={args.nvml_lib_path}") @@ -181,6 +154,100 @@ def _get_cxx11_abi(): class FbgemmGpuInstaller(PipInstall): """FBGEMM_GPU PIP Installer""" + @classmethod + def extract_package_name(cls, args_package_name: str) -> str: + package_name: str = "" + + if "BUILD_FROM_NOVA" in os.environ: + nova_flag = os.getenv("BUILD_FROM_NOVA") + print(f"[SETUP.PY] BUILD_FROM_NOVA={nova_flag}") + + # The package name is the same for all build variants in Nova + package_name = "fbgemm_gpu" + + if str(nova_flag) != "0": + # Skip build clean and build wheel steps in Nova workflow since + # they are done in pre-script + print("[SETUP.PY] Build from Nova detected... exiting.") + sys.exit(0) + + else: + package_name = args_package_name + + print(f"[SETUP.PY] Extracted the package name: '{package_name}'") + return package_name + + @classmethod + def extract_variant_version(cls, variant: str) -> str: + variant_version: str = "" + + if variant == "cpu": + variant_version = "+cpu" + elif variant == "cuda": + set_cuda_environment_variables() + if torch.version.cuda is not None: + cuda_version = torch.version.cuda.split(".") + variant_version = f"+cu{cuda_version[0]}{cuda_version[1]}" + else: + sys.exit( + "[SETUP.PY] Installed PyTorch variant is not CUDA; cannot determine the CUDA version!" + ) + elif variant == "rocm": + if torch.version.hip is not None: + rocm_version = torch.version.hip.split(".") + variant_version = f"+rocm{rocm_version[0]}.{rocm_version[1]}" + else: + sys.exit( + "[SETUP.PY] Installed PyTorch variant is not ROCm; cannot determine the ROCm version!" + ) + else: + sys.exit( + f"[SETUP.PY] Unrecognized build variant variant '{variant}'; cannot proceed with FBGEMM_GPU build!" + ) + + if "BUILD_FROM_NOVA" not in os.environ: + # If not building from Nova, use the fbgemm_gpu- + # PyPI does not accept version+xx in the name convention. + print("[SETUP.PY] Not building FBGEMM_GPU from Nova.") + variant_version = "" + + print(f"[SETUP.PY] Extracted the package variant+version: '{variant_version}'") + return variant_version + + @classmethod + def generate_package_version(cls, package_name: str, variant_version: str): + print("[SETUP.PY] Generating the package version ...") + + if "nightly" in package_name: + # Use date stamp for nightly versions + print( + "[SETUP.PY] Package is for NIGHTLY; using timestamp for the versioning" + ) + today = date.today() + version = f"{today.year}.{today.month}.{today.day}" + + elif "test" in package_name and "BUILD_FROM_NOVA" not in os.environ: + # Use random numbering for test versions + print( + "[SETUP.PY] Package is for TEST: using random number for the versioning" + ) + version = (f"0.0.{random.randint(0, 1000)}",) + + else: + # Use git tag / branch / commit info to generate a PEP-440-compliant version string + print( + "[SETUP.PY] Package is for RELEASE: using git info for the versioning" + ) + print( + f"[SETUP.PY] TAG: {gitversion.get_tag()}, BRANCH: {gitversion.get_branch()}, SHA: {gitversion.get_sha()}" + ) + # Remove the local version identifier, if any (e.g. 0.4.0rc0.post0+git.6a63116c.dirty => 0.4.0rc0.post0) + # Then remove post0 (keep postN for N > 0) (e.g. 0.4.0rc0.post0 => 0.4.0rc0) + version = re.sub(".post0$", "", gitversion.version_from_git().split("+")[0]) + version = str(version) + variant_version + print(f"[SETUP.PY] Setting the full package version string: {version}") + return version + @classmethod def generate_version_file(cls, package_version: str) -> None: with open("fbgemm_gpu/_fbgemm_gpu_version.py", "w") as file: @@ -262,41 +329,23 @@ def run(self): def main(argv: List[str]) -> None: # Handle command line args before passing to main setup() method. args, unknown = parse_args(argv) - print("Parsed Arguments: ", args) + print(f"[SETUP.PY] Parsed Arguments: {args}") if len(unknown) != 0 and (len(unknown) != 1 or unknown[0] != "clean"): - print("Unknown Arguments: ", unknown) + print(f"[SETUP.PY] Unknown Arguments: {unknown}") - if args.cpu_only: - version_variant = "+cpu" - else: - set_cuda_environment_variables() - if torch.version.cuda is not None: - cuda_version = torch.version.cuda.split(".") - version_variant = "+cu" + str(cuda_version[0]) + str(cuda_version[1]) - else: - # rocm or other gpus - to be specified if we offcially support them - version_variant = "" - - # Skip Nova build steps since it will be done in pre-script - if "BUILD_FROM_NOVA" in os.environ: - build_from_nova = os.getenv("BUILD_FROM_NOVA") - print("build_from_nova", build_from_nova) - # Package name is the same for all variants in Nova - package_name = "fbgemm_gpu" - if str(build_from_nova) != "0": - # Skip build clean and build wheel steps in Nova workflow since they are done in pre-script - print("Build from Nova detected... exiting") - sys.exit(0) - else: - # If not building from Nova, use the fbgemm_gpu- - # PyPi does not accept version+xx in the name convention. - version_variant = "" - package_name = args.package_name # Repair command line args for setup. sys.argv = [sys.argv[0]] + unknown - # Determine the package version - package_version = generate_package_version(args.package_name, version_variant) + # Extract the package name + package_name = FbgemmGpuInstaller.extract_package_name(args.package_name) + + # Extract the variant version, e.g. cpu, cu121, rocm5.6 + variant_version = FbgemmGpuInstaller.extract_variant_version(args.package_variant) + + # Generate the full package version string + package_version = FbgemmGpuInstaller.generate_package_version( + package_name, variant_version + ) # Generate the version file FbgemmGpuInstaller.generate_version_file(package_version)