Wheel file updates for FeaturizerLibrary data (#3640)

setup.py

@@ -7,14 +7,15 @@ from setuptools import setup, find_packages, Extension
 from distutils import log as logger
 from distutils.command.build_ext import build_ext as _build_ext
 from glob import glob
-from os import path, getcwd, environ, remove
-from shutil import copyfile
+from os import path, getcwd, environ, remove, walk, makedirs, listdir
+from shutil import copyfile, copytree, rmtree
 import platform
 import subprocess
 import sys
 import datetime
 
 nightly_build = False
+featurizers_build = False
 package_name = 'onnxruntime'
 wheel_name_suffix = None
 
@@ -47,6 +48,10 @@ elif '--use_nuphar' in sys.argv:
     package_name = 'onnxruntime-nuphar'
     sys.argv.remove('--use_nuphar')
 
+elif '--use_featurizers' in sys.argv:
+    featurizers_build = True
+    sys.argv.remove('--use_featurizers')
+
 if '--nightly_build' in sys.argv:
     package_name = 'ort-nightly'
     nightly_build = True
@@ -179,8 +184,6 @@ examples = [path.join('datasets', x) for x in examples_names]
 
 # Extra files such as EULA and ThirdPartyNotices
 extra = ["LICENSE", "ThirdPartyNotices.txt", "Privacy.md"]
-if package_name == 'onnxruntime-nuphar':
-  extra.extend([path.join('nuphar', 'NUPHAR_CACHE_VERSION')])
 
 # Description
 README = path.join(getcwd(), "docs/python/README.rst")
@@ -192,6 +195,41 @@ if not path.exists(README):
 with open(README) as f:
     long_description = f.read()
 
+packages = [
+    'onnxruntime',
+    'onnxruntime.backend',
+    'onnxruntime.capi',
+    'onnxruntime.datasets',
+    'onnxruntime.tools',
+]
+
+package_data = {}
+data_files = []
+
+if package_name == 'onnxruntime-nuphar':
+    packages += ["onnxruntime.nuphar"]
+    extra += [path.join('nuphar', 'NUPHAR_CACHE_VERSION')]
+
+if featurizers_build:
+    # Copy the featurizer data from its current directory into the onnx runtime directory so that the
+    # content can be included as module data.
+    featurizer_source_dir = path.join("external", "FeaturizersLibrary", "Data")
+    assert path.isdir(featurizer_source_dir), featurizer_source_dir
+
+    featurizer_dest_dir = path.join("onnxruntime", "FeaturizersLibrary", "Data")
+    if path.isdir(featurizer_dest_dir):
+        rmtree(featurizer_dest_dir)
+
+    for item in listdir(featurizer_source_dir):
+        this_featurizer_source_fullpath = path.join(featurizer_source_dir)
+        assert path.isdir(this_featurizer_source_fullpath), this_featurizer_source_fullpath
+
+        copytree(this_featurizer_source_fullpath, featurizer_dest_dir)
+
+        packages.append("{}.{}".format(featurizer_dest_dir.replace(path.sep, "."), item))
+        package_data[packages[-1]] = listdir(path.join(featurizer_dest_dir, item))
+
+package_data["onnxruntime"] = data + examples + extra
 
 version_number = ''
 with open('VERSION_NUMBER') as f:
@@ -234,16 +272,10 @@ setup(
     author_email='onnx@microsoft.com',
     cmdclass=cmd_classes,
     license="MIT License",
-    packages=['onnxruntime',
-              'onnxruntime.backend',
-              'onnxruntime.capi',
-              'onnxruntime.datasets',
-              'onnxruntime.tools',
-              ] + (['onnxruntime.nuphar'] if package_name == 'onnxruntime-nuphar' else []),
+    packages=packages,
     ext_modules=ext_modules,
-    package_data={
-        'onnxruntime': data + examples + extra,
-    },
+    package_data=package_data,
+    data_files=data_files,
     py_modules=python_modules_list,
     install_requires=install_requires,
     entry_points= {

tools/ci_build/build.py

@@ -1179,7 +1179,8 @@ def nuphar_run_python_tests(build_dir, configs):
 def build_python_wheel(
         source_dir, build_dir, configs, use_cuda, use_ngraph, use_dnnl,
         use_tensorrt, use_openvino, use_nuphar, wheel_name_suffix, use_acl,
-        nightly_build=False):
+        nightly_build=False,
+        featurizers_build=False):
     for config in configs:
         cwd = get_config_build_dir(build_dir, config)
         if is_windows():
@@ -1204,6 +1205,8 @@ def build_python_wheel(
             args.append('--wheel_name_suffix={}'.format(wheel_name_suffix))
         elif use_acl:
             args.append('--use_acl')
+        elif featurizers_build:
+            args.append("--use_featurizers")
 
         run_subprocess(args, cwd=cwd)
 
@@ -1537,6 +1540,7 @@ def main():
                 args.wheel_name_suffix,
                 args.use_acl,
                 nightly_build=nightly_build,
+                featurizers_build=args.use_featurizers,
             )
 
     if args.gen_doc and (args.build or args.test):

tools/ci_build/github/azure-pipelines/linux-ci-pipeline.yml

@@ -3,7 +3,7 @@ jobs:
   parameters:
     AgentPool : 'Linux-CPU'
     JobName: 'Linux_CI_Dev'
-    BuildCommand: 'tools/ci_build/github/linux/run_dockerbuild.sh -o ubuntu16.04 -d cpu -r $(Build.BinariesDirectory) -x "--use_mklml --use_llvm --use_nuphar --use_dnnl --use_tvm --use_featurizers --build_wheel --build_java --enable_language_interop_ops"'
+    BuildCommand: 'tools/ci_build/github/linux/run_dockerbuild.sh -o ubuntu16.04 -d cpu -r $(Build.BinariesDirectory) -x "--use_mklml --use_llvm --use_nuphar --use_dnnl --use_tvm --build_wheel --build_java --enable_language_interop_ops --use_featurizers --wheel_name_suffix featurizers"'
     DoNugetPack:  'false'
     ArtifactName: 'drop-linux'
     TimeoutInMinutes: 120