Merge pull request payu-org#376 from jo-basevi/375-date-based-restart…

…-MOM-mixin

Fix date-based restart pruning breaking on empty restart dirs

payu/experiment.py

@@ -1009,6 +1009,10 @@ def get_restarts_to_prune(self,
                           'model. To use integer based restart pruning, '
                           'set restart_freq to an integer value.')
                     raise
+                except FileNotFoundError as e:
+                    print(f'payu: warning: Ignoring {restart} from date-based '
+                          f'restart pruning. Error: {e}')
+                    continue
                 except Exception:
                     print('payu: error: Error parsing restart directory ',
                           f'{restart} for a datetime to prune restarts.')

payu/models/fms.py

@@ -17,8 +17,6 @@
 from itertools import count
 import fnmatch
 
-import cftime
-
 from payu.models.model import Model
 from payu import envmod
 from payu.fsops import required_libs
@@ -255,31 +253,3 @@ def archive(self, **kwargs):
 
     def collate(self):
         fms_collate(self)
-
-    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)"""
-        # Check for ocean_solo.res file
-        ocean_solo_path = os.path.join(restart_path, 'ocean_solo.res')
-        if not os.path.exists(ocean_solo_path):
-            raise NotImplementedError(
-                'Cannot find ocean_solo.res file, which is required for '
-                'date-based restart pruning')
-
-        with open(ocean_solo_path, 'r') as ocean_solo:
-            lines = ocean_solo.readlines()
-
-        calendar_int = int(lines[0].split()[0])
-        cftime_calendars = {
-            1: "360_day",
-            2: "julian",
-            3: "proleptic_gregorian",
-            4: "noleap"
-        }
-        calendar = cftime_calendars[calendar_int]
-
-        last_date_line = lines[-1].split()
-        date_values = [int(i) for i in last_date_line[:6]]
-        year, month, day, hour, minute, second = date_values
-        return cftime.datetime(year, month, day, hour, minute, second,
-                               calendar=calendar)

payu/models/mom.py

@@ -11,10 +11,11 @@
 import f90nml
 
 from payu.models.fms import Fms
+from payu.models.mom_mixin import MomMixin
 from payu.fsops import mkdir_p, make_symlink
 
 
-class Mom(Fms):
+class Mom(MomMixin, Fms):
 
     def __init__(self, expt, name, config):
 

payu/models/mom6.py

@@ -16,6 +16,7 @@
 
 # Local
 from payu.models.fms import Fms
+from payu.models.mom_mixin import MomMixin
 
 
 def mom6_add_parameter_files(model):
@@ -41,7 +42,7 @@ def mom6_add_parameter_files(model):
             model.config_files.extend(filenames)
 
 
-class Mom6(Fms):
+class Mom6(MomMixin, Fms):
     """Interface to GFDL's MOM6 ocean model."""
 
     def __init__(self, expt, name, config):

payu/models/mom_mixin.py

@@ -0,0 +1,40 @@
+"""Mixin class for MOM and MOM6 drivers
+
+:copyright: Copyright 2011 Marshall Ward, see AUTHORS for details
+:license: Apache License, Version 2.0, see LICENSE for details
+"""
+
+import os
+
+import cftime
+
+
+class MomMixin:
+
+    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)"""
+        # Check for ocean_solo.res file
+        ocean_solo_path = os.path.join(restart_path, 'ocean_solo.res')
+        if not os.path.exists(ocean_solo_path):
+            raise FileNotFoundError(
+                'Cannot find ocean_solo.res file, which is required for '
+                'date-based restart pruning')
+
+        with open(ocean_solo_path, 'r') as ocean_solo:
+            lines = ocean_solo.readlines()
+
+        calendar_int = int(lines[0].split()[0])
+        cftime_calendars = {
+            1: "360_day",
+            2: "julian",
+            3: "proleptic_gregorian",
+            4: "noleap"
+        }
+        calendar = cftime_calendars[calendar_int]
+
+        last_date_line = lines[-1].split()
+        date_values = [int(i) for i in last_date_line[:6]]
+        year, month, day, hour, minute, second = date_values
+        return cftime.datetime(year, month, day, hour, minute, second,
+                               calendar=calendar)

test/test_prune_restarts.py

@@ -10,8 +10,8 @@
 from test.common import config as config_orig
 from test.common import write_config
 from test.common import make_all_files
-from test.common import remove_expt_archive_dirs
-from test.models.test_mom import make_ocean_restart_dir
+from test.common import remove_expt_archive_dirs, make_expt_archive_dir
+from test.models.test_mom_mixin import make_ocean_restart_dir
 
 verbose = True
 
@@ -222,3 +222,38 @@ def test_prune_restarts(restarts,
     ]
 
     assert restarts_to_prune_indices == expected_restart_indices
+
+
+def test_prune_restarts_ignores_empty_restart_dirs():
+    # Test that in date-based restart pruning ignores any empty restart dirs
+
+    write_test_config(restart_freq='1YS')
+
+    # Create restart files
+    restart_datetimes = [(4, "1903-01-01 00:00:00"),
+                         (5, "1903-06-01 00:00:00"),
+                         (6, "1905-01-01 00:00:00")]
+    for index, datetime in restart_datetimes:
+        make_ocean_restart_dir(start_dt="1900-01-01 00:00:00",
+                               run_dt=datetime,
+                               calendar=4,
+                               restart_index=index,
+                               additional_path='ocean')
+
+    # Create empty restart dirs
+    make_expt_archive_dir(type='restart', index=0)
+    make_expt_archive_dir(type='restart', index=2)
+
+    with cd(ctrldir):
+        lab = payu.laboratory.Laboratory(lab_path=str(labdir))
+        expt = payu.experiment.Experiment(lab, reproduce=False)
+
+        # Function to test
+        restarts_to_prune = expt.get_restarts_to_prune()
+
+    # Extract out index
+    restarts_to_prune_indices = [
+        int(restart.lstrip('restart')) for restart in restarts_to_prune
+    ]
+
+    assert restarts_to_prune_indices == [5]