wrfrun 0.1.8__py3-none-any.whl → 0.2.0__py3-none-any.whl

wrfrun/run.py

@@ -4,15 +4,15 @@
 
 import sys
 import threading
-from os.path import abspath, dirname, exists
+from os.path import abspath, dirname
 from typing import Optional, Tuple, Union
 
-from .core import ExecConfigRecorder, WRFRUNConfig, WRFRunBasicError, WRFRunServer, WRFRunServerHandler, replay_config_generator, stop_server
+from .core import ExecConfigRecorder, WRFRUNConfig, init_wrfrun_config, WRFRunBasicError, WRFRunServer, WRFRunServerHandler, replay_config_generator, stop_server
 from .data import prepare_wps_input_data
-from .model import clear_model_logs, plot_domain_area
-from .pbs import in_pbs, prepare_pbs_script
+from .model import clear_model_logs, generate_domain_area
+from .scheduler import in_job_scheduler, prepare_scheduler_script
 from .utils import call_subprocess, logger, logger_add_file_handler
-from .workspace import prepare_workspace
+from .workspace import prepare_workspace, check_workspace
 
 
 def confirm_model_area():
@@ -20,9 +20,9 @@ def confirm_model_area():
     Ask user to check domain area.
 
     """
-    plot_domain_area()
+    generate_domain_area()
 
-    if not in_pbs():
+    if not in_job_scheduler():
         # ask user
         logger.warning(f"Check the domain image, is it right?")
         answer = input("Is it right? [y/N]: ")
@@ -41,16 +41,29 @@ class WRFRun:
     _instance = None
     _initialized = False
 
-    def __init__(self, config_file: str, init_workspace=True, start_server=False, pbs_mode=False, prepare_wps_data=False, wps_data_area: Optional[Tuple[int, int, int, int]] = None):
+    def __init__(
+        self, config_file: str, init_workspace=True,
+        start_server=False, submit_job=False, prepare_wps_data=False,
+        wps_data_area: Optional[Tuple[int, int, int, int]] = None,
+        skip_domain_confirm=False
+    ):
         """
         WRFRun, a context class to achieve some goals before and after running WRF, like save a copy of config file, start and close WRFRunServer.
 
         :param config_file: ``wrfrun`` config file's path.
+        :type config_file: str
         :param init_workspace: If True, clean old files in workspace and re-create it.
+        :type init_workspace: bool
         :param start_server: Whether to start WRFRunServer, defaults to True.
-        :param pbs_mode: If commit this task to the PBS system, defaults to True.
+        :type start_server: bool
+        :param submit_job: If commit this task to the PBS system, defaults to True.
+        :type submit_job: bool
         :param prepare_wps_data: If True, download input datas for WPS first.
+        :type prepare_wps_data: bool
         :param wps_data_area: If ``prepare_wps_data==True``, you need to give the area range of input data so download function can download data from ERA5.
