Enable registering external custom op schemas on Linux (#8889)

* Use manylinux instead of Ubuntu to run external custom ops build pipeline.

orttraining/orttraining/test/external_custom_ops/setup.py

@@ -1,6 +1,6 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
-
+import sys
 import os
 import subprocess
 from setuptools import setup, Extension
@@ -9,11 +9,14 @@ from subprocess import CalledProcessError
 import pybind11
 import onnx
 import onnxruntime
+
+
 class CMakeExtension(Extension):
     def __init__(self, name, sourcedir=""):
         Extension.__init__(self, name, sources=[])
         self.sourcedir = os.path.abspath(sourcedir)
 
+
 class CMakeBuild(build_ext):
     def build_extension(self, ext):
         extdir = os.path.abspath(os.path.dirname(self.get_ext_fullpath(ext.name)))
@@ -22,23 +25,27 @@ class CMakeBuild(build_ext):
             os.makedirs(self.build_temp)
         subprocess.check_call(
             ["cmake",
+             "-DPYBIND11_PYTHON_VERSION={}.{}.{}".format(sys.version_info.major, sys.version_info.minor,
+                                                         sys.version_info.micro),
              "-Dpybind11_DIR={}".format(pybind11.get_cmake_dir()),
-             "-DONNX_INCLUDE={}".format(os.path.dirname(os.path.dirname(onnx.__file__))),
+             "-DONNX_INCLUDE={}".format(
+                 os.path.dirname(os.path.dirname(onnx.__file__))),
              "-DONNXRUNTIME_EXTERNAL_INCLUDE={}".format(os.path.join(os.path.join(os.path.dirname(onnxruntime.__file__),
-                                                                                               "external"), "include")),
+                                                                                  "external"), "include")),
              "-DCMAKE_LIBRARY_OUTPUT_DIRECTORY={}".format(extdir),
              ext.sourcedir], cwd=self.build_temp)
         subprocess.check_call(
             ["cmake", "--build", "."], cwd=self.build_temp
         )
 
+
 setup(name='orttraining_external_custom_ops',
-    version='0.1',
-    author='',
-    author_email='',
-    description='External custom ops example',
-    long_description='',
-    ext_modules=[CMakeExtension('orttrainng_external_custom_ops')],
-    cmdclass=dict(build_ext=CMakeBuild),
-    zip_safe=False
-)
+      version='0.1',
+      author='',
+      author_email='',
+      description='External custom ops example',
+      long_description='',
+      ext_modules=[CMakeExtension('orttrainng_external_custom_ops')],
+      cmdclass=dict(build_ext=CMakeBuild),
+      zip_safe=False
+      )

tools/ci_build/build.py

@@ -588,11 +588,6 @@ def is_ubuntu_1604():
     return dist == 'Ubuntu' and ver.startswith('16.04')
 
 
-def is_ubuntu():
-    dist, _ = get_linux_distro()
-    return dist == 'Ubuntu'
-
-
 def get_config_build_dir(build_dir, config):
     # build directory per configuration
     return os.path.join(build_dir, config)
@@ -2215,8 +2210,8 @@ def main():
                         args.eager_customop_header, args.eager_customop_module, True)
 
             gen_ort_ops()
