wip: small changes

.github/workflows/update_regression_baseline.yml

@@ -33,7 +33,7 @@ jobs:
       run: |
         git config --global user.name 'GitHub Action'
         git config --global user.email '[email protected]'
-        UPDATE_BASELINE=1 pytest tests/test_regression.py
+        NEW_BASELINE=1 pytest tests/test_regression.py
         
     - name: Commit and push if baseline changed
       if: github.event.pull_request.base.ref == 'main'

tests/test_regression.py

@@ -30,14 +30,16 @@
 
 
 pytestmark = pytest.mark.regression  # mark all tests as regression tests in this file
-UPDATE_BASELINE = (
-    os.environ["UPDATE_BASELINE"] if "UPDATE_BASELINE" in os.environ else 0
-)
+NEW_BASELINE = os.environ["NEW_BASELINE"] if "NEW_BASELINE" in os.environ else 0
+dirname = os.path.dirname(__file__)
+fpath_baselines = os.path.join(dirname, "regression_test_baselines.json")
+fpath_results = os.path.join(dirname, "regression_test_results.json")
+
 tolerance = 0.2
 
 baselines = {}
-if os.path.exists("tests/regression_test_baselines.json"):
-    with open("tests/regression_test_baselines.json") as f:
+if os.path.exists(fpath_baselines):
+    with open(fpath_baselines) as f:
         baselines = json.load(f)
 
 
@@ -63,9 +65,15 @@ def generate_regression_report(base_results, new_results):
             base_time = base_runtimes.get(key)
             diff = None if base_time is None else ((new_time - base_time) / base_time)
 
-            emoji = "🆕" if diff is None else "⚪"
-            emoji = "🔴" if diff > tolerance else emoji
-            emoji = "🟢" if diff < 0 else emoji
+            emoji = ""
+            if diff is None:
+                emoji = "🆕"
+            elif diff > tolerance:
+                emoji = "🔴"
+            elif diff < 0:
+                emoji = "🟢"
+            else:
+                emoji = "⚪"
 
             time_str = (
                 f"({new_time:.6f}s)"
@@ -113,20 +121,16 @@ def test_wrapper(**kwargs):
             key = generate_unique_key(header)
 
             runs = []
-            num_iters = self.baseline_iters if UPDATE_BASELINE else self.test_iters
+            num_iters = self.baseline_iters if NEW_BASELINE else self.test_iters
             for _ in range(num_iters):
                 runtimes = func(**kwargs)
                 runs.append(runtimes)
             runtimes = {k: np.mean([d[k] for d in runs]) for k in runs[0]}
 
-            fpath = (
-                "tests/regression_test_results.json"
-                if not UPDATE_BASELINE
-                else "tests/regression_test_baselines.json"
-            )
+            fpath = fpath_results if not NEW_BASELINE else fpath_baselines
             append_to_json(fpath, header["test_name"], header["input_kwargs"], runtimes)
 
-            if not UPDATE_BASELINE:
+            if not NEW_BASELINE:
                 assert key in baselines, f"No basline found for {header}"
                 func_baselines = baselines[key]["runtimes"]
                 for key, baseline in func_baselines.items():
@@ -232,4 +236,4 @@ def simulate(params):
     start_time = time.time()
     _ = jitted_simulate(params).block_until_ready()
     runtimes["run_time"] = time.time() - start_time
-    return runtimes
+    return runtimes  # @compare_to_baseline decorator will compare this to the baseline