Add job to build stub

Author: mthrok

.github/workflows/_build_linux_cuda.yml

@@ -90,6 +90,70 @@ jobs:
           "${python_exe}" -m pip install -r ./packaging/requirements.txt
           "${python_exe}" -m twine check --strict ~/package/*.whl
 
+  check-stub:
+    if: ${{ inputs.run-test == 'true' }}
+    needs: ["build"]
+    runs-on: ${{ inputs.machine }}
+    strategy:
+      fail-fast: false
+      matrix:
+        ffmpeg-version: ["8.0"]
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          persist-credentials: false
+
+      - uses: actions/download-artifact@v4
+        with:
+          name: "${{ env.ARTIFACT }}"
+          path: ~/package
+
+      - uses: conda-incubator/setup-miniconda@v3
+        with:
+          python-version: "${{ inputs.python-version }}"
+          conda-remove-defaults: "true"
+
+      - name: Generate stub files
+        run: |
+          set -x
+
+          pip install $(find "${HOME}/package" -name '*.whl' -depth -maxdepth 1) nanobind
+
+
+          cu_ver="${{ inputs.cuda-version }}"
+          cu_ver="${cu_ver//[.]/}"  # Remove dots
+          pip install numpy torch --index-url "https://download.pytorch.org/whl/cu${cu_ver}"
+
+          # Install FFmpeg
+          # Note: Somehow FFmepg 5.1 does not install libiconv so specifying it here
+          conda install -q -c conda-forge "ffmpeg==${{ matrix.ffmpeg-version }}" libiconv
+
+          # Sort out library paths
+          # 1. For some reason, the FFmpeg dynamic libraries are not found.
+          export "LD_LIBRARY_PATH=${CONDA_PREFIX}/lib"
+          # 2. CUDA runtime PyPI package is installed as part of PyTorch installation.
+          # However, it is installed in site-packages dir, which is not in PATH by default.
+          # (PyTorch has a special way to load it, but we don't do that)
+          cudart_dir="$(
+            python -c 'import os.path,nvidia.cuda_runtime;print(os.path.join(nvidia.cuda_runtime.__path__[0], "lib"))'
+          )"
+          export "LD_LIBRARY_PATH=${LD_LIBRARY_PATH}:${cudart_dir}"
+
+
+          python scripts/generate_stub.py -o ~/stubs
+          cp ~/stubs/* src/spdl/io/lib/
+          git diff --exit-code
+
+      - uses: actions/upload-artifact@v4
+        name: Upload up-to-date stub files
+        if: failure()
+        with:
+          name: "stub"
+          path: ~/stubs
+          if-no-files-found: error
+          retention-days: 7
+          overwrite: true
+
   check-package-dependency:
     if: ${{ inputs.run-test == 'true' }}
     needs: ["build"]

.github/workflows/packaging.yml

@@ -65,4 +65,4 @@ jobs:
       build_container: "pytorch/manylinux2_28-builder:cuda${{ matrix.cuda-version }}"
       python-version: "${{ matrix.python-version }}"
       cuda-version: "${{ matrix.cuda-version }}"
-      run-test: "${{ matrix.python-version == '3.11'}}"
+      run-test: "${{ matrix.python-version == '3.12'}}"

scripts/generate_stub.py

@@ -0,0 +1,41 @@
+#!/usr/bin/env python3
+
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""Generate interface files for extension modules."""
+
+import logging
+from argparse import ArgumentParser
+from pathlib import Path
+
+from nanobind.stubgen import StubGen
+from spdl.io.lib import _import_libspdl, _import_libspdl_cuda
+
+
+def _parse_args():
+    parser = ArgumentParser(description=__doc__)
+    parser.add_argument("--output-dir", "-o", required=True, type=Path)
+    return parser.parse_args()
+
+
+def _generate(mod, output):
+    sg = StubGen(mod)
+    sg.put(mod)
+    with open(output, "w") as f:
+        f.write(sg.get())
+
+
+def _main():
+    args = _parse_args()
+    logging.basicConfig(level=logging.DEBUG)
+    args.output_dir.mkdir(exist_ok=True, parents=True)
+    _generate(_import_libspdl(), args.output_dir / "_libspdl.pyi")
+    _generate(_import_libspdl_cuda(), args.output_dir / "_libspdl_cuda.pyi")
+
+
+if __name__ == "__main__":
+    _main()