Packaging pipeline now builds with PythonOp (aka running autograd.Function) (#8652)

This PR disable UTs in training's package pipelines 
for building packages with PythonOp (torch.autograd.Function).

cmake/onnxruntime_pyop.cmake

@@ -6,6 +6,8 @@ add_dependencies(onnxruntime_pyop onnxruntime_graph)
 onnxruntime_add_include_to_target(onnxruntime_pyop onnxruntime_common onnxruntime_graph onnxruntime_framework onnx onnx_proto ${PROTOBUF_LIB} flatbuffers)
 target_include_directories(onnxruntime_pyop PRIVATE ${ONNXRUNTIME_ROOT} ${eigen_INCLUDE_DIRS})
 onnxruntime_add_include_to_target(onnxruntime_pyop Python::Module Python::NumPy)
-target_link_libraries(onnxruntime_pyop PRIVATE Python::Python)
-
-
+if (TARGET Python::Python)
+  target_link_libraries(onnxruntime_pyop PRIVATE Python::Python)
+else()
+  target_link_libraries(onnxruntime_pyop PRIVATE Python::Module)
+endif()

cmake/onnxruntime_session.cmake

@@ -39,3 +39,6 @@ endif()
 if (onnxruntime_ENABLE_TRAINING OR onnxruntime_ENABLE_TRAINING_OPS)
   target_include_directories(onnxruntime_session PRIVATE ${ORTTRAINING_ROOT})
 endif()
+if (onnxruntime_ENABLE_TRAINING_TORCH_INTEROP)
+  onnxruntime_add_include_to_target(onnxruntime_session Python::Module) 
+endif()

cmake/onnxruntime_training.cmake

@@ -71,7 +71,7 @@ if (onnxruntime_BUILD_UNIT_TESTS)
   add_dependencies(onnxruntime_training_runner ${onnxruntime_EXTERNAL_DEPENDENCIES} onnx onnxruntime_providers)
 
   if (onnxruntime_ENABLE_TRAINING_TORCH_INTEROP)
-    target_link_libraries(onnxruntime_training_runner PRIVATE Python::Python)
+    target_link_libraries(onnxruntime_training_runner PRIVATE Python::Module)
   endif()
 
   onnxruntime_add_include_to_target(onnxruntime_training_runner onnxruntime_training onnxruntime_framework onnxruntime_common onnx onnx_proto ${PROTOBUF_LIB} onnxruntime_training flatbuffers)
@@ -127,7 +127,7 @@ if (onnxruntime_BUILD_UNIT_TESTS)
   )
 
   if (onnxruntime_ENABLE_TRAINING_TORCH_INTEROP)
-    list(APPEND ONNXRUNTIME_LIBS Python::Python)
+    list(APPEND ONNXRUNTIME_LIBS Python::Module)
   endif()
 
   list(APPEND ONNXRUNTIME_LIBS

cmake/onnxruntime_unittests.cmake

@@ -693,7 +693,7 @@ if (onnxruntime_USE_ROCM)
   target_include_directories(onnxruntime_test_all PRIVATE  ${onnxruntime_ROCM_HOME}/hipfft/include ${onnxruntime_ROCM_HOME}/include ${onnxruntime_ROCM_HOME}/hiprand/include ${onnxruntime_ROCM_HOME}/rocrand/include ${CMAKE_CURRENT_BINARY_DIR}/amdgpu/onnxruntime ${CMAKE_CURRENT_BINARY_DIR}/amdgpu/orttraining)
 endif()
 if (onnxruntime_ENABLE_TRAINING_TORCH_INTEROP)
-  target_link_libraries(onnxruntime_test_all PRIVATE Python::Python)
+  target_link_libraries(onnxruntime_test_all PRIVATE Python::Module)
 endif()
 if (onnxruntime_BUILD_WEBASSEMBLY)
   set_target_properties(onnxruntime_test_all PROPERTIES LINK_DEPENDS ${TEST_SRC_DIR}/wasm/dump-test-result-in-nodejs.js)

tools/ci_build/build.py

@@ -160,6 +160,8 @@ def parse_arguments():
         "--enable_training", action='store_true', help="Enable training in ORT.")
     parser.add_argument(
         "--enable_training_ops", action='store_true', help="Enable training ops in inference graph.")
