Fixes release pipeline, adds separate build flow to include stubs

- Updates permissions for the release pipeline so it is allowed to push new
     releases.

- Adds a separate wheel building flow so that we can include stub files
     (see RELEASE.md for details).

.github/workflows/tripy-release.yml

@@ -8,7 +8,7 @@ on:
 
 # Sets permissions of the GITHUB_TOKEN to allow deployment to GitHub Pages
 permissions:
-  contents: read
+  contents: write
   pages: write
   id-token: write
 
@@ -31,29 +31,31 @@ jobs:
       options: --gpus all
     steps:
     - uses: actions/checkout@v4
+
     - uses: actions/configure-pages@v5
 
+    - name: build-package
+      run: |
+        cd /tripy/
+        python3 -m build . -w -n
+
     - name: build-docs
       run: |
         cd /tripy/
         python3 docs/generate_rsts.py
         sphinx-build build/doc_sources build/docs -c docs/ -j 4 -W -n
         cp docs/packages.html build/docs/
 
+    - name: Release
+      uses: softprops/action-gh-release@v2
+      with:
+        generate_release_notes: true
+        files: /tripy/dist/tripy-*.whl
+
     - uses: actions/upload-pages-artifact@v3
       with:
         path: "/tripy/build/docs"
+
     - name: Deploy to GitHub Pages
       id: deployment
       uses: actions/deploy-pages@v4
-
-    - name: build-package
-      run: |
-        cd /tripy/
-        python3 -m build .
-
-    - name: Release
-      uses: softprops/action-gh-release@v2
-      with:
-        generate_release_notes: true
-        files: /tripy/dist/tripy-*.whl

tripy/Dockerfile

@@ -25,7 +25,7 @@ RUN apt-get update && \
 COPY .lldbinit /root/
 COPY pyproject.toml /tripy/pyproject.toml
 
-RUN pip install build .[docs,dev,test] \
+RUN pip install build .[docs,dev,test,build] \
     -f https://nvidia.github.io/TensorRT-Incubator/packages.html \
     --extra-index-url https://download.pytorch.org/whl \
     --extra-index-url https://pypi.nvidia.com

tripy/README.md

@@ -47,7 +47,7 @@ To get the latest changes in the repository, you can build Tripy wheels from sou
 2. From the [`tripy` root directory](.), run:
 
     ```bash
-    python3 -m build .
+    python3 -m build . -w
     ```
 
 3. Install the wheel, which should have been created in the `dist/` directory.

tripy/RELEASE.md

@@ -21,3 +21,13 @@ This document explains how to release a new version of Tripy.
 
     This should trigger our release pipeline, which will build and deploy
     the documentation and create a GitHub release with the wheel.
+
+
+## Additional Notes
+
+The public build instructions use the usual isolated build flow to create the
+wheels. However, in our release pipelines, we would also like to include stub
+files in the generated wheels. `stubgen`, which we use in [setup.py](./setup.py)
+requires all dependencies of of the package to be installed for the `--inspect-mode`
+option to work. Obviously, this would be a bit circular for external users, but it
+is easy for us to do in the development container which already includes the dependencies.

tripy/examples/nanogpt/model.py

@@ -17,10 +17,8 @@
 
 import math
 from dataclasses import dataclass
-from typing import Optional
 
 import tripy as tp
-from tripy import utils
 
 
 @dataclass

tripy/pyproject.toml

@@ -21,13 +21,24 @@ Documentation = "https://nvidia.github.io/TensorRT-Incubator/"
 
 
 [build-system]
-requires = ["setuptools>=45", "wheel", "mypy==1.11.0"]
+requires = [
+  "setuptools>=45",
+  "wheel",
+  # For stubgen:
+  "mypy==1.11.0",
+]
 build-backend = "setuptools.build_meta"
 
 [project.optional-dependencies]
 dev = [
   "pre-commit==3.6.0",
 ]
