MPAS-Dev · xylar · Aug 13, 2021 · Jul 27, 2021 · Jul 27, 2021 · Aug 3, 2021
diff --git a/compass/__main__.py b/compass/__main__.py
@@ -2,9 +2,10 @@
 
 import sys
 import argparse
+import os
 
 import compass
-from compass import list, setup, clean, suite, run
+from compass import list, setup, clean, suite, run, cache
 
 
 def main():
@@ -40,11 +41,18 @@ def main():
 
     args = parser.parse_args(sys.argv[1:2])
 
+    # only allow the "compass cache" command if we're on Anvil or Chrysalis
+    allow_cache = ('COMPASS_MACHINE' in os.environ and
+                   os.environ['COMPASS_MACHINE'] in ['anvil', 'chrysalis'])
+
     commands = {'list': list.main,
                 'setup': setup.main,
                 'clean': clean.main,
                 'suite': suite.main,
                 'run': run.main}
+    if allow_cache:
+        commands['cache'] = cache.main
+
     if args.command not in commands:
         print('Unrecognized command {}'.format(args.command))
         parser.print_help()

diff --git a/compass/cache.py b/compass/cache.py
@@ -0,0 +1,135 @@
+import argparse
+import json
+import sys
+from datetime import datetime
+import os
+from importlib import resources
+import configparser
+import shutil
+import pickle
+
+from compass.config import add_config
+
+
+def update_cache(step_paths, date_string=None, dry_run=False):
+    """
+    Cache one or more compass output files for use in a cached variant of the
+    test case or step
+
+    Parameters
+    ----------
+    step_paths : list of str
+        The relative path of the original (uncached) steps from the base work
+        directory
+
+    date_string : str, optional
+        The datestamp (YYMMDD) to use on the files.  Default is today's date.
+
+    dry_run : bool, optional
+        Whether this is a dry run (producing the json file but not copying
+        files to the LCRC server)
+    """
+    if 'COMPASS_MACHINE' not in os.environ:
+        machine = None
+        invalid = True
+    else:
+        machine = os.environ['COMPASS_MACHINE']
+        invalid = machine not in ['anvil', 'chrysalis']
+
+    if invalid:
+        raise ValueError('You must cache files from either Anvil or Chrysalis')
+
+    config = configparser.ConfigParser(
+        interpolation=configparser.ExtendedInterpolation())
+    add_config(config, 'compass.machines', '{}.cfg'.format(machine))
+
+    if date_string is None:
+        date_string = datetime.now().strftime("%y%m%d")
+
+    # make a dictionary with MPAS cores as keys, and lists of steps as values
+    steps = dict()
+    for path in step_paths:
+        with open(f'{path}/step.pickle', 'rb') as handle:
+            _, step = pickle.load(handle)
+
+        mpas_core = step.mpas_core.name
+
+        if mpas_core in steps:
+            steps[mpas_core].append(step)
+        else:
+            steps[mpas_core] = [step]
+
+    # now, iterate over cores and steps
+    for mpas_core in steps:
+        database_root = config.get('paths', f'{mpas_core}_database_root')
+        cache_root = f'{database_root}/compass_cache'
+
+        package = f'compass.{mpas_core}'
+        try:
+            with open(f'{mpas_core}_cached_files.json') as data_file:
+                cached_files = json.load(data_file)
+        except FileNotFoundError:
+            # we don't have a local version of the file yet, let's see if
+            # there's a remote one for this MPAS core
+            try:
+                with resources.path(package, 'cached_files.json') as path:
+                    with open(path) as data_file:
+                        cached_files = json.load(data_file)
+            except FileNotFoundError:
+                # no cached files yet for this core
+                cached_files = dict()
+
+        for step in steps[mpas_core]:
+            # load the step from its pickle file
+
+            step_path = step.path
+
+            for output in step.outputs:
+                output = os.path.basename(output)
+                out_filename = os.path.join(step_path, output)
+                # remove the MPAS core from the file path
+                target = out_filename[len(mpas_core)+1:]
+                path, ext = os.path.splitext(target)
+                target = f'{path}.{date_string}{ext}'
+                cached_files[out_filename] = target
+
+                print(out_filename)
+                print(f'  ==> {target}')
+                output_path = f'{cache_root}/{target}'
+                print(f'  copy to: {output_path}')
+                print()
+                if not dry_run:
+                    directory = os.path.dirname(output_path)
+                    try:
+                        os.makedirs(directory)
+                    except FileExistsError:
+                        pass
+                    shutil.copyfile(out_filename, output_path)
+
+        out_filename = f'{mpas_core}_cached_files.json'
+        with open(out_filename, 'w') as data_file:
+            json.dump(cached_files, data_file, indent=4)
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description='Cache the output files from one or more steps for use in '
+                    'a cached variant of the step',
+        prog='compass cache')
+    parser.add_argument("-i", "--orig_steps", nargs='+', dest="orig_steps",
+                        type=str,
+                        help="The relative path of the original (uncached) "
+                             "steps from the base work directory",
+                        metavar="STEP")
+    parser.add_argument("-d", "--date_string", dest="date_string", type=str,
+                        help="The datestamp (YYMMDD) to use on the files.  "
+                             "Default is today's date.",
+                        metavar="DATE")
+    parser.add_argument("-r", "--dry_run", dest="dry_run",
+                        help="Whether this is a dry run (producing the json "
+                             "file but not copying files to the LCRC server).",
+                        action="store_true")
+
+    args = parser.parse_args(sys.argv[2:])
+    update_cache(step_paths=args.orig_steps, date_string=args.date_string,
+                 dry_run=args.dry_run)
diff --git a/compass/mpas_core.py b/compass/mpas_core.py
@@ -1,3 +1,7 @@
+from importlib import resources
+import json
+
+
 class MpasCore:
     """
     The base class for housing all the tests for a given MPAS core, such as
@@ -10,6 +14,11 @@ class MpasCore:
 
     test_groups : dict
         A dictionary of test groups for the MPAS core with their names as keys
+
+    cached_files : dict
+        A dictionary that maps from output file names in test cases to cached
+        files in the ``compass_cache`` database for the MPAS core.  These
+        file mappings are read in from ``cached_files.json`` in the MPAS core.
     """
 
     def __init__(self, name):
