From 332baa32f6ec9022e9c570a4b244e28731b4bda9 Mon Sep 17 00:00:00 2001
From: Marius Brehler <marius.brehler@amd.com>
Date: Fri, 13 Dec 2024 19:36:18 +0100
Subject: [PATCH] [sharktank] Extend to test with Python 3.12 (#693)

Progress on https://github.com/nod-ai/shark-ai/issues/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.
---
 .github/workflows/ci-sharktank.yml            | 23 +++++++++++++------
 .../paged_llama_attention_block_test.py       |  4 ++++
 2 files changed, 20 insertions(+), 7 deletions(-)

diff --git a/.github/workflows/ci-sharktank.yml b/.github/workflows/ci-sharktank.yml
index f169eb2b2..5ee2fbcae 100644
--- 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
index d5cb6863d..bbb52f235 100644
--- 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