+        :type wps_data_area: tuple
+        :param skip_domain_confirm: If ``True``, skip domain confirm.
+        :type skip_domain_confirm: bool
         :return:
         """
         if self._initialized:
@@ -63,17 +76,18 @@ class WRFRun:
         self._ip = ""
         self._port = -1
 
-        self._pbs_mode = pbs_mode
+        self._submit_job = submit_job
         self._init_workspace = init_workspace
         self._prepare_wps_data = prepare_wps_data
         self._wps_data_area = wps_data_area
+        self._skip_domain_confirm = skip_domain_confirm
         self._entry_file_path = abspath(sys.argv[0])
         self._entry_file_dir_path = dirname(self._entry_file_path)
         self._replay_configs = None
 
         # make sure we can read the config file, because sometimes the user may run the Python script in a different path.
         abs_config_path = f"{self._entry_file_dir_path}/{config_file}"
-        WRFRUNConfig.load_wrfrun_config(abs_config_path)
+        init_wrfrun_config(abs_config_path)
 
         self._record_output_path = None
         self._record_include_data = False
@@ -89,35 +103,33 @@ class WRFRun:
 
     def __enter__(self):
         # check workspace
-        check_list = [WRFRUNConfig.WPS_WORK_PATH, WRFRUNConfig.WRF_WORK_PATH, WRFRUNConfig.WRFDA_WORK_PATH]
-        check_list = [WRFRUNConfig.parse_resource_uri(x) for x in check_list]
-        for _path in check_list:
-            if not exists(_path) and not self._init_workspace:
-                logger.info(f"Force re-create workspace because it is broken.")
-                self._init_workspace = True
-                break
+        if not check_workspace():
+            logger.info(f"Force re-create workspace because it is broken.")
+            self._init_workspace = True
 
         # here is the condition we need to initialize workspace:
-        # 1. pbs_mode = True and init_workspace = True, do prepare_workspace before committing the task to the PBS system.
-        # 2. pbs_mode = False and init_workspace = True, do prepare_workspace.
-        if self._pbs_mode and not in_pbs():
+        # 1. submit_job = True and init_workspace = True, do prepare_workspace before submitting the task to job scheduler.
+        # 2. submit_job = False and init_workspace = True, do prepare_workspace.
+        if self._submit_job and not in_job_scheduler():
             if self._init_workspace:
                 prepare_workspace()
 
             # ask user before commit the task
-            confirm_model_area()
+            if not self._skip_domain_confirm:
+                confirm_model_area()
 
-            prepare_pbs_script(self._entry_file_path)
+            prepare_scheduler_script(self._entry_file_path)
 
             call_subprocess(["qsub", f"{self._entry_file_dir_path}/run.sh"])
             logger.info(f"Work has been submit to PBS system")
             exit(0)
 
-        elif not self._pbs_mode:
+        elif not self._submit_job:
             if self._init_workspace:
                 prepare_workspace()
 
-            confirm_model_area()
+            if not self._skip_domain_confirm:
+                confirm_model_area()
 
         # save a copy of config to the output path
         WRFRUNConfig.save_wrfrun_config(f"{WRFRUNConfig.WRFRUN_OUTPUT_PATH}/config.toml")
@@ -139,7 +151,7 @@ class WRFRun:
             else:
                 prepare_wps_input_data(self._wps_data_area)
 
-        logger.info(r"Enter wrfrun context")
+        logger.debug(r"Enter wrfrun context")
 
         return self
 
@@ -157,7 +169,7 @@ class WRFRun:
 
         clear_model_logs()
 
-        logger.info(r"Exit wrfrun context")
+        logger.debug(r"Exit wrfrun context")
 
     def _start_wrfrun_server(self):
         """

wrfrun/scheduler/__init__.py

@@ -0,0 +1,35 @@
+"""
+wrfrun.scheduler
+################
+
+``wrfrun`` provides functions to help users take care of the job scheduler.
+
+Submodules
+**********
+
+======================================= ===========================================================
+:doc:`env </api/scheduler.env>`         Functions to manage environment variables in job scheduler.
+:doc:`lsf </api/scheduler.lsf>`         Scheduler interface for LSF job scheduler.
+:doc:`pbs </api/scheduler.pbs>`         Scheduler interface for PBS job scheduler.
+:doc:`script </api/scheduler.script>`   Generate shell scripts for job scheduler.
+:doc:`slurm </api/scheduler.slurm>`     Scheduler interface for Slurm job scheduler.
+:doc:`utils </api/scheduler.utils>`     Utility functions.
+======================================= ===========================================================
+
+.. toctree::
+    :maxdepth: 1
+    :hidden:
+
+    env <scheduler.env>
+    lsf <scheduler.lsf>
+    pbs <scheduler.pbs>
+    slurm <scheduler.slurm>
+    utils <scheduler.utils>
+"""
+
+from .env import *
+from .lsf import *
+from .pbs import *
+from .script import *
+from .slurm import *
+from .utils import *

wrfrun/scheduler/env.py

