build: Improves python packaging infrastructure

.github/workflows/build-and-test.yml

@@ -42,7 +42,6 @@ jobs:
           mkdir -p /core/build
           cd /core/build
           cmake -DCMAKE_INSTALL_PREFIX:PATH=`pwd`/install -DTRITON_CORE_HEADERS_ONLY=OFF ..
-          export TRITON_PYBIND="_c/triton_bindings.cpython-310-x86_64-linux-gnu.so"
           make -j8
 
       - name: Run tests with pytest

pyproject.toml

@@ -24,6 +24,53 @@
 # (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
+[project]
+name = "tritonserver"
+authors = [{ name = "NVIDIA Inc.", email = "[email protected]" }]
+description = "Triton Inference Server In-Process Python API"
+license = { file = "LICENSE.txt" }
+dynamic = ["version"]
+dependencies = ["numpy<2"]
+classifiers = [
+    "Development Status :: 5 - Production/Stable",
+    "Intended Audience :: Developers",
+    "Intended Audience :: Science/Research",
+    "Intended Audience :: Information Technology",
+    "Topic :: Scientific/Engineering",
+    "Topic :: Scientific/Engineering :: Image Recognition",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+    "Topic :: Software Development :: Libraries",
+    "Topic :: Utilities",
+    "License :: OSI Approved :: BSD License",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.10",
+    "Environment :: Console",
+    "Natural Language :: English",
+    "Operating System :: OS Independent",
+]
+
+[tool.setuptools]
+include-package-data = true
+
+[tool.setuptools.package-data]
+tritonserver = ["_c/triton_bindings.*.so"]
+
+[build-system]
+requires = [
+    "setuptools==75.3.0",
+    "wheel==0.44.0",
+    # For stubgen:
+    "mypy==1.11.0",
+    "numpy<2",
+]
+build-backend = "setuptools.build_meta"
+
+[project.optional-dependencies]
+GPU = ["cupy-cuda12x"]
+test = ["pytest"]
+all = ["tritonserver[GPU]", "tritonserver[test]"]
+
+
 [tool.codespell]
 # note: pre-commit passes explicit lists of files here, which this skip file list doesn't override -
 # this is only to allow you to run codespell interactively

python/CMakeLists.txt

@@ -30,6 +30,7 @@ add_subdirectory(tritonserver)
 file(WRITE ${CMAKE_CURRENT_BINARY_DIR}/TRITON_VERSION ${TRITON_VERSION})
 configure_file(../LICENSE LICENSE.txt COPYONLY)
 configure_file(setup.py setup.py @ONLY)
