Add date-based restart pruning to OM3 driver (#562)

* Update payu/models/cesm_cmeps.py to provide date-based restart pruning

docs/source/config.rst

@@ -217,9 +217,9 @@ configuration.
    For example, ``restart_freq: 10YS`` would save earliest restart of the year,
    10 years from the last permanently archived restart's datetime.
 
-   Please note that currently, only ACCESS-OM2, MOM5 and MOM6 models support
-   date-based restart frequency, as it depends on the payu model driver being
-   able to parse restarts files for a datetime.
+   Please note that currently, only ACCESS-ESM1.5, ACCESS-OM2, ACCESS-OM3, MOM5 
+   , MOM6 and UM7 models support date-based restart frequency, as it depends on the payu 
+   model driver being able to parse restarts files for a datetime.
 
 ``restart_history``
     Specifies how many of the most recent restart files to retain regardless of 

payu/models/cesm_cmeps.py

@@ -13,6 +13,7 @@
 import os
 import re
 import shutil
+import cftime
 from warnings import warn
 
 from payu.fsops import mkdir_p, make_symlink
@@ -23,6 +24,14 @@
 NUOPC_CONFIG = "nuopc.runconfig"
 NUOPC_RUNSEQ = "nuopc.runseq"
 
+# mapping of runconfig to cftime calendars:
+# these match the supported calendars in CMEPS
+# https://github.com/ESCOMP/CMEPS/blob/bc29792d76c16911046dbbfcfc7f4c3ae89a6f00/cesm/driver/ensemble_driver.F90#L196
+CFTIME_CALENDARS = {
+    "NO_LEAP" : "noleap" ,
+    "GREGORIAN" : "proleptic_gregorian"
+}
+
 # Add as needed
 component_info = {
     "mom": {
@@ -330,6 +339,48 @@ def collate(self):
         else:
             super().collate()
 
+    def get_restart_datetime(self, restart_path):
+        """Given a restart path, parse the restart files and
+        return a cftime datetime (for date-based restart pruning)
+        Supports noleap and proleptic gregorian calendars"""
+
+        # Check for rpointer.cpl file
+        rpointer_path = os.path.join(restart_path, 'rpointer.cpl')
+        if not os.path.exists(rpointer_path):
+            raise FileNotFoundError(
+                'Cannot find rpointer.cpl file, which is required for '
+                'date-based restart pruning')
+
+        with open(rpointer_path, 'r') as ocean_solo:
+            lines = ocean_solo.readlines()
+            # example lines would be access-om3.cpl.r.1900-01-02-00000.nc
+
+        date_str = lines[0].split('.')[3]
+        year, month, day, seconds = date_str.split('-')
+
+        # convert seconds into hours, mins, seconds
+        seconds = int(seconds)
+        hour = seconds // 3600 ; min = (seconds % 3600) // 60 ; sec = seconds % 60
+
+        # TODO: change to self.control_path if https://github.com/payu-org/payu/issues/509 is implemented
+        self.get_runconfig(self.expt.control_path)
+        run_calendar = self.runconfig.get("CLOCK_attributes", "calendar")
+
+        try:
+            cf_cal = CFTIME_CALENDARS[run_calendar]
+        except KeyError as e:
+            raise RuntimeError(
+                f"Unsupported calendar for restart pruning: {run_calendar}, in {NUOPC_CONFIG}."
+                f" Try {' or '.join(CFTIME_CALENDARS.keys())}"
+            ) from e
+            return False
+
+        return cftime.datetime(
+            int(year), int(month), int(day), 
+            hour, min, sec, 
+            calendar  = cf_cal
+        )
+
 
 class AccessOm3(CesmCmeps):
 

test/models/access-om3/test_access_om3.py

@@ -5,10 +5,12 @@
 import pytest
 
 import payu
+import cftime
 
 from test.common import cd, tmpdir, ctrldir, labdir, workdir, write_config, config_path
 from test.common import config as config_orig
 from test.common import make_inputs, make_exe
+from test.common import list_expt_archive_dirs, make_expt_archive_dir, remove_expt_archive_dirs
 
 MODEL = 'access-om3'
 
@@ -312,3 +314,111 @@ def test__setup_checks_bad_io(pio_numiotasks, pio_stride):
             model._setup_checks()
 
     teardown_cmeps_config()
+
+
+# test restart datetime pruning
+
+def make_restart_dir(start_dt):
+    """Create restart directory with rpointer.cpl file"""
+    # Create restart directory
+    restart_path = make_expt_archive_dir(type='restart')
+
+    rpath = os.path.join(restart_path, "rpointer.cpl")
+    with open(rpath, "w") as rpointer_file:
+        rpointer_file.write(
+            f"access-om3.cpl.r.{start_dt}.nc"
+        )
+
+@pytest.mark.parametrize(
+    "start_dt, calendar, cmeps_calendar, expected_cftime",
+    [
+        (
+            "0001-01-01-00000",
+            "proleptic_gregorian",
+            "GREGORIAN",
+            cftime.datetime(1, 1, 1, calendar="proleptic_gregorian")
+        ),
+        (
+            "9999-12-31-86399",
+            "proleptic_gregorian",
+            "GREGORIAN",
+            cftime.datetime(9999, 12, 31,23,59,59, calendar="proleptic_gregorian")
+        ),
+        (
+            "1900-02-01-00000",
+            "noleap",
+            "NO_LEAP",
+            cftime.datetime(1900, 2, 1, calendar="noleap")
+        ),
+    ])
+@pytest.mark.filterwarnings("error")
+def test_get_restart_datetime(start_dt, calendar, cmeps_calendar, expected_cftime):
+
+    cmeps_config(1)
+
+    make_restart_dir(start_dt)
+
+    test_runconf = {
+        "CLOCK_attributes": {
+            "calendar":cmeps_calendar
+        }
+    }
+
+    with cd(ctrldir):
+        lab = payu.laboratory.Laboratory(lab_path=str(labdir))
+        expt = payu.experiment.Experiment(lab, reproduce=False)
+
+        model = expt.model
+        model.get_runconfig = lambda a : (True)         # mock reading runconf from file
+        model.runconfig = MockRunConfig(test_runconf)
+
+        print(model.runconfig.get("CLOCK_attributes","calendar"))
+
+    restart_path = list_expt_archive_dirs()[0]
+    parsed_run_dt = expt.model.get_restart_datetime(restart_path)
+    assert parsed_run_dt == expected_cftime
+
+    teardown_cmeps_config()
+    remove_expt_archive_dirs(type='restart')
+
+@pytest.mark.parametrize(
+    "start_dt, calendar, cmeps_calendar, expected_cftime",
+    [
+        (
+            "1900-02-01-00000",
+            "julian",
+            "JULIAN",
+            cftime.datetime(1900, 2, 1, calendar="julian")
+        ),
+    ])
+def test_get_restart_datetime_badcal(start_dt, calendar, cmeps_calendar, expected_cftime):
+
+    cmeps_config(1)
+
+    make_restart_dir(start_dt)
+
+    test_runconf = {
+        "CLOCK_attributes": {
+            "calendar":cmeps_calendar
+        }
+    }
+
+    with cd(ctrldir):
+        lab = payu.laboratory.Laboratory(lab_path=str(labdir))
+        expt = payu.experiment.Experiment(lab, reproduce=False)
+
+        model = expt.model
+        model.get_runconfig = lambda a : (True)         # mock reading runconf from file
+        model.runconfig = MockRunConfig(test_runconf)
+
+        print(model.runconfig.get("CLOCK_attributes","calendar"))
+
+    restart_path = list_expt_archive_dirs()[0]
+    with pytest.raises(
+            RuntimeError, match="Unsupported calendar"
+        ):
+        expt.model.get_restart_datetime(restart_path)
+
+    teardown_cmeps_config()
+    remove_expt_archive_dirs(type='restart')
+