[Test] Rework PR Gate

.github/workflows/template_unit_tests_on_gpu.yml → .github/workflows/action_gpu_unit_tests.yml

@@ -6,29 +6,32 @@ on:
       os:
         required: true
         type: string
-      python:
-        required: true
-        type: string
-      model:
+      python-version:
         required: true
         type: string
 
 jobs:
   unit-tests-on-gpu:
+    strategy:
+      fail-fast: false # Don't cancel all on first failure
+      matrix:
+        model:
+          - "gpt2gpu"
+          - "phi2gpu"
+          - "hfllama_7b_gpu"
     runs-on: ${{ inputs.os }}
 
     steps:
       - uses: actions/checkout@v4
       - name: Install Rust
-        uses: actions-rs/toolchain@v1
-        with:
-          profile: minimal
-          toolchain: stable
-          override: true
-      - name: Set up Python ${{ inputs.python }}
+        shell: bash
+        run: |
+           curl https://sh.rustup.rs -sSf | sh -s -- -y --default-toolchain 1.75.0
+           echo "$HOME/.cargo/bin" >> $GITHUB_PATH
+      - name: Set up Python ${{ inputs.python-version }}
         uses: actions/setup-python@v5
         with:
-          python-version: ${{ inputs.python }}
+          python-version: ${{ inputs.python-version }}
       - name: Show GPUs
         run: |
           nvidia-smi
@@ -54,10 +57,10 @@ jobs:
       - name: Check GPU available
         run: |
           python -c "import torch; assert torch.cuda.is_available()"
-      - name: Run tests (except server) for ${{ inputs.model }}
+      - name: Run tests (except server) for ${{ matrix.model }}
         run: |
           pytest -vv --cov=guidance --cov-report=xml --cov-report=term-missing \
-            --selected_model ${{ inputs.model }} \
+            --selected_model ${{ matrix.model }} \
             -m "not (server or needs_credentials)" \
             ./tests/
       - name: Upload coverage reports to Codecov

.github/workflows/unit_tests.yml → .github/workflows/action_plain_unit_tests.yml

@@ -1,33 +1,40 @@
-name: Unit tests
+name: plain-unit-tests
 
 on:
-  push:
-    branches: [main]
-  pull_request:
-    branches: [main]
-  workflow_dispatch:
+  workflow_call:
+    inputs:
+      os:
+        required: true
+        type: string
+      python-version:
+        required: true
+        type: string
 
-jobs:
-  unit-tests:
 
+jobs:
+  run-plain-unit-tests:
     strategy:
+      fail-fast: false # Don't cancel all on first failure
       matrix:
-        os: [ubuntu-latest, windows-latest, macos-12]
-        python-version: ["3.8", "3.9", "3.10", "3.11", "3.12"]
         model:
           - "gpt2cpu"
           - "phi2cpu"
           # - "transformers_mistral_7b" See Issue 713
           - "hfllama7b"
           - "hfllama_mistral_7b"
 
-    runs-on: ${{ matrix.os }}
+    runs-on: ${{ inputs.os }}
     steps:
       - uses: actions/checkout@v4
-      - name: Set up Python ${{ matrix.python-version }}
+      - name: Install Rust
+        shell: bash
+        run: |
+           curl https://sh.rustup.rs -sSf | sh -s -- -y --default-toolchain 1.75.0
+           echo "$HOME/.cargo/bin" >> $GITHUB_PATH
+      - name: Set up Python ${{ inputs.python-version }}
         uses: actions/setup-python@v5
         with:
-          python-version: ${{ matrix.python-version }}
+          python-version: ${{ inputs.python-version }}
       - name: Install dependencies
         shell: bash
         run: |

.github/workflows/server_tests.yml → .github/workflows/action_server_tests.yml

@@ -1,26 +1,25 @@
 name: Server tests
 
 on:
-  push:
-    branches: [main]
-  pull_request:
-    branches: [main]
-  workflow_dispatch:
+  workflow_call:
+    inputs:
+      os:
+        required: false
+        type: string
+        default: ubuntu-latest
+      python-version:
+        required: true
+        type: string
 
 jobs:
   server-tests:
-    strategy:
-      matrix:
-        os: [ubuntu-latest]
-        python-version: ["3.8", "3.9", "3.10", "3.11", "3.12"]
-
-    runs-on: ${{ matrix.os }}
+    runs-on: ${{ inputs.os }}
     steps:
       - uses: actions/checkout@v4
-      - name: Set up Python ${{ matrix.python-version }}
+      - name: Set up Python ${{ inputs.python-version }}
         uses: actions/setup-python@v5
         with:
