Add checkpointer tests to CircleCI

.circleci/config.yml

@@ -97,6 +97,7 @@ commands:
         description: "make target to run"
         type: enum
         enum:
+          - "test_savepoint"
           - "savepoint_tests"
           - "savepoint_tests_mpi"
           - "physics_savepoint_tests"
@@ -231,7 +232,7 @@ jobs:
   driver_savepoints_mpi:
     machine:
       image: ubuntu-2004:202111-02
-      resource_class: large
+      resource_class: xlarge
     parameters:
       backend:
         description: "gt4py backend"
@@ -402,7 +403,7 @@ jobs:
   test_mpi_54rank:
     docker:
       - image: cimg/python:3.8
-    resource_class: large
+    resource_class: xlarge
     working_directory: ~/repo
     steps:
       - checkout
@@ -416,6 +417,20 @@ jobs:
             . venv/bin/activate
             mpirun -n 54 --oversubscribe --mca btl_vader_single_copy_mechanism none python3 -m mpi4py -m pytest tests/mpi_54rank
 
+  test_savepoint:
+    machine:
+      image: ubuntu-2004:202111-02
+      resource_class: large
+    environment:
+      GOOGLE_APPLICATION_CREDENTIALS: /tmp/key.json
+    steps:
+      - checkout
+      - make_savepoints:
+          backend: numpy
+          experiment: c12_6ranks_standard
+          target: test_savepoint
+          num_ranks: 6
+
   test_notebooks:
     docker:
       - image: gcr.io/vcm-ml/pace_notebook_examples
@@ -528,6 +543,12 @@ workflows:
           filters:
             tags:
               only: /^v.*/
+      - test_savepoint:
+          context:
+            - GCLOUD_ENCODED_KEY
+          filters:
+            tags:
+              only: /^v.*/
       - test_util:
           filters:
             tags:

Makefile

@@ -103,7 +103,7 @@ test_util:
 		$(MAKE) -C util test; \
 	fi
 
-savepoint_tests: build
+savepoint_tests: build  ## dycore-only savepoint tests
 	TARGET=dycore $(MAKE) get_test_data
 	$(CONTAINER_CMD) $(CONTAINER_FLAGS) bash -c "$(SAVEPOINT_SETUP) && cd $(PACE_PATH) && pytest --data_path=$(EXPERIMENT_DATA_RUN)/dycore/ $(TEST_ARGS) $(FV3CORE_THRESH_ARGS) $(PACE_PATH)/fv3core/tests/savepoint"
 
@@ -130,6 +130,10 @@ physics_savepoint_tests_mpi: build
 test_main: build
 	$(CONTAINER_CMD) $(CONTAINER_FLAGS) bash -c "$(SAVEPOINT_SETUP) && cd $(PACE_PATH) && pytest $(TEST_ARGS) $(PACE_PATH)/tests/main"
 
+test_savepoint:  ## top level savepoint tests
+	TARGET=dycore $(MAKE) get_test_data
+	$(CONTAINER_CMD) $(CONTAINER_FLAGS) bash -c "$(SAVEPOINT_SETUP) && cd $(PACE_PATH) && $(MPIRUN_CALL) python -m pytest --data_path=$(EXPERIMENT_DATA_RUN)/dycore/ $(TEST_ARGS) $(PACE_PATH)/tests/savepoint"
+
 test_mpi_54rank:
 	mpirun -n 54 $(MPIRUN_ARGS) python3 -m mpi4py -m pytest tests/mpi_54rank
 

driver/examples/configs/baroclinic_c12_orch_cpu.yaml

@@ -13,7 +13,7 @@ performance_config:
 nx_tile: 12
 nz: 79
 dt_atmos: 225
-minutes: 1
+minutes: 5
 layout:
   - 1
   - 1

fv3core/pace/fv3core/stencils/fv_dynamics.py

@@ -22,7 +22,7 @@
 from pace.fv3core.stencils.neg_adj3 import AdjustNegativeTracerMixingRatio
 from pace.fv3core.stencils.remapping import LagrangianToEulerian
 from pace.stencils.c2l_ord import CubedToLatLon
-from pace.util import Timer
+from pace.util import X_DIM, Y_DIM, Z_INTERFACE_DIM, Timer
 from pace.util.grid import DampingCoefficients, GridData
 from pace.util.mpi import MPI
 from pace.util.quantity import Quantity
@@ -179,6 +179,34 @@ def __init__(
             dace_compiletime_args=["state", "tag"],
         )
 
+        orchestrate(
+            obj=self,
+            config=stencil_factory.config.dace_config,
+            method_to_orchestrate="_checkpoint_remapping_in",
+            dace_compiletime_args=[
+                "state",
+            ],
+        )
+
+        orchestrate(
+            obj=self,
+            config=stencil_factory.config.dace_config,
+            method_to_orchestrate="_checkpoint_remapping_out",
+            dace_compiletime_args=["state"],
+        )
+
+        orchestrate(
+            obj=self,
+            config=stencil_factory.config.dace_config,
+            method_to_orchestrate="_checkpoint_tracer_advection_in",
+            dace_compiletime_args=["state"],
+        )
+        orchestrate(
+            obj=self,
+            config=stencil_factory.config.dace_config,
+            method_to_orchestrate="_checkpoint_tracer_advection_out",
+            dace_compiletime_args=["state"],
+        )
         # nested and stretched_grid are options in the Fortran code which we
         # have not implemented, so they are hard-coded here.
         self.call_checkpointer = checkpointer is not None
