wrfrun 0.1.7__py3-none-any.whl → 0.1.9__py3-none-any.whl

wrfrun/res/{config.toml.template → config/wrf.template.toml}

@@ -1,46 +1,5 @@
-# Config template of wrfrun.
-# Path of the directory which contains input data.
-input_data_path = ""
-
-# Path of the directory to store all outputs.
-output_path = "./outputs"
-log_path = "./logs"
-
-# wrfrun can launch a socket server during NWP execution to report simulation progress.
-# To enable the server, you need to configure the IP address and port on which it will listen to.
-server_host = "localhost"
-server_port = 54321
-
-# How many cores you will use.
-# Note that if you use a job scheduler (like PBS), this value means the number of cores each node you use.
-core_num = 36
-
-
-[job_scheduler]
-# Job scheduler settings.
-# How many nodes you will use.
-node_num = 1
-
-# Custom environment settings
-env_settings = {}
-
-# Path of the python interpreter that will be used to run wrfrun.
-# You can also give its name only.
-python_interpreter = "/usr/bin/python3"     # or just "python3"
-
-
-[model]
-# Model debug level
-debug_level = 100
-
-# ################################################### Only settings above is necessary ###########################################
-
-
-# ####################################### You can give more settings about the NWP you will use ##################################
-
-[model.wrf]
-# Config for WRF
-# WRF model path
+# Config for WRF model.
+# WRF model path.
 wps_path = '/path/to/your/WPS/folder'
 wrf_path = '/path/to/your/WRF/folder'
 # WRFDA is optional.
@@ -59,8 +18,10 @@ user_wrfda_namelist = ''
 # If you make a restart run?
 restart_mode = false
 
+# debug level for WRF model
+debug_level = 100
 
-[model.wrf.time]
+[time]
 # Advance time config for WRF
 # Set the start and end date. It will be used for all domains.
 # You can also provide all the dates as a list, with each date for the corresponding domain.
@@ -87,7 +48,7 @@ parent_time_step_ratio = [1, 3, 4]
 restart_interval = -1
 
 
-[model.wrf.domain]
+[domain]
 # Advance domain config for WRF.
 # Set domain number.
 domain_num = 3
@@ -109,7 +70,9 @@ dx = 9000
 dy = 9000
 
 # Projection.
-map_proj = { name = 'lambert', truelat1 = 34.0, truelat2 = 40.0}
+map_proj = 'lambert'
+truelat1 = 34.0
+truelat2 = 40.0
 
 # Central point of the first area.
 ref_lat = 37.0
@@ -117,7 +80,7 @@ ref_lon = 120.5
 stand_lon = 120.5
 
 
-[model.wrf.scheme]
+[scheme]
 # Advance physics scheme config for WRF.
 # To loop up the nickname for all physics schemes, please see: https://wrfrun.syize.cn.
 # Option contains many other settings related to the scheme.

wrfrun/res/run.template.sh

@@ -0,0 +1,10 @@
+#!/bin/bash
+{SCHEDULER_SETTINGS}
+
+ulimit -s unlimited
+
+{ENV_SETTINGS}
+
+cd {WORK_PATH}
+
+{WORK_COMMAND}

wrfrun/res/scheduler/lsf.template

@@ -0,0 +1,5 @@
+#BSUB -q {QUEUE_NAME}
+#BSUB -o {STDOUT_LOG_PATH}
+#BSUB -e {STDERR_LOG_PATH}
+#BSUB -n {CORE_NUM}
+#BSUB -J wrfrun

wrfrun/res/{job_scheduler → scheduler}/pbs.template

@@ -1,4 +1,4 @@
-#PBS -q batch
+#PBS -q {QUEUE_NAME}
 #PBS -N wrfrun
 #PBS -o {STDOUT_LOG_PATH}
 #PBS -e {STDERR_LOG_PATH}

wrfrun/res/{job_scheduler → scheduler}/slurm.template

@@ -3,4 +3,5 @@
 #SBATCH -n{CORE_NUM}                %core num
 #SBATCH --time=9999:00:00           %maximum run time
 #SBATCH --output={STDOUT_LOG_PATH}  %stdout output log path
-#SBATCH --error={STDERR_LOG_PATH}   %error output log path
+#SBATCH --error={STDERR_LOG_PATH}   %error output log path
+#SBATCH -p {QUEUE_NAME}

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 .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,14 +41,14 @@ 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):
         """
         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.
         :param init_workspace: If True, clean old files in workspace and re-create it.
         :param start_server: Whether to start WRFRunServer, defaults to True.
-        :param pbs_mode: If commit this task to the PBS system, defaults to True.
+        :param submit_job: If commit this task to the PBS system, defaults to True.
         :param prepare_wps_data: If True, download input datas for WPS first.
         :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.
         :return:
@@ -63,7 +63,7 @@ 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
@@ -89,31 +89,27 @@ 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()
 
-            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()
 
@@ -139,7 +135,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 +153,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,47 @@
+"""
+wrfrun.scheduler.lsf
+####################
+
+Scheduler interface for LSF system.
+
+.. autosummary::
+    :toctree: generated/
+
+    lsf_generate_settings
+"""
+
+from wrfrun.core import WRFRUNConfig
+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
+    """
+    # 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,48 @@
+"""
+wrfrun.scheduler.pbs
+####################
+
+Scheduler interface for PBS system.
+
+.. autosummary::
+    :toctree: generated/
+
+    pbs_generate_settings
+"""
+
+from wrfrun.core import WRFRUNConfig
+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
+    """
+    # 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,70 @@
+from os.path import exists, abspath, dirname
+
+from wrfrun import WRFRUNConfig
+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
+    """
+    # 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,48 @@
+"""
+wrfrun.scheduler.slurm
+######################
+
+Scheduler interface for Slurm system.
+
+.. autosummary::
+    :toctree: generated/
+
+    slurm_generate_settings
+"""
+
+from wrfrun.core import WRFRUNConfig
+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
+    """
+    # 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 import WRFRUNConfig
+
+
+def get_core_num() -> int:
+    """
+    Get core number.
+
+    :return: Core number.
+    :rtype: int
+    """
+    return WRFRUNConfig["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 *