-          python-version: ${{ matrix.python-version }}
+          python-version: ${{ inputs.python-version }}
       - name: Install dependencies
         shell: bash
         run: |

.github/workflows/workflow-pr-gate.yml

@@ -0,0 +1,156 @@
+name: Pull Request Gate
+
+on:
+  push:
+    branches: [main]
+  pull_request:
+    branches: [main]
+  workflow_dispatch:
+
+jobs:
+  # First Stage =======================================================================
+  # Linting and basic CPU-based tests
+
+  linting-black:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      # https://black.readthedocs.io/en/stable/integrations/github_actions.html
+      - uses: psf/black@stable
+        with:
+          options: "--diff" # Remove this to start enforcement
+
+  linting-mypy:
+    runs-on: ubuntu-latest
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v4
+      - name: Setup Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: 3.8
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip
+          python -m pip install .[all,test]
+      - name: Run mypy
+        run: |
+          python -m mypy guidance
+
+  bare-install:
+    strategy:
+      matrix:
+        os: [ubuntu-latest, windows-latest, macos-12]
+        python-version: ["3.8", "3.9", "3.10", "3.11", "3.12"]
+    runs-on: ${{ matrix.os }}
+    steps:
+      - uses: actions/checkout@v4
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python-version }}
+      - name: Minimal install
+        run: |
+          pip install -e .
+      - name: Attempt import
+        run: |
+          python -c "import guidance"
+
+
+  unit-tests-linux-python-latest:
+    uses: ./.github/workflows/action_plain_unit_tests.yml
+    with:
+      os: ubuntu-latest
+      python-version: "3.12"
+
+  end-stage-1:
+    needs:
+      - linting-black
+      - linting-mypy
+      - bare-install
+      - unit-tests-linux-python-latest
+    name: End Stage 1
+    runs-on: ubuntu-latest
+    steps:
+    - run: echo "End of Stage 1"
+
+# Second stage =======================================================================
+# Remaining CPU-based tests on Linux, plus server tests and basic GPU testing
+
+
+  unit-tests-gpu-python-latest:
+    needs: end-stage-1
+    uses: ./.github/workflows/action_gpu_unit_tests.yml
+    with:
+      os: gpu-runner
+      python-version: "3.12"
+
+  unit-tests-linux-python-other:
+    needs: end-stage-1
+    strategy:
+      fail-fast: false # Don't cancel all on first failure
+      matrix:
+        python-version: ["3.8", "3.9", "3.10", "3.11"]
+    uses: ./.github/workflows/action_plain_unit_tests.yml
+    with:
+      os: ubuntu-latest
+      python-version: ${{ matrix.python-version }}
+
+  server-tests:
+    needs: end-stage-1
+    strategy:
+      fail-fast: false # Don't cancel all on first failure
+      matrix:
+        python-version: ["3.8", "3.9", "3.10", "3.11", "3.12"]
+    uses: ./.github/workflows/action_server_tests.yml
+    with:
+      os: ubuntu-latest
+      python-version: ${{ matrix.python-version }}
+
+
+
+  end-stage-2:
+    needs:
+      - unit-tests-linux-python-other
+      - unit-tests-gpu-python-latest
+    name: End Stage 2
+    runs-on: ubuntu-latest
+    steps:
+    - run: echo "End of Stage 1"
+
+# Third Stage ==============================================================
+# Windows and MacOS, plus other GPU Linux tests
+
+  unit-tests-win:
+    needs: end-stage-2
+    strategy:
+      fail-fast: false # Don't cancel all on first failure
+      matrix:
+        python-version: ["3.8", "3.9", "3.10", "3.11", "3.12"]
+    uses: ./.github/workflows/action_plain_unit_tests.yml
+    with:
+      os: windows-latest
+      python-version: ${{ matrix.python-version }}
+
+  unit-tests-mac:
+    needs: end-stage-2
+    strategy:
+      fail-fast: false # Don't cancel all on first failure
+      matrix:
+        python-version: ["3.8", "3.9", "3.10", "3.11", "3.12"]
+    uses: ./.github/workflows/action_plain_unit_tests.yml
+    with:
+      os: macos-12
+      python-version: ${{ matrix.python-version }}
+
+  unit-tests-gpu-python-others:
+    needs: end-stage-2
+    strategy:
+      fail-fast: false # Don't cancel all on first failure
+      matrix:
+        python-version: ["3.8", "3.9", "3.10", "3.11"]
+    uses: ./.github/workflows/action_gpu_unit_tests.yml
+    with:
+      os: gpu-runner
+      python-version: ${{ matrix.python-version }}
+