refine release binary build for 3.0 API (#1336)

chensuyue · pre-commit-ci[bot] · web-flow · commit 2582361aee60 · 2023-10-24T11:58:09.000+08:00
Signed-off-by: chensuyue &lt;suyue.chen@intel.com&gt;
Co-authored-by: pre-commit-ci[bot] &lt;66853113+pre-commit-ci[bot]@users.noreply.github.com&gt;
diff --git a/setup.cfg b/setup.cfg
@@ -0,0 +1,4 @@
+[metadata]
+license_files =
+   LICENSE
+   third-party-programs.txt
diff --git a/setup.py b/setup.py
@@ -17,68 +17,126 @@ def fetch_requirements(path):
 except Exception as error:
     assert False, "Error: Could not open '%s' due %s\n" % (filepath, error)
 
-neural_insights = False
-if "neural_insights" in sys.argv:
-    neural_insights = True
-    sys.argv.remove("neural_insights")
+PKG_INSTALL_CFG = {
+    "neural_compressor_2x": {
+        "project_name": "neural_compressor",
+        "include_packages": find_packages(
+            include=["neural_compressor", "neural_compressor.*", "neural_coder", "neural_coder.*"],
+            exclude=[
+                "neural_compressor.template",
+                "neural_compressor.common",
+                "neural_compressor.common.*",
+                "neural_compressor.torch",
+                "neural_compressor.torch.*",
+                "neural_compressor.tensorflow",
+                "neural_compressor.tensorflow.*",
+            ],
+        ),
+        "package_data": {"": ["*.yaml"]},
+        "install_requires": fetch_requirements("requirements.txt"),
+        "extras_require": {
+            "pt": [f"neural_compressor_3x_pt=={__version__}"],
+            "tf": [f"neural_compressor_3x_tf=={__version__}"],
+            "ort": [f"neural_compressor_3x_ort=={__version__}"],
+        },
+    },
+    "neural_compressor_3x_pt": {
+        "project_name": "neural_compressor_3x_pt",
+        "include_packages": find_packages(
+            include=[
+                "neural_compressor.common",
+                "neural_compressor.common.*",
+                "neural_compressor.version.py",
+                "neural_compressor.torch",
+                "neural_compressor.torch.*",
+            ],
+        ),
+    },
+    "neural_compressor_3x_tf": {
+        "project_name": "neural_compressor_3x_tf",
+        "include_packages": find_packages(
+            include=[
+                "neural_compressor.common",
+                "neural_compressor.common.*",
+                "neural_compressor.version.py",
+                "neural_compressor.tensorflow",
+                "neural_compressor.tensorflow.*",
+            ],
+        ),
+    },
+    "neural_compressor_3x_ort": {
+        "project_name": "neural_compressor_3x_ort",
+        "include_packages": find_packages(
+            include=[
+                "neural_compressor.common",
+                "neural_compressor.common.*",
+                "neural_compressor.version.py",
+                "neural_compressor.onnxrt",
+                "neural_compressor.onnxrt.*",
+            ],
+        ),
+    },
+    "neural_insights": {
+        "project_name": "neural_insights",
+        "include_packages": find_packages(include=["neural_insights", "neural_insights.*"], exclude=["test.*", "test"]),
+        "package_data": {
+            "neural_insights": [
+                "bin/*",
+                "*.yaml",
+                "web/app/*.*",
+                "web/app/static/css/*",
+                "web/app/static/js/*",
+                "web/app/static/media/*",
+                "web/app/icons/*",
+            ]
+        },
+        "install_requires": fetch_requirements("neural_insights/requirements.txt"),
+        "entry_points": {"console_scripts": ["neural_insights = neural_insights.bin.neural_insights:execute"]},
+    },
+    "neural_solution": {
+        "project_name": "neural_solution",
+        "include_packages": find_packages(include=["neural_solution", "neural_solution.*"]),
+        "package_data": {
+            "neural_solution": [
+                "scripts/*.*",
+                "frontend/*.json",
+            ]
+        },
+        "install_requires": fetch_requirements("neural_solution/requirements.txt"),
+        "entry_points": {"console_scripts": ["neural_solution = neural_solution.bin.neural_solution:exec"]},
+    },
+}
 
-neural_solution = False
-if "neural_solution" in sys.argv:
-    neural_solution = True
-    sys.argv.remove("neural_solution")
 
-# define include packages
-include_packages = find_packages(
-    include=["neural_compressor", "neural_compressor.*", "neural_coder", "neural_coder.*"],
-    exclude=["neural_compressor.template"],
-)
-neural_insights_packages = find_packages(include=["neural_insights", "neural_insights.*"], exclude=["test.*", "test"])
-neural_solution_packages = find_packages(include=["neural_solution", "neural_solution.*"])
+if __name__ == "__main__":
+    cfg_key = "neural_compressor_2x"
+    if "neural_insights" in sys.argv:
+        sys.argv.remove("neural_insights")
+        cfg_key = "neural_insights"
 
-# define package data
-package_data = {"": ["*.yaml"]}
-neural_insights_data = {
-    "neural_insights": [
-        "bin/*",
-        "*.yaml",
-        "web/app/*.*",
-        "web/app/static/css/*",
-        "web/app/static/js/*",
-        "web/app/static/media/*",
-        "web/app/icons/*",
-    ]
-}
-neural_solution_data = {
-    "neural_solution": [
-        "scripts/*.*",
-        "frontend/*.json",
-    ]
-}
+    if "neural_solution" in sys.argv:
+        sys.argv.remove("neural_solution")
+        cfg_key = "neural_solution"
 
-# define install requirements
-install_requires_list = fetch_requirements("requirements.txt")
-neural_insights_requires = fetch_requirements("neural_insights/requirements.txt")
-neural_solution_requires = fetch_requirements("neural_solution/requirements.txt")
+    if "pt" in sys.argv:
+        sys.argv.remove("pt")
+        cfg_key = "neural_compressor_3x_pt"
 
-# define entry points
-entry_points = {}
+    if "tf" in sys.argv:
+        sys.argv.remove("tf")
+        cfg_key = "neural_compressor_3x_tf"
 
-if neural_insights:
-    project_name = "neural_insights"
-    package_data = neural_insights_data
-    install_requires_list = neural_insights_requires
-    include_packages = neural_insights_packages
-    entry_points = {"console_scripts": ["neural_insights = neural_insights.bin.neural_insights:execute"]}
-elif neural_solution:
-    project_name = "neural_solution"
-    package_data = neural_solution_data
-    install_requires_list = neural_solution_requires
-    include_packages = neural_solution_packages
-    entry_points = {"console_scripts": ["neural_solution = neural_solution.bin.neural_solution:exec"]}
-else:
-    project_name = "neural_compressor"
+    if "ort" in sys.argv:
+        sys.argv.remove("ort")
+        cfg_key = "neural_compressor_3x_ort"
+
+    project_name = PKG_INSTALL_CFG[cfg_key].get("project_name")
+    include_packages = PKG_INSTALL_CFG[cfg_key].get("include_packages") or {}
+    package_data = PKG_INSTALL_CFG[cfg_key].get("package_data") or {}
+    install_requires = PKG_INSTALL_CFG[cfg_key].get("install_requires") or []
+    entry_points = PKG_INSTALL_CFG[cfg_key].get("entry_points") or {}
+    extras_require = PKG_INSTALL_CFG[cfg_key].get("extras_require") or {}
 
-if __name__ == "__main__":
     setup(
         name=project_name,
         version=__version__,
@@ -87,15 +145,17 @@ def fetch_requirements(path):
         description="Repository of Intel® Neural Compressor",
         long_description=open("README.md", "r", encoding="utf-8").read(),
         long_description_content_type="text/markdown",
-        keywords="quantization, auto-tuning, post-training static quantization, post-training dynamic quantization, quantization-aware training",
+        keywords="quantization,auto-tuning,post-training static quantization,"
+        "post-training dynamic quantization,quantization-aware training",
         license="Apache 2.0",
         url="https://github.com/intel/neural-compressor",
         packages=include_packages,
         include_package_data=True,
         package_data=package_data,
-        install_requires=install_requires_list,
+        install_requires=install_requires,
         entry_points=entry_points,
-        python_requires=">=3.6.0",
+        extras_require=extras_require,
+        python_requires=">=3.7.0",
         classifiers=[
             "Intended Audience :: Science/Research",
             "Programming Language :: Python :: 3",
