From 9171a74d729e00a180f4c312c3b1efaca9d77070 Mon Sep 17 00:00:00 2001 From: Roman Cattaneo <1116746+romanc@users.noreply.github.com> Date: Fri, 27 Jun 2025 09:59:26 +0200 Subject: [PATCH 1/6] Re-enable caching tests --- tests/dsl/test_caches.py | 113 ++++++++++++++++++++++----------------- 1 file changed, 64 insertions(+), 49 deletions(-) diff --git a/tests/dsl/test_caches.py b/tests/dsl/test_caches.py index d550c123..8f25152d 100644 --- a/tests/dsl/test_caches.py +++ b/tests/dsl/test_caches.py @@ -1,5 +1,7 @@ import os import shutil +import sys +from pathlib import Path import pytest from gt4py.cartesian import config as gt_config @@ -16,11 +18,12 @@ from ndsl.comm.mpi import MPI from ndsl.dsl.dace.orchestration import orchestrate from ndsl.dsl.gt4py import PARALLEL, Field, computation, interval +from ndsl.dsl.stencil import CompareToNumpyStencil, FrozenStencil def _make_storage( func, - grid_indexing, + grid_indexing: GridIndexing, stencil_config: StencilConfig, *, dtype=float, @@ -39,7 +42,9 @@ def _stencil(inp: Field[float], out: Field[float], scalar: float): out = inp -def _build_stencil(backend, orchestrated: DaCeOrchestration): +def _build_stencil( + backend: str, orchestrated: DaCeOrchestration +) -> tuple[FrozenStencil | CompareToNumpyStencil, GridIndexing, StencilConfig]: # Make stencil and verify it ran grid_indexing = GridIndexing( domain=(5, 5, 5), @@ -65,7 +70,7 @@ def _build_stencil(backend, orchestrated: DaCeOrchestration): class OrchestratedProgram: - def __init__(self, backend, orchestration): + def __init__(self, backend, orchestration: DaCeOrchestration): self.stencil, grid_indexing, stencil_config = _build_stencil( backend, orchestration ) @@ -77,52 +82,49 @@ def __call__(self): self.stencil(self.inp, self.out, self.inp[0, 0, 0]) -@pytest.mark.parametrize( - "backend", - [ - pytest.param("dace:cpu"), - ], -) +@pytest.mark.parametrize("backend", [pytest.param("dace:cpu")]) @pytest.mark.skipif( - MPI is not None, reason="relocatibility checked with a one-rank setup" + MPI.COMM_WORLD.Get_size() > 1, reason="relocatibility checked with a one-rank setup" ) -def test_relocatability_orchestration(backend): - original_root_directory = gt_config.cache_settings["root_path"] - working_dir = str(os.getcwd()) +def test_relocatability_orchestration(backend: str) -> None: + original_root_directory: str = gt_config.cache_settings["root_path"] + cwd = Path.cwd() # Compile on default p0 = OrchestratedProgram(backend, DaCeOrchestration.BuildAndRun) p0() - assert os.path.exists( - f"{working_dir}/.gt_cache_FV3_A/dacecache/" - "test_caches_OrchestratedProgam___call__", - ) or os.path.exists( - f"{working_dir}/.gt_cache_FV3_A/dacecache/OrchestratedProgam___call__", + expected_cache_dir = ( + cwd + / ".gt_cache_FV3_A" + / "dacecache" + / "test_caches_OrchestratedProgram___call__" ) + assert expected_cache_dir.exists() # Compile in another directory - custom_path = f"{working_dir}/.my_cache_path" + custom_path = cwd / ".my_cache_path" gt_config.cache_settings["root_path"] = custom_path p1 = OrchestratedProgram(backend, DaCeOrchestration.BuildAndRun) p1() - assert os.path.exists( - f"{custom_path}/.gt_cache_FV3_A/dacecache/" - "test_caches_OrchestratedProgam___call__", - ) or os.path.exists( - f"{working_dir}/.gt_cache_FV3_A/dacecache/OrchestratedProgam___call__", + expected_cache_dir = ( + custom_path + / ".gt_cache_FV3_A" + / "dacecache" + / "test_caches_OrchestratedProgram___call__" ) + assert expected_cache_dir.exists() # Check relocability by copying the second cache directory, # changing the path of gt_config.cache_settings and trying to Run on it - relocated_path = f"{working_dir}/.my_relocated_cache_path" + relocated_path = cwd / ".my_relocated_cache_path" shutil.copytree(custom_path, relocated_path, dirs_exist_ok=True) gt_config.cache_settings["root_path"] = relocated_path p2 = OrchestratedProgram(backend, DaCeOrchestration.Run) p2() # Generate a file exists error to check for bad path - bogus_path = "./nope/notatall/nothappening" + bogus_path = "./nope/not_at_all/not_happening" gt_config.cache_settings["root_path"] = bogus_path with pytest.raises(RuntimeError): OrchestratedProgram(backend, DaCeOrchestration.Run) @@ -131,49 +133,62 @@ def test_relocatability_orchestration(backend): gt_config.cache_settings["root_path"] = original_root_directory -@pytest.mark.parametrize( - "backend", - [ - pytest.param("dace:cpu"), - ], -) +@pytest.mark.parametrize("backend", [pytest.param("dace:cpu")]) @pytest.mark.skipif( - MPI is not None, reason="relocatibility checked with a one-rank setup" + MPI.COMM_WORLD.Get_size() > 1, reason="relocatibility checked with a one-rank setup" ) -def test_relocatability(backend: str): +def test_relocatability(backend: str) -> None: # Restore original dir name gt_config.cache_settings["dir_name"] = os.environ.get( "GT_CACHE_DIR_NAME", f".gt_cache_{MPI.COMM_WORLD.Get_rank():06}" ) + cwd = Path.cwd() backend_sanitized = backend.replace(":", "") + python_version = f"py{sys.version_info[0]}{sys.version_info[1]}" + expected_cache_path = ( + cwd + / ".gt_cache_000000" + / f"{python_version}_1013" + / f"{backend_sanitized}" + / "test_caches" + / "_stencil" + ) # Compile on default p0 = OrchestratedProgram(backend, DaCeOrchestration.Python) p0() - assert os.path.exists( - f"./.gt_cache_000000/py38_1013/{backend_sanitized}/test_caches/_stencil/" - ) + assert expected_cache_path.exists() # Compile in another directory - custom_path = "./.my_cache_path" + custom_path = cwd / ".my_cache_path" + expected_cache_path = ( + custom_path + / ".gt_cache_000000" + / f"{python_version}_1013" + / f"{backend_sanitized}" + / "test_caches" + / "_stencil" + ) gt_config.cache_settings["root_path"] = custom_path p1 = OrchestratedProgram(backend, DaCeOrchestration.Python) p1() - assert os.path.exists( - f"{custom_path}/.gt_cache_000000/py38_1013/{backend_sanitized}" - "/test_caches/_stencil/" - ) + assert expected_cache_path.exists() - # Check relocability by copying the second cache directory, + # Check relocability by copying the first cache directory, # changing the path of gt_config.cache_settings and trying to Run on it - relocated_path = "./.my_relocated_cache_path" - shutil.copytree("./.gt_cache_000000", relocated_path, dirs_exist_ok=True) + relocated_path = cwd / ".my_relocated_cache_path" + shutil.copytree(cwd / ".gt_cache_000000", relocated_path, dirs_exist_ok=True) gt_config.cache_settings["root_path"] = relocated_path + expected_cache_path = ( + relocated_path + / ".gt_cache_000000" + / f"{python_version}_1013" + / f"{backend_sanitized}" + / "test_caches" + / "_stencil" + ) p2 = OrchestratedProgram(backend, DaCeOrchestration.Python) p2() - assert os.path.exists( - f"{relocated_path}/.gt_cache_000000/py38_1013/{backend_sanitized}" - "/test_caches/_stencil/" - ) + assert expected_cache_path.exists() From a261f633348d3dedeb925abab85c6a1d7110d18d Mon Sep 17 00:00:00 2001 From: Roman Cattaneo <1116746+romanc@users.noreply.github.com> Date: Fri, 27 Jun 2025 10:08:49 +0200 Subject: [PATCH 2/6] Type hints in tests/dsl/compilation_config --- tests/dsl/test_compilation_config.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/tests/dsl/test_compilation_config.py b/tests/dsl/test_compilation_config.py index fa323b06..326eb222 100644 --- a/tests/dsl/test_compilation_config.py +++ b/tests/dsl/test_compilation_config.py @@ -30,7 +30,7 @@ def test_safety_checks(): ) def test_check_communicator_valid( size: int, use_minimal_caching: bool, run_mode: RunMode -): +) -> None: partitioner = CubedSpherePartitioner( TilePartitioner((int(sqrt(size / 6)), int((sqrt(size / 6))))) ) @@ -50,7 +50,7 @@ def test_check_communicator_valid( ) def test_check_communicator_invalid( nx: int, ny: int, use_minimal_caching: bool, run_mode: RunMode -): +) -> None: partitioner = CubedSpherePartitioner(TilePartitioner((nx, ny))) comm = NullComm(rank=0, total_ranks=nx * ny * 6) cubed_sphere_comm = CubedSphereCommunicator(comm, partitioner) @@ -61,7 +61,7 @@ def test_check_communicator_invalid( config.check_communicator(cubed_sphere_comm) -def test_get_decomposition_info_from_no_comm(): +def test_get_decomposition_info_from_no_comm() -> None: config = CompilationConfig() ( computed_rank, @@ -86,7 +86,7 @@ def test_get_decomposition_info_from_no_comm(): ) def test_get_decomposition_info_from_comm( rank: int, size: int, is_compiling: bool, equivalent: int -): +) -> None: partitioner = CubedSpherePartitioner( TilePartitioner((int(sqrt(size / 6)), int(sqrt(size / 6)))) ) @@ -123,8 +123,8 @@ def test_get_decomposition_info_from_comm( ], ) def test_determine_compiling_equivalent( - rank, size, minimal_caching, run_mode, equivalent -): + rank: int, size: int, minimal_caching: bool, run_mode: RunMode, equivalent: int +) -> None: config = CompilationConfig(use_minimal_caching=minimal_caching, run_mode=run_mode) partitioner = CubedSpherePartitioner( TilePartitioner((sqrt(size / 6), sqrt(size / 6))) @@ -138,7 +138,7 @@ def test_determine_compiling_equivalent( ) -def test_as_dict(): +def test_as_dict() -> None: config = CompilationConfig() asdict = config.as_dict() assert asdict["backend"] == "numpy" @@ -151,7 +151,7 @@ def test_as_dict(): assert len(asdict) == 7 -def test_from_dict(): +def test_from_dict() -> None: specification_dict = {} config = CompilationConfig.from_dict(specification_dict) assert config.backend == "numpy" From 7f6b0ba8452a7c9f9bc9f734f04cc6567db814ad Mon Sep 17 00:00:00 2001 From: Roman Cattaneo <1116746+romanc@users.noreply.github.com> Date: Fri, 27 Jun 2025 10:42:03 +0200 Subject: [PATCH 3/6] more test cleanups --- tests/dsl/test_caches.py | 21 +------ tests/dsl/test_dace_config.py | 25 ++++---- tests/dsl/test_skip_passes.py | 2 +- tests/dsl/test_stencil.py | 23 ++----- tests/dsl/test_stencil_config.py | 100 +++++++++++------------------- tests/dsl/test_stencil_factory.py | 32 ++++++---- tests/dsl/test_stencil_wrapper.py | 65 ++++++++----------- 7 files changed, 106 insertions(+), 162 deletions(-) diff --git a/tests/dsl/test_caches.py b/tests/dsl/test_caches.py index 8f25152d..677bc59a 100644 --- a/tests/dsl/test_caches.py +++ b/tests/dsl/test_caches.py @@ -19,22 +19,7 @@ from ndsl.dsl.dace.orchestration import orchestrate from ndsl.dsl.gt4py import PARALLEL, Field, computation, interval from ndsl.dsl.stencil import CompareToNumpyStencil, FrozenStencil - - -def _make_storage( - func, - grid_indexing: GridIndexing, - stencil_config: StencilConfig, - *, - dtype=float, - aligned_index=(0, 0, 0), -): - return func( - backend=stencil_config.compilation_config.backend, - shape=grid_indexing.domain, - dtype=dtype, - aligned_index=aligned_index, - ) +from tests.dsl import utils def _stencil(inp: Field[float], out: Field[float], scalar: float): @@ -75,8 +60,8 @@ def __init__(self, backend, orchestration: DaCeOrchestration): backend, orchestration ) orchestrate(obj=self, config=stencil_config.dace_config) - self.inp = _make_storage(ones, grid_indexing, stencil_config, dtype=float) - self.out = _make_storage(empty, grid_indexing, stencil_config, dtype=float) + self.inp = utils.make_storage(ones, grid_indexing, stencil_config, dtype=float) + self.out = utils.make_storage(empty, grid_indexing, stencil_config, dtype=float) def __call__(self): self.stencil(self.inp, self.out, self.inp[0, 0, 0]) diff --git a/tests/dsl/test_dace_config.py b/tests/dsl/test_dace_config.py index c044cb16..e89b69c2 100644 --- a/tests/dsl/test_dace_config.py +++ b/tests/dsl/test_dace_config.py @@ -1,6 +1,7 @@ import unittest.mock from ndsl import CubedSpherePartitioner, DaceConfig, DaCeOrchestration, TilePartitioner +from ndsl.comm.partitioner import Partitioner from ndsl.dsl.dace.dace_config import _determine_compiling_ranks from ndsl.dsl.dace.orchestration import orchestrate, orchestrate_function @@ -11,8 +12,8 @@ """ -def test_orchestrate_function_calls_dace(): - def foo(): +def test_orchestrate_function_calls_dace() -> None: + def foo() -> None: pass dace_config = DaceConfig( @@ -29,8 +30,8 @@ def foo(): assert mock_call_sdfg.call_args.args[0].f == foo -def test_orchestrate_function_does_not_call_dace(): - def foo(): +def test_orchestrate_function_does_not_call_dace() -> None: + def foo() -> None: pass dace_config = DaceConfig( @@ -46,7 +47,7 @@ def foo(): assert not mock_call_sdfg.called -def test_orchestrate_calls_dace(): +def test_orchestrate_calls_dace() -> None: dace_config = DaceConfig( communicator=None, backend="gtc:dace", @@ -54,10 +55,10 @@ def test_orchestrate_calls_dace(): ) class A: - def __init__(self): + def __init__(self) -> None: orchestrate(obj=self, config=dace_config, method_to_orchestrate="foo") - def foo(self): + def foo(self) -> None: pass with unittest.mock.patch( @@ -68,7 +69,7 @@ def foo(self): assert mock_call_sdfg.called -def test_orchestrate_does_not_call_dace(): +def test_orchestrate_does_not_call_dace() -> None: dace_config = DaceConfig( communicator=None, backend="gtc:dace", @@ -76,10 +77,10 @@ def test_orchestrate_does_not_call_dace(): ) class A: - def __init__(self): + def __init__(self) -> None: orchestrate(obj=self, config=dace_config, method_to_orchestrate="foo") - def foo(self): + def foo(self) -> None: pass with unittest.mock.patch( @@ -90,14 +91,14 @@ def foo(self): assert not mock_call_sdfg.called -def test_orchestrate_distributed_build(): +def test_orchestrate_distributed_build() -> None: dummy_dace_config = DaceConfig( communicator=None, backend="gtc:dace", orchestration=DaCeOrchestration.BuildAndRun, ) - def _does_compile(rank, partitioner) -> bool: + def _does_compile(rank: int, partitioner: Partitioner) -> bool: dummy_dace_config.layout = partitioner.layout dummy_dace_config.rank_size = partitioner.layout[0] * partitioner.layout[1] * 6 dummy_dace_config.my_rank = rank diff --git a/tests/dsl/test_skip_passes.py b/tests/dsl/test_skip_passes.py index 22b840cb..6ffb0c95 100644 --- a/tests/dsl/test_skip_passes.py +++ b/tests/dsl/test_skip_passes.py @@ -23,7 +23,7 @@ def stencil_definition(a: FloatField): a = 0.0 -def test_skip_passes_becomes_oir_pipeline(): +def test_skip_passes_becomes_oir_pipeline() -> None: backend = "numpy" dace_config = DaceConfig(None, backend) config = StencilConfig( diff --git a/tests/dsl/test_stencil.py b/tests/dsl/test_stencil.py index 5348f346..19742b49 100644 --- a/tests/dsl/test_stencil.py +++ b/tests/dsl/test_stencil.py @@ -2,25 +2,10 @@ from ndsl import CompilationConfig, GridIndexing, StencilConfig, StencilFactory from ndsl.dsl.gt4py import PARALLEL, Field, computation, interval +from tests.dsl import utils -def _make_storage( - func, - grid_indexing, - stencil_config: StencilConfig, - *, - dtype=float, - aligned_index=(0, 0, 0), -): - return func( - backend=stencil_config.compilation_config.backend, - shape=grid_indexing.domain, - dtype=dtype, - aligned_index=aligned_index, - ) - - -def test_timing_collector(): +def test_timing_collector() -> None: grid_indexing = GridIndexing( domain=(5, 5, 5), n_halo=2, @@ -46,8 +31,8 @@ def func(inp: Field[float], out: Field[float]): build_report = stencil_factory.build_report(key="parse_time") assert "func" in build_report - inp = _make_storage(ones, grid_indexing, stencil_config, dtype=float) - out = _make_storage(empty, grid_indexing, stencil_config, dtype=float) + inp = utils.make_storage(ones, grid_indexing, stencil_config, dtype=float) + out = utils.make_storage(empty, grid_indexing, stencil_config, dtype=float) test(inp, out) exec_report = stencil_factory.exec_report() diff --git a/tests/dsl/test_stencil_config.py b/tests/dsl/test_stencil_config.py index 7e6b4da3..89498695 100644 --- a/tests/dsl/test_stencil_config.py +++ b/tests/dsl/test_stencil_config.py @@ -14,7 +14,7 @@ def test_same_config_equal( validate_args: bool, format_source: bool, compare_to_numpy: bool, -): +) -> None: dace_config = DaceConfig( communicator=None, backend=backend, @@ -30,7 +30,6 @@ def test_same_config_equal( compare_to_numpy=compare_to_numpy, dace_config=dace_config, ) - assert config == config same_config = StencilConfig( compilation_config=CompilationConfig( @@ -46,19 +45,14 @@ def test_same_config_equal( assert config == same_config -@pytest.mark.parametrize("validate_args", [True]) -@pytest.mark.parametrize("device_sync", [False]) -@pytest.mark.parametrize("rebuild", [True]) -@pytest.mark.parametrize("format_source", [True]) -@pytest.mark.parametrize("compare_to_numpy", [True]) def test_different_backend_not_equal( - backend: str, - rebuild: bool, - validate_args: bool, - format_source: bool, - device_sync: bool, - compare_to_numpy: bool, -): + backend: str = "numpy", + rebuild: bool = True, + validate_args: bool = True, + format_source: bool = True, + device_sync: bool = False, + compare_to_numpy: bool = True, +) -> None: dace_config = DaceConfig( communicator=None, backend=backend, @@ -77,7 +71,7 @@ def test_different_backend_not_equal( different_config = StencilConfig( compilation_config=CompilationConfig( - backend="fakebackend", + backend="fake_backend", rebuild=rebuild, validate_args=validate_args, format_source=format_source, @@ -89,19 +83,14 @@ def test_different_backend_not_equal( assert config != different_config -@pytest.mark.parametrize("validate_args", [True]) -@pytest.mark.parametrize("device_sync", [False]) -@pytest.mark.parametrize("rebuild", [True]) -@pytest.mark.parametrize("format_source", [True]) -@pytest.mark.parametrize("compare_to_numpy", [True]) def test_different_rebuild_not_equal( - backend: str, - rebuild: bool, - validate_args: bool, - format_source: bool, - device_sync: bool, - compare_to_numpy: bool, -): + backend: str = "numpy", + rebuild: bool = True, + validate_args: bool = True, + format_source: bool = True, + device_sync: bool = False, + compare_to_numpy: bool = True, +) -> None: dace_config = DaceConfig( communicator=None, backend=backend, @@ -132,18 +121,13 @@ def test_different_rebuild_not_equal( assert config != different_config -@pytest.mark.parametrize("validate_args", [True]) -@pytest.mark.parametrize("device_sync", [False]) -@pytest.mark.parametrize("rebuild", [True]) -@pytest.mark.parametrize("format_source", [True]) -@pytest.mark.parametrize("compare_to_numpy", [True]) def test_different_device_sync_not_equal( - rebuild: bool, - validate_args: bool, - format_source: bool, - device_sync: bool, - compare_to_numpy: bool, -): + rebuild: bool = True, + validate_args: bool = True, + format_source: bool = True, + device_sync: bool = False, + compare_to_numpy: bool = True, +) -> None: dace_config = DaceConfig( communicator=None, backend="gt:gpu", @@ -174,19 +158,14 @@ def test_different_device_sync_not_equal( assert config != different_config -@pytest.mark.parametrize("validate_args", [True]) -@pytest.mark.parametrize("device_sync", [False]) -@pytest.mark.parametrize("rebuild", [True]) -@pytest.mark.parametrize("format_source", [True]) -@pytest.mark.parametrize("compare_to_numpy", [True]) def test_different_validate_args_not_equal( - backend: str, - rebuild: bool, - validate_args: bool, - format_source: bool, - device_sync: bool, - compare_to_numpy: bool, -): + backend: str = "numpy", + rebuild: bool = True, + validate_args: bool = True, + format_source: bool = True, + device_sync: bool = False, + compare_to_numpy: bool = True, +) -> None: dace_config = DaceConfig( None, backend, @@ -217,19 +196,14 @@ def test_different_validate_args_not_equal( assert config != different_config -@pytest.mark.parametrize("validate_args", [True]) -@pytest.mark.parametrize("device_sync", [False]) -@pytest.mark.parametrize("rebuild", [True]) -@pytest.mark.parametrize("format_source", [True]) -@pytest.mark.parametrize("compare_to_numpy", [True]) def test_different_format_source_not_equal( - backend: str, - rebuild: bool, - validate_args: bool, - format_source: bool, - device_sync: bool, - compare_to_numpy: bool, -): + backend: str = "numpy", + rebuild: bool = True, + validate_args: bool = True, + format_source: bool = True, + device_sync: bool = False, + compare_to_numpy: bool = True, +) -> None: dace_config = DaceConfig(communicator=None, backend=backend) config = StencilConfig( compilation_config=CompilationConfig( @@ -265,7 +239,7 @@ def test_different_compare_to_numpy_not_equal( format_source: bool = True, rebuild: bool = True, validate_args: bool = False, -): +) -> None: dace_config = DaceConfig(communicator=None, backend=backend) config = StencilConfig( compilation_config=CompilationConfig( diff --git a/tests/dsl/test_stencil_factory.py b/tests/dsl/test_stencil_factory.py index 65bf1cf2..ce9de962 100644 --- a/tests/dsl/test_stencil_factory.py +++ b/tests/dsl/test_stencil_factory.py @@ -13,7 +13,10 @@ from ndsl.dsl.gt4py import PARALLEL, computation, horizontal, interval, region from ndsl.dsl.gt4py_utils import make_storage_from_shape from ndsl.dsl.stencil import CompareToNumpyStencil, get_stencils_with_varied_bounds -from ndsl.dsl.typing import FloatField +from ndsl.dsl.typing import Field, FloatField + + +BACKENDS = ["numpy", "dace:cpu"] def copy_stencil(q_in: FloatField, q_out: FloatField): @@ -36,7 +39,7 @@ def add_1_in_region_stencil(q_in: FloatField, q_out: FloatField): q_out = q_in + 1.0 -def setup_data_vars(backend: str): +def setup_data_vars(backend: str) -> tuple[Field, Field]: shape = (7, 7, 3) q = make_storage_from_shape(shape, backend=backend) q[:] = 1.0 @@ -68,7 +71,8 @@ def get_stencil_factory(backend: str) -> StencilFactory: return StencilFactory(config=config, grid_indexing=indexing) -def test_get_stencils_with_varied_bounds(backend: str): +@pytest.mark.parametrize("backend", BACKENDS) +def test_get_stencils_with_varied_bounds(backend: str) -> None: origins = [(2, 2, 0), (1, 1, 0)] domains = [(1, 1, 3), (2, 2, 3)] factory = get_stencil_factory(backend) @@ -87,7 +91,8 @@ def test_get_stencils_with_varied_bounds(backend: str): np.testing.assert_array_equal(q.data, q_ref.data) -def test_get_stencils_with_varied_bounds_and_regions(backend: str): +@pytest.mark.parametrize("backend", BACKENDS) +def test_get_stencils_with_varied_bounds_and_regions(backend: str) -> None: factory = get_stencil_factory(backend) origins = [(3, 3, 0), (2, 2, 0)] domains = [(1, 1, 3), (2, 2, 3)] @@ -107,7 +112,8 @@ def test_get_stencils_with_varied_bounds_and_regions(backend: str): np.testing.assert_array_equal(q_orig.data, q_ref.data) -def test_stencil_vertical_bounds(backend: str): +@pytest.mark.parametrize("backend", BACKENDS) +def test_stencil_vertical_bounds(backend: str) -> None: factory = get_stencil_factory(backend) origins = [(3, 3, 0), (2, 2, 1)] domains = [(1, 1, 3), (2, 2, 4)] @@ -124,9 +130,11 @@ def test_stencil_vertical_bounds(backend: str): assert "k_end" in stencils[1].externals and stencils[1].externals["k_end"] == 4 +@pytest.mark.parametrize("backend", BACKENDS) @pytest.mark.parametrize("enabled", [True, False]) -def test_stencil_factory_numpy_comparison_from_dims_halo(enabled: bool): - backend = "numpy" +def test_stencil_factory_numpy_comparison_from_dims_halo( + backend: str, enabled: bool +) -> None: dace_config = DaceConfig(communicator=None, backend=backend) config = StencilConfig( compilation_config=CompilationConfig( @@ -159,9 +167,11 @@ def test_stencil_factory_numpy_comparison_from_dims_halo(enabled: bool): assert isinstance(stencil, FrozenStencil) +@pytest.mark.parametrize("backend", BACKENDS) @pytest.mark.parametrize("enabled", [True, False]) -def test_stencil_factory_numpy_comparison_from_origin_domain(enabled: bool): - backend = "numpy" +def test_stencil_factory_numpy_comparison_from_origin_domain( + backend: str, enabled: bool +) -> None: dace_config = DaceConfig(communicator=None, backend=backend) config = StencilConfig( compilation_config=CompilationConfig( @@ -192,8 +202,8 @@ def test_stencil_factory_numpy_comparison_from_origin_domain(enabled: bool): assert isinstance(stencil, FrozenStencil) -def test_stencil_factory_numpy_comparison_runs_without_exceptions(): - backend = "numpy" +@pytest.mark.parametrize("backend", BACKENDS) +def test_stencil_factory_numpy_comparison_runs_without_exceptions(backend: str) -> None: dace_config = DaceConfig(communicator=None, backend=backend) config = StencilConfig( compilation_config=CompilationConfig( diff --git a/tests/dsl/test_stencil_wrapper.py b/tests/dsl/test_stencil_wrapper.py index 94ea894c..df91f5c0 100644 --- a/tests/dsl/test_stencil_wrapper.py +++ b/tests/dsl/test_stencil_wrapper.py @@ -24,7 +24,7 @@ def get_stencil_config( backend: str, orchestration: DaCeOrchestration = DaCeOrchestration.Python, **kwargs, -): +) -> StencilConfig: dace_config = DaceConfig(None, backend=backend, orchestration=orchestration) config = StencilConfig( compilation_config=CompilationConfig( @@ -104,7 +104,7 @@ def __init__(self, axes): ), ], ) -def test_compute_field_origins(field_info, origin, field_origins): +def test_compute_field_origins(field_info, origin, field_origins) -> None: result = FrozenStencil._compute_field_origins(field_info, origin) assert result == field_origins @@ -115,16 +115,13 @@ def copy_stencil(q_in: FloatField, q_out: FloatField): @pytest.mark.parametrize("validate_args", [True, False]) -@pytest.mark.parametrize("device_sync", [False]) -@pytest.mark.parametrize("rebuild", [False]) -@pytest.mark.parametrize("format_source", [False]) def test_copy_frozen_stencil( - backend: str, - rebuild: bool, validate_args: bool, - format_source: bool, - device_sync: bool, -): + backend: str = "numpy", + rebuild: bool = False, + format_source: bool = False, + device_sync: bool = False, +) -> None: config = get_stencil_config( backend=backend, rebuild=rebuild, @@ -147,15 +144,12 @@ def test_copy_frozen_stencil( np.testing.assert_array_equal(q_in, q_out) -@pytest.mark.parametrize("device_sync", [False]) -@pytest.mark.parametrize("rebuild", [False]) -@pytest.mark.parametrize("format_source", [False]) def test_frozen_stencil_raises_if_given_origin( - backend: str, - rebuild: bool, - format_source: bool, - device_sync: bool, -): + backend: str = "numpy", + rebuild: bool = False, + format_source: bool = False, + device_sync: bool = False, +) -> None: # only guaranteed when validating args config = get_stencil_config( backend=backend, @@ -177,14 +171,11 @@ def test_frozen_stencil_raises_if_given_origin( stencil(q_in, q_out, origin=(0, 0, 0)) -@pytest.mark.parametrize("device_sync", [False]) -@pytest.mark.parametrize("rebuild", [False]) -@pytest.mark.parametrize("format_source", [False]) def test_frozen_stencil_raises_if_given_domain( - backend: str, - rebuild: bool, - format_source: bool, - device_sync: bool, + backend: str = "numpy", + rebuild: bool = False, + format_source: bool = False, + device_sync: bool = False, ): # only guaranteed when validating args config = get_stencil_config( @@ -212,11 +203,11 @@ def test_frozen_stencil_raises_if_given_domain( [[False, False, False, False], [True, False, False, False]], ) def test_frozen_stencil_kwargs_passed_to_init( - backend: str, rebuild: bool, validate_args: bool, format_source: bool, device_sync: bool, + backend: str = "numpy", ): config = get_stencil_config( backend=backend, @@ -255,9 +246,9 @@ def field_after_parameter_stencil(q_in: FloatField, param: float, q_out: FloatFi q_out = param * q_in -def test_frozen_field_after_parameter(backend): +def test_frozen_field_after_parameter() -> None: config = get_stencil_config( - backend=backend, + backend="numpy", rebuild=False, validate_args=False, format_source=False, @@ -273,13 +264,11 @@ def test_frozen_field_after_parameter(backend): @pytest.mark.parametrize("backend", ("numpy", "cuda")) -@pytest.mark.parametrize("rebuild", [True]) -@pytest.mark.parametrize("validate_args", [True]) def test_backend_options( backend: str, - rebuild: bool, - validate_args: bool, -): + rebuild: bool = True, + validate_args: bool = True, +) -> None: expected_options = { "numpy": { "backend": "numpy", @@ -307,7 +296,7 @@ def get_mock_quantity(): return unittest.mock.MagicMock(spec=Quantity) -def test_convert_quantities_to_storage_no_args(): +def test_convert_quantities_to_storage_no_args() -> None: args = [] kwargs = {} _convert_quantities_to_storage(args, kwargs) @@ -315,7 +304,7 @@ def test_convert_quantities_to_storage_no_args(): assert len(kwargs) == 0 -def test_convert_quantities_to_storage_one_arg_quantity(): +def test_convert_quantities_to_storage_one_arg_quantity() -> None: quantity = get_mock_quantity() args = [quantity] kwargs = {} @@ -325,7 +314,7 @@ def test_convert_quantities_to_storage_one_arg_quantity(): assert len(kwargs) == 0 -def test_convert_quantities_to_storage_one_kwarg_quantity(): +def test_convert_quantities_to_storage_one_kwarg_quantity() -> None: quantity = get_mock_quantity() args = [] kwargs = {"val": quantity} @@ -335,7 +324,7 @@ def test_convert_quantities_to_storage_one_kwarg_quantity(): assert kwargs["val"] == quantity.data -def test_convert_quantities_to_storage_one_arg_nonquantity(): +def test_convert_quantities_to_storage_one_arg_nonquantity() -> None: non_quantity = unittest.mock.MagicMock(spec=tuple) args = [non_quantity] kwargs = {} @@ -345,7 +334,7 @@ def test_convert_quantities_to_storage_one_arg_nonquantity(): assert len(kwargs) == 0 -def test_convert_quantities_to_storage_one_kwarg_non_quantity(): +def test_convert_quantities_to_storage_one_kwarg_non_quantity() -> None: non_quantity = unittest.mock.MagicMock(spec=tuple) args = [] kwargs = {"val": non_quantity} From 820640396daf1f3994d8b2499b8b5b58dc2ee6c7 Mon Sep 17 00:00:00 2001 From: Roman Cattaneo <1116746+romanc@users.noreply.github.com> Date: Fri, 27 Jun 2025 10:48:31 +0200 Subject: [PATCH 4/6] simplify --- tests/dsl/test_caches.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/dsl/test_caches.py b/tests/dsl/test_caches.py index 677bc59a..5e1afd95 100644 --- a/tests/dsl/test_caches.py +++ b/tests/dsl/test_caches.py @@ -67,7 +67,7 @@ def __call__(self): self.stencil(self.inp, self.out, self.inp[0, 0, 0]) -@pytest.mark.parametrize("backend", [pytest.param("dace:cpu")]) +@pytest.mark.parametrize("backend", ["dace:cpu"]) @pytest.mark.skipif( MPI.COMM_WORLD.Get_size() > 1, reason="relocatibility checked with a one-rank setup" ) @@ -118,7 +118,7 @@ def test_relocatability_orchestration(backend: str) -> None: gt_config.cache_settings["root_path"] = original_root_directory -@pytest.mark.parametrize("backend", [pytest.param("dace:cpu")]) +@pytest.mark.parametrize("backend", ["dace:cpu"]) @pytest.mark.skipif( MPI.COMM_WORLD.Get_size() > 1, reason="relocatibility checked with a one-rank setup" ) From a000a9b6e09e13e5d15a7da0c2ceafd49dee3ae4 Mon Sep 17 00:00:00 2001 From: Roman Cattaneo <1116746+romanc@users.noreply.github.com> Date: Fri, 27 Jun 2025 13:12:30 +0200 Subject: [PATCH 5/6] More robust caches tests. --- tests/dsl/test_caches.py | 101 +++++++++++++++++++++++---------------- tests/dsl/utils.py | 17 +++++++ 2 files changed, 78 insertions(+), 40 deletions(-) create mode 100644 tests/dsl/utils.py diff --git a/tests/dsl/test_caches.py b/tests/dsl/test_caches.py index 5e1afd95..2aecfbb3 100644 --- a/tests/dsl/test_caches.py +++ b/tests/dsl/test_caches.py @@ -1,4 +1,3 @@ -import os import shutil import sys from pathlib import Path @@ -22,6 +21,26 @@ from tests.dsl import utils +@pytest.fixture +def tmp_cache_root(tmpdir): + original_root = gt_config.cache_settings["root_path"] + gt_config.cache_settings["root_path"] = tmpdir + + yield tmpdir + + # restore original cache settings + gt_config.cache_settings["root_path"] = original_root + + +@pytest.fixture +def restore_cache_dir(): + cache_dir = gt_config.cache_settings["dir_name"] + + yield + + gt_config.cache_settings["dir_name"] = cache_dir + + def _stencil(inp: Field[float], out: Field[float], scalar: float): with computation(PARALLEL), interval(...): out = inp @@ -67,33 +86,34 @@ def __call__(self): self.stencil(self.inp, self.out, self.inp[0, 0, 0]) -@pytest.mark.parametrize("backend", ["dace:cpu"]) @pytest.mark.skipif( MPI.COMM_WORLD.Get_size() > 1, reason="relocatibility checked with a one-rank setup" ) -def test_relocatability_orchestration(backend: str) -> None: - original_root_directory: str = gt_config.cache_settings["root_path"] - cwd = Path.cwd() - +def test_relocatability_orchestration(restore_cache_dir) -> None: # Compile on default - p0 = OrchestratedProgram(backend, DaCeOrchestration.BuildAndRun) + p0 = OrchestratedProgram("dace:cpu", DaCeOrchestration.BuildAndRun) p0() + expected_cache_dir = ( - cwd + Path.cwd() / ".gt_cache_FV3_A" / "dacecache" / "test_caches_OrchestratedProgram___call__" ) assert expected_cache_dir.exists() - # Compile in another directory - custom_path = cwd / ".my_cache_path" - gt_config.cache_settings["root_path"] = custom_path +@pytest.mark.skipif( + MPI.COMM_WORLD.Get_size() > 1, reason="relocatibility checked with a one-rank setup" +) +def test_relocatability_orchestration_tmpdir(restore_cache_dir, tmp_cache_root) -> None: + # Compile in temporary directory that is only available in this test session. + backend = "dace:cpu" p1 = OrchestratedProgram(backend, DaCeOrchestration.BuildAndRun) p1() + expected_cache_dir = ( - custom_path + tmp_cache_root / ".gt_cache_FV3_A" / "dacecache" / "test_caches_OrchestratedProgram___call__" @@ -102,8 +122,8 @@ def test_relocatability_orchestration(backend: str) -> None: # Check relocability by copying the second cache directory, # changing the path of gt_config.cache_settings and trying to Run on it - relocated_path = cwd / ".my_relocated_cache_path" - shutil.copytree(custom_path, relocated_path, dirs_exist_ok=True) + relocated_path = tmp_cache_root / ".my_relocated_cache_path" + shutil.copytree(tmp_cache_root, relocated_path, dirs_exist_ok=False) gt_config.cache_settings["root_path"] = relocated_path p2 = OrchestratedProgram(backend, DaCeOrchestration.Run) p2() @@ -114,59 +134,62 @@ def test_relocatability_orchestration(backend: str) -> None: with pytest.raises(RuntimeError): OrchestratedProgram(backend, DaCeOrchestration.Run) - # Restore cache settings - gt_config.cache_settings["root_path"] = original_root_directory - -@pytest.mark.parametrize("backend", ["dace:cpu"]) @pytest.mark.skipif( MPI.COMM_WORLD.Get_size() > 1, reason="relocatibility checked with a one-rank setup" ) -def test_relocatability(backend: str) -> None: - # Restore original dir name - gt_config.cache_settings["dir_name"] = os.environ.get( - "GT_CACHE_DIR_NAME", f".gt_cache_{MPI.COMM_WORLD.Get_rank():06}" - ) +def test_relocatability(restore_cache_dir) -> None: + # Compile on default + backend = "dace:cpu" + p0 = OrchestratedProgram(backend, DaCeOrchestration.Python) + p0() - cwd = Path.cwd() backend_sanitized = backend.replace(":", "") python_version = f"py{sys.version_info[0]}{sys.version_info[1]}" expected_cache_path = ( - cwd + Path.cwd() / ".gt_cache_000000" / f"{python_version}_1013" / f"{backend_sanitized}" / "test_caches" / "_stencil" ) - - # Compile on default - p0 = OrchestratedProgram(backend, DaCeOrchestration.Python) - p0() assert expected_cache_path.exists() + +@pytest.mark.skipif( + MPI.COMM_WORLD.Get_size() > 1, reason="relocatibility checked with a one-rank setup" +) +def test_relocatability_tmpdir(restore_cache_dir, tmp_cache_root) -> None: # Compile in another directory + backend = "dace:cpu" + p1 = OrchestratedProgram(backend, DaCeOrchestration.Python) + p1() - custom_path = cwd / ".my_cache_path" + backend_sanitized = backend.replace(":", "") + python_version = f"py{sys.version_info[0]}{sys.version_info[1]}" expected_cache_path = ( - custom_path + tmp_cache_root / ".gt_cache_000000" / f"{python_version}_1013" / f"{backend_sanitized}" / "test_caches" / "_stencil" ) - gt_config.cache_settings["root_path"] = custom_path - p1 = OrchestratedProgram(backend, DaCeOrchestration.Python) - p1() assert expected_cache_path.exists() # Check relocability by copying the first cache directory, # changing the path of gt_config.cache_settings and trying to Run on it - relocated_path = cwd / ".my_relocated_cache_path" - shutil.copytree(cwd / ".gt_cache_000000", relocated_path, dirs_exist_ok=True) + relocated_path = tmp_cache_root / ".my_relocated_cache_path" + shutil.copytree( + tmp_cache_root / ".gt_cache_000000", relocated_path, dirs_exist_ok=False + ) gt_config.cache_settings["root_path"] = relocated_path - expected_cache_path = ( + + p2 = OrchestratedProgram(backend, DaCeOrchestration.Python) + p2() + + relocated_cache_path = ( relocated_path / ".gt_cache_000000" / f"{python_version}_1013" @@ -174,6 +197,4 @@ def test_relocatability(backend: str) -> None: / "test_caches" / "_stencil" ) - p2 = OrchestratedProgram(backend, DaCeOrchestration.Python) - p2() - assert expected_cache_path.exists() + assert relocated_cache_path.exists() diff --git a/tests/dsl/utils.py b/tests/dsl/utils.py new file mode 100644 index 00000000..5e8535e5 --- /dev/null +++ b/tests/dsl/utils.py @@ -0,0 +1,17 @@ +from ndsl import GridIndexing, StencilConfig + + +def make_storage( + func, + grid_indexing: GridIndexing, + stencil_config: StencilConfig, + *, + dtype=float, + aligned_index=(0, 0, 0), +): + return func( + backend=stencil_config.compilation_config.backend, + shape=grid_indexing.domain, + dtype=dtype, + aligned_index=aligned_index, + ) From d100394175bbad28e3aeb20ba7bc405b4395891c Mon Sep 17 00:00:00 2001 From: Roman Cattaneo <1116746+romanc@users.noreply.github.com> Date: Fri, 27 Jun 2025 13:25:30 +0200 Subject: [PATCH 6/6] MPI cleanup --- ndsl/stencils/testing/conftest.py | 2 +- ndsl/stencils/testing/test_translate.py | 4 ++-- tests/mpi/test_mpi_all_reduce_sum.py | 16 +++++++--------- tests/mpi/test_mpi_halo_update.py | 24 +++++++++--------------- tests/mpi/test_mpi_mock.py | 4 +--- tests/test_decomposition.py | 4 ++-- 6 files changed, 22 insertions(+), 32 deletions(-) diff --git a/ndsl/stencils/testing/conftest.py b/ndsl/stencils/testing/conftest.py index 3e31d191..4b62dc5f 100644 --- a/ndsl/stencils/testing/conftest.py +++ b/ndsl/stencils/testing/conftest.py @@ -376,7 +376,7 @@ def parallel_savepoint_cases( def pytest_generate_tests(metafunc): backend = metafunc.config.getoption("backend") - if MPI is not None and MPI.COMM_WORLD.Get_size() > 1: + if MPI.COMM_WORLD.Get_size() > 1: if metafunc.function.__name__ == "test_parallel_savepoint": generate_parallel_stencil_tests(metafunc, backend=backend) elif metafunc.function.__name__ == "test_sequential_savepoint": diff --git a/ndsl/stencils/testing/test_translate.py b/ndsl/stencils/testing/test_translate.py index 02d00420..894b8ebc 100644 --- a/ndsl/stencils/testing/test_translate.py +++ b/ndsl/stencils/testing/test_translate.py @@ -140,7 +140,7 @@ def _get_thresholds(compute_function, input_data) -> None: @pytest.mark.sequential @pytest.mark.skipif( - MPI is not None and MPI.COMM_WORLD.Get_size() > 1, + MPI.COMM_WORLD.Get_size() > 1, reason="Running in parallel with mpi", ) def test_sequential_savepoint( @@ -293,7 +293,7 @@ def get_tile_communicator(comm, layout): @pytest.mark.parallel @pytest.mark.skipif( - MPI is None or MPI.COMM_WORLD.Get_size() == 1, + MPI.COMM_WORLD.Get_size() == 1, reason="Not running in parallel with mpi", ) def test_parallel_savepoint( diff --git a/tests/mpi/test_mpi_all_reduce_sum.py b/tests/mpi/test_mpi_all_reduce_sum.py index c7c864cf..6cab1023 100644 --- a/tests/mpi/test_mpi_all_reduce_sum.py +++ b/tests/mpi/test_mpi_all_reduce_sum.py @@ -15,14 +15,14 @@ @pytest.fixture def layout(): - if MPI is not None: - size = MPI.COMM_WORLD.Get_size() - ranks_per_tile = size // 6 - ranks_per_edge = int(ranks_per_tile**0.5) - return (ranks_per_edge, ranks_per_edge) - else: + if MPI is None: return (1, 1) + size = MPI.COMM_WORLD.Get_size() + ranks_per_tile = size // 6 + ranks_per_edge = int(ranks_per_tile**0.5) + return (ranks_per_edge, ranks_per_edge) + @pytest.fixture(params=[0.1, 1.0]) def edge_interior_ratio(request): @@ -47,9 +47,7 @@ def communicator(cube_partitioner): ) -@pytest.mark.skipif( - MPI is None, reason="mpi4py is not available or pytest was not run in parallel" -) +@pytest.mark.skipif(MPI is None, reason="pytest is not run in parallel") def test_all_reduce(communicator): backends = ["dace:cpu", "gt:cpu_kfirst", "numpy"] diff --git a/tests/mpi/test_mpi_halo_update.py b/tests/mpi/test_mpi_halo_update.py index c9b35ecc..76aca71e 100644 --- a/tests/mpi/test_mpi_halo_update.py +++ b/tests/mpi/test_mpi_halo_update.py @@ -37,14 +37,14 @@ def dtype(numpy): @pytest.fixture def layout(): - if MPI is not None: - size = MPI.COMM_WORLD.Get_size() - ranks_per_tile = size // 6 - ranks_per_edge = int(ranks_per_tile**0.5) - return (ranks_per_edge, ranks_per_edge) - else: + if MPI is None: return (1, 1) + size = MPI.COMM_WORLD.Get_size() + ranks_per_tile = size // 6 + ranks_per_edge = int(ranks_per_tile**0.5) + return (ranks_per_edge, ranks_per_edge) + @pytest.fixture def nz(): @@ -281,9 +281,7 @@ def depth_quantity( return quantity -@pytest.mark.skipif( - MPI is None, reason="mpi4py is not available or pytest was not run in parallel" -) +@pytest.mark.skipif(MPI is None, reason="pytest is not run in parallel") def test_depth_halo_update( depth_quantity, communicator, @@ -332,9 +330,7 @@ def zeros_quantity(dims, units, origin, extent, shape, numpy, dtype): return quantity -@pytest.mark.skipif( - MPI is None, reason="mpi4py is not available or pytest was not run in parallel" -) +@pytest.mark.skipif(MPI is None, reason="pytest is not run in parallel") def test_zeros_halo_update( zeros_quantity, communicator, @@ -371,9 +367,7 @@ def test_zeros_halo_update( ) -@pytest.mark.skipif( - MPI is None, reason="mpi4py is not available or pytest was not run in parallel" -) +@pytest.mark.skipif(MPI is None, reason="pytest is not run in parallel") def test_zeros_vector_halo_update( zeros_quantity, communicator, diff --git a/tests/mpi/test_mpi_mock.py b/tests/mpi/test_mpi_mock.py index 1e4113fd..0da2f21b 100644 --- a/tests/mpi/test_mpi_mock.py +++ b/tests/mpi/test_mpi_mock.py @@ -293,9 +293,7 @@ def dummy_results(worker_function, dummy_list, numpy): return result_list -@pytest.mark.skipif( - MPI is None, reason="mpi4py is not available or pytest was not run in parallel" -) +@pytest.mark.skipif(MPI is None, reason="pytest is not run in parallel") def test_worker(comm, dummy_results, mpi_results, numpy): comm.barrier() # synchronize the test "dots" across ranks if comm.Get_rank() == 0: diff --git a/tests/test_decomposition.py b/tests/test_decomposition.py index bf7363e3..2160e23e 100644 --- a/tests/test_decomposition.py +++ b/tests/test_decomposition.py @@ -70,8 +70,8 @@ def test_build_cache_path( @pytest.mark.skipif( - MPI is None or MPI.COMM_WORLD.Get_size() != 6, - reason="mpi4py is not available or pytest was not run in parallel", + MPI is None, + reason="pytest is not run in parallel", ) def test_unblock_waiting_tiles(): comm = MPI.COMM_WORLD