Implemented a new way of recording properties that works even if test is skipped

.github/workflows/build-and-test.yml

@@ -250,7 +250,7 @@ jobs:
       run: |
         export LD_LIBRARY_PATH="/opt/ttmlir-toolchain/lib/:${{ steps.strings.outputs.install-output-dir }}/lib:${LD_LIBRARY_PATH}"
         source venv/activate
-        pytest ./tests \
+        pytest ./tests/jax \
                -m "${{ inputs.test_mark }}" \
                --junitxml=${{ steps.strings.outputs.test_report_path }} \
                2>&1 | tee pytest.log

tests/conftest.py

@@ -2,101 +2,67 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
-from datetime import datetime
-from enum import Enum
-from typing import Callable
-
 import pytest
 
 
-class RecordProperties(Enum):
-    """Properties we can record."""
-
-    # Timestamp of test start.
-    START_TIMESTAMP = "start_timestamp"
-    # Timestamp of test end.
-    END_TIMESTAMP = "end_timestamp"
-    # Frontend or framework used to run the test.
-    FRONTEND = "frontend"
-    # Kind of operation. e.g. eltwise.
-    OP_KIND = "op_kind"
-    # Name of the operation in the framework. e.g. torch.conv2d.
-    FRAMEWORK_OP_NAME = "framework_op_name"
-    # Name of the operation. e.g. ttir.conv2d.
-    OP_NAME = "op_name"
-    # Name of the model in which this op appears.
-    MODEL_NAME = "model_name"
-
-
-@pytest.fixture(scope="function", autouse=True)
-def record_test_timestamp(record_property: Callable):
+def pytest_configure(config: pytest.Config):
     """
-    Autouse fixture used to capture execution time of a test.
-
-    Parameters:
-    ----------
-    record_property: Callable
-        A pytest built-in function used to record test metadata, such as custom
-        properties or additional information about the test execution.
-
-    Yields:
-    -------
-    Callable
-        The `record_property` callable, allowing tests to add additional properties if
-        needed.
-
-
-    Example:
-    --------
-    ```
-    def test_model(fixture1, fixture2, ..., record_tt_xla_property):
-        record_tt_xla_property("key", value)
-
-        # Test logic...
-    ```
+    Registers custom pytest marker `record_properties(key1=val1, key2=val2, ...)`.
+
+    Allowed keys are ["test_category", "jax_op_name", "op_name", "model_name"].
+        - `test_category`: one of ["op_test", "graph_test", "model_test", "multichip_test", "other"]
+        - `jax_op_name`: name of the operation in jax, e.g. `jax.numpy.exp`
+        - `shlo_op_name`: name of the matching stablehlo operation
+        - `model_name`: name of the model under test (if recorded from a model test, or op
+          under test comes from some model and we want to note that in the report)
+        - `run_mode`: one of ["inference", "training"]. Only exists for model tests.
+
+    These are used to tag the function under test with properties which will be dumped
+    to the final XML test report. These reports get picked up by other CI workflows and
+    are used to display state of tests on a dashboard.
     """
-    start_timestamp = datetime.strftime(datetime.now(), "%Y-%m-%dT%H:%M:%S%z")
-    record_property(RecordProperties.START_TIMESTAMP.value, start_timestamp)
+    config.addinivalue_line(
+        "markers",
+        "record_properties(key_value_pairs): Record custom properties for the test",
+    )
 
-    # Run the test.
-    yield
 
-    end_timestamp = datetime.strftime(datetime.now(), "%Y-%m-%dT%H:%M:%S%z")
-    record_property(RecordProperties.END_TIMESTAMP.value, end_timestamp)
-
-
-@pytest.fixture(scope="function", autouse=True)
-def record_tt_xla_property(record_property: Callable):
+def pytest_collection_modifyitems(items):
     """
-    Autouse fixture that automatically records some test properties for each test
-    function.
-
-    It also yields back callable which can be explicitly used in tests to record
-    additional properties.
-
-    Example:
-
-    ```
-    def test_model(fixture1, fixture2, ..., record_tt_xla_property):
-        record_tt_xla_property("key", value)
-
-        # Test logic...
-    ```
-
-    Parameters:
-    ----------
-    record_property: Callable
-        A pytest built-in function used to record test metadata, such as custom
-        properties or additional information about the test execution.
-
-    Yields:
-    -------
-    Callable
-        The `record_property` callable, allowing tests to add additional properties if
-        needed.
+    Pytest hook to process the custom marker and attach recorder properties to the test.
     """