+    parser.add_argument(
+        "--enable_training_torch_interop", action='store_true', help="Enable training kernels interop with torch.")
     parser.add_argument(
         "--disable_nccl", action='store_true', help="Disable Nccl.")
     parser.add_argument(
@@ -763,6 +765,7 @@ def generate_build_tree(cmake_path, source_dir, build_dir, cuda_home, cudnn_home
         "-Donnxruntime_ENABLE_NVTX_PROFILE=" + ("ON" if args.enable_nvtx_profile else "OFF"),
         "-Donnxruntime_ENABLE_TRAINING=" + ("ON" if args.enable_training else "OFF"),
         "-Donnxruntime_ENABLE_TRAINING_OPS=" + ("ON" if args.enable_training_ops else "OFF"),
+        "-Donnxruntime_ENABLE_TRAINING_TORCH_INTEROP=" + ("ON" if args.enable_training_torch_interop else "OFF"),
         # Enable advanced computations such as AVX for some traininig related ops.
         "-Donnxruntime_ENABLE_CPU_FP16_OPS=" + ("ON" if args.enable_training else "OFF"),
         "-Donnxruntime_USE_NCCL=" + ("OFF" if args.disable_nccl else "ON"),

tools/ci_build/github/azure-pipelines/orttraining-py-packaging-pipeline-torch190-cuda102.yml

@@ -3,7 +3,7 @@ trigger: none
 stages:
 - template: templates/py-packaging-training-cuda-stage.yml
   parameters:
-    build_py_parameters: --enable_training --update --build
+    build_py_parameters: --enable_training --update --build 
     torch_version: '1.9.0'
     cuda_version: '10.2'
     gcc_version: 8

tools/ci_build/github/azure-pipelines/templates/py-packaging-training-cuda-stage.yml

@@ -128,8 +128,9 @@ stages:
                   --build_wheel \
                   --enable_onnx_tests \
                   ${{ parameters.build_py_parameters }} \
-                  --cmake_extra_defines CMAKE_CUDA_HOST_COMPILER=/opt/rh/devtoolset-$(GccVersion)/root/usr/bin/cc 'CMAKE_CUDA_ARCHITECTURES=${{ parameters.cmake_cuda_architectures }}' \
-                  --use_cuda --cuda_version=$(CudaVersion) --cuda_home=/usr/local/cuda-$(CudaVersion) --cudnn_home=/usr/local/cuda-$(CudaVersion) ;
+                  --cmake_extra_defines CMAKE_CUDA_HOST_COMPILER=/opt/rh/devtoolset-$(GccVersion)/root/usr/bin/cc 'CMAKE_CUDA_ARCHITECTURES=${{ parameters.cmake_cuda_architectures }}' onnxruntime_BUILD_UNIT_TESTS=OFF \
+                  --use_cuda --cuda_version=$(CudaVersion) --cuda_home=/usr/local/cuda-$(CudaVersion) --cudnn_home=/usr/local/cuda-$(CudaVersion) \
+                  --enable_training_torch_interop ;
           workingDirectory: $(Build.SourcesDirectory)
 
       - task: CmdLine@2
@@ -149,7 +150,7 @@ stages:
               --volume /bert_data:/bert_data \
               --volume /hf_models_cache:/hf_models_cache \
               onnxruntimetraininggpubuild \
-                bash -c " $(PythonManylinuxDir)/bin/python3 -m pip install /build/Release/dist/$whlfilename && $(PythonManylinuxDir)/bin/python3 -m onnxruntime.training.ortmodule.torch_cpp_extensions.install && $(PythonManylinuxDir)/bin/python3 /build/Release/launch_test.py --cmd_line_with_args 'python orttraining_ortmodule_tests.py --mnist /mnist --bert_data /bert_data/hf_data/glue_data/CoLA/original/raw --transformers_cache /hf_models_cache/huggingface/transformers' --cwd /build/Release " ;
+                bash -c " $(PythonManylinuxDir)/bin/python3 -m pip install /build/Release/dist/$whlfilename && $(PythonManylinuxDir)/bin/python3 -m onnxruntime.training.ortmodule.torch_cpp_extensions.install " ;
           workingDirectory: $(Build.SourcesDirectory)
 
       - task: CopyFiles@2