Plugins for zppy (#402)

Implementation of plugins for zppy

---------

Co-authored-by: Ryan Forsyth <forsyth2@llnl.gov>

tests/test_sections.py

@@ -81,6 +81,7 @@ def test_sections(self):
             "output": "OUTPUT",
             "nodes": 1,
             "partition": "SHORT",
+            "plugins": [],
             "qos": "regular",
             "templateDir": "zppy/templates",
             "vars": "FSNTOA,FLUT,FSNT,FLNT,FSNS,FLNS,SHFLX,QFLX,TAUX,TAUY,PRECC,PRECL,PRECSC,PRECSL,TS,TREFHT,CLDTOT,CLDHGH,CLDMED,CLDLOW,U",
@@ -132,6 +133,7 @@ def test_sections(self):
             "nodes": 1,
             "output": "OUTPUT",
             "partition": "SHORT",
+            "plugins": [],
             "qos": "regular",
             "subsection": None,
             "templateDir": "zppy/templates",
@@ -179,6 +181,7 @@ def test_sections(self):
             "nodes": 4,
             "output": "OUTPUT",
             "partition": "SHORT",
+            "plugins": [],
             "qos": "regular",
             "subsection": None,
             "templateDir": "zppy/templates",
@@ -227,6 +230,7 @@ def test_subsections(self):
             "nodes": 1,
             "output": "OUTPUT",
             "partition": "SHORT",
+            "plugins": [],
             "qos": "regular",
             "templateDir": "zppy/templates",
             "vars": "FSNTOA,FLUT,FSNT,FLNT,FSNS,FLNS,SHFLX,QFLX,TAUX,TAUY,PRECC,PRECL,PRECSC,PRECSL,TS,TREFHT,CLDTOT,CLDHGH,CLDMED,CLDLOW,U",
@@ -296,6 +300,7 @@ def test_subsections(self):
             "nodes": 1,
             "output": "OUTPUT",
             "partition": "SHORT",
+            "plugins": [],
             "qos": "regular",
             "subsection": "ts_grid1",
             "templateDir": "zppy/templates",
@@ -331,6 +336,7 @@ def test_subsections(self):
             "nodes": 1,
             "output": "OUTPUT",
             "partition": "SHORT",
+            "plugins": [],
             "qos": "regular",
             "subsection": "ts_grid2",
             "templateDir": "zppy/templates",
@@ -392,6 +398,7 @@ def test_subsections(self):
             "nodes": 4,
             "output": "OUTPUT",
             "partition": "SHORT",
+            "plugins": [],
             "qos": "regular",
             "subsection": "climo_grid1",
             "templateDir": "zppy/templates",
@@ -422,6 +429,7 @@ def test_subsections(self):
             "nodes": 4,
             "output": "OUTPUT",
             "partition": "LONG",
+            "plugins": [],
             "qos": "regular",
             "subsection": "climo_grid2",
             "templateDir": "zppy/templates",

zppy/__main__.py

@@ -1,5 +1,7 @@
 import argparse
 import errno
+import importlib
+import io
 import os
 from typing import List
 
@@ -40,6 +42,38 @@ def main():  # noqa: C901
     # Read configuration file and validate it
     default_config = os.path.join(templateDir, "default.ini")
     user_config = ConfigObj(args.config, configspec=default_config)
+
+    # Load all external plugins. Build a list.
+    plugins = []
+    if "plugins" in user_config["default"].keys():
+        for plugin_name in user_config["default"]["plugins"]:
+            # Load plugin module
+            try:
+                plugin_module = importlib.import_module(plugin_name)
+            except BaseException:
+                raise ValueError(
+                    "Could not load external zppy plugin module {}".format(plugin_name)
+                )
+            # Path
+            plugin_path = plugin_module.__path__[0]
+            # Add to list
+            plugins.append(
+                {"name": plugin_name, "module": plugin_module, "path": plugin_path}
+            )
+
+    # Read configuration files again, this time including all plugins
+    with open(default_config) as f:
+        default = f.read()
+    for plugin in plugins:
+        # Read plugin 'default.ini' if it exists
+        plugin_default_file = os.path.join(plugin["path"], "templates/default.ini")
+        print(plugin_default_file)
+        if os.path.isfile(plugin_default_file):
+            with open(plugin_default_file) as f:
+                default += "\n" + f.read()
+    user_config = ConfigObj(args.config, configspec=io.StringIO(default))
+
+    # Handle 'campaign' option
     if "campaign" in user_config["default"]:
         campaign = user_config["default"]["campaign"]
     else:
@@ -56,6 +90,8 @@ def main():  # noqa: C901
         config = campaign_config
     else:
         config = user_config
+
+    # Validate
     _validate_config(config)
 
     # Add templateDir to config
@@ -169,6 +205,15 @@ def main():  # noqa: C901
     # ilamb tasks
     existing_bundles = ilamb(config, scriptDir, existing_bundles, job_ids_file)
 
+    # zppy external plugins
+    for plugin in plugins:
+        # Get plugin module function
+        plugin_func = getattr(plugin["module"], plugin["name"])
+        # Call plugin
+        existing_bundles = plugin_func(
+            plugin["path"], config, scriptDir, existing_bundles, job_ids_file
+        )
+
     # Submit bundle jobs
     for b in existing_bundles:
         skip = checkStatus(b.bundle_status)

zppy/templates/default.ini

@@ -41,6 +41,8 @@ nodes = integer(default=1)
 output = string
 # The partition of the machine to run on
 partition = string(default="")
+# External zppy plugin modules
+plugins = force_list(default=list())
 # Quality of service
 qos = string(default="regular")
 # scriptDir -- NOTE: this parameter is created internally

zppy/utils.py

@@ -4,6 +4,7 @@
 import stat
 import time
 from subprocess import PIPE, Popen
+from typing import Any, Dict, List
 
 # -----------------------------------------------------------------------------
 # Process specified section and its sub-sections to build list of tasks
@@ -14,7 +15,14 @@
 
 def getTasks(config, section_name):
 
-    tasks = []
+    # mypy: resolves error: Need type annotation for "tasks" (hint: "tasks: List[<type>] = ...")
+    tasks: List[Dict[str, Any]] = []
+
+    # Sanity check
+    # flake8: resolves E713 test for membership should be 'not in'
+    if section_name not in config:
+        print('WARNING: Skipping section not found = "%s"' % (section_name))
+        return tasks
 
     # List of sub-sections
     sub_section_names = config[section_name].sections