Merge pull request #254 from CABLE-LSM/251-use-jinja-for-job-script-r…

…endering

251 use jinja for job script rendering

benchcab/data/pbs_jobscript.j2

@@ -0,0 +1,19 @@
+#!/bin/bash
+#PBS -l wd
+#PBS -l ncpus={{ncpus}}
+#PBS -l mem={{mem}}
+#PBS -l walltime={{walltime}}
+#PBS -q normal
+#PBS -P {{project}}
+#PBS -j oe
+#PBS -m e
+#PBS -l storage={{storage}}
+
+module purge
+{% for module in modules -%}
+module load {{module}}
+{% endfor %}
+set -ev
+
+{{benchcab_path}} fluxsite-run-tasks --config={{config_path}}{{verbose_flag}}
+{% if skip_bitwise_cmp == False %}{{benchcab_path}} fluxsite-bitwise-cmp --config={{config_path}}{{verbose_flag}}{% endif %}

benchcab/data/test/pbs_jobscript_default.sh

@@ -0,0 +1,20 @@
+#!/bin/bash
+#PBS -l wd
+#PBS -l ncpus=18
+#PBS -l mem=30GB
+#PBS -l walltime=6:00:00
+#PBS -q normal
+#PBS -P tm70
+#PBS -j oe
+#PBS -m e
+#PBS -l storage=gdata/ks32+gdata/hh5+gdata/wd9
+
+module purge
+module load foo
+module load bar
+module load baz
+
+set -ev
+
+/absolute/path/to/benchcab fluxsite-run-tasks --config=/path/to/config.yaml
+/absolute/path/to/benchcab fluxsite-bitwise-cmp --config=/path/to/config.yaml

benchcab/data/test/pbs_jobscript_skip_bitwise.sh

@@ -0,0 +1,19 @@
+#!/bin/bash
+#PBS -l wd
+#PBS -l ncpus=18
+#PBS -l mem=30GB
+#PBS -l walltime=6:00:00
+#PBS -q normal
+#PBS -P tm70
+#PBS -j oe
+#PBS -m e
+#PBS -l storage=gdata/ks32+gdata/hh5+gdata/wd9
+
+module purge
+module load foo
+module load bar
+module load baz
+
+set -ev
+
+/absolute/path/to/benchcab fluxsite-run-tasks --config=/path/to/config.yaml

benchcab/data/test/pbs_jobscript_verbose.sh

@@ -0,0 +1,20 @@
+#!/bin/bash
+#PBS -l wd
+#PBS -l ncpus=18
+#PBS -l mem=30GB
+#PBS -l walltime=6:00:00
+#PBS -q normal
+#PBS -P tm70
+#PBS -j oe
+#PBS -m e
+#PBS -l storage=gdata/ks32+gdata/hh5+gdata/wd9
+
+module purge
+module load foo
+module load bar
+module load baz
+
+set -ev
+
+/absolute/path/to/benchcab fluxsite-run-tasks --config=/path/to/config.yaml -v
+/absolute/path/to/benchcab fluxsite-bitwise-cmp --config=/path/to/config.yaml -v

benchcab/data/test/template.j2

@@ -0,0 +1 @@
+This is a template. {{myarg}}

benchcab/fluxsite.py

