Set logger in Benchcab object (#281)

Author: SeanBryan51

src/benchcab/benchcab.py

@@ -4,6 +4,7 @@
 """Contains the benchcab application class."""
 
 import grp
+import logging
 import os
 import sys
 from pathlib import Path
@@ -16,7 +17,6 @@
 from benchcab.environment_modules import EnvironmentModules, EnvironmentModulesInterface
 from benchcab.internal import get_met_forcing_file_names
 from benchcab.model import Model
-from benchcab.utils import get_logger
 from benchcab.utils.fs import mkdir, next_path
 from benchcab.utils.pbs import render_job_script
 from benchcab.utils.repo import create_repo
@@ -59,26 +59,37 @@ def __init__(
         self._spatial_tasks: list[spatial.SpatialTask] = []
 
         # Get the logger object
-        self.logger = get_logger()
+        self._logger: Optional[logging.Logger] = None
         self._set_environment()
 
+    def set_logger(self, logger: logging.Logger):
+        """Set logger."""
+        self._logger = logger
+
+    def _get_logger(self) -> logging.Logger:
+        if self._logger is None:
+            msg = "Logger uninitialised."
+            raise RuntimeError(msg)
+        return self._logger
+
     def _set_environment(self):
         """Sets environment variables on current user environment."""
         # Prioritize system binaries over externally set $PATHs (#220)
         os.environ["PATH"] = f"{':'.join(internal.SYSTEM_PATHS)}:{os.environ['PATH']}"
 
     def _validate_environment(self, project: str, modules: list):
         """Performs checks on current user environment."""
+        logger = self._get_logger()
         if not self.validate_env:
             return
 
         if "gadi.nci" not in internal.NODENAME:
-            self.logger.error("benchcab is currently implemented only on Gadi")
+            logger.error("benchcab is currently implemented only on Gadi")
             sys.exit(1)
 
         namelist_dir = Path(internal.NAMELIST_DIR)
         if not namelist_dir.exists():
-            self.logger.error(
+            logger.error(
                 "Cannot find 'namelists' directory in current working directory"
             )
             sys.exit(1)
@@ -101,7 +112,7 @@ def _validate_environment(self, project: str, modules: list):
 
         for modname in modules:
             if not self.modules_handler.module_is_avail(modname):
-                self.logger.error(f"Module ({modname}) is not available.")
+                logger.error(f"Module ({modname}) is not available.")
                 sys.exit(1)
 
         system_paths = os.getenv("PATH").split(":")[: len(internal.SYSTEM_PATHS)]
@@ -119,14 +130,12 @@ def _validate_environment(self, project: str, modules: list):
         for site_id in all_site_ids:
             paths = list(internal.MET_DIR.glob(f"{site_id}*"))
             if not paths:
-                self.logger.error(
-                    f"Failed to infer met file for site id '{site_id}' in "
-                )
-                self.logger.error(f"{internal.MET_DIR}.")
+                logger.error(f"Failed to infer met file for site id '{site_id}' in ")
+                logger.error(f"{internal.MET_DIR}.")
 
                 sys.exit(1)
             if len(paths) > 1:
-                self.logger.error(
+                logger.error(
                     f"Multiple paths infered for site id: '{site_id}' in {internal.MET_DIR}."
                 )
                 sys.exit(1)
@@ -174,18 +183,17 @@ def validate_config(self, config_path: str):
 
     def fluxsite_submit_job(self, config_path: str, skip: list[str]) -> None:
         """Submits the PBS job script step in the fluxsite test workflow."""
+        logger = self._get_logger()
         config = self._get_config(config_path)
         self._validate_environment(project=config["project"], modules=config["modules"])
         if self.benchcab_exe_path is None:
             msg = "Path to benchcab executable is undefined."
             raise RuntimeError(msg)
 
         job_script_path = Path(internal.QSUB_FNAME)
-        self.logger.info(
-            "Creating PBS job script to run fluxsite tasks on compute nodes"
-        )
+        logger.info("Creating PBS job script to run fluxsite tasks on compute nodes")
 
-        self.logger.info(f"job_script_path = {job_script_path}")
+        logger.info(f"job_script_path = {job_script_path}")
 
         with job_script_path.open("w", encoding="utf-8") as file:
             contents = render_job_script(
@@ -204,92 +212,97 @@ def fluxsite_submit_job(self, config_path: str, skip: list[str]) -> None:
                 capture_output=True,
             )
         except CalledProcessError as exc:
-            self.logger.error("when submitting job to NCI queue, details to follow")
-            self.logger.error(exc.output)
+            logger.error("when submitting job to NCI queue, details to follow")
+            logger.error(exc.output)
             raise
 
-        self.logger.info(f"PBS job submitted: {proc.stdout.strip()}")
-        self.logger.info("CABLE log file for each task is written to:")
-        self.logger.info(f"{internal.FLUXSITE_DIRS['LOG']}/<task_name>_log.txt")
-        self.logger.info("The CABLE standard output for each task is written to:")
-        self.logger.info(f"{internal.FLUXSITE_DIRS['TASKS']}/<task_name>/out.txt")
-        self.logger.info("The NetCDF output for each task is written to:")
-        self.logger.info(f"{internal.FLUXSITE_DIRS['OUTPUT']}/<task_name>_out.nc")
+        logger.info(f"PBS job submitted: {proc.stdout.strip()}")
+        logger.info("CABLE log file for each task is written to:")
+        logger.info(f"{internal.FLUXSITE_DIRS['LOG']}/<task_name>_log.txt")
+        logger.info("The CABLE standard output for each task is written to:")
+        logger.info(f"{internal.FLUXSITE_DIRS['TASKS']}/<task_name>/out.txt")
+        logger.info("The NetCDF output for each task is written to:")
+        logger.info(f"{internal.FLUXSITE_DIRS['OUTPUT']}/<task_name>_out.nc")
 
     def checkout(self, config_path: str):
         """Endpoint for `benchcab checkout`."""
+        logger = self._get_logger()
         config = self._get_config(config_path)
         self._validate_environment(project=config["project"], modules=config["modules"])
 
         mkdir(internal.SRC_DIR, exist_ok=True)
 
-        self.logger.info("Checking out repositories...")
+        logger.info("Checking out repositories...")
         rev_number_log = ""
 
         for model in self._get_models(config):
             try:
                 model.repo.checkout()
             except Exception:
                 msg = "Try using `benchcab clean realisations` first"
-                self.logger.error(
+                logger.error(
                     "Model checkout failed, probably due to existing realisation name"
                 )
                 raise FileExistsError(msg)
             rev_number_log += f"{model.name}: {model.repo.get_revision()}\n"
 
         rev_number_log_path = next_path("rev_number-*.log")
-        self.logger.info(f"Writing revision number info to {rev_number_log_path}")
+        logger.info(f"Writing revision number info to {rev_number_log_path}")
         with rev_number_log_path.open("w", encoding="utf-8") as file:
             file.write(rev_number_log)
 
     def build(self, config_path: str, mpi=False):
         """Endpoint for `benchcab build`."""
+        logger = self._get_logger()
         config = self._get_config(config_path)
         self._validate_environment(project=config["project"], modules=config["modules"])
 
         for repo in self._get_models(config):
             if repo.build_script:
 
-                self.logger.info("Compiling CABLE using custom build script for")
-                self.logger.info(f"realisation {repo.name}")
+                logger.info("Compiling CABLE using custom build script for")
+                logger.info(f"realisation {repo.name}")
                 repo.custom_build(modules=config["modules"])
 
             else:
                 build_mode = "serial and MPI" if mpi else "serial"
-                self.logger.info(
+                logger.info(
                     f"Compiling CABLE {build_mode} for realisation {repo.name}..."
                 )
                 repo.build(modules=config["modules"], mpi=mpi)
-            self.logger.info(f"Successfully compiled CABLE for realisation {repo.name}")
+            logger.info(f"Successfully compiled CABLE for realisation {repo.name}")
 
     def fluxsite_setup_work_directory(self, config_path: str):
         """Endpoint for `benchcab fluxsite-setup-work-dir`."""
+        logger = self._get_logger()
         config = self._get_config(config_path)
         self._validate_environment(project=config["project"], modules=config["modules"])
 
-        self.logger.info("Setting up run directory tree for fluxsite tests...")
+        logger.info("Setting up run directory tree for fluxsite tests...")
         setup_fluxsite_directory_tree()
-        self.logger.info("Setting up tasks...")
+        logger.info("Setting up tasks...")
         for task in self._get_fluxsite_tasks(config):
             task.setup_task()
-        self.logger.info("Successfully setup fluxsite tasks")
+        logger.info("Successfully setup fluxsite tasks")
 
     def fluxsite_run_tasks(self, config_path: str):
         """Endpoint for `benchcab fluxsite-run-tasks`."""
+        logger = self._get_logger()
         config = self._get_config(config_path)
         self._validate_environment(project=config["project"], modules=config["modules"])
         tasks = self._get_fluxsite_tasks(config)
 
-        self.logger.info("Running fluxsite tasks...")
+        logger.info("Running fluxsite tasks...")
         if config["fluxsite"]["multiprocess"]:
             ncpus = config["fluxsite"]["pbs"]["ncpus"]
             fluxsite.run_tasks_in_parallel(tasks, n_processes=ncpus)
         else:
             fluxsite.run_tasks(tasks)
-        self.logger.info("Successfully ran fluxsite tasks")
+        logger.info("Successfully ran fluxsite tasks")
 
     def fluxsite_bitwise_cmp(self, config_path: str):
         """Endpoint for `benchcab fluxsite-bitwise-cmp`."""
+        logger = self._get_logger()
         config = self._get_config(config_path)
         self._validate_environment(project=config["project"], modules=config["modules"])
 
@@ -302,13 +315,13 @@ def fluxsite_bitwise_cmp(self, config_path: str):
             self._get_fluxsite_tasks(config)
         )
 
-        self.logger.info("Running comparison tasks...")
+        logger.info("Running comparison tasks...")
         if config["fluxsite"]["multiprocess"]:
             ncpus = config["fluxsite"]["pbs"]["ncpus"]
             run_comparisons_in_parallel(comparisons, n_processes=ncpus)
         else:
             run_comparisons(comparisons)
-        self.logger.info("Successfully ran comparison tasks")
+        logger.info("Successfully ran comparison tasks")
 
     def fluxsite(self, config_path: str, no_submit: bool, skip: list[str]):
         """Endpoint for `benchcab fluxsite`."""
@@ -331,28 +344,30 @@ def clean(self, config_path: str, clean_option: str):
 
     def spatial_setup_work_directory(self, config_path: str):
         """Endpoint for `benchcab spatial-setup-work-dir`."""
+        logger = self._get_logger()
         config = self._get_config(config_path)
         self._validate_environment(project=config["project"], modules=config["modules"])
 
-        self.logger.info("Setting up run directory tree for spatial tests...")
+        logger.info("Setting up run directory tree for spatial tests...")
         setup_spatial_directory_tree()
-        self.logger.info("Setting up tasks...")
+        logger.info("Setting up tasks...")
         try:
             payu_config = config["spatial"]["payu"]["config"]
         except KeyError:
             payu_config = None
         for task in self._get_spatial_tasks(config):
             task.setup_task(payu_config=payu_config)
-        self.logger.info("Successfully setup spatial tasks")
+        logger.info("Successfully setup spatial tasks")
 
     def spatial_run_tasks(self, config_path: str):
         """Endpoint for `benchcab spatial-run-tasks`."""
+        logger = self._get_logger()
         config = self._get_config(config_path)
         self._validate_environment(project=config["project"], modules=config["modules"])
 
-        self.logger.info("Running spatial tasks...")
+        logger.info("Running spatial tasks...")
         spatial.run_tasks(tasks=self._get_spatial_tasks(config))
-        self.logger.info("Successfully dispatched payu jobs")
+        logger.info("Successfully dispatched payu jobs")
 
     def spatial(self, config_path: str, skip: list):
         """Endpoint for `benchcab spatial`."""

src/benchcab/main.py

@@ -8,13 +8,15 @@
 from benchcab.utils import get_logger
 
 
-def parse_and_dispatch(parser):
+def parse_and_dispatch(parser, app):
     """Parse arguments for the script and dispatch to the correct function.
 
     Args:
     ----
     parser : argparse.ArgumentParser
         Parser object.
+    app : Benchcab
+        Benchcab application instance.
 
     """
     args = vars(parser.parse_args(sys.argv[1:] if sys.argv[1:] else ["-h"]))
@@ -26,7 +28,8 @@ def parse_and_dispatch(parser):
     _ = args.pop("verbose")
 
     # We just need to instantiate this with the desired level
-    get_logger(level=log_level)
+    logger = get_logger(level=log_level)
+    app.set_logger(logger)
 
     func = args.pop("func")
     func(**args)
@@ -39,7 +42,7 @@ def main():
     """
     app = Benchcab(benchcab_exe_path=shutil.which(sys.argv[0]))
     parser = generate_parser(app)
-    parse_and_dispatch(parser)
+    parse_and_dispatch(parser, app)
 
 
 if __name__ == "__main__":