diff --git a/zppy/__main__.py b/zppy/__main__.py index 366f5320..8a7afac2 100644 --- a/zppy/__main__.py +++ b/zppy/__main__.py @@ -12,6 +12,7 @@ from mache import MachineInfo from validate import Validator +from zppy.budget import budget from zppy.bundle import Bundle, predefined_bundles from zppy.climo import climo from zppy.e3sm_diags import e3sm_diags @@ -271,6 +272,9 @@ def _launch_scripts(config: ConfigObj, script_dir, job_ids_file, plugins) -> Non # ilamb tasks existing_bundles = ilamb(config, script_dir, existing_bundles, job_ids_file) + # budget tasks + existing_bundles = budget(config, script_dir, existing_bundles, job_ids_file) + # pcmdi_diags tasks existing_bundles = pcmdi_diags(config, script_dir, existing_bundles, job_ids_file) diff --git a/zppy/budget.py b/zppy/budget.py new file mode 100644 index 00000000..30ba7268 --- /dev/null +++ b/zppy/budget.py @@ -0,0 +1,79 @@ +from typing import Any, Dict, List, Tuple + +from configobj import ConfigObj + +from zppy.bundle import handle_bundles +from zppy.utils import ( + check_status, + get_file_names, + get_tasks, + get_years, + initialize_template, + make_executable, + submit_script, + write_settings_file, +) + + +# ----------------------------------------------------------------------------- +def budget(config: ConfigObj, script_dir: str, existing_bundles, job_ids_file): + + template, _ = initialize_template(config, "budget.bash") + + # --- List of budget tasks --- + tasks: List[Dict[str, Any]] = get_tasks(config, "budget") + if len(tasks) == 0: + return existing_bundles + + # --- Generate and submit budget scripts --- + for c in tasks: + + dependencies: List[str] = [] + + # Loop over year sets + year_sets: List[Tuple[int, int]] = get_years(c["years"]) + for s in year_sets: + c["year1"] = s[0] + c["year2"] = s[1] + if ("last_year" in c.keys()) and (c["year2"] > c["last_year"]): + continue # Skip this year set + c["scriptDir"] = script_dir + + prefix = f"budget_{c['year1']:04d}-{c['year2']:04d}" + print(prefix) + c["prefix"] = prefix + bash_file, settings_file, status_file = get_file_names(script_dir, prefix) + skip: bool = check_status(status_file) + if skip: + continue + # Create script + with open(bash_file, "w") as f: + f.write(template.render(**c)) + make_executable(bash_file) + c["dependencies"] = dependencies + write_settings_file(settings_file, c, s) + export = "NONE" + existing_bundles = handle_bundles( + c, + bash_file, + export, + dependFiles=dependencies, + existing_bundles=existing_bundles, + ) + if not c["dry_run"]: + if c["bundle"] == "": + # Submit job + submit_script( + bash_file, + status_file, + export, + job_ids_file, + dependFiles=dependencies, + fail_on_dependency_skip=c["fail_on_dependency_skip"], + ) + else: + print(f"...adding to bundle {c['bundle']}") + + print(f" environment_commands={c['environment_commands']}") + + return existing_bundles diff --git a/zppy/defaults/default.ini b/zppy/defaults/default.ini index 01e12c3d..9173ff55 100755 --- a/zppy/defaults/default.ini +++ b/zppy/defaults/default.ini @@ -179,6 +179,19 @@ input_files = string(default="eam.h2") # If not set, zppy will attempt to infer this value from the topography file. res = string(default="") +[budget] +# Budget types to analyze: water, heat (comma-separated) +budget_types = string(default="water,heat") +# Generate HTML plots with interactive visualizations +output_html = boolean(default=True) +# Generate ASCII summary tables +output_ascii = boolean(default=True) + + [[__many__]] + budget_types = string(default=None) + output_html = boolean(default=None) + output_ascii = boolean(default=None) + [e3sm_diags] # See https://e3sm-project.github.io/e3sm_diags/_build/html/master/available-parameters.html backend = string(default="mpl") diff --git a/zppy/templates/budget.bash b/zppy/templates/budget.bash new file mode 100644 index 00000000..7ba4a649 --- /dev/null +++ b/zppy/templates/budget.bash @@ -0,0 +1,147 @@ +#!/bin/bash +{% include 'inclusions/slurm_header.bash' %} +{% include 'inclusions/boilerplate.bash' %} +set -e +{{ environment_commands }} +set +e + +# Generate water and energy budget analysis plots +################################################################################ + +# Create temporary workdir +hash=`mktemp --dry-run -d XXXX` +workdir=tmp.{{ prefix }}.${id}.${hash} +mkdir ${workdir} +cd ${workdir} +echo "Working in temporary directory: ${workdir}" + +# Execute budget analysis using external CLI tool +echo "Running budget analysis for years {{ year1 }} to {{ year2 }}..." + +# Construct log path from zppy parameters +LOG_PATH="{{ input }}/{{ case }}/archive/logs" + +zi-budget-analysis \ + --log_path "${LOG_PATH}" \ + --start_year {{ year1 }} \ + --end_year {{ year2 }} \ + --budget_types {{ budget_types }} \ +{%- if output_html %} + --output_html \ +{%- endif %} +{%- if output_ascii %} + --output_ascii \ +{%- endif %} + --output_dir . + +if [ $? != 0 ]; then + cd {{ scriptDir }} + echo 'ERROR (1)' > {{ prefix }}.status + exit 1 +fi + +# Copy results to case directory +echo +echo "===== COPY RESULTS TO CASE DIRECTORY =====" +echo + +case_dir={{ output }} +dest_dir=${case_dir}/post/budget/{{ year1 }}-{{ year2 }} +mkdir -p ${dest_dir} +if [ $? != 0 ]; then + cd {{ scriptDir }} + echo 'ERROR (2)' > {{ prefix }}.status + exit 2 +fi + +# Copy all generated files +cp *.html ${dest_dir}/ 2>/dev/null || true +cp *.txt ${dest_dir}/ 2>/dev/null || true +cp *.png ${dest_dir}/ 2>/dev/null || true + +echo "Results copied to ${dest_dir}/" + +# Copy output to web server +echo +echo "===== COPY FILES TO WEB SERVER =====" +echo + +# Create web directory +www={{ www }} +case={{ case }} +web_dir=${www}/${case}/budget/{{ year1 }}-{{ year2 }} +mkdir -p ${web_dir} +if [ $? != 0 ]; then + cd {{ scriptDir }} + echo 'ERROR (3)' > {{ prefix }}.status + exit 3 +fi + +{% if machine in ['pm-cpu', 'pm-gpu'] %} +# For NERSC, make sure it is world readable +f=`realpath ${web_dir}` +while [[ $f != "/" ]] +do + owner=`stat --format '%U' $f` + if [ "${owner}" = "${USER}" ]; then + chgrp e3sm $f + chmod go+rx $f + fi + f=$(dirname $f) +done +{% endif %} + +# Copy HTML files for web viewing +cp *.html ${web_dir}/ 2>/dev/null || true +if [ $? != 0 ]; then + cd {{ scriptDir }} + echo 'ERROR (4)' > {{ prefix }}.status + exit 4 +fi + +{% if machine in ['pm-cpu', 'pm-gpu'] %} +# For NERSC, change permissions of new files +pushd ${web_dir} +chgrp -R e3sm *.html 2>/dev/null || true +chmod -R go+rX,go-w *.html 2>/dev/null || true +popd +{% endif %} + +{% if machine in ['anvil', 'chrysalis'] %} +# For LCRC, change permissions of new files +pushd ${web_dir} +chmod -R go+rX,go-w *.html 2>/dev/null || true +popd +{% endif %} + +echo "Web files copied to ${web_dir}/" + +# Clean up temporary workdir +echo +echo "===== CLEANUP TEMPORARY DIRECTORY =====" +echo + +cd {{ scriptDir }} +if [[ "${debug,,}" != "true" ]]; then + rm -rf ${workdir} + if [ $? = 0 ]; then + echo "Successfully cleaned up ${workdir}" + else + echo "Warning: Failed to clean up ${workdir}" + fi +else + echo "Debug mode: keeping temporary directory ${workdir}" +fi + +################################################################################ +# Update status file and exit +{% raw %} +ENDTIME=$(date +%s) +ELAPSEDTIME=$(($ENDTIME - $STARTTIME)) +{% endraw %} +echo ============================================== +echo "Elapsed time: $ELAPSEDTIME seconds" +echo ============================================== +rm -f {{ prefix }}.status +echo 'OK' > {{ prefix }}.status +exit 0