Turn main unit test & lint on PR, logger clean up [NASA:Update]

.github/workflows/lint.yml

@@ -0,0 +1,27 @@
+name: "Lint"
+on:
+  pull_request:
+      types: [opened, synchronize, reopened, ready_for_review, labeled, unlabeled]
+
+jobs:
+  lint:
+    runs-on: ubuntu-latest
+    steps:
+        - name: Checkout Pace repository
+          uses: actions/checkout@v3.5.2
+          with:
+            submodules: 'recursive'
+        - name: Step Python 3.8.12
+          uses: actions/setup-python@v4.6.0
+          with:
+            python-version: '3.8.12'
+        - name: Install OpenMPI for gt4py
+          run: |
+            sudo apt-get install libopenmpi-dev
+        - name: Install Python packages
+          run: |
+            python -m pip install --upgrade pip
+            pip install -r requirements_dev.txt -r requirements_lint.txt
+        - name: Run lint via pre-commit
+          run: |
+            pre-commit run --all-files

.github/workflows/main_unit_tests.yml

@@ -0,0 +1,27 @@
+name: "Main unit tests"
+on:
+  pull_request:
+      types: [opened, synchronize, reopened, ready_for_review, labeled, unlabeled]
+
+jobs:
+  main_unit_tests:
+    runs-on: ubuntu-latest
+    steps:
+        - name: Checkout Pace repository
+          uses: actions/checkout@v3.5.2
+          with:
+            submodules: 'recursive'
+        - name: Step Python 3.8.12
+          uses: actions/setup-python@v4.6.0
+          with:
+            python-version: '3.8.12'
+        - name: Install OpenMPI for gt4py
+          run: |
+            sudo apt-get install libopenmpi-dev
+        - name: Install Python packages
+          run: |
+            python -m pip install --upgrade pip
+            pip install -r requirements_dev.txt
+        - name: Run all main tests
+          run: |
+            pytest -x tests/main

README.md

