From b476d21c52f2354ad1c36ac72dcb5dbf5a9a2398 Mon Sep 17 00:00:00 2001
From: Jo Basevi <johanna.basevi@edu.anu.au>
Date: Fri, 22 Sep 2023 10:32:57 +1000
Subject: [PATCH] Refactor set_counter() and some bug fixes

- Fix to payu crashing with empty restart files/dirs
- Add default values of LD_LIBRARY_PATHS
- Raise exceptions for missing model executables
---
 payu/experiment.py   | 73 ++++++++++++++++++++++----------------------
 payu/models/model.py |  4 +++
 2 files changed, 41 insertions(+), 36 deletions(-)

diff --git a/payu/experiment.py b/payu/experiment.py
index 970a1f48..0746799e 100644
--- a/payu/experiment.py
+++ b/payu/experiment.py
@@ -13,6 +13,7 @@
 import errno
 import getpass
 import os
+import re
 import resource
 import sys
 import shlex
@@ -172,41 +173,40 @@ def set_counters(self):
 
         # Initialize counter if unset
         if self.counter is None:
-            # TODO: this logic can probably be streamlined
-            try:
-                restart_dirs = [d for d in os.listdir(self.archive_path)
-                                if d.startswith('restart')]
-            except EnvironmentError as exc:
-                if exc.errno == errno.ENOENT:
-                    restart_dirs = None
-                else:
-                    raise
-
-            # First test for restarts
-            if restart_dirs:
-                self.counter = 1 + max([int(d.lstrip('restart'))
-                                        for d in restart_dirs
-                                        if d.startswith('restart')])
+            # Check for restart index
+            max_restart_index = self.max_output_index(output_type="restart")
+            if max_restart_index:
+                self.counter = 1 + max_restart_index
             else:
-                # repeat runs do not generate restart files, so check outputs
-                try:
-                    output_dirs = [d for d in os.listdir(self.archive_path)
-                                   if d.startswith('output')]
-                except EnvironmentError as exc:
-                    if exc.errno == errno.ENOENT:
-                        output_dirs = None
-                    else:
-                        raise
-
-                # First test for restarts
-                # Now look for output directories
-                if output_dirs:
-                    self.counter = 1 + max([int(d.lstrip('output'))
-                                            for d in output_dirs
-                                            if d.startswith('output')])
+                # Now look for output directories,
+                # as repeat runs do not generate restart files.
+                max_output_index = self.max_output_index()
+                if max_output_index:
+                    self.counter = 1 + max_output_index
                 else:
                     self.counter = 0
 
+    def max_output_index(self, output_type="output"):
+        """Given a output directory type (output or restart), 
+        return the maximum index of output directories found"""
+        try:
+            output_dirs = self.list_output_dirs(output_type)
+        except EnvironmentError as exc:
+            if exc.errno == errno.ENOENT:
+                output_dirs = None
+            else:
+                raise
+        
+        if output_dirs and len(output_dirs):
+            return max([int(d.lstrip(output_type))
+                                for d in output_dirs])
+    
+    def list_output_dirs(self, output_type="output"):
+        """Return a list of restart or output directories in archive"""
+        naming_pattern = re.compile(fr"^{output_type}[0-9][0-9][0-9]$")
+        return [d for d in os.listdir(self.archive_path)
+                            if naming_pattern.match(d)]       
+
     def set_stacksize(self, stacksize):
 
         if stacksize == 'unlimited':
@@ -749,8 +749,7 @@ def archive(self):
                                           default_restart_history)
 
         # Remove any outdated restart files
-        prior_restart_dirs = [d for d in os.listdir(self.archive_path)
-                              if d.startswith('restart')]
+        prior_restart_dirs = self.list_output_dirs(output_type="restart")
 
         for res_dir in prior_restart_dirs:
 
@@ -766,10 +765,12 @@ def archive(self):
                     shutil.rmtree(res_path)
 
         # Ensure dynamic library support for subsequent python calls
-        ld_libpaths = os.environ['LD_LIBRARY_PATH']
+        ld_libpaths = os.environ.get('LD_LIBRARY_PATH', None)
         py_libpath = sysconfig.get_config_var('LIBDIR')
-        if py_libpath not in ld_libpaths.split(':'):
-            os.environ['LD_LIBRARY_PATH'] = ':'.join([py_libpath, ld_libpaths])
+        if ld_libpaths is None:
+            os.environ['LD_LIBRARY_PATH'] = py_libpath
+        elif py_libpath not in ld_libpaths.split(':'):
+            os.environ['LD_LIBRARY_PATH'] = f'{py_libpath}:{ld_libpaths}'
 
         collate_config = self.config.get('collate', {})
         collating = collate_config.get('enable', True)
diff --git a/payu/models/model.py b/payu/models/model.py
index ac6643d2..87daccab 100644
--- a/payu/models/model.py
+++ b/payu/models/model.py
@@ -283,6 +283,10 @@ def setup(self):
 
         # Make symlink to executable in work directory
         if self.exec_path:
+            # Check whether executable path exists
+            if not os.path.isfile(self.exec_path):
+                raise FileNotFoundError(f'Executable not found on path: {self.exec_path}')
+
             # If have exe manifest this implies exe reproduce is True. Do not
             # want to overwrite exe manifest in this case
             if not self.expt.manifest.have_manifest['exe']: