From 68a402e739abd596e6fa8d05a3e185626aa1b669 Mon Sep 17 00:00:00 2001
From: Baiju Meswani <bmeswani@microsoft.com>
Date: Thu, 2 Feb 2023 11:32:41 -0800
Subject: [PATCH] Add support for python 3.10 for onnxruntime-training cuda and
 cpu (#14100)

---
 .../orttraining-py-packaging-pipeline-cpu.yml              | 2 ++
 .../templates/py-packaging-training-cuda-stage.yml         | 7 +++++++
 .../docker/Dockerfile.manylinux2014_training_cuda11_6      | 6 +++++-
 tools/ci_build/github/linux/docker/scripts/install_rust.sh | 5 +++++
 4 files changed, 19 insertions(+), 1 deletion(-)
 create mode 100755 tools/ci_build/github/linux/docker/scripts/install_rust.sh

diff --git a/tools/ci_build/github/azure-pipelines/orttraining-py-packaging-pipeline-cpu.yml b/tools/ci_build/github/azure-pipelines/orttraining-py-packaging-pipeline-cpu.yml
index e3dcdd16ff..b424028260 100644
--- a/tools/ci_build/github/azure-pipelines/orttraining-py-packaging-pipeline-cpu.yml
+++ b/tools/ci_build/github/azure-pipelines/orttraining-py-packaging-pipeline-cpu.yml
@@ -35,6 +35,8 @@ stages:
             PythonVersion: '3.8'
           Python39:
             PythonVersion: '3.9'
+          Python310:
+            PythonVersion: '3.10'
 
       steps:
       - task: CmdLine@2
diff --git a/tools/ci_build/github/azure-pipelines/templates/py-packaging-training-cuda-stage.yml b/tools/ci_build/github/azure-pipelines/templates/py-packaging-training-cuda-stage.yml
index c32192bc05..ae940c839c 100644
--- a/tools/ci_build/github/azure-pipelines/templates/py-packaging-training-cuda-stage.yml
+++ b/tools/ci_build/github/azure-pipelines/templates/py-packaging-training-cuda-stage.yml
@@ -100,6 +100,13 @@ stages:
             CudaVersion: ${{ parameters.cuda_version }}
             GccVersion: ${{ parameters.gcc_version }}
             UploadWheel: ${{ parameters.upload_wheel }}
+          Python310:
+            PythonVersion: '3.10'
+            TorchVersion: ${{ parameters.torch_version }}
+            OpsetVersion: ${{ parameters.opset_version }}
+            CudaVersion: ${{ parameters.cuda_version }}
+            GccVersion: ${{ parameters.gcc_version }}
+            UploadWheel: ${{ parameters.upload_wheel }}
 
       steps:
       - task: CmdLine@2
diff --git a/tools/ci_build/github/linux/docker/Dockerfile.manylinux2014_training_cuda11_6 b/tools/ci_build/github/linux/docker/Dockerfile.manylinux2014_training_cuda11_6
index fc7a6bc827..15e946cc97 100644
--- a/tools/ci_build/github/linux/docker/Dockerfile.manylinux2014_training_cuda11_6
+++ b/tools/ci_build/github/linux/docker/Dockerfile.manylinux2014_training_cuda11_6
@@ -174,7 +174,11 @@ ADD scripts /tmp/scripts
 RUN cd /tmp/scripts && \
     /tmp/scripts/manylinux/install_centos.sh && \
     /tmp/scripts/install_os_deps.sh -d gpu $INSTALL_DEPS_EXTRA_ARGS && \
-    /tmp/scripts/install_ninja.sh && \
+    /tmp/scripts/install_rust.sh
+
+ENV PATH="/root/.cargo/bin/:$PATH"
+
+RUN /tmp/scripts/install_ninja.sh && \
     /tmp/scripts/install_python_deps.sh -d gpu -v 11.6 -p $PYTHON_VERSION -h $TORCH_VERSION $INSTALL_DEPS_EXTRA_ARGS && \
     rm -rf /tmp/scripts
 
diff --git a/tools/ci_build/github/linux/docker/scripts/install_rust.sh b/tools/ci_build/github/linux/docker/scripts/install_rust.sh
new file mode 100755
index 0000000000..03b219ea39
--- /dev/null
+++ b/tools/ci_build/github/linux/docker/scripts/install_rust.sh
@@ -0,0 +1,5 @@
+#!/bin/bash
+set -e
+
+echo "Installing rust from https://sh.rustup.rs non interatively."
+curl https://sh.rustup.rs -sSf | sh -s -- -y