@@ -0,0 +1,44 @@
+"""
+wrfrun.scheduler.env
+####################
+
+Functions to manage environment variables in job scheduler.
+
+.. autosummary::
+    :toctree: generated/
+
+    in_job_scheduler
+
+How Does wrfrun Check If It Is In A Job Scheduler?
+**************************************************
+
+If you submit your task through ``wrfrun``, that is,
+set ``submit_job = True`` in :class:`WRFRun <wrfrun.run.WRFRun>`,
+an environment variable called ``WRFRUN_ENV_JOB_SCHEDULER`` will be set.
+``wrfrun`` will determine if it is in a job scheduler by checking if ``WRFRUN_ENV_JOB_SCHEDULER`` appears in environment.
+
+If you submit your task by your own,
+it is recommended that add ``WRFRUN_ENV_JOB_SCHEDULER`` to the environment,
+which can ensure ``wrfrun`` works properly in the job scheduler.
+"""
+
+from os import environ
+
+
+def in_job_scheduler() -> bool:
+    """
+    Check if ``wrfrun`` runs in a job scheduler task.
+
+    This function checks the environment variable ``WRFRUN_ENV_JOB_SCHEDULER``
+    to determine if ``wrfrun`` is running in a job scheduler task.
+
+    :return: ``True`` if in a job scheduler task, else ``False``.
+    :rtype: bool
+    """
+    if "WRFRUN_ENV_JOB_SCHEDULER" in environ:
+        return True
+    else:
+        return False
+
+
+__all__ = ["in_job_scheduler"]

wrfrun/scheduler/lsf.py

@@ -0,0 +1,49 @@
+"""
+wrfrun.scheduler.lsf
+####################
+
+Scheduler interface for LSF system.
+
+.. autosummary::
+    :toctree: generated/
+
+    lsf_generate_settings
+"""
+
+from wrfrun.core import get_wrfrun_config
+from wrfrun.res import SCHEDULER_LSF_TEMPLATE
+from .utils import get_core_num
+
+
+def lsf_generate_settings(scheduler_config: dict) -> str:
+    """
+    This function generate bash settings for LSF job scheduler.
+
+    :return: Generated settings.
+    :rtype: str
+    """
+    WRFRUNConfig = get_wrfrun_config()
+
+    # get log path and job scheduler config
+    log_path = WRFRUNConfig.get_log_path()
+
+    # get scheduler configs
+    stdout_log_path = f"{log_path}/lsf.log"
+    stderr_log_path = f"{log_path}/lsf.err"
+    node_num = scheduler_config["node_num"]
+    queue_name = scheduler_config["queue_name"]
+    core_num = get_core_num()
+
+    template_path = WRFRUNConfig.parse_resource_uri(SCHEDULER_LSF_TEMPLATE)
+    with open(template_path, "r") as f:
+        template = f.read()
+
+    return template.format(
+        STDOUT_LOG_PATH=stdout_log_path,
+        STDERR_LOG_PATH=stderr_log_path,
+        CORE_NUM=core_num * node_num,
+        QUEUE_NAME=queue_name,
+    )
+
+
+__all__ = ["lsf_generate_settings"]

wrfrun/scheduler/pbs.py

