perf: make test suite snappier and clean up some testing things

.github/workflows/benchmarks.yaml

@@ -20,10 +20,14 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        python-version: ["3.12",]
+        python-version: ["3.12"]
 
     steps:
       - uses: actions/checkout@v6
+        with:
+          fetch-depth: 0
+          submodules: recursive
+          fetch-tags: true
 
       - name: Set up Python ${{ matrix.python-version }}
         uses: actions/setup-python@v6
@@ -35,11 +39,10 @@ jobs:
           python -m pip install --upgrade pip
           python -m pip install pytest pytest-codspeed
           python -m pip install .
-          git submodule update --init --recursive
 
       - name: Run benchmarks
         uses: CodSpeedHQ/action@v4
         with:
           token: ${{ secrets.CODSPEED_TOKEN }}
           run: pytest -vvs --codspeed
-          mode: instrumentation
+          mode: simulation

.github/workflows/python_package.yaml

@@ -21,6 +21,9 @@ jobs:
       fail-fast: false
       matrix:
         python-version: ["3.12"]
+        group: [1, 2, 3, 4]
+    env:
+      NUM_SPLITS: 4
 
     steps:
       - uses: actions/checkout@v6
@@ -37,18 +40,67 @@ jobs:
       - name: Install dependencies
         run: |
           python -m pip install --upgrade pip
-          python -m pip install pytest pytest-codspeed pytest-randomly
+          python -m pip install pytest pytest-codspeed pytest-randomly pytest-split
           python -m pip install .
 
+      - name: Restore test durations
+        uses: actions/cache/restore@cdf6c1fa76f9f475f3d7449005a359c84ca0f306  # v5.0.3
+        with:
+          path: .test_durations
+          key: test-durations-${{ github.ref }}-${{ github.sha }}
+          restore-keys: |
+            test-durations-${{ github.ref }}-
+            test-durations-
+
+      - name: Post-process test durations
+        run: |
+          if [[ -f .test_durations ]]; then
+            cp .test_durations .test_durations.${{ matrix.group }}
+            ls -lah .test_durations*
+            echo " "
+            cat .test_durations*
+          fi
+
       - name: Test with pytest
         run: |
-          pytest -vv --durations=100 --randomly-seed=42
+          pytest \
+            -vv \
+            --durations=100 \
+            --randomly-seed=42 \
+            --splits ${NUM_SPLITS} --group ${{ matrix.group }} \
+            --store-durations \
+            --durations-path=.test_durations.${{ matrix.group }} \
+            --splitting-algorithm least_duration \
+            --clean-durations
+
+      - name: Upload test durations
+        uses: actions/upload-artifact@v6
+        with:
+          name: test-durations-${{ matrix.group }}
+          path: .test_durations.${{ matrix.group }}
+          include-hidden-files: true
 
   build-status:
     needs: build
     runs-on: ubuntu-latest
 
     steps:
+      - name: Download test duration artifacts
+        uses: actions/download-artifact@v7
+        with:
+          pattern: test-durations-*
+
+      - name: Combine test durations
+        run: |
+          # see https://stackoverflow.com/a/71416016/1745538
+          jq 'reduce inputs as $i (.; . + $i)' test-durations-*/.test_durations.* > .test_durations
+
+      - name: Save test durations
+        uses: actions/cache/save@v5
+        with:
+          path: .test_durations
+          key: test-durations-${{ github.ref }}-${{ github.sha }}
+
       - name: check status
         run: |
           echo "Builds all passed!"

.pre-commit-config.yaml

@@ -13,7 +13,7 @@ repos:
   - repo: https://github.com/astral-sh/ruff-pre-commit
     rev: v0.14.14
     hooks:
-      - id: ruff
+      - id: ruff-check
         args: [ --fix ]
         exclude: tests/GalSim/|tests/Coord/|tests/jax/galsim/|dev/notebooks/
       - id: ruff-format

tests/conftest.py

@@ -18,6 +18,9 @@
 # Identify the path to this current file
 test_directory = os.path.dirname(os.path.abspath(__file__))
 
