[sharktank] Extend to test with Python 3.12 (#693)

Progress on #357 Extends to test with Python 3.12 and newer `torch` versions on Linux. For `torch.compile` support on Python 3.12, torch ge 2.4.0 is required. Testing with torch 2.3.0 is therefore skipped when running Python 3.12. To limit the number of jobs, tests are currently only expanded for Linux, whereas Windows (and unpinning NumPy) should be handled in a follow up. The error logged for `testExportNondecomposed` when failing with torch 2.5.1 slows down the CI by more than 20 minutes, therefore the test is now skipped instead of marked xfail.
nod-ai · Dec 13, 2024 · 332baa3 · 332baa3
1 parent 0bec8f1
commit 332baa3
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 7 deletions.
diff --git a/.github/workflows/ci-sharktank.yml b/.github/workflows/ci-sharktank.yml
@@ -23,11 +23,20 @@ concurrency:
 
 jobs:
   test:
-    name: "Unit Tests and Type Checking"
+    name: "Unit Tests (${{ matrix.os }}, ${{ matrix.python-version }}, ${{ matrix.torch-version }})"
     strategy:
       matrix:
-        version: [3.11]
-        os: [ubuntu-24.04, windows-2022]
+        python-version: ["3.11", "3.12"]
+        torch-version: ["2.3.0", "2.4.1", "2.5.1"]
+        os: [ubuntu-24.04]
+        include:
+          - os: windows-2022
+            python-version: "3.11"
+            torch-version: "2.3.0"
+        exclude:
+          - python-version: "3.12"
+             # `torch.compile` requires torch>=2.4.0 for Python 3.12+
+            torch-version: "2.3.0"
       fail-fast: false
     runs-on: ${{matrix.os}}
     defaults:
@@ -42,7 +51,7 @@ jobs:
         id: setup_python
         uses: actions/setup-python@0b93645e9fea7318ecaed2b359559ac225c90a2b # v5.3.0
         with:
-          python-version: ${{matrix.version}}
+          python-version: ${{matrix.python-version}}
 
       - name: Cache Pip Packages
         uses: actions/cache@6849a6489940f00c2f30c0fb92c6274307ccb58a # v4.1.2
@@ -57,7 +66,7 @@ jobs:
           # Note: We install in three steps in order to satisfy requirements
           # from non default locations first. Installing the PyTorch CPU
           # wheels saves multiple minutes and a lot of bandwidth on runner setup.
-          pip install --no-compile -r pytorch-cpu-requirements.txt
+          pip install --no-compile --pre --index-url https://download.pytorch.org/whl/test/cpu torch==${{matrix.torch-version}}+cpu
           pip install --no-compile -r requirements.txt -r sharktank/requirements-tests.txt -e sharktank/
 
           # Install nightly IREE packages.
@@ -77,7 +86,7 @@ jobs:
     name: "Data-dependent Tests"
     strategy:
       matrix:
-        version: [3.11]
+        python-version: [3.11]
         runs-on: [llama-mi300x-3]
       fail-fast: false
     runs-on: ${{matrix.runs-on}}
@@ -94,7 +103,7 @@ jobs:
         id: setup_python
         uses: actions/setup-python@0b93645e9fea7318ecaed2b359559ac225c90a2b # v5.3.0
         with:
-          python-version: ${{matrix.version}}
+          python-version: ${{matrix.python-version}}
 
       - name: Create Python venv
         run: python -m venv ${VENV_DIR}

diff --git a/sharktank/tests/layers/paged_llama_attention_block_test.py b/sharktank/tests/layers/paged_llama_attention_block_test.py
@@ -124,6 +124,10 @@ def forward(self, h, seq_block_ids, cache_state):
         torch.__version__ >= (2, 4),
         reason="https://github.com/nod-ai/shark-ai/issues/684",
     )
+    @pytest.mark.skipif(
+        torch.__version__ >= (2, 5),
+        reason="https://github.com/nod-ai/shark-ai/issues/684, error slows down CI",
+    )
     def testExportNondecomposed(self):
         dtype = torch.float32