@@ -0,0 +1,50 @@
+"""
+wrfrun.scheduler.pbs
+####################
+
+Scheduler interface for PBS system.
+
+.. autosummary::
+    :toctree: generated/
+
+    pbs_generate_settings
+"""
+
+from wrfrun.core import get_wrfrun_config
+from wrfrun.res import SCHEDULER_PBS_TEMPLATE
+from .utils import get_core_num
+
+
+def pbs_generate_settings(scheduler_config: dict) -> str:
+    """
+    This function generate bash settings for PBS job scheduler.
+
+    :return: Generated settings.
+    :rtype: str
+    """
+    WRFRUNConfig = get_wrfrun_config()
+
+    # get log path and job scheduler config
+    log_path = WRFRUNConfig.get_log_path()
+
+    # get scheduler configs
+    stdout_log_path = f"{log_path}/pbs.log"
+    stderr_log_path = f"{log_path}/pbs.err"
+    node_num = scheduler_config["node_num"]
+    queue_name = scheduler_config["queue_name"]
+    core_num = get_core_num()
+
+    template_path = WRFRUNConfig.parse_resource_uri(SCHEDULER_PBS_TEMPLATE)
+    with open(template_path, "r") as f:
+        template = f.read()
+
+    return template.format(
+        STDOUT_LOG_PATH=stdout_log_path,
+        STDERR_LOG_PATH=stderr_log_path,
+        NODE_NUM=node_num,
+        CORE_NUM=core_num,
+        QUEUE_NAME=queue_name,
+    )
+
+
+__all__ = ["pbs_generate_settings"]

wrfrun/scheduler/script.py

@@ -0,0 +1,72 @@
+from os.path import abspath, dirname, exists
+
+from wrfrun import get_wrfrun_config
+from wrfrun.res import RUN_SH_TEMPLATE
+from wrfrun.utils import logger
+from .lsf import lsf_generate_settings
+from .pbs import pbs_generate_settings
+from .slurm import slurm_generate_settings
+
+
+def prepare_scheduler_script(main_file_path: str):
+    """
+    Prepare the bash script to be submitted to job scheduler.
+
+    :param main_file_path: Path of the main entry file.
+    :type main_file_path: str
+    """
+    WRFRUNConfig = get_wrfrun_config()
+
+    # check main file path
+    if not exists(main_file_path):
+        logger.error(f"Wrong path of main entry file: {main_file_path}")
+        raise FileNotFoundError(f"Wrong path of main entry file: {main_file_path}")
+
+    # get absolute path of main entry file's parent directory
+    dir_path = abspath(dirname(main_file_path))
+
+    scheduler_configs = WRFRUNConfig.get_job_scheduler_config()
+
+    # generate scheduler settings
+    match scheduler_configs["job_scheduler"]:
+        case "lsf":
+            scheduler_settings = lsf_generate_settings(scheduler_configs)
+
+        case "pbs":
+            scheduler_settings = pbs_generate_settings(scheduler_configs)
+
+        case "slurm":
+            scheduler_settings = slurm_generate_settings(scheduler_configs)
+
+        case _:
+            logger.error(f"Unknown scheduler name: {scheduler_configs['job_scheduler']}")
+            raise ValueError(f"Unknown scheduler name: {scheduler_configs['job_scheduler']}")
+
+    # generate environment settings
+    env_settings = 'export WRFRUN_ENV_JOB_SCHEDULER=1\n'
+    if len(scheduler_configs["env_settings"]) > 0:
+        for key in scheduler_configs["env_settings"]:
+            env_settings += f"export {key}={scheduler_configs['env_settings'][key]}\n"
+
+    # generate command
+    exec_cmd = f"{scheduler_configs['python_interpreter']} {main_file_path}"
+
+    # generate shell script
+    shell_template_path = WRFRUNConfig.parse_resource_uri(RUN_SH_TEMPLATE)
+    with open(f"{dir_path}/run.sh", "w") as f:
+
+        with open(shell_template_path, "r") as f_template:
+            template = f_template.read()
+
+        template = template.format(
+            SCHEDULER_SETTINGS=scheduler_settings,
+            ENV_SETTINGS=env_settings,
+            WORK_COMMAND=exec_cmd,
+        )
+
+        f.write(template)
+
+    logger.info(f"Job scheduler script written to {dir_path}/run.sh")
+
+
+__all__ = ["prepare_scheduler_script"]

wrfrun/scheduler/slurm.py

