[ROCm] ONNX Runtime training rocm package for ADO (microsoft#17683)

- we will publish the onnxruntime-training-rocm package on ADO feeds. The onnxruntime-training package will solely be for cuda. - Add new pipeline for onnxruntime-training-rocm ADO feeds https://aiinfra.visualstudio.com/Lotus/_build?definitionId=1278. Only package with latest rocm version is publish to ADO.
kleiti · Mar 22, 2024 · eb7cbd7 · eb7cbd7
1 parent b36efed
commit eb7cbd7
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 11 deletions.
diff --git a/setup.py b/setup.py
@@ -520,6 +520,10 @@ def finalize_options(self):
                 # Training CPU package for ADO feeds is called onnxruntime-training-cpu
                 package_name = "onnxruntime-training-cpu"
 
+            if rocm_version:
+                # Training ROCM package for ADO feeds is called onnxruntime-training-rocm
+                package_name = "onnxruntime-training-rocm"
+
 if package_name == "onnxruntime-tvm":
     packages += ["onnxruntime.providers.tvm"]
 

diff --git a/tools/ci_build/github/azure-pipelines/orttraining-py-packaging-pipeline-rocm.yml b/tools/ci_build/github/azure-pipelines/orttraining-py-packaging-pipeline-rocm.yml
@@ -9,44 +9,52 @@ resources:
     ref: 5eda9aded5462201e6310105728d33016e637ea7
 
 stages:
-- stage: Python_Packaging
+- stage: "Python_Packaging_ROCm57_Release"
   jobs:
   - template: templates/rocm.yml
     parameters:
       PythonVersion: '3.8'
-      RocmVersion: '5.6'
+      RocmVersion: '5.7'
   - template: templates/rocm.yml
     parameters:
       PythonVersion: '3.9'
-      RocmVersion: '5.6'
+      RocmVersion: '5.7'
   - template: templates/rocm.yml
     parameters:
       PythonVersion: '3.10'
-      RocmVersion: '5.6'
+      RocmVersion: '5.7'
+
+- stage: "Python_Packaging_ROCm57_RelWithDebInfo"
+  condition: ne(variables['ORT_DISABLE_PYTHON_PACKAGE_LOCAL_VERSION'], 'true')
+  jobs:
   - template: templates/rocm.yml
     parameters:
       PythonVersion: '3.8'
       RocmVersion: '5.7'
+      BuildConfig: 'RelWithDebInfo'
   - template: templates/rocm.yml
     parameters:
       PythonVersion: '3.9'
       RocmVersion: '5.7'
+      BuildConfig: 'RelWithDebInfo'
   - template: templates/rocm.yml
     parameters:
       PythonVersion: '3.10'
       RocmVersion: '5.7'
+      BuildConfig: 'RelWithDebInfo'
+
+- stage: "Python_Packaging_ROCm56_Release"
+  condition: ne(variables['ORT_DISABLE_PYTHON_PACKAGE_LOCAL_VERSION'], 'true')
+  jobs:
   - template: templates/rocm.yml
     parameters:
       PythonVersion: '3.8'
-      RocmVersion: '5.7'
-      BuildConfig: 'RelWithDebInfo'
+      RocmVersion: '5.6'
   - template: templates/rocm.yml
     parameters:
       PythonVersion: '3.9'
-      RocmVersion: '5.7'
-      BuildConfig: 'RelWithDebInfo'
+      RocmVersion: '5.6'
   - template: templates/rocm.yml
     parameters:
       PythonVersion: '3.10'
-      RocmVersion: '5.7'
-      BuildConfig: 'RelWithDebInfo'
+      RocmVersion: '5.6'
diff --git a/tools/ci_build/github/azure-pipelines/templates/rocm.yml b/tools/ci_build/github/azure-pipelines/templates/rocm.yml
@@ -73,6 +73,7 @@ jobs:
           --entrypoint $(PythonManylinuxDir)/bin/python3 \
           -e NIGHTLY_BUILD \
           -e BUILD_BUILDNUMBER \
+          -e ORT_DISABLE_PYTHON_PACKAGE_LOCAL_VERSION \
           --user onnxruntimedev \
           onnxruntimetrainingrocmbuild-rocm${{ parameters.RocmVersion }} \
             /onnxruntime_src/tools/ci_build/build.py \
@@ -131,7 +132,7 @@ jobs:
       python3 tools/ci_build/upload_python_package_to_azure_storage.py \
           --python_wheel_path ${files[0]} \
           --final_storage
-    condition: and(succeeded(), eq(variables['DRY_RUN'], '0'))
+    condition: and(ne(variables['ORT_DISABLE_PYTHON_PACKAGE_LOCAL_VERSION'], 'true'), and(succeeded(), eq(variables['DRY_RUN'], '0')))
     displayName: 'Upload Rocm wheel to release repository'
 
   - template: component-governance-component-detection-steps.yml