-    # Record default properties for tt-xla.
-    record_property(RecordProperties.FRONTEND.value, "tt-xla")
 
-    # Run the test.
-    yield record_property
+    def validate_keys(keys):
+        valid_keys = [
+            "test_category",
+            "jax_op_name",
+            "shlo_op_name",
+            "model_name",
+            "run_mode",
+        ]
+
+        if not all(key in valid_keys for key in keys):
+            raise KeyError(
+                f"Invalid keys found in 'record_properties' marker: {', '.join(keys)}. "
+                f"Allowed keys are: {', '.join(valid_keys)}"
+            )
+
+    for item in items:
+        # Add some test metadata in a 'tags' dictionary.
+        tags = {"test_name": item.originalname, "specific_test_case": item.name}
+
+        # Look for the custom marker.
+        properties_marker = item.get_closest_marker(name="record_properties")
+
+        if properties_marker:
+            # Extract the key-value pairs passed to the marker.
+            properties: dict = properties_marker.kwargs
+            # Validate that only allowed keys are used.
+            validate_keys(properties.keys())
+
+            # Tag them.
+            for key, value in properties.items():
+                tags[key] = value
+
+        # Attach the tags dictionary as a single property.
+        item.user_properties.append(("tags", tags))

tests/jax/graphs/test_MLP_regression.py

@@ -17,6 +17,7 @@ def comparison_config() -> ComparisonConfig:
 
 @pytest.mark.push
 @pytest.mark.nightly
