Add requested changes from code review

benchcab/benchcab.py

@@ -2,7 +2,7 @@
 
 import sys
 
-from benchcab.job_script import create_job_script, submit_job
+from benchcab.job_script import submit_job
 from benchcab.bench_config import read_config
 from benchcab.benchtree import setup_fluxnet_directory_tree, setup_src_dir
 from benchcab.build_cable import default_build, custom_build
@@ -144,14 +144,13 @@ def fluxnet(self):
             if "fluxnet-bitwise-cmp" not in self.args.skip:
                 self.fluxnet_bitwise_cmp()
         else:
-            create_job_script(
+            submit_job(
                 project=self.config["project"],
                 config_path=self.args.config,
                 modules=self.config["modules"],
                 verbose=self.args.verbose,
                 skip_bitwise_cmp="fluxnet-bitwise-cmp" in self.args.skip,
             )
-            submit_job()
             print(
                 "The CABLE log file for each task is written to "
                 f"{SITE_LOG_DIR}/<task_name>_log.txt"

benchcab/build_cable.py

@@ -3,12 +3,12 @@
 import os
 import contextlib
 import stat
-import subprocess
 import shlex
 import shutil
 import pathlib
 
 from benchcab import internal
+from benchcab.utils import subprocess
 from benchcab import environment_modules
 
 
@@ -23,7 +23,7 @@ def chdir(newdir: pathlib.Path):
         os.chdir(prevdir)
 
 
-def patch_build_script(file_path):
+def remove_module_lines(file_path):
     """Remove lines from `file_path` that call the environment modules package."""
     with open(file_path, "r", encoding="utf-8") as file:
         contents = file.read()
@@ -58,39 +58,28 @@ def default_build(branch_name: str, modules: list, verbose=False):
     tmp_script_path = default_script_path.parent / "tmp-build3.sh"
 
     if verbose:
-        print(f"  Copying {default_script_path} to {tmp_script_path}")
+        print(f"Copying {default_script_path} to {tmp_script_path}")
     shutil.copy(default_script_path, tmp_script_path)
+
     if verbose:
-        print(f"  chmod +x {tmp_script_path}")
+        print(f"chmod +x {tmp_script_path}")
     tmp_script_path.chmod(tmp_script_path.stat().st_mode | stat.S_IEXEC)
 
     if verbose:
         print(
-            f"  Patching {tmp_script_path.name}: remove lines that call "
+            f"Modifying {tmp_script_path.name}: remove lines that call "
             "environment modules"
         )
-    patch_build_script(tmp_script_path)
-
-    if verbose:
-        print("  Loading modules: " + " ".join(modules))
-    environment_modules.module_load(*modules)
-
-    with chdir(default_script_path.parent):
-        cmd = f"./{tmp_script_path.name}" + (" mpi" if internal.MPI else "")
-        if verbose:
-            print(f"  {cmd}")
-        subprocess.run(
-            cmd,
-            shell=True,
-            check=True,
-            stdout=None if verbose else subprocess.DEVNULL,
-            stderr=subprocess.STDOUT,
+    remove_module_lines(tmp_script_path)
+
+    with chdir(default_script_path.parent), environment_modules.load(
+        modules, verbose=verbose
+    ):
+        subprocess.run_cmd(
+            f"./{tmp_script_path.name}" + (" mpi" if internal.MPI else ""),
+            verbose=verbose,
         )
 
-    if verbose:
-        print("  Unloading modules: " + " ".join(modules))
-    environment_modules.module_unload(*modules)
-
 
 def custom_build(config_build_script: str, branch_name: str, verbose=False):
     """Build CABLE with a script provided in configuration file"""
@@ -103,13 +92,4 @@ def custom_build(config_build_script: str, branch_name: str, verbose=False):
     )
 
     with chdir(build_script_path.parent):
-        cmd = f"./{build_script_path.name}"
-        if verbose:
-            print(f"  {cmd}")
-        subprocess.run(
-            cmd,
-            shell=True,
-            check=True,
-            stdout=None if verbose else subprocess.DEVNULL,
-            stderr=subprocess.STDOUT,
-        )
+        subprocess.run_cmd(f"./{build_script_path.name}", verbose=verbose)

benchcab/environment_modules.py

@@ -1,6 +1,7 @@
 """Contains a wrapper around the environment modules API."""
 
 import sys
+import contextlib
 
 sys.path.append("/opt/Modules/v4.3.0/init")
 try:
@@ -18,6 +19,20 @@ class EnvironmentModulesError(Exception):
     """Custom exception class for environment modules errors."""
 
 
+@contextlib.contextmanager
+def load(modules, verbose=False):
+    """Context manager for loading and unloading modules."""
+    if verbose:
+        print("Loading modules: " + " ".join(modules))
+    module_load(*modules)
+    try:
+        yield
+    finally:
+        if verbose:
+            print("Unloading modules: " + " ".join(modules))
+        module_unload(*modules)
+
+
 def module_is_avail(*args: str):
     """Wrapper around `module is-avail modulefile...`"""
     return module("is-avail", *args)

benchcab/get_cable.py

@@ -1,10 +1,10 @@
 """A module containing functions for checking out CABLE repositories."""
 
-import subprocess
 from typing import Union
 from pathlib import Path
 
 from benchcab import internal
+from benchcab.utils import subprocess
 
 
 def next_path(path_pattern, sep="-"):
@@ -32,27 +32,20 @@ def next_path(path_pattern, sep="-"):
 
 def svn_info_show_item(path: Union[Path, str], item: str) -> str:
     """A wrapper around `svn info --show-item <item> <path>`."""
-    cmd = f"svn info --show-item {item} {path}"
-    out = subprocess.run(cmd, shell=True, capture_output=True, text=True, check=True)
-    return out.stdout.strip()
+    proc = subprocess.run_cmd(
+        f"svn info --show-item {item} {path}", capture_output=True
+    )
+    return proc.stdout.strip()
 
 
 def checkout_cable_auxiliary(verbose=False) -> Path:
     """Checkout CABLE-AUX."""
 
     cable_aux_dir = Path(internal.CWD / internal.CABLE_AUX_DIR)
 
-    cmd = f"svn checkout {internal.CABLE_SVN_ROOT}/branches/Share/CABLE-AUX {cable_aux_dir}"
-
-    if verbose:
-        print(cmd)
-
-    subprocess.run(
-        cmd,
-        shell=True,
-        check=True,
-        stdout=None if verbose else subprocess.DEVNULL,
-        stderr=subprocess.STDOUT,
+    subprocess.run_cmd(
+        f"svn checkout {internal.CABLE_SVN_ROOT}/branches/Share/CABLE-AUX {cable_aux_dir}",
+        verbose=verbose,
     )
 
     revision = svn_info_show_item(cable_aux_dir, "revision")
@@ -92,16 +85,7 @@ def checkout_cable(branch_config: dict, verbose=False) -> Path:
     path_to_repo = Path(internal.CWD, internal.SRC_DIR, branch_config["name"])
     cmd += f" {internal.CABLE_SVN_ROOT}/{branch_config['path']} {path_to_repo}"
 
-    if verbose:
-        print(cmd)
-
-    subprocess.run(
-        cmd,
-        shell=True,
-        check=True,
-        stdout=None if verbose else subprocess.DEVNULL,
-        stderr=subprocess.STDOUT,
-    )
+    subprocess.run_cmd(cmd, verbose=verbose)
 
     revision = svn_info_show_item(path_to_repo, "revision")
     print(f"Successfully checked out {branch_config['name']} at revision {revision}")

benchcab/job_script.py

@@ -1,10 +1,10 @@
 """Contains functions for job script creation and submission on Gadi."""
 
-import os
-import subprocess
+from subprocess import CalledProcessError
 from pathlib import Path
 
 from benchcab import internal
+from benchcab.utils import subprocess
 
 
 def get_local_storage_flag(path: Path) -> str:
@@ -16,19 +16,18 @@ def get_local_storage_flag(path: Path) -> str:
     raise RuntimeError("Current directory structure unknown on Gadi.")
 
 
-def create_job_script(
+def submit_job(
     project: str,
     config_path: str,
     modules: list,
     verbose=False,
     skip_bitwise_cmp=False,
 ):
-    """Creates a job script that executes all computationally expensive commands.
-
-    Executed commands are:
-    - benchcab fluxnet-run-tasks
-    - benchcab fluxnet-bitwise-cmp
+    """Submits a PBS job that executes all computationally expensive commands.
 
+    This includes things such as running CABLE and running bitwise comparison jobs
+    between model output files.
+    The PBS job script is written to the current working directory as a side effect.
     """
 
     job_script_path = internal.CWD / internal.QSUB_FNAME
@@ -73,18 +72,12 @@ def create_job_script(
 """
         )
 
-
-def submit_job():
-    """Submits the job script specified by `QSUB_FNAME`."""
-
-    job_script_path = internal.CWD / internal.QSUB_FNAME
-    cmd = f"qsub {job_script_path}"
     try:
-        proc = subprocess.run(
-            cmd, shell=True, check=True, capture_output=True, text=True
+        proc = subprocess.run_cmd(
+            f"qsub {job_script_path}", capture_output=True, verbose=verbose
         )
         print(f"PBS job submitted: {proc.stdout.strip()}")
-    except subprocess.CalledProcessError as exc:
+    except CalledProcessError as exc:
         print("Error when submitting job to NCI queue")
         print(exc.stderr)
         raise

benchcab/task.py

@@ -3,18 +3,19 @@
 
 import os
 import shutil
-import subprocess
 import multiprocessing
 import queue
 import dataclasses
 from pathlib import Path
 from typing import TypeVar, Dict, Any
+from subprocess import CalledProcessError
 
 import flatdict
 import netCDF4
 import f90nml
 
 from benchcab import internal
+from benchcab.utils import subprocess
 import benchcab.get_cable
 
 
@@ -251,12 +252,12 @@ def run_cable(self, verbose=False):
         exe_path = task_dir / internal.CABLE_EXE
         nml_path = task_dir / internal.CABLE_NML
         stdout_path = task_dir / internal.CABLE_STDOUT_FILENAME
-        cmd = f"{exe_path} {nml_path} > {stdout_path} 2>&1"
+
         try:
-            if verbose:
-                print(f"  {cmd}")
-            subprocess.run(cmd, shell=True, check=True)
-        except subprocess.CalledProcessError as exc:
+            subprocess.run_cmd(
+                f"{exe_path} {nml_path}", output_file=stdout_path, verbose=verbose
+            )
+        except CalledProcessError as exc:
             print(f"Error: CABLE returned an error for task {task_name}")
             raise CableError from exc
 
@@ -445,16 +446,17 @@ def run_comparison(task_a: Task, task_b: Task, verbose=False):
         print(
             f"Comparing files {task_a_output.name} and {task_b_output.name} bitwise..."
         )
-    cmd = f"nccmp -df {task_a_output} {task_b_output} 2>&1"
-    if verbose:
-        print(f"  {cmd}")
-    proc = subprocess.run(cmd, shell=True, check=False, capture_output=True, text=True)
-    if proc.returncode != 0:
+    try:
+        subprocess.run_cmd(
+            f"nccmp -df {task_a_output} {task_b_output}",
+            capture_output=True,
+            verbose=verbose,
+        )
+        print(f"Success: files {task_a_output.name} {task_b_output.name} are identical")
+    except CalledProcessError as exc:
         with open(output_file, "w", encoding="utf-8") as file:
-            file.write(proc.stdout)
+            file.write(exc.stdout)
         print(
             f"Failure: files {task_a_output.name} {task_b_output.name} differ. "
             f"Results of diff have been written to {output_file}"
         )
-    else:
-        print(f"Success: files {task_a_output.name} {task_b_output.name} are identical")

benchcab/utils/subprocess.py

@@ -0,0 +1,35 @@
+"""A module containing utility functions that wraps around the `subprocess` module."""
+
+import subprocess
+import contextlib
+import pathlib
+from typing import Any, Optional
+
+
+def run_cmd(
+    cmd: str,
+    capture_output: bool = False,
+    output_file: Optional[pathlib.Path] = None,
+    verbose: bool = False,
+) -> subprocess.CompletedProcess:
+    """Helper function that wraps around `subprocess.run()`"""
+
+    kwargs: Any = {}
+    with contextlib.ExitStack() as stack:
+        if capture_output:
+            kwargs["capture_output"] = True
+            kwargs["text"] = True
+        else:
+            if output_file:
+                kwargs["stdout"] = stack.enter_context(
+                    output_file.open("w", encoding="utf-8")
+                )
+            else:
+                kwargs["stdout"] = None if verbose else subprocess.DEVNULL
+            kwargs["stderr"] = subprocess.STDOUT
+
+        if verbose:
+            print(cmd)
+        proc = subprocess.run(cmd, shell=True, check=True, **kwargs)
+
+    return proc