+build = [
+  "setuptools>=45",
+  "wheel",
+  # For stubgen:
+  "mypy==1.11.0",
+]
 doc_test_common = [
   "torch==2.4.0+cu121",
   "numpy==1.25.0",

tripy/setup.py

@@ -33,7 +33,16 @@ def generate_stubs(self):
         package_name = self.distribution.metadata.name
         try:
             subprocess.run(
-                ["stubgen", "-m", package_name, "--include-docstrings", "--inspect-mode", "-o", self.build_lib],
+                [
+                    "stubgen",
+                    "-m",
+                    package_name,
+                    "--include-docstrings",
+                    "--inspect-mode",
+                    "--export-less",
+                    "-o",
+                    self.build_lib,
+                ],
                 check=True,
             )
         except subprocess.CalledProcessError as e:

tripy/tests/test_packaging.py

@@ -12,16 +12,62 @@
 # 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.
-from tests import helper
-import tempfile
 import glob
 import os
-import tripy as tp
+import subprocess as sp
+import tempfile
+
 import pytest
+from tests import helper
+
+import tripy as tp
+
+
+def check_wheel(virtualenv, wheel_dir):
+    wheel_paths = glob.glob(os.path.join(wheel_dir, "*.whl"))
+    assert len(wheel_paths) == 1, "Expected exactly one wheel file."
+
+    wheel_path = wheel_paths[0]
+    assert os.path.basename(wheel_path) == f"tripy-{tp.__version__}-py3-none-any.whl"
+
+    virtualenv.run(
+        [
+            virtualenv.python,
+            "-m",
+            "pip",
+            "install",
+            wheel_path,
+            # Needed to install MLIR-TRT:
+            "-f",
+            "https://nvidia.github.io/TensorRT-Incubator/packages.html",
+            # For some reason, using the cache causes pip not to correctly install TRT
+            # if this test is run multiple times in succession.
+            "--no-cache-dir",
+        ]
+    )
+
+    assert "tripy" in virtualenv.installed_packages()
+    tripy_pkg = virtualenv.installed_packages()["tripy"]
+    assert tripy_pkg.version == tp.__version__
+
+    # Check that we only package things we actually want.
+    # If tests are packaged, they'll end up in a higher-level directory.
+    assert not os.path.exists(os.path.join(tripy_pkg.source_path, "tests"))
+
+    # Lastly check we can actually import it and run a simple sanity test:
+    virtualenv.run(
+        [
+            virtualenv.python,
+            "-c",
+            "import tripy as tp; x = tp.ones((5,), dtype=tp.int32); assert x.tolist() == [1] * 5",
+        ]
+    )
+    return tripy_pkg
 
 
 @pytest.mark.l1
-def test_wheel_packaging_and_install(virtualenv):
+def test_isolated_wheel_packaging_and_install(virtualenv):
+    # Tests wheel packaging and installation as an external user would do it.
     with helper.raises(Exception, "returned non-zero exit status"):
         virtualenv.run([virtualenv.python, "-c", "import tripy"])
 
@@ -30,41 +76,16 @@ def test_wheel_packaging_and_install(virtualenv):
     with tempfile.TemporaryDirectory() as tmp:
         virtualenv.run([virtualenv.python, "-m", "build", ".", "-o", tmp], cwd=helper.ROOT_DIR)
 
-        wheel_paths = glob.glob(os.path.join(tmp, "*.whl"))
-        assert len(wheel_paths) == 1, "Expected exactly one wheel file."
-
-        wheel_path = wheel_paths[0]
-        assert os.path.basename(wheel_path) == f"tripy-{tp.__version__}-py3-none-any.whl"
-
-        virtualenv.run(
-            [
-                virtualenv.python,
-                "-m",
-                "pip",
-                "install",
-                wheel_path,
-                # Needed to install MLIR-TRT:
-                "-f",
-                "https://nvidia.github.io/TensorRT-Incubator/packages.html",
-                # For some reason, using the cache causes pip not to correctly install TRT
-                # if this test is run multiple times in succession.
-                "--no-cache-dir",
-            ]
-        )
-
-        assert "tripy" in virtualenv.installed_packages()
-        tripy_pkg = virtualenv.installed_packages()["tripy"]
-        assert tripy_pkg.version == tp.__version__
-
-        # Check that we only package things we actually want.
-        # If tests are packaged, they'll end up in a higher-level directory.
-        assert not os.path.exists(os.path.join(tripy_pkg.source_path, "tests"))
-
-        # Lastly check we can actually import it and run a simple sanity test:
-        virtualenv.run(
-            [
-                virtualenv.python,
-                "-c",
-                "import tripy as tp; x = tp.ones((5,), dtype=tp.int32); assert x.tolist() == [1] * 5",
-            ]
-        )
+        check_wheel(virtualenv, tmp)
+
+
+@pytest.mark.l1
+def test_container_wheel_packaging_and_install(virtualenv):
+    # Tests wheel packaging as we do it in the development container.
+    with tempfile.TemporaryDirectory() as tmp:
+        sp.run(["python3", "-m", "build", ".", "-n", "-o", tmp], cwd=helper.ROOT_DIR)
+
+        tripy_pkg = check_wheel(virtualenv, tmp)
+
+        # Check that stubs are included
+        assert os.path.exists(os.path.join(tripy_pkg.source_path, tp.__name__, "__init__.pyi"))