@@ -44,7 +44,7 @@ def deep_update(mapping: Dict[KeyType, Any], *updating_mappings: Dict[KeyType, A
     -------
     Dict[KeyType, Any]
         Updated mapping.
-    
+
     """
     updated_mapping = mapping.copy()
     for updating_mapping in updating_mappings:

benchcab/utils/__init__.py

@@ -11,8 +11,10 @@
 import sys
 from importlib import resources
 from pathlib import Path
+from typing import Union
 
 import yaml
+from jinja2 import BaseLoader, Environment
 
 # List of one-argument decoding functions.
 PACKAGE_DATA_DECODERS = dict(json=json.loads, yml=yaml.safe_load)
@@ -30,14 +32,19 @@ def get_installed_root() -> Path:
     return Path(resources.files("benchcab"))
 
 
-def load_package_data(filename: str) -> dict:
+def load_package_data(filename: str) -> Union[str, dict]:
     """Load data out of the installed package data directory.
 
     Parameters
     ----------
     filename : str
         Filename of the file to load out of the data directory.
 
+    Returns
+    -------
+    str or dict
+        String or dictionary, depending on format of data read.
+
     """
     # Work out the encoding of requested file.
     ext = filename.split(".")[-1]
@@ -48,10 +55,54 @@ def load_package_data(filename: str) -> dict:
     # Extract from the installations data directory.
     raw = pkgutil.get_data("benchcab", os.path.join("data", filename)).decode("utf-8")
 
+    # If there is no explicit decoder, just return the raw text
+    if ext not in PACKAGE_DATA_DECODERS.keys():
+        return raw
+
     # Decode and return.
     return PACKAGE_DATA_DECODERS[ext](raw)
 
 
+def interpolate_string_template(template, **kwargs):
+    """Interpolate a string template with kwargs.
+
+    Parameters
+    ----------
+    template : str
+        Template string to interpolate over.
+    **kwargs :
+        Keyword arguments to interpolate into the string.
+
+    Returns
+    -------
+    str
+        Interpolated string.
+
+    """
+    _template = Environment(loader=BaseLoader()).from_string(template)
+    return _template.render(**kwargs)
+
+
+def interpolate_file_template(template_file, **kwargs):
+    """Interpolate kwargs directly into a j2 template file from the data directory.
+
+    Parameters
+    ----------
+    template_file : str
+        Filepath slug in the benchcab data directory.
+    **kwargs :
+        Keyword arguments to interpolate into the file.
+
+    Returns
+    -------
+    str
+        Interpolated template string.
+
+    """
+    template = load_package_data(template_file)
+    return interpolate_string_template(template, **kwargs)
+
+
 def get_logger(name="benchcab", level="debug"):
     """Get a logger instance.
 

benchcab/utils/pbs.py

@@ -5,6 +5,8 @@
 
 from typing import TypedDict
 
+from benchcab.utils import interpolate_file_template
+
 
 class PBSConfig(TypedDict):
     """Default parameters for PBS runs via benchcab."""
@@ -29,29 +31,20 @@ def render_job_script(
     This includes things such as running CABLE and running bitwise comparison jobs
     between model output files.
     """
-    module_load_lines = "\n".join(
-        f"module load {module_name}" for module_name in modules
-    )
-    verbose_flag = "-v" if verbose else ""
+    verbose_flag = " -v" if verbose else ""
     storage_flags = ["gdata/ks32", "gdata/hh5", "gdata/wd9", *pbs_config["storage"]]
-    return f"""#!/bin/bash
-#PBS -l wd
-#PBS -l ncpus={pbs_config["ncpus"]}
-#PBS -l mem={pbs_config["mem"]}
-#PBS -l walltime={pbs_config["walltime"]}
-#PBS -q normal
-#PBS -P {project}
-#PBS -j oe
-#PBS -m e
-#PBS -l storage={'+'.join(storage_flags)}
-
-module purge
-{module_load_lines}
-
-set -ev
-
-{benchcab_path} fluxsite-run-tasks --config={config_path} {verbose_flag}
-{'' if skip_bitwise_cmp else f'''
-{benchcab_path} fluxsite-bitwise-cmp --config={config_path} {verbose_flag}
-''' }
-"""
+
+    context = dict(
+        modules=modules,
+        verbose_flag=verbose_flag,
+        ncpus=pbs_config["ncpus"],
+        mem=pbs_config["mem"],
+        walltime=pbs_config["walltime"],
+        project=project,
+        storage="+".join(storage_flags),
+        benchcab_path=benchcab_path,
+        config_path=config_path,
+        skip_bitwise_cmp=skip_bitwise_cmp,
+    )
+
+    return interpolate_file_template("pbs_jobscript.j2", **context)

tests/test_pbs.py

@@ -1,6 +1,7 @@
 """`pytest` tests for `utils/pbs.py`."""
 
 from benchcab import internal
+from benchcab.utils import load_package_data
 from benchcab.utils.pbs import render_job_script
 
 
@@ -15,31 +16,7 @@ def test_default_job_script(self):
             modules=["foo", "bar", "baz"],
             pbs_config=internal.FLUXSITE_DEFAULT_PBS,
             benchcab_path="/absolute/path/to/benchcab",
-        ) == (
-            f"""#!/bin/bash
-#PBS -l wd
-#PBS -l ncpus={internal.FLUXSITE_DEFAULT_PBS["ncpus"]}
-#PBS -l mem={internal.FLUXSITE_DEFAULT_PBS["mem"]}
-#PBS -l walltime={internal.FLUXSITE_DEFAULT_PBS["walltime"]}
-#PBS -q normal
-#PBS -P tm70
-#PBS -j oe
-#PBS -m e
-#PBS -l storage=gdata/ks32+gdata/hh5+gdata/wd9
-
-module purge
-module load foo
-module load bar
-module load baz
-
-set -ev
-
-/absolute/path/to/benchcab fluxsite-run-tasks --config=/path/to/config.yaml 
-
-/absolute/path/to/benchcab fluxsite-bitwise-cmp --config=/path/to/config.yaml 
-
-"""
-        )
+        ) == load_package_data("test/pbs_jobscript_default.sh")
 
     def test_verbose_flag_added_to_command_line_arguments(self):
         """Success case: test verbose flag is added to command line arguments."""
@@ -50,31 +27,7 @@ def test_verbose_flag_added_to_command_line_arguments(self):
             pbs_config=internal.FLUXSITE_DEFAULT_PBS,
             verbose=True,
             benchcab_path="/absolute/path/to/benchcab",
-        ) == (
-            f"""#!/bin/bash
-#PBS -l wd
-#PBS -l ncpus={internal.FLUXSITE_DEFAULT_PBS["ncpus"]}
-#PBS -l mem={internal.FLUXSITE_DEFAULT_PBS["mem"]}
-#PBS -l walltime={internal.FLUXSITE_DEFAULT_PBS["walltime"]}
-#PBS -q normal
-#PBS -P tm70
-#PBS -j oe
-#PBS -m e
-#PBS -l storage=gdata/ks32+gdata/hh5+gdata/wd9
-
-module purge
-module load foo
-module load bar
-module load baz
-
-set -ev
-
-/absolute/path/to/benchcab fluxsite-run-tasks --config=/path/to/config.yaml -v
-
-/absolute/path/to/benchcab fluxsite-bitwise-cmp --config=/path/to/config.yaml -v
-
-"""
-        )
+        ) == load_package_data("test/pbs_jobscript_verbose.sh")
 
     def test_skip_bitwise_comparison_step(self):
         """Success case: skip fluxsite-bitwise-cmp step."""
@@ -85,64 +38,4 @@ def test_skip_bitwise_comparison_step(self):
             pbs_config=internal.FLUXSITE_DEFAULT_PBS,
             skip_bitwise_cmp=True,
             benchcab_path="/absolute/path/to/benchcab",
-        ) == (
-            f"""#!/bin/bash
-#PBS -l wd
-#PBS -l ncpus={internal.FLUXSITE_DEFAULT_PBS["ncpus"]}
-#PBS -l mem={internal.FLUXSITE_DEFAULT_PBS["mem"]}
-#PBS -l walltime={internal.FLUXSITE_DEFAULT_PBS["walltime"]}
-#PBS -q normal
-#PBS -P tm70
-#PBS -j oe
-#PBS -m e
-#PBS -l storage=gdata/ks32+gdata/hh5+gdata/wd9
-
-module purge
-module load foo
-module load bar
-module load baz
-
-set -ev
-
-/absolute/path/to/benchcab fluxsite-run-tasks --config=/path/to/config.yaml 
-
-"""
-        )
-
-    def test_pbs_config_parameters(self):
-        """Success case: specify parameters in pbs_config."""
-        assert render_job_script(
-            project="tm70",
-            config_path="/path/to/config.yaml",
-            modules=["foo", "bar", "baz"],
-            skip_bitwise_cmp=True,
-            benchcab_path="/absolute/path/to/benchcab",
-            pbs_config={
-                "ncpus": 4,
-                "mem": "16GB",
-                "walltime": "00:00:30",
-                "storage": ["gdata/foo"],
-            },
-        ) == (
-            """#!/bin/bash
-#PBS -l wd
-#PBS -l ncpus=4
-#PBS -l mem=16GB
-#PBS -l walltime=00:00:30
-#PBS -q normal
-#PBS -P tm70
-#PBS -j oe
-#PBS -m e
-#PBS -l storage=gdata/ks32+gdata/hh5+gdata/wd9+gdata/foo
-
-module purge
-module load foo
-module load bar
-module load baz
-
-set -ev
-
-/absolute/path/to/benchcab fluxsite-run-tasks --config=/path/to/config.yaml 
-
-"""
-        )
+        ) == load_package_data("test/pbs_jobscript_skip_bitwise.sh")