@@ -0,0 +1,50 @@
+"""
+wrfrun.scheduler.slurm
+######################
+
+Scheduler interface for Slurm system.
+
+.. autosummary::
+    :toctree: generated/
+
+    slurm_generate_settings
+"""
+
+from wrfrun.core import get_wrfrun_config
+from wrfrun.res import SCHEDULER_SLURM_TEMPLATE
+from .utils import get_core_num
+
+
+def slurm_generate_settings(scheduler_config: dict) -> str:
+    """
+    This function generate bash settings for Slurm job scheduler.
+
+    :return: Generated settings.
+    :rtype: str
+    """
+    WRFRUNConfig = get_wrfrun_config()
+
+    # get log path and job scheduler config
+    log_path = WRFRUNConfig.get_log_path()
+
+    # get scheduler configs
+    stdout_log_path = f"{log_path}/slurm.log"
+    stderr_log_path = f"{log_path}/slurm.err"
+    node_num = scheduler_config["node_num"]
+    queue_name = scheduler_config["queue_name"]
+    core_num = get_core_num()
+
+    template_path = WRFRUNConfig.parse_resource_uri(SCHEDULER_SLURM_TEMPLATE)
+    with open(template_path, "r") as f:
+        template = f.read()
+
+    return template.format(
+        STDOUT_LOG_PATH=stdout_log_path,
+        STDERR_LOG_PATH=stderr_log_path,
+        CORE_NUM=core_num,
+        NODE_NUM=node_num,
+        QUEUE_NAME=queue_name,
+    )
+
+
+__all__ = ["slurm_generate_settings"]

wrfrun/scheduler/utils.py

@@ -0,0 +1,14 @@
+from wrfrun.core import get_wrfrun_config
+
+
+def get_core_num() -> int:
+    """
+    Get core number.
+
+    :return: Core number.
+    :rtype: int
+    """
+    return get_wrfrun_config()["core_num"]
+
+
+__all__ = ["get_core_num"]

wrfrun/utils.py

@@ -1,7 +1,7 @@
 import logging
 import subprocess
 from datetime import datetime
-from os import chdir, getcwd, makedirs
+from os import chdir, getcwd, makedirs, environ
 from os.path import exists
 from shutil import rmtree
 from time import time
@@ -45,7 +45,12 @@ def set_logger(logger_list: List[str], logger_level: Optional[Dict] = None):
                 
 # init wrfrun logger
 logger = logging.getLogger("wrfrun")
-set_logger(["wrfrun", ], {"wrfrun": logging.DEBUG})
+# check environment variables and set logger level
+if "WRFRUN_DEBUG_MODE" in environ and environ["WRFRUN_DEBUG_MODE"]:
+    _logger_level = logging.DEBUG
+else:
+    _logger_level = logging.INFO
+set_logger(["wrfrun", ], {"wrfrun": _logger_level})
 
 
 def unify_logger_format():
@@ -218,7 +223,7 @@ def check_subprocess_status(status: subprocess.CompletedProcess):
         logger.error(f"====== ====== ======")
 
         # raise error
-        raise RuntimeError
+        raise RuntimeError(f"Failed to exec command: '{command}'. Please check the log above.")
 
 
 def call_subprocess(command: list[str], work_path: Optional[str] = None, print_output=False):

wrfrun/workspace/__init__.py

@@ -0,0 +1,38 @@
+"""
+wrfrun.workspace
+################
+
+Prepare ``workspace`` for ``wrfrun`` and numerical models.
+
+Submodules
+**********
+
+=================================    ===========================================================
+:doc:`core </api/workspace.core>`    Core functions of this submodule.
+:doc:`wrf </api/workspace.wrf>`      Functions to prepare workspace for WPS/WRF model.
+=================================    ===========================================================
+
+Workspace
+*********
+
+``workspace`` is a collection of several directories where ``wrfrun``, extensions and numerical model works.
+These directories and their purpose are listed below.
+
+===================================         ===========================================================
+Director Path                               Purpose
+===================================         ===========================================================
+``/tmp/wrfrun``                             Store temporary files.
+``$HOME/.config/wrfrun``                    Main work directory.
+``$HOME/.config/wrfrun/replay``             Work directory for :doc:`replay <wrfrun.core.replay>`.
+``$HOME/.config/wrfrun/model``              Work directory for numerical models.
+===================================         ===========================================================
+
+.. toctree::
+    :maxdepth: 1
+    :hidden:
+
+    core <workspace.core>
+    wrf <workspace.wrf>
+"""
+
+from .core import *