+# make the output directory for the test output
+os.makedirs("output", exist_ok=True)
+
 # Loading which tests to run
 with open(os.path.join(test_directory, "galsim_tests_config.yaml"), "r") as f:
     test_config = yaml.safe_load(f)

tests/jax/test_benchmarks.py

@@ -33,7 +33,7 @@ def _run():
     return tr.dt
 
 
-@pytest.mark.parametrize("kind", ["compile", "run"])
+@pytest.mark.parametrize("kind", ["run"])
 @pytest.mark.parametrize(
     "conserve_dc", [True, False], ids=["conserve_dc", "no_conserve_dc"]
 )
@@ -53,7 +53,7 @@ def test_benchmarks_lanczos_interp(benchmark, kind, conserve_dc, method):
     print(f"time: {dt:0.4g} ms", end=" ")
 
 
-@pytest.mark.parametrize("kind", ["compile", "run"])
+@pytest.mark.parametrize("kind", ["run"])
 def test_benchmarks_interpolated_image(benchmark, kind):
     gal = jgs.Gaussian(fwhm=1.2)
     im_gal = gal.drawImage(nx=32, ny=32, scale=0.2)
@@ -174,7 +174,7 @@ def _run_spergel_bench_conv(gsmod):
 _run_spergel_bench_conv_jit = jax.jit(partial(_run_spergel_bench_conv, jgs))
 
 
