diff --git a/setup.py b/setup.py index 7f22fb9de903..796d91184961 100644 --- a/setup.py +++ b/setup.py @@ -266,6 +266,8 @@ def run(self): local_version = None enable_training = parse_arg_remove_boolean(sys.argv, '--enable_training') +default_training_package_device = parse_arg_remove_boolean(sys.argv, '--default_training_package_device') + if enable_training: packages.extend(['onnxruntime.training', 'onnxruntime.training.amp', @@ -280,13 +282,16 @@ def run(self): # this is needed immediately by pytorch/ort so that the user is able to # install an onnxruntime training package with matching torch cuda version. package_name = 'onnxruntime-training' - if cuda_version: - # removing '.' to make local Cuda version number in the same form as Pytorch. - local_version = '+cu' + cuda_version.replace('.', '') - if rocm_version: - # removing '.' to make Cuda version number in the same form as Pytorch. - rocm_version = rocm_version.replace('.', '') - local_version = '+rocm' + rocm_version + + # we want put default training packages to pypi. pypi does not accept package with a local version. + if not default_training_package_device: + if cuda_version: + # removing '.' to make local Cuda version number in the same form as Pytorch. + local_version = '+cu' + cuda_version.replace('.', '') + if rocm_version: + # removing '.' to make Cuda version number in the same form as Pytorch. + rocm_version = rocm_version.replace('.', '') + local_version = '+rocm' + rocm_version package_data = {} diff --git a/tools/ci_build/build.py b/tools/ci_build/build.py index f601a779e30b..b36d74558328 100644 --- a/tools/ci_build/build.py +++ b/tools/ci_build/build.py @@ -1536,7 +1536,8 @@ def run_nodejs_tests(nodejs_binding_dir): def build_python_wheel( source_dir, build_dir, configs, use_cuda, cuda_version, use_rocm, rocm_version, use_dnnl, use_tensorrt, use_openvino, use_nuphar, use_vitisai, use_acl, use_armnn, use_dml, - wheel_name_suffix, enable_training, nightly_build=False, featurizers_build=False, use_ninja=False): + wheel_name_suffix, enable_training, nightly_build=False, default_training_package_device=False, + featurizers_build=False, use_ninja=False): for config in configs: cwd = get_config_build_dir(build_dir, config) if is_windows() and not use_ninja: @@ -1558,6 +1559,8 @@ def build_python_wheel( # Any combination of the following arguments can be applied if nightly_build: args.append('--nightly_build') + if default_training_package_device: + args.append('--default_training_package_device') if featurizers_build: args.append("--use_featurizers") if wheel_name_suffix: @@ -2087,6 +2090,7 @@ def main(): if args.build: if args.build_wheel: nightly_build = bool(os.getenv('NIGHTLY_BUILD') == '1') + default_training_package_device = bool(os.getenv('DEFAULT_TRAINING_PACKAGE_DEVICE') == '1') build_python_wheel( source_dir, build_dir, @@ -2106,6 +2110,7 @@ def main(): args.wheel_name_suffix, args.enable_training, nightly_build=nightly_build, + default_training_package_device=default_training_package_device, featurizers_build=args.use_featurizers, use_ninja=(args.cmake_generator == 'Ninja') ) diff --git a/tools/ci_build/github/azure-pipelines/templates/py-packaging-stage.yml b/tools/ci_build/github/azure-pipelines/templates/py-packaging-stage.yml index bb121d6b88e5..21e997905d87 100644 --- a/tools/ci_build/github/azure-pipelines/templates/py-packaging-stage.yml +++ b/tools/ci_build/github/azure-pipelines/templates/py-packaging-stage.yml @@ -545,6 +545,7 @@ stages: --volume $HOME/.onnx:/home/onnxruntimedev/.onnx \ -e NVIDIA_VISIBLE_DEVICES=all \ -e NIGHTLY_BUILD \ + -e DEFAULT_TRAINING_PACKAGE_DEVICE \ -e BUILD_BUILDNUMBER \ onnxruntimetraininggpubuild \ $(PythonManylinuxDir)/bin/python3 /onnxruntime_src/tools/ci_build/build.py \