+@pytest.mark.record_properties(test_category="graph_test")
 @pytest.mark.parametrize(
     ["W1", "b1", "W2", "b2", "X", "y"],
     [

tests/jax/graphs/test_activation_functions.py

@@ -10,6 +10,7 @@
 
 @pytest.mark.push
 @pytest.mark.nightly
+@pytest.mark.record_properties(test_category="graph_test")
 @pytest.mark.parametrize("x_shape", [(32, 32), (64, 64)])
 def test_relu(x_shape: tuple):
     """Test ReLU activation function."""

tests/jax/graphs/test_example_graph.py

@@ -17,6 +17,7 @@ def example_graph(x: jax.Array, y: jax.Array) -> jax.Array:
 
 @pytest.mark.push
 @pytest.mark.nightly
+@pytest.mark.record_properties(test_category="graph_test")
 @pytest.mark.parametrize(
     ["x_shape", "y_shape"],
     [

tests/jax/graphs/test_linear_transformation.py

@@ -10,6 +10,7 @@
 
 @pytest.mark.push
 @pytest.mark.nightly
+@pytest.mark.record_properties(test_category="graph_test")
 @pytest.mark.parametrize(
     ["x_shape", "y_shape", "bias_shape"],
     [

tests/jax/graphs/test_simple_gradient.py

@@ -9,6 +9,7 @@
 
 @pytest.mark.push
 @pytest.mark.nightly
+@pytest.mark.record_properties(test_category="graph_test")
 @pytest.mark.parametrize("x_shape", [(32, 32), (64, 64)])
 def test_simple_gradient(x_shape: tuple):
     def simple_gradient(x: jax.Array):

tests/jax/graphs/test_simple_regression.py

@@ -9,6 +9,7 @@
 
 @pytest.mark.push
 @pytest.mark.nightly
+@pytest.mark.record_properties(test_category="graph_test")
 @pytest.mark.parametrize(
     ["weights", "bias", "X", "y"], [[(1, 2), (1, 1), (2, 1), (1, 1)]]
 )

tests/jax/graphs/test_softmax.py

@@ -9,6 +9,7 @@
 
 @pytest.mark.push
 @pytest.mark.nightly
+@pytest.mark.record_properties(test_category="graph_test")
 @pytest.mark.parametrize(
     ["x_shape", "axis"],
     [

tests/jax/models/albert/v2/base/test_albert_base.py

@@ -2,11 +2,9 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
-from typing import Callable
-
 import pytest
 from infra import RunMode
-from utils import record_model_test_properties, runtime_fail
+from utils import runtime_fail
 
 from ..tester import AlbertV2Tester
 
@@ -31,6 +29,11 @@ def training_tester() -> AlbertV2Tester:
 
 
 @pytest.mark.nightly
+@pytest.mark.record_properties(
+    test_category="model_test",
+    model_name=MODEL_NAME,
+    run_mode=RunMode.INFERENCE.value,
+)
 @pytest.mark.xfail(
     reason=(
         runtime_fail(
@@ -39,21 +42,16 @@ def training_tester() -> AlbertV2Tester:
         )
     )
 )
-def test_flax_albert_v2_base_inference(
-    inference_tester: AlbertV2Tester,
-    record_tt_xla_property: Callable,
-):
-    record_model_test_properties(record_tt_xla_property, MODEL_NAME)
-
+def test_flax_albert_v2_base_inference(inference_tester: AlbertV2Tester):
     inference_tester.test()
 
 
 @pytest.mark.nightly
+@pytest.mark.record_properties(
+    test_category="model_test",
+    model_name=MODEL_NAME,
+    run_mode=RunMode.TRAINING.value,
+)
 @pytest.mark.skip(reason="Support for training not implemented")
-def test_flax_albert_v2_base_training(
-    training_tester: AlbertV2Tester,
-    record_tt_xla_property: Callable,
-):
-    record_model_test_properties(record_tt_xla_property, MODEL_NAME)
-
+def test_flax_albert_v2_base_training(training_tester: AlbertV2Tester):
     training_tester.test()

tests/jax/models/albert/v2/large/test_albert_large.py

@@ -2,11 +2,9 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
-from typing import Callable
-
 import pytest
 from infra import RunMode
-from utils import record_model_test_properties, runtime_fail
+from utils import runtime_fail
 
 from ..tester import AlbertV2Tester
 
@@ -31,6 +29,11 @@ def training_tester() -> AlbertV2Tester:
 
 
 @pytest.mark.nightly
+@pytest.mark.record_properties(
+    test_category="model_test",
+    model_name=MODEL_NAME,
+    run_mode=RunMode.INFERENCE.value,
+)
 @pytest.mark.xfail(
     reason=(
         runtime_fail(
@@ -39,21 +42,16 @@ def training_tester() -> AlbertV2Tester:
         )
     )
 )
-def test_flax_albert_v2_large_inference(
-    inference_tester: AlbertV2Tester,
-    record_tt_xla_property: Callable,
-):
-    record_model_test_properties(record_tt_xla_property, MODEL_NAME)
-
+def test_flax_albert_v2_large_inference(inference_tester: AlbertV2Tester):
     inference_tester.test()
 
 
 @pytest.mark.nightly
+@pytest.mark.record_properties(
+    test_category="model_test",
+    model_name=MODEL_NAME,
+    run_mode=RunMode.TRAINING.value,
+)
 @pytest.mark.skip(reason="Support for training not implemented")
-def test_flax_albert_v2_large_training(
-    training_tester: AlbertV2Tester,
-    record_tt_xla_property: Callable,
-):
-    record_model_test_properties(record_tt_xla_property, MODEL_NAME)
-
+def test_flax_albert_v2_large_training(training_tester: AlbertV2Tester):
     training_tester.test()

tests/jax/models/albert/v2/xlarge/test_albert_xlarge.py

@@ -2,11 +2,9 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
-from typing import Callable
-
 import pytest
 from infra import RunMode
-from utils import record_model_test_properties, runtime_fail
+from utils import runtime_fail
 
 from ..tester import AlbertV2Tester
 
@@ -31,6 +29,11 @@ def training_tester() -> AlbertV2Tester:
 
 
 @pytest.mark.nightly
+@pytest.mark.record_properties(
+    test_category="model_test",
+    model_name=MODEL_NAME,
+    run_mode=RunMode.INFERENCE.value,
+)
 @pytest.mark.xfail(
     reason=(
         runtime_fail(
@@ -39,21 +42,16 @@ def training_tester() -> AlbertV2Tester:
         )
     )
 )
-def test_flax_albert_v2_xlarge_inference(
-    inference_tester: AlbertV2Tester,
-    record_tt_xla_property: Callable,
-):
-    record_model_test_properties(record_tt_xla_property, MODEL_NAME)
-
+def test_flax_albert_v2_xlarge_inference(inference_tester: AlbertV2Tester):
     inference_tester.test()
 
 
 @pytest.mark.nightly
+@pytest.mark.record_properties(
+    test_category="model_test",
+    model_name=MODEL_NAME,
+    run_mode=RunMode.TRAINING.value,
+)
 @pytest.mark.skip(reason="Support for training not implemented")
-def test_flax_albert_v2_xlarge_training(
-    training_tester: AlbertV2Tester,
-    record_tt_xla_property: Callable,
-):
-    record_model_test_properties(record_tt_xla_property, MODEL_NAME)
-
+def test_flax_albert_v2_xlarge_training(training_tester: AlbertV2Tester):
     training_tester.test()