From f9e7ba74b314cfe2764569312f767139d07ead91 Mon Sep 17 00:00:00 2001 From: Supadchaya Puangpontip Date: Fri, 15 Sep 2023 14:41:59 -0700 Subject: [PATCH] Change name convention for PyPI uploads (#2023) Summary: Pull Request resolved: https://github.com/pytorch/FBGEMM/pull/2023 PyPI does not accept version to be specified with `+` (e.g., `+cpu` or `+cu118`). Reviewed By: sryap Differential Revision: D49329679 fbshipit-source-id: 1bc0c81ae15fdf157b54b2ca9226d67dacaade8e --- fbgemm_gpu/setup.py | 25 ++++++++++++++----------- 1 file changed, 14 insertions(+), 11 deletions(-) diff --git a/fbgemm_gpu/setup.py b/fbgemm_gpu/setup.py index b0d7b56b3..fba4846ca 100644 --- a/fbgemm_gpu/setup.py +++ b/fbgemm_gpu/setup.py @@ -266,6 +266,17 @@ def main(argv: List[str]) -> None: if len(unknown) != 0 and (len(unknown) != 1 or unknown[0] != "clean"): print("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") @@ -277,18 +288,10 @@ def main(argv: List[str]) -> None: 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 - - 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: - version_variant = "" - # Repair command line args for setup. sys.argv = [sys.argv[0]] + unknown