Switch to pytest (#44)

* move to pytest; factor out operator registration into context class

* fixes

* fix

* format

* fix application test

* remove test

* fix

* make all operators share the same default context

* unify tests into single function (incorporate Erika's feedback)

* reduce iterations of application tests to one for now due to CI inconsistency

* oops

Author: andrej

.github/actions/build/action.yaml

@@ -6,13 +6,17 @@ name: "Run test suite"
 inputs:
   test_flags:
     required: false
-    default: "--log tests.log --csv-latest tests_latest.csv"
+    default: ""
   env_name:
     required: false
     default: "ci_env"
-  testing_dir:
+  extensive:
     required: false
-    default: "testing"
+    default: "false"
+  test_suite:
+    required: false
+    default: "operators"
+    description: "Which tests to run: 'operators', 'applications', or 'both'"
 
 runs:
   using: "composite"
@@ -23,4 +27,23 @@ runs:
       shell: "bash"
       run: |
         source ${{ inputs.env_name }}/bin/activate
-        ./scripts/run_tests.py ${{ inputs.test_flags }} ${{ inputs.testing_dir }}
+        
+        case "${{ inputs.test_suite }}" in
+          operators)
+            TEST_PATH="operators/"
+            ;;
+          applications)
+            TEST_PATH="applications/"
+            ;;
+          *)
+            echo "Invalid test_suite: ${{ inputs.test_suite }}"
+            exit 1
+            ;;
+        esac
+        
+        if [ "${{ inputs.extensive }}" = "true" ]; then
+          pytest $TEST_PATH --csv-output=tests_latest.csv ${{ inputs.test_flags }}
+        else
+          pytest -m "not extensive" $TEST_PATH --csv-output=tests_latest.csv ${{ inputs.test_flags }}
+        fi
+

.github/actions/test/action.yaml

@@ -32,18 +32,14 @@ jobs:
         with:
           env_name: ci_env
 
-      - name: Build
-        uses: ./.github/actions/build
-        with:
-          env_name: ci_env
-          extensive: "true"
-
       - name: Test
         id: test
         continue-on-error: true
         uses: ./.github/actions/test
         with:
           env_name: ci_env
+          extensive: "true"
+          test_suite: "operators"
 
       - name: Commit test results
         id: commit_results

.github/workflows/extensive.yml

@@ -31,18 +31,14 @@ jobs:
         with:
           env_name: ci_env
 
-      - name: Build
-        uses: ./.github/actions/build
-        with:
-          env_name: ci_env
-          extensive: "false"
-
       - name: Test
         id: test
         continue-on-error: true
         uses: ./.github/actions/test
         with:
           env_name: ci_env
+          extensive: "false"
+          test_suite: "operators"
 
       - name: Commit test results
         id: commit_results

.github/workflows/small.yml

@@ -38,28 +38,14 @@ jobs:
           source ci_env/bin/activate
           pip install -r requirements_examples.txt
 
-      - name: Generate application tests
-        shell: bash
-        env:
-          HOME: /workspace
-        run: |
-          source ci_env/bin/activate
-          set -euxo pipefail
-          mkdir -p testing
-          cd testing
-          python3 ../applications/discover_application_tests.py --output-dir .
-          echo "Application tests generated:"
-          ls -la
-
       - name: Run examples
         id: test
         continue-on-error: true
-        shell: bash
-        env:
-          HOME: /workspace
-        run: |
-          source ci_env/bin/activate
-          ./scripts/run_tests.py --log tests.log --csv-latest tests_latest.csv testing
+        uses: ./.github/actions/test
+        with:
+          env_name: ci_env
+          test_suite: "applications"
+          test_flags: "--iterations 1"
 
       - name: Commit example test results
         uses: ./.github/actions/commit_results

.github/workflows/test-examples.yml

@@ -135,14 +135,20 @@ All available operators can be found in `operators`. These each contain:
 >       `source /opt/xilinx/xrt/setup.sh`
 >       `source /path/to/ironenv/bin/activate`
 
-To build and test all the operators, first generate a list of all test cases, then run them:
-``` python
-mkdir testing && cd testing
-../operators/common/discover_tests.py
-../scripts/run_tests.py --iter 1
+To build and test all the operators:
+``` bash
+pytest operators/ -m "not extensive"
 ``` 
 
-You can select a single test to run using the `--select` flag.
+To run the extensive test suite:
+``` bash
+pytest operators/
+```
+
+To run a specific operator's tests:
+``` bash
+pytest operators/axpy/
+```
 
 ### Git Hooks (Optional but Recommended)
 

README.md

@@ -299,8 +299,8 @@ def inference(
 
     logging.info("Preparing AIE operators...")
     # At this point the model is fully described (operators and their dimensions and how to compile them)
-    AIEOperatorBase.compile_all_operators()
-    AIEOperatorBase.prepare_runtime()
+    AIEOperatorBase.get_default_context().compile_all()
+    AIEOperatorBase.get_default_context().prepare_runtime()
     logging.info("AIE operator preparation completed.")
     print(f"Starting text generation...")
     print(f"Generating {num_tokens} tokens...")

applications/discover_application_tests.py

@@ -2,18 +2,50 @@
 # SPDX-FileCopyrightText: Copyright (C) 2025 Advanced Micro Devices, Inc. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+import subprocess
+import pytest
 from pathlib import Path
 
-# Get the directory containing this test file
-test_dir = Path(__file__).parent.absolute()
+test_dir = Path(__file__).parent
+weights_dir = Path("/srv")
 
-run = f"python3 {test_dir}/inference.py /srv/llama3.2-1b/model.safetensors /srv/llama3.2-1b/tokenizer.model --prompt_len 2048 --num_tokens 40"
 
-checks = []
+def generate_test_params():
+    prompt_lengths = [2048]
+    num_tokens_list = [40]
 
-metrics = [
-    ("Total", r"  Total time: (?P<metric>[\d\.e\+-]+) seconds"),
-    ("TTFT", r"  Prefill time: (?P<metric>[\d\.e\+-]+) seconds"),
-    ("TPS", r"  Tokens per second: (?P<metric>[\d\.e\+-]+)"),
-    ("Num Tokens", r"  Tokens generated: (?P<metric>[\d\.e\+-]+)"),
-]
+    params = []
+    names = []
+    for prompt_len in prompt_lengths:
+        for num_tokens in num_tokens_list:
+            params.append((prompt_len, num_tokens))
+            names.append(f"llama_3.2_1b_prompt_{prompt_len}_tokens_{num_tokens}")
+    return params, names
+
+
+params, names = generate_test_params()
+
+
+@pytest.mark.metrics(
+    TTFT=r"Prefill time: (?P<value>[\d\.e\+-]+) seconds",
+    TPS=r"Tokens per second: (?P<value>[\d\.e\+-]+)",
+    Num_Tokens=r"Tokens generated: (?P<value>[\d\.e\+-]+)",
+)
+@pytest.mark.parametrize("prompt_len,num_tokens", params, ids=names)
+def test_llama_3_2_1b(prompt_len, num_tokens):
+    command = f"python3 {test_dir}/inference.py {weights_dir}/llama3.2-1b/model.safetensors {weights_dir}/llama3.2-1b/tokenizer.model --prompt_len {prompt_len} --num_tokens {num_tokens}"
+
+    result = subprocess.run(
+        command,
+        cwd=test_dir,
+        shell=True,
+        capture_output=True,
+        text=True,
+        timeout=300,
+    )
+
+    assert (
+        result.returncode == 0
+    ), f"Command failed with return code {result.returncode}\nStderr: {result.stderr}"
+
+    print(result.stdout)