@@ -11,6 +11,7 @@ Pace is an implementation of the FV3GFS / SHiELD atmospheric model developed by
 Full Sphinx documentation can be found at [https://ai2cm.github.io/pace/](https://ai2cm.github.io/pace/).
 
 **WARNING** This repo is under active development - supported features and procedures can change rapidly and without notice.
+
 ## Quickstart - bare metal
 
 ### Build
@@ -27,10 +28,13 @@ export BOOST_ROOT=BOOST/ROOT/boost_1_79_0
 ```
 
 When cloning Pace you will need to update the repository's submodules as well:
+
 ```shell
 git clone --recursive https://github.com/ai2cm/pace.git
 ```
+
 or if you have already cloned the repository:
+
 ```
 git submodule update --init --recursive
 ```
@@ -43,6 +47,7 @@ source venv_name/bin/activate
 ```
 
 Inside of your pace `venv` or conda environment pip install the Python requirements, GT4Py, and Pace:
+
 ```shell
 pip3 install -r requirements_dev.txt -c constraints.txt
 ```
@@ -52,6 +57,7 @@ Shell scripts to install Pace on specific machines such as Gaea can be found in
 ### Run
 
 With the environment activated, you can run an example baroclinic test case with the following command:
+
 ```shell
 mpirun -n 6 python3 -m pace.driver.run driver/examples/configs/baroclinic_c12.yaml
 
@@ -61,20 +67,33 @@ mpirun -n 6 --oversubscribe python3 -m pace.driver.run driver/examples/configs/b
 
 After the run completes, you will see an output direcotry `output.zarr`. An example to visualize the output is provided in `driver/examples/plot_output.py`. See the [driver example](driver/examples/README.md) section for more details.
 
+### Environment variable configuration
+
+- `PACE_CONSTANTS`: Pace is bundled with various constants (see _util/pace/util/constants.py_).
+  - `FV3DYCORE` NOAA's FV3 dynamical core constants (original port)
+  - `GFS` Constant as defined in NOAA GFS
+  - `GEOS`  Constant as defined in GEOS v13
+- `PACE_FLOAT_PRECISION`: default precision of the field & scalars in the numerics. Default to 64.
+- `PACE_LOGLEVEL`: logging level to display (DEBUG, INFO, WARNING, ERROR, CRITICAL). Default to INFO.
+
 ## Quickstart - Docker
+
 ### Build
 
 While it is possible to install and build pace bare-metal, we can ensure all system libraries are installed with the correct versions by using a Docker container to test and develop pace.
 
 First, you will need to update the git submodules so that any dependencies are cloned and at the correct version:
+
 ```shell
 git submodule update --init --recursive
 ```
 
 Then build the `pace` docker image at the top level.
+
 ```shell
 make build
 ```
+
 ### Run
 
 ```shell
@@ -94,7 +113,6 @@ This git repository is laid out as a mono-repo, containing multiple independent
 
 ![Graph of interdependencies of Pace modules, generated from dependences.dot](./dependencies.svg)
 
-
 ## ML emulation
 
 An example of integration of an ML model replacing the microphysics parametrization is available on the `feature/microphysics-emulator` branch.

driver/pace/driver/run.py

@@ -1,58 +1,15 @@
 import dataclasses
 import gc
-import logging
 from typing import Optional
 
 import click
 import yaml
 
-from pace.util.mpi import MPI
+from pace.util import pace_log, AVAILABLE_LOG_LEVELS
 
 from .driver import Driver, DriverConfig
 
 
-logger = logging.getLogger(__name__)
-
-
-log_levels = {
-    "info": logging.INFO,
-    "debug": logging.DEBUG,
-    "warning": logging.WARNING,
-    "error": logging.ERROR,
-    "critical": logging.CRITICAL,
-}
-
-
-def configure_logging(log_rank: Optional[int], log_level: str):
-    """
-    Configure logging for the driver.
-
-    Args:
-        log_rank: rank to log from, or 'all' to log to all ranks,
-            forced to 'all' if running without MPI
-        log_level: log level to use
-    """
-    level = log_levels[log_level.lower()]
-    if MPI is None:
-        logging.basicConfig(
-            level=level,
-            format="%(asctime)s [%(levelname)s] %(name)s:%(message)s",
-            handlers=[logging.StreamHandler()],
-            datefmt="%Y-%m-%d %H:%M:%S",
-        )
-    else:
-        if log_rank is None or int(log_rank) == MPI.COMM_WORLD.Get_rank():
-            logging.basicConfig(
-                level=level,
-                format=(
-                    f"%(asctime)s [%(levelname)s] (rank {MPI.COMM_WORLD.Get_rank()}) "
-                    "%(name)s:%(message)s"
-                ),
-                handlers=[logging.StreamHandler()],
-                datefmt="%Y-%m-%d %H:%M:%S",
-            )
-
-
 @click.command()
 @click.argument(
     "CONFIG_PATH",
@@ -75,12 +32,15 @@ def command_line(config_path: str, log_rank: Optional[int], log_level: str):
 
     CONFIG_PATH is the path to a DriverConfig yaml file.
     """
-    configure_logging(log_rank=log_rank, log_level=log_level)
-    logger.info("loading DriverConfig from yaml")
+    level = AVAILABLE_LOG_LEVELS[log_level.lower()]
+    pace_log.setLevel(level)
+    pace_log.info("loading DriverConfig from yaml")
     with open(config_path, "r") as f:
         config = yaml.safe_load(f)
         driver_config = DriverConfig.from_dict(config)
-    logging.info(f"DriverConfig loaded: {yaml.dump(dataclasses.asdict(driver_config))}")
+    pace_log.info(
+        f"DriverConfig loaded: {yaml.dump(dataclasses.asdict(driver_config))}"
+    )
     main(driver_config=driver_config)
 
 

dsl/pace/dsl/dace/utils.py

@@ -28,7 +28,7 @@ def __init__(self, config: DaceConfig, label: str):
 
     @classmethod
     def log(cls, prefix: str, message: str):
-        pace_log.info(f"{prefix} {message}")
+        pace_log.debug(f"{prefix} {message}")
 
     @classmethod
     def default_prefix(cls, config: DaceConfig) -> str:

fv3core/pace/fv3core/initialization/geos_wrapper.py

@@ -132,9 +132,12 @@ def __init__(
         self._allocate_output_dir()
 
         pace_log.info(
-            "GEOS-Wrapper with: \n"
+            "Pace GEOS wrapper initialized: \n"
             f"  dt     : {self.dycore_state.bdt}\n"
             f"  bridge : {self._fortran_mem_space} > {self._pace_mem_space}\n"
+            f"  backend: {backend}\n"
+            f"  orchestration: {self._is_orchestrated}\n"
+            f"  sizer  : {sizer.nx}x{sizer.ny}x{sizer.nz} (halo: {sizer.n_halo})"
         )
 
     def _critical_path(self):
@@ -173,7 +176,6 @@ def __call__(
         cyd: np.ndarray,
         diss_estd: np.ndarray,
     ) -> Tuple[Dict[str, np.ndarray], Dict[str, List[float]]]:
-
         with self.perf_collector.timestep_timer.clock("numpy-to-dycore"):
             self.dycore_state = self._put_fortran_data_in_dycore(
                 u,
@@ -246,7 +248,6 @@ def _put_fortran_data_in_dycore(
         cyd: np.ndarray,
         diss_estd: np.ndarray,
     ) -> fv3core.DycoreState:
-
         isc = self._grid_indexing.isc
         jsc = self._grid_indexing.jsc
         iec = self._grid_indexing.iec + 1
@@ -315,7 +316,6 @@ def _put_fortran_data_in_dycore(
         return state
 
     def _prep_outputs_for_geos(self) -> Dict[str, np.ndarray]:
-
         output_dict = self.output_dict
         isc = self._grid_indexing.isc
         jsc = self._grid_indexing.jsc

fv3core/pace/fv3core/stencils/d_sw.py

@@ -932,7 +932,7 @@ def make_quantity():
             )
         )
 
-        if (self._d_con > 1.e-5) or (self._do_stochastic_ke_backscatter):
+        if (self._d_con > 1.0e-5) or (self._do_stochastic_ke_backscatter):
             self._accumulate_heat_source_and_dissipation_estimate_stencil = (
                 stencil_factory.from_dims_halo(
                     func=accumulate_heat_source_and_dissipation_estimate,
@@ -1254,11 +1254,11 @@ def __call__(
             self._column_namelist["d_con"],
         )
 
-        if (self._d_con > 1.e-5) or (self._do_stochastic_ke_backscatter):
+        if (self._d_con > 1.0e-5) or (self._do_stochastic_ke_backscatter):
             self._accumulate_heat_source_and_dissipation_estimate_stencil(
                 self._tmp_heat_s, heat_source, self._tmp_diss_e, diss_est
             )
-        
+
         self._update_u_and_v_stencil(
             self._tmp_ut,
             self._tmp_vt,

fv3core/pace/fv3core/stencils/fv_dynamics.py

@@ -7,7 +7,6 @@
 import pace.dsl.gt4py_utils as utils
 import pace.fv3core.stencils.moist_cv as moist_cv
 import pace.util
-import pace.util.constants as constants
 from pace.dsl.dace.orchestration import dace_inhibitor, orchestrate
 from pace.dsl.dace.wrapped_halo_exchange import WrappedHaloUpdater
 from pace.dsl.stencil import StencilFactory
@@ -21,7 +20,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 X_DIM, Y_DIM, Z_INTERFACE_DIM, Timer
+from pace.util import X_DIM, Y_DIM, Z_INTERFACE_DIM, Timer, constants
 from pace.util.grid import DampingCoefficients, GridData
 from pace.util.logging import pace_log
 from pace.util.mpi import MPI

tests/main/fv3core/test_init_from_geos.py

@@ -7,7 +7,6 @@
 
 
 def test_geos_wrapper():
-
     namelist_dict = {
         "stencil_config": {
             "compilation_config": {
@@ -82,7 +81,12 @@ def test_geos_wrapper():
     comm = NullComm(rank=0, total_ranks=6, fill_value=0.0)
     backend = "numpy"
 
-    wrapper = fv3core.GeosDycoreWrapper(namelist, comm, backend)
+    wrapper = fv3core.GeosDycoreWrapper(
+        namelist=namelist,
+        comm=comm,
+        backend=backend,
+        bdt=namelist_dict["dt_atmos"],
+    )
     nhalo = 3
     shape_centered = (
         namelist["nx_tile"] + 2 * nhalo,
@@ -191,31 +195,33 @@ def test_geos_wrapper():
     )
     diss_estd = np.ones(shape_centered)
 
-    output_dict = wrapper(
-        u,
-        v,
-        w,
-        delz,
-        pt,
-        delp,
-        q,
-        ps,
-        pe,
-        pk,
-        peln,
-        pkz,
-        phis,
-        q_con,
-        omga,
-        ua,
-        va,
-        uc,
-        vc,
-        mfxd,
-        mfyd,
-        cxd,
-        cyd,
-        diss_estd,
+    timings = {}
+    output_dict, timings = wrapper(
+        timings=timings,
+        u=u,
+        v=v,
+        w=w,
+        delz=delz,
+        pt=pt,
+        delp=delp,
+        q=q,
+        ps=ps,
+        pe=pe,
+        pk=pk,
+        peln=peln,
+        pkz=pkz,
+        phis=phis,
+        q_con=q_con,
+        omga=omga,
+        ua=ua,
+        va=va,
+        uc=uc,
+        vc=vc,
+        mfxd=mfxd,
+        mfyd=mfyd,
+        cxd=cxd,
+        cyd=cyd,
+        diss_estd=diss_estd,
     )
 
     assert isinstance(output_dict["u"], np.ndarray)

util/HISTORY.md

@@ -5,6 +5,7 @@ latest
 ------
 
 - Added f32 support to halo exchange data transformation
+- Use one single logger, from logging.py
 
 v0.10.0
 -------

util/pace/util/__init__.py

@@ -54,7 +54,7 @@
 from .initialization import GridSizer, QuantityFactory, SubtileGridSizer
 from .io import read_state, write_state
 from .local_comm import LocalComm
-from .logging import pace_log
+from .logging import pace_log, AVAILABLE_LOG_LEVELS
 from .monitor import Monitor, NetCDFMonitor, ZarrMonitor
 from .mpi import MPIComm
 from .namelist import Namelist, NamelistDefaults