@@ -26,6 +35,9 @@ def __init__(self, name):
         # test groups are added with add_test_groups()
         self.test_groups = dict()
 
+        self.cached_files = dict()
+        self._read_cached_files()
+
     def add_test_group(self, test_group):
         """
         Add a test group to the MPAS core
@@ -36,3 +48,16 @@ def add_test_group(self, test_group):
             the test group to add
         """
         self.test_groups[test_group.name] = test_group
+
+    def _read_cached_files(self):
+        """ Read in the dictionary of cached files from cached_files.json """
+
+        package = f'compass.{self.name}'
+        filename = 'cached_files.json'
+        try:
+            with resources.path(package, filename) as path:
+                with open(path) as data_file:
+                    self.cached_files = json.load(data_file)
+        except FileNotFoundError:
+            # no cached files for this core
+            pass
diff --git a/compass/ocean/cached_files.json b/compass/ocean/cached_files.json
@@ -0,0 +1,69 @@
+{
+    "ocean/global_ocean/QU240/mesh/mesh/culled_mesh.nc": "global_ocean/QU240/mesh/mesh/culled_mesh.210809.nc",
+    "ocean/global_ocean/QU240/mesh/mesh/culled_graph.info": "global_ocean/QU240/mesh/mesh/culled_graph.210809.info",
+    "ocean/global_ocean/QU240/mesh/mesh/critical_passages_mask_final.nc": "global_ocean/QU240/mesh/mesh/critical_passages_mask_final.210809.nc",
+    "ocean/global_ocean/QU240/PHC/init/initial_state/initial_state.nc": "global_ocean/QU240/PHC/init/initial_state/initial_state.210809.nc",
+    "ocean/global_ocean/QU240/PHC/init/initial_state/init_mode_forcing_data.nc": "global_ocean/QU240/PHC/init/initial_state/init_mode_forcing_data.210809.nc",
+    "ocean/global_ocean/QU240/PHC/init/initial_state/graph.info": "global_ocean/QU240/PHC/init/initial_state/graph.210809.info",
+    "ocean/global_ocean/QUwISC240/mesh/mesh/culled_mesh.nc": "global_ocean/QUwISC240/mesh/mesh/culled_mesh.210809.nc",
+    "ocean/global_ocean/QUwISC240/mesh/mesh/culled_graph.info": "global_ocean/QUwISC240/mesh/mesh/culled_graph.210809.info",
+    "ocean/global_ocean/QUwISC240/mesh/mesh/critical_passages_mask_final.nc": "global_ocean/QUwISC240/mesh/mesh/critical_passages_mask_final.210809.nc",
+    "ocean/global_ocean/QUwISC240/PHC/init/initial_state/initial_state.nc": "global_ocean/QUwISC240/PHC/init/initial_state/initial_state.210809.nc",
+    "ocean/global_ocean/QUwISC240/PHC/init/initial_state/init_mode_forcing_data.nc": "global_ocean/QUwISC240/PHC/init/initial_state/init_mode_forcing_data.210809.nc",
+    "ocean/global_ocean/QUwISC240/PHC/init/initial_state/graph.info": "global_ocean/QUwISC240/PHC/init/initial_state/graph.210809.info",
+    "ocean/global_ocean/QUwISC240/PHC/init/ssh_adjustment/adjusted_init.nc": "global_ocean/QUwISC240/PHC/init/ssh_adjustment/adjusted_init.210809.nc",
+    "ocean/global_ocean/EC30to60/mesh/mesh/culled_mesh.nc": "global_ocean/EC30to60/mesh/mesh/culled_mesh.210809.nc",
+    "ocean/global_ocean/EC30to60/mesh/mesh/culled_graph.info": "global_ocean/EC30to60/mesh/mesh/culled_graph.210809.info",
+    "ocean/global_ocean/EC30to60/mesh/mesh/critical_passages_mask_final.nc": "global_ocean/EC30to60/mesh/mesh/critical_passages_mask_final.210809.nc",
+    "ocean/global_ocean/EC30to60/PHC/init/initial_state/initial_state.nc": "global_ocean/EC30to60/PHC/init/initial_state/initial_state.210809.nc",
+    "ocean/global_ocean/EC30to60/PHC/init/initial_state/init_mode_forcing_data.nc": "global_ocean/EC30to60/PHC/init/initial_state/init_mode_forcing_data.210809.nc",
+    "ocean/global_ocean/EC30to60/PHC/init/initial_state/graph.info": "global_ocean/EC30to60/PHC/init/initial_state/graph.210809.info",
+    "ocean/global_ocean/WC14/mesh/mesh/culled_mesh.nc": "global_ocean/WC14/mesh/mesh/culled_mesh.210809.nc",
+    "ocean/global_ocean/WC14/mesh/mesh/culled_graph.info": "global_ocean/WC14/mesh/mesh/culled_graph.210809.info",
+    "ocean/global_ocean/WC14/mesh/mesh/critical_passages_mask_final.nc": "global_ocean/WC14/mesh/mesh/critical_passages_mask_final.210809.nc",
+    "ocean/global_ocean/WC14/PHC/init/initial_state/initial_state.nc": "global_ocean/WC14/PHC/init/initial_state/initial_state.210809.nc",
+    "ocean/global_ocean/WC14/PHC/init/initial_state/init_mode_forcing_data.nc": "global_ocean/WC14/PHC/init/initial_state/init_mode_forcing_data.210809.nc",
+    "ocean/global_ocean/WC14/PHC/init/initial_state/graph.info": "global_ocean/WC14/PHC/init/initial_state/graph.210809.info",
+    "ocean/global_ocean/ECwISC30to60/mesh/mesh/culled_mesh.nc": "global_ocean/ECwISC30to60/mesh/mesh/culled_mesh.210809.nc",
+    "ocean/global_ocean/ECwISC30to60/mesh/mesh/culled_graph.info": "global_ocean/ECwISC30to60/mesh/mesh/culled_graph.210809.info",
+    "ocean/global_ocean/ECwISC30to60/mesh/mesh/critical_passages_mask_final.nc": "global_ocean/ECwISC30to60/mesh/mesh/critical_passages_mask_final.210809.nc",
+    "ocean/global_ocean/ECwISC30to60/PHC/init/initial_state/initial_state.nc": "global_ocean/ECwISC30to60/PHC/init/initial_state/initial_state.210809.nc",
+    "ocean/global_ocean/ECwISC30to60/PHC/init/initial_state/init_mode_forcing_data.nc": "global_ocean/ECwISC30to60/PHC/init/initial_state/init_mode_forcing_data.210809.nc",
+    "ocean/global_ocean/ECwISC30to60/PHC/init/ssh_adjustment/adjusted_init.nc": "global_ocean/ECwISC30to60/PHC/init/ssh_adjustment/adjusted_init.210809.nc",
+    "ocean/global_ocean/ECwISC30to60/PHC/init/initial_state/graph.info": "global_ocean/ECwISC30to60/PHC/init/initial_state/graph.210809.info",
+    "ocean/global_ocean/SOwISC12to60/mesh/mesh/culled_mesh.nc": "global_ocean/SOwISC12to60/mesh/mesh/culled_mesh.210810.nc",
+    "ocean/global_ocean/SOwISC12to60/mesh/mesh/culled_graph.info": "global_ocean/SOwISC12to60/mesh/mesh/culled_graph.210810.info",
+    "ocean/global_ocean/SOwISC12to60/mesh/mesh/critical_passages_mask_final.nc": "global_ocean/SOwISC12to60/mesh/mesh/critical_passages_mask_final.210810.nc",
+    "ocean/global_ocean/SOwISC12to60/PHC/init/initial_state/initial_state.nc": "global_ocean/SOwISC12to60/PHC/init/initial_state/initial_state.210810.nc",
+    "ocean/global_ocean/SOwISC12to60/PHC/init/initial_state/init_mode_forcing_data.nc": "global_ocean/SOwISC12to60/PHC/init/initial_state/init_mode_forcing_data.210810.nc",
+    "ocean/global_ocean/SOwISC12to60/PHC/init/ssh_adjustment/adjusted_init.nc": "global_ocean/SOwISC12to60/PHC/init/ssh_adjustment/adjusted_init.210810.nc",
+    "ocean/global_ocean/SOwISC12to60/PHC/init/initial_state/graph.info": "global_ocean/SOwISC12to60/PHC/init/initial_state/graph.210810.info",
+    "ocean/global_convergence/cosine_bell/QU60/mesh/mesh.nc": "global_convergence/cosine_bell/QU60/mesh/mesh.210803.nc",
+    "ocean/global_convergence/cosine_bell/QU60/mesh/graph.info": "global_convergence/cosine_bell/QU60/mesh/graph.210803.info",
+    "ocean/global_convergence/cosine_bell/QU60/init/namelist.ocean": "global_convergence/cosine_bell/QU60/init/namelist.210803.ocean",
+    "ocean/global_convergence/cosine_bell/QU60/init/initial_state.nc": "global_convergence/cosine_bell/QU60/init/initial_state.210803.nc",
+    "ocean/global_convergence/cosine_bell/QU90/mesh/mesh.nc": "global_convergence/cosine_bell/QU90/mesh/mesh.210803.nc",
+    "ocean/global_convergence/cosine_bell/QU90/mesh/graph.info": "global_convergence/cosine_bell/QU90/mesh/graph.210803.info",
+    "ocean/global_convergence/cosine_bell/QU90/init/namelist.ocean": "global_convergence/cosine_bell/QU90/init/namelist.210803.ocean",
+    "ocean/global_convergence/cosine_bell/QU90/init/initial_state.nc": "global_convergence/cosine_bell/QU90/init/initial_state.210803.nc",
+    "ocean/global_convergence/cosine_bell/QU120/mesh/mesh.nc": "global_convergence/cosine_bell/QU120/mesh/mesh.210803.nc",
+    "ocean/global_convergence/cosine_bell/QU120/mesh/graph.info": "global_convergence/cosine_bell/QU120/mesh/graph.210803.info",
+    "ocean/global_convergence/cosine_bell/QU120/init/namelist.ocean": "global_convergence/cosine_bell/QU120/init/namelist.210803.ocean",
+    "ocean/global_convergence/cosine_bell/QU120/init/initial_state.nc": "global_convergence/cosine_bell/QU120/init/initial_state.210803.nc",
+    "ocean/global_convergence/cosine_bell/QU180/mesh/mesh.nc": "global_convergence/cosine_bell/QU180/mesh/mesh.210803.nc",
+    "ocean/global_convergence/cosine_bell/QU180/mesh/graph.info": "global_convergence/cosine_bell/QU180/mesh/graph.210803.info",
+    "ocean/global_convergence/cosine_bell/QU180/init/namelist.ocean": "global_convergence/cosine_bell/QU180/init/namelist.210803.ocean",
+    "ocean/global_convergence/cosine_bell/QU180/init/initial_state.nc": "global_convergence/cosine_bell/QU180/init/initial_state.210803.nc",
+    "ocean/global_convergence/cosine_bell/QU210/mesh/mesh.nc": "global_convergence/cosine_bell/QU210/mesh/mesh.210803.nc",
+    "ocean/global_convergence/cosine_bell/QU210/mesh/graph.info": "global_convergence/cosine_bell/QU210/mesh/graph.210803.info",
+    "ocean/global_convergence/cosine_bell/QU210/init/namelist.ocean": "global_convergence/cosine_bell/QU210/init/namelist.210803.ocean",
+    "ocean/global_convergence/cosine_bell/QU210/init/initial_state.nc": "global_convergence/cosine_bell/QU210/init/initial_state.210803.nc",
+    "ocean/global_convergence/cosine_bell/QU240/mesh/mesh.nc": "global_convergence/cosine_bell/QU240/mesh/mesh.210803.nc",
+    "ocean/global_convergence/cosine_bell/QU240/mesh/graph.info": "global_convergence/cosine_bell/QU240/mesh/graph.210803.info",
+    "ocean/global_convergence/cosine_bell/QU240/init/namelist.ocean": "global_convergence/cosine_bell/QU240/init/namelist.210803.ocean",
+    "ocean/global_convergence/cosine_bell/QU240/init/initial_state.nc": "global_convergence/cosine_bell/QU240/init/initial_state.210803.nc",
+    "ocean/global_convergence/cosine_bell/QU150/mesh/mesh.nc": "global_convergence/cosine_bell/QU150/mesh/mesh.210803.nc",
+    "ocean/global_convergence/cosine_bell/QU150/mesh/graph.info": "global_convergence/cosine_bell/QU150/mesh/graph.210803.info",
+    "ocean/global_convergence/cosine_bell/QU150/init/namelist.ocean": "global_convergence/cosine_bell/QU150/init/namelist.210803.ocean",
+    "ocean/global_convergence/cosine_bell/QU150/init/initial_state.nc": "global_convergence/cosine_bell/QU150/init/initial_state.210803.nc"
+}
diff --git a/compass/ocean/suites/cosine_bell_cached_init.txt b/compass/ocean/suites/cosine_bell_cached_init.txt
@@ -0,0 +1,4 @@
+ocean/global_convergence/cosine_bell
+  cached: QU60_mesh QU60_init QU90_mesh QU90_init QU120_mesh QU120_init
+  cached: QU150_mesh QU150_init QU180_mesh QU180_init QU210_mesh QU210_init
+  cached: QU240_mesh QU240_init
diff --git a/compass/ocean/suites/nightly.txt b/compass/ocean/suites/nightly.txt
@@ -22,6 +22,12 @@ ocean/global_ocean/QU240/EN4_1900/performance_test
 ocean/global_ocean/QU240/PHC_BGC/init
 ocean/global_ocean/QU240/PHC_BGC/performance_test
 
+ocean/global_ocean/QUwISC240/mesh
+  cached
+ocean/global_ocean/QUwISC240/PHC/init
+  cached
+ocean/global_ocean/QUwISC240/PHC/performance_test
+
 ocean/ice_shelf_2d/5km/z-star/restart_test
 ocean/ice_shelf_2d/5km/z-level/restart_test