[Py] Add NGRAPH_ONNX_IMPORT_ENABLE env var to setup.py (#1687)

Add support for an environment variable `NGRAPH_ONNX_IMPORT_ENABLE` to setup.py
This variable controls whether the Python API is build with or without support for ONNX import functionality built into core nGraph.

.ci/onnx/jenkins/prepare_environment.sh

@@ -74,6 +74,7 @@ function build_ngraph() {
     fi
     export PYBIND_HEADERS_PATH="${ngraph_directory}/ngraph/python/pybind11"
     export NGRAPH_CPP_BUILD_PATH="${ngraph_directory}/ngraph_dist"
+    export NGRAPH_ONNX_IMPORT_ENABLE="TRUE"
     python3 setup.py bdist_wheel
     # Clean build artifacts
     rm -rf "${ngraph_directory}/ngraph/python/_pyngraph.cpython* ${ngraph_directory}/ngraph/python/build"

.ci/travis/run_test.sh

@@ -17,11 +17,11 @@ if [ ${TRAVIS_OS_NAME} == "linux" ]; then
     fi
 
     if [ ${TASK} == "python2_test" ]; then
-        docker run -w '/root/ngraph/python' test_ngraph tox -e py27
+        docker run -w '/root/ngraph/python' -e NGRAPH_ONNX_IMPORT_ENABLE=TRUE test_ngraph tox -e py27
     fi
 
     if [ ${TASK} == "python3_test" ]; then
-        docker run -w '/root/ngraph/python' test_ngraph tox -e py3
+        docker run -w '/root/ngraph/python' -e NGRAPH_ONNX_IMPORT_ENABLE=TRUE test_ngraph tox -e py3
     fi
 
 fi

.ci/travis/ubuntu/Dockerfile

@@ -47,4 +47,4 @@ ENV PYBIND_HEADERS_PATH /root/ngraph/python/pybind11
 RUN df -k
 
 # Test nGraph and nGraph Python API
-CMD cd /root/ngraph/build && make unit-test-check ; cd /root/ngraph/python && tox
+CMD cd /root/ngraph/build && make unit-test-check ; cd /root/ngraph/python && NGRAPH_ONNX_IMPORT_ENABLE=TRUE tox

python/ngraph/impl/onnx_import/__init__.py

@@ -32,7 +32,7 @@ else:
     flags = sys.getdlopenflags() | ctypes.RTLD_GLOBAL
 sys.setdlopenflags(flags)
 
-from _pyngraph.onnx_import import load_onnx_model
-from _pyngraph.onnx_import import load_onnx_model_file
-from _pyngraph.onnx_import import import_onnx_function
-from _pyngraph.onnx_import import import_onnx_function_file
+from _pyngraph_onnx_import import load_onnx_model
+from _pyngraph_onnx_import import load_onnx_model_file
+from _pyngraph_onnx_import import import_onnx_function
+from _pyngraph_onnx_import import import_onnx_function_file

python/pyngraph/pyngraph.cpp

@@ -22,7 +22,6 @@
 #include "pyngraph/function.hpp"
 #include "pyngraph/node.hpp"
 #include "pyngraph/node_vector.hpp"
-#include "pyngraph/onnx_import/onnx_import.hpp"
 #include "pyngraph/ops/op.hpp"
 #include "pyngraph/ops/regmodule_pyngraph_op.hpp"
 #include "pyngraph/ops/util/regmodule_pyngraph_op_util.hpp"
@@ -50,10 +49,6 @@ PYBIND11_MODULE(_pyngraph, m)
     regmodule_pyngraph_types(m);
     regclass_pyngraph_Function(m);
     regclass_pyngraph_Serializer(m);
-    py::module m_onnx_import = m.def_submodule("onnx_import",
-                                               "Package ngraph.impl.onnx_import "
-                                               "that wraps ngraph::onnx_import");
-    regmodule_pyngraph_onnx_import(m_onnx_import);
     py::module m_op = m.def_submodule("op", "Package ngraph.impl.op that wraps ngraph::op");
     regclass_pyngraph_op_Op(m_op);
     regmodule_pyngraph_op_util(m_op);

python/pyngraph/pyngraph_onnx_import.cpp

+//*****************************************************************************
+// Copyright 2017-2018 Intel Corporation
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//*****************************************************************************
+#include <pybind11/pybind11.h>
+
+#include "pyngraph/onnx_import/onnx_import.hpp"
+
+namespace py = pybind11;
+
+PYBIND11_MODULE(_pyngraph_onnx_import, m)
+{
+    m.doc() = "Package ngraph.impl that wraps ngraph::onnx_import";
+    regmodule_pyngraph_onnx_import(m);
+}

python/setup.py

@@ -126,7 +126,6 @@ sources = ['pyngraph/function.cpp',
            'pyngraph/serializer.cpp',
            'pyngraph/node.cpp',
            'pyngraph/node_vector.cpp',
-           'pyngraph/onnx_import/onnx_import.cpp',
            'pyngraph/shape.cpp',
            'pyngraph/strides.cpp',
            'pyngraph/coordinate_diff.cpp',
@@ -219,6 +218,16 @@ sources = ['pyngraph/function.cpp',
            'pyngraph/types/regmodule_pyngraph_types.cpp',
            ]
 
+package_dir={'ngraph': PYNGRAPH_SOURCE_DIR + "/ngraph",
+             'ngraph.utils': PYNGRAPH_SOURCE_DIR + "/ngraph/utils",
+             'ngraph.impl': PYNGRAPH_SOURCE_DIR + "/ngraph/impl",
+             'ngraph.impl.op': PYNGRAPH_SOURCE_DIR + "/ngraph/impl/op",
+             'ngraph.impl.op.util': PYNGRAPH_SOURCE_DIR + "/ngraph/impl/op/util",
+             'ngraph.impl.passes': PYNGRAPH_SOURCE_DIR + "/ngraph/impl/passes",
+             'ngraph.impl.runtime': PYNGRAPH_SOURCE_DIR + "/ngraph/impl/runtime"}
+packages = ['ngraph', 'ngraph.utils', 'ngraph.impl', 'ngraph.impl.op',
+            'ngraph.impl.op.util', 'ngraph.impl.passes', 'ngraph.impl.runtime']
+
 sources = [PYNGRAPH_SOURCE_DIR + "/" + source for source in sources]
 
 include_dirs = [PYNGRAPH_SOURCE_DIR,
@@ -251,6 +260,27 @@ ext_modules = [Extension(
               ]
 
 
+if(os.environ.get('NGRAPH_ONNX_IMPORT_ENABLE') == 'TRUE'):
+    onnx_sources = ['pyngraph/pyngraph_onnx_import.cpp']
+    onnx_sources.append('pyngraph/onnx_import/onnx_import.cpp')
+    onnx_sources = [PYNGRAPH_SOURCE_DIR + "/" + source for source in onnx_sources]
+
+    package_dir['ngraph.impl.onnx_import'] = PYNGRAPH_SOURCE_DIR + "/ngraph/impl/onnx_import"
+    packages.append('ngraph.impl.onnx_import')
+
+    ext_modules.append(Extension(
+                           '_pyngraph_onnx_import',
+                           sources = onnx_sources,
+                           include_dirs = include_dirs,
+                           define_macros = [("VERSION_INFO", __version__)],
+                           library_dirs = library_dirs,
+                           libraries = libraries,
+                           extra_link_args = extra_link_args,
+                           language = "c++",
+                           )
+                      )
+
+
 class BuildExt(build_ext):
     """
     A custom build extension for adding compiler-specific options.
@@ -295,17 +325,8 @@ setup(
     description='Python wrapper for ngraph',
     long_description='',
     ext_modules=ext_modules,
-    package_dir={'ngraph': PYNGRAPH_SOURCE_DIR + "/ngraph",
-                 'ngraph.utils': PYNGRAPH_SOURCE_DIR + "/ngraph/utils",
-                 'ngraph.impl': PYNGRAPH_SOURCE_DIR + "/ngraph/impl",
-                 'ngraph.impl.onnx_import': PYNGRAPH_SOURCE_DIR + "/ngraph/impl/onnx_import",
-                 'ngraph.impl.op': PYNGRAPH_SOURCE_DIR + "/ngraph/impl/op",
-                 'ngraph.impl.op.util': PYNGRAPH_SOURCE_DIR + "/ngraph/impl/op/util",
-                 'ngraph.impl.passes': PYNGRAPH_SOURCE_DIR + "/ngraph/impl/passes",
-                 'ngraph.impl.runtime': PYNGRAPH_SOURCE_DIR + "/ngraph/impl/runtime"},
-    packages = ['ngraph', 'ngraph.utils', 'ngraph.impl', 'ngraph.impl.onnx_import',
-                'ngraph.impl.op', 'ngraph.impl.op.util', 'ngraph.impl.passes',
-                'ngraph.impl.runtime'],
+    package_dir=package_dir,
+    packages = packages,
     cmdclass={'build_ext': BuildExt},
     data_files = data_files,
     install_requires = requirements,

python/tox.ini

@@ -9,6 +9,7 @@ deps =
   -rtest_requirements.txt
 setenv =
   NGRAPH_CPP_BUILD_PATH = {env:NGRAPH_CPP_BUILD_PATH:{homedir}/ngraph_dist}
+  NGRAPH_ONNX_IMPORT_ENABLE = {env:NGRAPH_ONNX_IMPORT_ENABLE:"FALSE"}
   LD_LIBRARY_PATH = {env:LD_LIBRARY_PATH:{homedir}/ngraph_dist/lib}
   DYLD_LIBRARY_PATH = {env:DYLD_LIBRARY_PATH:{homedir}/ngraph_dist/lib}
   PYBIND_HEADERS_PATH = {env:PYBIND_HEADERS_PATH:}