add environment variable to control default training package's local version

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 = {}

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,8 @@ 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 +2111,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')
  )

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 \