@@ -234,7 +262,11 @@ def __init__(
 
         # Build advection stencils
         self.tracer_advection = tracer_2d_1l.TracerAdvection(
-            stencil_factory, tracer_transport, self.grid_data, comm, self.tracers
+            stencil_factory,
+            tracer_transport,
+            self.grid_data,
+            comm,
+            self.tracers,
         )
         self._ak = grid_data.ak
         self._bk = grid_data.bk
@@ -316,6 +348,7 @@ def __init__(
             NQ,
             self._pfull,
             tracers=self.tracers,
+            checkpointer=checkpointer,
         )
 
         full_xyz_spec = grid_indexing.get_quantity_halo_spec(
@@ -353,6 +386,93 @@ def _checkpoint_fvdynamics(self, state: DycoreState, tag: str):
                 qvapor=state.qvapor,
             )
 
+    def _checkpoint_remapping_in(
+        self,
+        state: DycoreState,
+    ):
+        if self.call_checkpointer:
+            self.checkpointer(
+                "Remapping-In",
+                pt=state.pt,
+                delp=state.delp,
+                delz=state.delz,
+                peln=state.peln.transpose(
+                    [X_DIM, Z_INTERFACE_DIM, Y_DIM]
+                ),  # [x, z, y] fortran data
+                u=state.u,
+                v=state.v,
+                w=state.w,
+                ua=state.ua,
+                va=state.va,
+                cappa=self._cappa,
+                pkz=state.pkz,
+                pk=state.pk,
+                pe=state.pe.transpose(
+                    [X_DIM, Z_INTERFACE_DIM, Y_DIM]
+                ),  # [x, z, y] fortran data
+                phis=state.phis,
+                te_2d=self._te0_2d,
+                ps=state.ps,
+                wsd=self._wsd,
+                omga=state.omga,
+                dp1=self._dp1,
+            )
+
+    def _checkpoint_remapping_out(
+        self,
+        state: DycoreState,
+    ):
+        if self.call_checkpointer:
+            self.checkpointer(
+                "Remapping-Out",
+                pt=state.pt,
+                delp=state.delp,
+                delz=state.delz,
+                peln=state.peln.transpose(
+                    [X_DIM, Z_INTERFACE_DIM, Y_DIM]
+                ),  # [x, z, y] fortran data
+                u=state.u,
+                v=state.v,
+                w=state.w,
+                cappa=self._cappa,
+                pkz=state.pkz,
+                pk=state.pk,
+                pe=state.pe.transpose(
+                    [X_DIM, Z_INTERFACE_DIM, Y_DIM]
+                ),  # [x, z, y] fortran data
+                te_2d=self._te0_2d,
+                omga=state.omga,
+                dp1=self._dp1,
+            )
+
+    def _checkpoint_tracer_advection_in(
+        self,
+        state: DycoreState,
+    ):
+        if self.call_checkpointer:
+            self.checkpointer(
+                "Tracer2D1L-In",
+                dp1=self._dp1,
+                mfxd=state.mfxd,
+                mfyd=state.mfyd,
+                cxd=state.cxd,
+                cyd=state.cyd,
+            )
+
+    def _checkpoint_tracer_advection_out(
+        self,
+        state: DycoreState,
+    ):
+        if self.call_checkpointer:
+            self.checkpointer(
+                "Tracer2D1L-Out",
+                dp1=self._dp1,
+                mfxd=state.mfxd,
+                mfyd=state.mfyd,
+                cxd=state.cxd,
+                cyd=state.cyd,
+            )
+
     def step_dynamics(
         self,
         state: DycoreState,
@@ -449,6 +569,7 @@ def _compute(self, state: DycoreState, timer: pace.util.Timer):
                 if __debug__:
                     log_on_rank_0("Remapping")
                 with timer.clock("Remapping"):
+                    self._checkpoint_remapping_in(state)
                     self._lagrangian_to_eulerian_obj(
                         self.tracer_storages,
                         state.pt,
@@ -483,6 +604,7 @@ def _compute(self, state: DycoreState, timer: pace.util.Timer):
                         self._timestep / self._k_split,
                         self._timestep,
                     )
+                    self._checkpoint_remapping_out(state)
                 if last_step:
                     da_min: float = self._get_da_min()
                     self.post_remap(
@@ -519,6 +641,7 @@ def _dyn(
             if __debug__:
                 log_on_rank_0("TracerAdvection")
             with timer.clock("TracerAdvection"):
+                self._checkpoint_tracer_advection_in(state)
                 self.tracer_advection(
                     tracers,
                     self._dp1,
@@ -528,6 +651,7 @@ def _dyn(
                     state.cyd,
                     self._timestep / self._k_split,
                 )
+                self._checkpoint_tracer_advection_out(state)
 
     def post_remap(
         self,

fv3core/pace/fv3core/stencils/remapping.py

@@ -1,4 +1,4 @@
-from typing import Dict
+from typing import Dict, Optional
 
 from gt4py.gtscript import (
     __INLINED,
@@ -15,6 +15,7 @@
 
 import pace.dsl.gt4py_utils as utils
 import pace.fv3core.stencils.moist_cv as moist_cv
+import pace.util
 from pace.dsl.dace.orchestration import orchestrate
 from pace.dsl.stencil import StencilFactory
 from pace.dsl.typing import FloatField, FloatFieldIJ, FloatFieldK
@@ -286,12 +287,17 @@ def __init__(
         nq,
         pfull,
         tracers: Dict[str, Quantity],
+        checkpointer: Optional[pace.util.Checkpointer] = None,
     ):
         orchestrate(
             obj=self,
             config=stencil_factory.config.dace_config,
             dace_compiletime_args=["tracers"],
         )
+        self._checkpointer = checkpointer
+        # this is only computed in init because Dace does not yet support
+        # this operation
+        self._call_checkpointer = checkpointer is not None
         grid_indexing = stencil_factory.grid_indexing
         if config.kord_tm >= 0:
             raise NotImplementedError("map ppm, untested mode where kord_tm >= 0")