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,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')
             )

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 \