From 4ed5a5b2a80a938d68193bb5a4c99f40a7cee93b Mon Sep 17 00:00:00 2001 From: Baiju Meswani Date: Fri, 16 Sep 2022 22:51:18 -0700 Subject: [PATCH] Disable local versions based on environment variable (#12997) --- setup.py | 30 +++++++++++++++++++----------- 1 file changed, 19 insertions(+), 11 deletions(-) diff --git a/setup.py b/setup.py index 13808a0190..d194f03658 100644 --- a/setup.py +++ b/setup.py @@ -518,17 +518,25 @@ if enable_training: # To support the package consisting of both openvino and training modules part of it package_name = "onnxruntime-training" - # we want put default training packages to pypi. pypi does not accept package with a local version. - if not default_training_package_device or nightly_build: - if cuda_version: - # removing '.' to make Cuda version number in the same form as Pytorch. - local_version = "+cu" + cuda_version.replace(".", "") - elif rocm_version: - # removing '.' to make Rocm version number in the same form as Pytorch. - local_version = "+rocm" + rocm_version.replace(".", "") - else: - # cpu version for documentation - local_version = "+cpu" + disable_local_version = environ.get("ORT_DISABLE_PYTHON_PACKAGE_LOCAL_VERSION", "0") + disable_local_version = ( + disable_local_version == "1" + or disable_local_version.lower() == "true" + or disable_local_version.lower() == "yes" + ) + # local version should be disabled for internal feeds. + if not disable_local_version: + # we want put default training packages to pypi. pypi does not accept package with a local version. + if not default_training_package_device or nightly_build: + if cuda_version: + # removing '.' to make Cuda version number in the same form as Pytorch. + local_version = "+cu" + cuda_version.replace(".", "") + elif rocm_version: + # removing '.' to make Rocm version number in the same form as Pytorch. + local_version = "+rocm" + rocm_version.replace(".", "") + else: + # cpu version for documentation + local_version = "+cpu" if package_name == "onnxruntime-nuphar": packages += ["onnxruntime.nuphar"]