-@pytest.mark.parametrize("kind", ["compile", "run"])
+@pytest.mark.parametrize("kind", ["run"])
 def test_benchmark_spergel_conv(benchmark, kind):
     dt = _run_benchmarks(
         benchmark, kind, lambda: _run_spergel_bench_conv_jit().block_until_ready()
@@ -190,7 +190,7 @@ def _run_spergel_bench_xvalue(gsmod):
 _run_spergel_bench_xvalue_jit = jax.jit(partial(_run_spergel_bench_xvalue, jgs))
 
 
-@pytest.mark.parametrize("kind", ["compile", "run"])
+@pytest.mark.parametrize("kind", ["run"])
 def test_benchmark_spergel_xvalue(benchmark, kind):
     dt = _run_benchmarks(
         benchmark, kind, lambda: _run_spergel_bench_xvalue_jit().block_until_ready()
@@ -206,7 +206,7 @@ def _run_spergel_bench_kvalue(gsmod):
 _run_spergel_bench_kvalue_jit = jax.jit(partial(_run_spergel_bench_kvalue, jgs))
 
 
-@pytest.mark.parametrize("kind", ["compile", "run"])
+@pytest.mark.parametrize("kind", ["run"])
 def test_benchmark_spergel_kvalue(benchmark, kind):
     dt = _run_benchmarks(
         benchmark, kind, lambda: _run_spergel_bench_kvalue_jit().block_until_ready()
@@ -219,7 +219,7 @@ def _run_spergel_bench_init():
     return jgs.Spergel(nu=-0.6, half_light_radius=3.4).scale_radius
 
 
-@pytest.mark.parametrize("kind", ["compile", "run"])
+@pytest.mark.parametrize("kind", ["run"])
 def test_benchmark_spergel_init(benchmark, kind):
     dt = _run_benchmarks(
         benchmark, kind, lambda: _run_spergel_bench_init().block_until_ready()
@@ -232,7 +232,7 @@ def _run_gaussian_bench_init():
     return jgs.Gaussian(half_light_radius=3.4).sigma
 
 
-@pytest.mark.parametrize("kind", ["compile", "run"])
+@pytest.mark.parametrize("kind", ["run"])
 def test_benchmark_gaussian_init(benchmark, kind):
     dt = _run_benchmarks(
         benchmark, kind, lambda: _run_gaussian_bench_init().block_until_ready()
@@ -247,7 +247,7 @@ def _run_benchmark_interpimage_flux_frac(img):
     return jgs.interpolatedimage._flux_frac(img.array, x, y, cenx, ceny)
 
 
-@pytest.mark.parametrize("kind", ["compile", "run"])
+@pytest.mark.parametrize("kind", ["run"])
 def test_benchmark_interpimage_flux_frac(benchmark, kind):
     obj = jgs.Gaussian(half_light_radius=0.9).shear(g1=0.1, g2=0.2)
     img = obj.drawImage(nx=55, ny=55, scale=0.2, method="no_pixel")
@@ -265,7 +265,7 @@ def _run_benchmark_rng_discard(rng):
     return rng._state.key
 
 
-@pytest.mark.parametrize("kind", ["compile", "run"])
+@pytest.mark.parametrize("kind", ["run"])
 def test_benchmark_rng_discard(benchmark, kind):
     rng = jgs.BaseDeviate(seed=42)
     dt = _run_benchmarks(
@@ -278,7 +278,7 @@ def _run_benchmark_invert_ab_noraise(u, v, ab):
     return jgs.fitswcs._invert_ab_noraise(u, v, ab)[0]
 
 
-@pytest.mark.parametrize("kind", ["compile", "run"])
+@pytest.mark.parametrize("kind", ["run"])
 def test_benchmark_invert_ab_noraise(benchmark, kind):
     u = jnp.arange(1000).astype(jnp.float64)
     v = jnp.arange(1000).astype(jnp.float64)
@@ -295,7 +295,7 @@ def _run_benchmark_moffat_init():
     return jgs.Moffat(beta=2.5, half_light_radius=0.6, trunc=1.2).scale_radius
 
 
-@pytest.mark.parametrize("kind", ["compile", "run"])
+@pytest.mark.parametrize("kind", ["run"])
 def test_benchmark_moffat_init(benchmark, kind):
     dt = _run_benchmarks(
         benchmark, kind, lambda: _run_benchmark_moffat_init().block_until_ready()
@@ -307,7 +307,7 @@ def _run_benchmark_spergel_calcfluxrad():
     return jgs.spergel.calculateFluxRadius(1e-10, 2.0)
 
 
-@pytest.mark.parametrize("kind", ["compile", "run"])
+@pytest.mark.parametrize("kind", ["run"])
 def test_benchmark_spergel_calcfluxrad(benchmark, kind):
     dt = _run_benchmarks(
         benchmark,

tests/jax/test_image_wrapping.py

@@ -1,7 +1,7 @@
 import jax
 import numpy as np
 import pytest
-from galsim_test_helpers import timer
+from galsim.utilities import timer
 
 import jax_galsim as galsim
 from jax_galsim.core.wrap_image import (

tests/jax/test_interpolant_jax.py

@@ -11,7 +11,8 @@
 import jax
 import numpy as np
 import pytest
-from galsim_test_helpers import assert_raises, timer
+from galsim.utilities import timer
+from numpy.testing import assert_raises
 from scipy.special import sici
 
 import jax_galsim as galsim

tests/jax/test_spergel_comp_galsim.py

@@ -113,7 +113,7 @@ def _run():
 
 
 @pytest.mark.parametrize("kind", ["compile", "run"])
-def test_spergel_comp_galsim_perf_conv(benchmark, kind):
+def test_spergel_comp_galsim_perf_conv(kind):
     dt = _run_time_test(kind, lambda: _run_spergel_bench_conv_jit().block_until_ready())
     print(f"\njax-galsim time: {dt:0.4g} ms")
 
@@ -125,7 +125,7 @@ def test_spergel_comp_galsim_perf_conv(benchmark, kind):
 
 
 @pytest.mark.parametrize("kind", ["compile", "run"])
-def test_spergel_comp_galsim_perf_kvalue(benchmark, kind):
+def test_spergel_comp_galsim_perf_kvalue(kind):
     dt = _run_time_test(
         kind, lambda: _run_spergel_bench_kvalue_jit().block_until_ready()
     )
@@ -139,7 +139,7 @@ def test_spergel_comp_galsim_perf_kvalue(benchmark, kind):
 
 
 @pytest.mark.parametrize("kind", ["compile", "run"])
-def test_spergel_comp_galsim_perf_xvalue(benchmark, kind):
+def test_spergel_comp_galsim_perf_xvalue(kind):
     dt = _run_time_test(
         kind, lambda: _run_spergel_bench_xvalue_jit().block_until_ready()
     )