wrfrun/workspace/core.py

@@ -0,0 +1,94 @@
+"""
+wrfrun.workspace.core
+#####################
+
+Core functions to prepare ``wrfrun`` workspace.
+
+.. autosummary::
+    :toctree: generated/
+
+    prepare_workspace
+"""
+
+from os.path import exists
+from shutil import rmtree
+
+from wrfrun.core import get_wrfrun_config
+from wrfrun.utils import check_path, logger
+from .wrf import check_wrf_workspace, prepare_wrf_workspace
+
+
+def prepare_workspace():
+    """
+    Initialize ``wrfrun`` workspace.
+
+    This function will check following paths,
+    and create them if them don't exist:
+
+    1. ``/tmp/wrfrun``
+    2. ``$HOME/.config/wrfrun``
+    3. ``$HOME/.config/wrfrun/replay``
+    4. ``$HOME/.config/wrfrun/model``
+
+    It will call other responding functions to initialize workspace for numerical models:
+
+    1. :doc:`WPS/WRF model </api/workspace.wrf>`
+    """
+    WRFRUNConfig = get_wrfrun_config()
+    logger.info(f"Initialize main workspace.")
+
+    wrfrun_temp_path = WRFRUNConfig.parse_resource_uri(WRFRUNConfig.WRFRUN_TEMP_PATH)
+    workspace_path = WRFRUNConfig.parse_resource_uri(WRFRUNConfig.WRFRUN_WORKSPACE_ROOT)
+    replay_work_path = WRFRUNConfig.parse_resource_uri(WRFRUNConfig.WRFRUN_WORKSPACE_REPLAY)
+    output_path = WRFRUNConfig.parse_resource_uri(WRFRUNConfig.WRFRUN_OUTPUT_PATH)
+
+    if exists(workspace_path):
+        logger.info(f"Remove old files in workspace.")
+        rmtree(workspace_path)
+
+    # check folder
+    check_path(wrfrun_temp_path)
+    check_path(replay_work_path)
+    check_path(output_path)
+
+    func_map = {
+        "wrf": prepare_wrf_workspace
+    }
+    model_configs = WRFRUNConfig["model"]
+
+    for model_name in model_configs:
+        func_map[model_name](model_configs[model_name])
+
+
+def check_workspace() -> bool:
+    """
+    Check if workspace exists.
+
+    :return: ``True`` if workspace exists, ``False`` otherwise.
+    :rtype: bool
+    """
+    WRFRUNConfig = get_wrfrun_config()
+
+    wrfrun_temp_path = WRFRUNConfig.parse_resource_uri(WRFRUNConfig.WRFRUN_TEMP_PATH)
+    workspace_path = WRFRUNConfig.parse_resource_uri(WRFRUNConfig.WRFRUN_WORKSPACE_ROOT)
+    replay_work_path = WRFRUNConfig.parse_resource_uri(WRFRUNConfig.WRFRUN_WORKSPACE_REPLAY)
+    output_path = WRFRUNConfig.parse_resource_uri(WRFRUNConfig.WRFRUN_OUTPUT_PATH)
+
+    flag = True
+    flag = flag & exists(wrfrun_temp_path) & exists(replay_work_path) & exists(output_path) & exists(workspace_path)
+
+    func_map = {
+        "wrf": check_wrf_workspace
+    }
+    model_configs = WRFRUNConfig["model"]
+
+    for model_name in model_configs:
+        if model_name == "debug_level":
+            continue
+
+        flag = flag & func_map[model_name](model_configs[model_name])
+
+    return True
+
+
+__all__ = ["prepare_workspace", "check_workspace"]