-        if args.enable_external_custom_op_schemas and not is_ubuntu():
-            raise BuildError("Registering external custom op schemas is only supported on Ubuntu.")
+        if args.enable_external_custom_op_schemas and not is_linux():
+            raise BuildError("Registering external custom op schemas is only supported on Linux.")
 
         generate_build_tree(
             cmake_path, source_dir, build_dir, cuda_home, cudnn_home, rocm_home, mpi_home, nccl_home,

tools/ci_build/github/azure-pipelines/orttraining-linux-external-custom-ops.yml

@@ -9,32 +9,38 @@ jobs:
     clean: true
     submodules: recursive
 
+  - template: templates/get-docker-image-steps.yml
+    parameters:
+      Dockerfile: tools/ci_build/github/linux/docker/Dockerfile.manylinux2014_cpu
+      Context: tools/ci_build/github/linux/docker
+      DockerBuildArgs: "--build-arg BUILD_UID=$( id -u )"
+      Repository: onnxruntimecpubuild
+
   - task: CmdLine@2
-    displayName: 'Docker build'
+    displayName: 'Create and start docker container'
     inputs:
       script: |
-        docker build --pull -t onnxruntime-ubuntu \
-        --build-arg BUILD_USER=onnxruntimedev \
-        --build-arg BUILD_UID=$(id -u) \
-        --build-arg OS_VERSION=20.04 \
-        --build-arg PYTHON_VERSION=3.8 -f Dockerfile.ubuntu .
-      workingDirectory: $(Build.SourcesDirectory)/tools/ci_build/github/linux/docker
+        docker run -it -d \
+          --name external_custom_ops_container \
+          --volume /data/onnx:/data/onnx:ro \
+          --volume $(Build.SourcesDirectory):/onnxruntime_src \
+          --volume $(Build.BinariesDirectory):/build \
+          --volume $HOME/.onnx:/home/onnxruntimedev/.onnx \
+          onnxruntimecpubuild \
+          /bin/bash
+      workingDirectory: $(Build.SourcesDirectory)
 
   - task: CmdLine@2
     displayName: 'ONNXRuntime build'
     inputs:
       script: |
         mkdir -p $HOME/.onnx
-        docker run --rm \
-          --volume /data/onnx:/data/onnx:ro \
-          --volume $(Build.SourcesDirectory):/onnxruntime_src \
-          --volume $(Build.BinariesDirectory):/build \
-          --volume $HOME/.onnx:/home/onnxruntimedev/.onnx \
+        docker exec -t \
           -e ALLOW_RELEASED_ONNX_OPSET_ONLY=0 \
           -e NIGHTLY_BUILD \
           -e BUILD_BUILDNUMBER \
-          onnxruntime-ubuntu \
-            /usr/bin/python3 /onnxruntime_src/tools/ci_build/build.py \
+          external_custom_ops_container \
+            /opt/python/cp38-cp38/bin/python3 /onnxruntime_src/tools/ci_build/build.py \
               --build_dir /build \
               --config Debug Release \
               --parallel \
@@ -49,43 +55,73 @@ jobs:
     inputs:
       script: |
          rm -rf $(Build.BinariesDirectory)/Release/onnxruntime $(Build.BinariesDirectory)/Release/pybind11
-         python3 -m pip install $(Build.BinariesDirectory)/Release/dist/*.whl
+         docker exec -t external_custom_ops_container \
+            /bin/bash -c "/opt/python/cp38-cp38/bin/python3 -m pip install /build/Release/dist/*.whl" 
+      workingDirectory: $(Build.BinariesDirectory)
+
+  - task: CmdLine@2
+    displayName: 'Install Pybind11'
+    inputs:
+      script: |
+         docker exec -t \
+          -w  /build/Debug/external_custom_ops \
+          external_custom_ops_container \
+          /opt/python/cp38-cp38/bin/python3 -m pip install pybind11
+      workingDirectory: $(Build.BinariesDirectory)
 
   - task: CmdLine@2
     displayName: 'Build and Install custom ops python package'
     inputs:
       script: |
-         python3 -m pip install pybind11
-         python3 -m pip install .
-      workingDirectory: $(Build.BinariesDirectory)/Release/external_custom_ops
+         docker exec -t \
+          -w  /build/Debug/external_custom_ops \
+          external_custom_ops_container \
+          /opt/python/cp38-cp38/bin/python3 -m pip install .
+      workingDirectory: $(Build.BinariesDirectory)
 
   - task: CmdLine@2
     displayName: 'Test using external custom op module'
     inputs:
       script: |
-         python3 test.py
-      workingDirectory: $(Build.BinariesDirectory)/Release/external_custom_ops
+         docker exec -t \
+          -w  /build/Debug/external_custom_ops \
+          external_custom_ops_container \
+          /opt/python/cp38-cp38/bin/python3 test.py
+      workingDirectory: $(Build.BinariesDirectory)
+
 
   - task: CmdLine@2
     displayName: 'Install Debug python package'
     inputs:
       script: |
          rm -rf $(Build.BinariesDirectory)/Debug/onnxruntime $(Build.BinariesDirectory)/Debug/pybind11
-         python3 -m pip uninstall -y onnxruntime-traning
-         python3 -m pip install $(Build.BinariesDirectory)/Debug/dist/*.whl
+         docker exec -t external_custom_ops_container \
+            /bin/bash -c "/opt/python/cp38-cp38/bin/python3 -m pip install /build/Debug/dist/*.whl" 
+      workingDirectory: $(Build.BinariesDirectory)
 
   - task: CmdLine@2
     displayName: 'Build and Install custom ops python package'
     inputs:
       script: |
-         python3 -m pip uninstall -y orttraining_external_custom_ops
-         python3 -m pip install .
-      workingDirectory: $(Build.BinariesDirectory)/Debug/external_custom_ops
+         docker exec -t \
+          -w  /build/Debug/external_custom_ops \
+          external_custom_ops_container \
+          /opt/python/cp38-cp38/bin/python3 -m pip install .
+      workingDirectory: $(Build.BinariesDirectory)
 
   - task: CmdLine@2
     displayName: 'Test using external custom op module'
     inputs:
       script: |
-         python3 test.py
-      workingDirectory: $(Build.BinariesDirectory)/Debug/external_custom_ops
+         docker exec -t \
+          -w  /build/Debug/external_custom_ops \
+          external_custom_ops_container \
+          /opt/python/cp38-cp38/bin/python3 test.py
+      workingDirectory: $(Build.BinariesDirectory)
 
+  - task: CmdLine@2
+    displayName: 'Stop and remove docker container'
+    inputs:
+      script: |
+         docker stop external_custom_ops_container
+         docker rm external_custom_ops_container