chore: fix the docgen job

Signed-off-by: Naren Dasan <[email protected]>
Signed-off-by: Naren Dasan <[email protected]>

Author: narendasan

.github/workflows/docgen.yml

@@ -12,75 +12,33 @@ jobs:
   build-docs:
     runs-on: ubuntu-20.04
     container:
-      image: ghcr.io/pytorch/tensorrt/docgen:latest
+      image: docker.io/pytorch/manylinux-builder:cuda12.1
       credentials:
         username: ${{ github.actor }}
         password: ${{ secrets.GITHUB_TOKEN }}
     steps:
-      - name: Reclaim space
-        run: |
-          rm -rf /usr/share/dotnet
-          rm -rf /opt/ghc
-          rm -rf "/usr/local/share/boost"
-          rm -rf /usr/local/cuda/cuda-*
-      - name: Install base deps
-        run: |
-          apt update
-          DEBIAN_FRONTEND=noninteractive apt install -y software-properties-common gcc git curl wget make zlib1g-dev bzip2 libbz2-dev lzma lzma-dev libreadline-dev libsqlite3-dev libssl-dev libffi-dev doxygen pandoc
-          git config --global --add safe.directory '*'
-      - name: Set up Python 3.10.12
-        uses: actions/setup-python@v4
-        with:
-          python-version: 3.10.12
       - uses: actions/checkout@v3
         with:
           ref: ${{github.head_ref}}
+      - name: Install base deps
+        run: |
+          ./packaging/pre_build_script.sh
       - name: Get HEAD SHA
         id: vars
         run: echo "sha=$(git rev-parse --short HEAD)" >> $GITHUB_OUTPUT
       - name: Get Bazel version
         id: bazel_info
         run: echo "version=$(cat .bazelversion)" >> $GITHUB_OUTPUT
-      - name: Install Bazel
-        run: |
-          wget -q https://github.com/bazelbuild/bazel/releases/download/${{ steps.bazel_info.outputs.version }}/bazel-${{ steps.bazel_info.outputs.version }}-linux-x86_64 -O /usr/bin/bazel
-          chmod a+x /usr/bin/bazel
-      - name: Install cudnn + tensorrt
-        run: |
-          apt-get update
-          wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/cuda-ubuntu2004.pin
-          mv cuda-ubuntu2004.pin /etc/apt/preferences.d/cuda-repository-pin-600
-          apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/7fa2af80.pub
-          apt-key adv --keyserver keyserver.ubuntu.com --recv-keys 536F8F1DE80F6A35
-          apt-key adv --keyserver keyserver.ubuntu.com --recv-keys A4B469963BF863CC
-          add-apt-repository "deb https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/ /"
-          apt-get update
-          apt-get install -y libcudnn8 libcudnn8-dev
-
-          apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/3bf863cc.pub
-          add-apt-repository "deb https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/ /"
-          apt-get update
-
-          apt-get install -y libnvinfer8 libnvinfer-plugin8 libnvinfer-dev libnvinfer-plugin-dev
-      - name: Install Torch
-        run: |
-          python3 -m pip install -r py/requirements.txt --user
       - name: Build Python Package
         run: |
-          cp toolchains/ci_workspaces/WORKSPACE.x86_64 WORKSPACE
-          cd py
-          python3 -m pip install pip==21.3.1
-          echo $(which python3)
-          echo $(python3 -c 'import site; print(site.getsitepackages()[0])')
-          mkdir -p /opt/circleci/.pyenv/versions/3.9.4/lib/python3.9/
-          ln -s $(python3 -c 'import site; print(site.getsitepackages()[0])') /opt/circleci/.pyenv/versions/3.9.4/lib/python3.9/site-packages
-          python3 setup.py install
-          cd ..
+          cp toolchains/ci_workspaces/WORKSPACE.x86_64.cu121.release.rhel WORKSPACE
+          python -m pip install pip<=23
+          python -m pip install --pre -e . --extra-index-url https://download.pytorch.org/whl/nightly/cu121
       - name: Generate New Docs
         run: |
           cd docsrc
-          python3 -m pip install -r requirements.txt
-          python3 -c "import torch_tensorrt; print(torch_tensorrt.__version__)"
+          python -m pip install -r requirements.txt
+          python -c "import torch_tensorrt; print(torch_tensorrt.__version__)"
           make html
           cd ..
       - uses: stefanzweifel/git-auto-commit-action@v4

packaging/pre_build_script.sh

@@ -0,0 +1,13 @@
+#!/bin/bash
+
+# Install dependencies
+TRT_VERSION=$(python3 -c "import versions; versions.tensorrt_version()")
+yum-config-manager --add-repo https://developer.download.nvidia.com/compute/cuda/repos/rhel7/x86_64/cuda-rhel7.repo
+yum check-update
+yum install -y ninja-build tensorrt-${TRT_VERSION}.*
+wget https://github.com/bazelbuild/bazelisk/releases/download/v1.17.0/bazelisk-linux-amd64 \
+    && mv bazelisk-linux-amd64 /usr/bin/bazel \
+    && chmod +x /usr/bin/bazel
+
+cp toolchains/ci_workspaces/WORKSPACE.x86_64.${VERSION_SUFFIX#*+}.release.rhel WORKSPACE
+export CI_BUILD=1