+configure_file(../pyproject.toml pyproject.toml COPYONLY)
 file(COPY test/  DESTINATION ./test/.)
 
 set(WHEEL_DEPENDS

python/build_wheel.py

@@ -108,17 +108,18 @@ def sed(pattern, replace, source, dest=None):
 
     shutil.copyfile("LICENSE.txt", os.path.join(FLAGS.whl_dir, "LICENSE.txt"))
     shutil.copyfile("setup.py", os.path.join(FLAGS.whl_dir, "setup.py"))
+    shutil.copyfile("pyproject.toml", os.path.join(FLAGS.whl_dir, "pyproject.toml"))
 
     os.chdir(FLAGS.whl_dir)
     print("=== Building wheel")
-    args = ["python3", "setup.py", "bdist_wheel"]
+    args = ["python3", "-m", "build"]
 
     wenv = os.environ.copy()
     wenv["VERSION"] = FLAGS.triton_version
     wenv["TRITON_PYBIND"] = PYBIND_LIB
     p = subprocess.Popen(args, env=wenv)
     p.wait()
-    fail_if(p.returncode != 0, "setup.py failed")
+    fail_if(p.returncode != 0, "Building wheel failed failed")
 
     cpdir("dist", FLAGS.dest_dir)
 

python/setup.py

@@ -25,90 +25,23 @@
 # (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
-import os
-import sys
-from itertools import chain
 
-from setuptools import find_packages, setup
+import subprocess
 
-if "--plat-name" in sys.argv:
-    PLATFORM_FLAG = sys.argv[sys.argv.index("--plat-name") + 1]
-else:
-    PLATFORM_FLAG = "any"
+from setuptools import setup
+from setuptools.command.build_py import build_py
 
-if "VERSION" not in os.environ:
-    raise Exception("envvar VERSION must be specified")
 
-VERSION = os.environ["VERSION"]
+class BuildPyCommand(build_py):
+    def run(self):
+        build_py.run(self)
+        # Generate stub files:
+        package_name = self.distribution.metadata.name
+        subprocess.run(
+            ["stubgen", "-p", f"{package_name}._c", "-o", f"{self.build_lib}"],
+            check=True,
+        )
 
-try:
-    from wheel.bdist_wheel import bdist_wheel as _bdist_wheel
 
-    class bdist_wheel(_bdist_wheel):
-        def finalize_options(self):
-            _bdist_wheel.finalize_options(self)
-            self.root_is_pure = False
-
-        def get_tag(self):
-            pyver, abi, plat = "py3", "none", PLATFORM_FLAG
-            return pyver, abi, plat
-
-except ImportError:
-    bdist_wheel = None
-
-this_directory = os.path.abspath(os.path.dirname(__file__))
-
-data_files = [
-    ("", ["LICENSE.txt"]),
-]
-
-# Type checking marker file indicating support for type checkers.
-# https://peps.python.org/pep-0561/
-# Type hints for c extension generated by mypy
-platform_package_data = [
-    os.environ["TRITON_PYBIND"],
-    "py.typed",
-    "_c/__init__.pyi",
-    "_c/triton_bindings.pyi",
-]
-
-gpu_extras = ["cupy-cuda12x"]
-test_extras = ["pytest"]
-all_extras = gpu_extras + test_extras
-
-setup(
-    name="tritonserver",
-    version=VERSION,
-    author="NVIDIA Inc.",
-    author_email="[email protected]",
-    description="Triton Inference Server In-Process Python API",
-    license="BSD",
-    url="https://developer.nvidia.com/nvidia-triton-inference-server",
-    classifiers=[
-        "Development Status :: 5 - Production/Stable",
-        "Intended Audience :: Developers",
-        "Intended Audience :: Science/Research",
-        "Intended Audience :: Information Technology",
-        "Topic :: Scientific/Engineering",
-        "Topic :: Scientific/Engineering :: Image Recognition",
-        "Topic :: Scientific/Engineering :: Artificial Intelligence",
-        "Topic :: Software Development :: Libraries",
-        "Topic :: Utilities",
-        "License :: OSI Approved :: BSD License",
-        "Programming Language :: Python :: 3",
-        "Programming Language :: Python :: 3.10",
-        "Programming Language :: Python :: 3.12",
-        "Environment :: Console",
-        "Natural Language :: English",
-        "Operating System :: OS Independent",
-    ],
-    packages=find_packages(),
-    package_data={
-        "": platform_package_data,
-    },
-    zip_safe=False,
-    cmdclass={"bdist_wheel": bdist_wheel},
-    data_files=data_files,
-    install_requires=["numpy<2"],
-    extras_require={"GPU": gpu_extras, "test": test_extras, "all": all_extras},
-)
+if __name__ == "__main__":
+    setup(cmdclass={"build_py": BuildPyCommand})

python/tritonserver/CMakeLists.txt

@@ -33,8 +33,6 @@ file(COPY __init__.py DESTINATION .)
 file(COPY py.typed DESTINATION .)
 # Copy the '__init__.py' for the '_c' module
 file(COPY _c/__init__.py DESTINATION ./_c/.)
-file(COPY _c/__init__.pyi DESTINATION ./_c/.)
-file(COPY _c/triton_bindings.pyi DESTINATION ./_c/.)
 # Find and copy _api modules
 file(GLOB PYTHON_MODULE_FILES ./_api/*.py)
 file(COPY ${PYTHON_MODULE_FILES} DESTINATION ./_api/.)
@@ -65,3 +63,4 @@ target_compile_features(python-bindings PRIVATE cxx_std_17)
 set_property(TARGET python-bindings PROPERTY OUTPUT_NAME triton_bindings)
 # Add Triton library default path in 'rpath' for runtime library lookup
 set_target_properties(python-bindings PROPERTIES BUILD_RPATH "$ORIGIN:/opt/tritonserver/lib")
+set_target_properties(python-bindings PROPERTIES LIBRARY_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR}/python/tritonserver/_c/)