Default Templates

This page describes the structure and design of the default task templates. These templates are used to write task and calculation scripts when writing task inputs (i.e., with Task.write_inputs()). All templates are written in Jinja.

run.sh.j2

run.sh.j2 is the default task script template for all autojob-defined TaskBase implementations. With this template, scheduler inputs are written, configuration files (/etc/profile and $HOME/.bash_profile) are sourced, and the task is executed. The resulting script is intended to be executed with Bash.

#!/bin/bash
{% if scheduler_inputs %}
{% for key, value in scheduler_inputs.items() %}
{% if value %}
#SBATCH {{ key }}={{ value }}
{% endif %}
{% endfor %}
{% endif %}

echo "\n... Setting up shell environment ...\n"
shopt -s expand_aliases

if test -e "/etc/profile"; then source "/etc/profile"; fi
if test -e "$HOME/.bash_profile"; then source "$HOME/.bash_profile"; fi

# run task
{% if calculation_inputs %}
python {{ calculation_inputs.calculation_script }}
{% else %}
cp {{ task_inputs.atoms_filename }} {{ settings.OUTPUT_ATOMS_FILE }}
{% endif %}

The first block (i.e., the block starting with {% if scheduler_inputs %}) prints out the scheduler inputs for the task. The block writes the scheduler inputs (if defined) in a format supported by the SLURM workload manager. Next, the expand_aliases shell option is enabled. (This is useful since aliases are often used to condense useful commands for use in task scripts. For example, it is common to define an alias for activating Python virtual environments.) This is followed by sourcing common Bash configuration files, /etc/profile and $HOME/.bash_profile. Finally, depending on whether the task has calculation inputs or not, the last block will either execute the calculation script with Python or copy the inputs file to the outputs file.

calc.sh.j2

calc.sh.j2 is a built-in task script template that leverages advanced features of autojob. This template extends run.sh.j2 be loading modules, calling custom commands (activate_env, debug-slurm, log-job), and then calling autojob CLI commands (autojob run, autojob harvest, autojob clean, autojob restart, and autojob advance).

#!/bin/bash
{% if scheduler_inputs %}
{% for key, value in scheduler_inputs.items() %}
{% if value %}
#SBATCH {{ key }}={{ value }}
{% endif %}
{% endfor %}
{% endif %}

echo "\n... Setting up shell environment ...\n"
shopt -s expand_aliases
module purge

if test -e "/etc/profile"; then source "/etc/profile"; fi
if test -e "$HOME/.bash_profile"; then source "$HOME/.bash_profile"; fi

unset LANG; ulimit -s unlimited; ulimit -a -S; ulimit -a -H
export LC_ALL="C"; export MKL_NUM_THREADS=1; export OMP_NUM_THREADS=1

module load python {{ calculator|lower }}
activate_env
debug-slurm

# run task
{% if calculation_inputs %}
autojob run --buffer=0.1 --time-source={{ task_inputs.task_script }} 'python {{ calculation_inputs.calculation_script }}'
exit_code=$?
autojob -vv harvest
{% endif %}
autojob clean
log-job

if [ "$exit_code" -eq 124 ]; then
autojob -vv restart -S; else
autojob -vv advance; fi

Lines 17-18 define common environment variables for running computational codes.

run.py.j2

run.sh.j2 is the default calculation script template for all autojob-defined Calculation subclasses. This Python script template loads the input Atoms, configures the ASE Calculator and Optimizer, calls the ccu function run_calculation(), and then runs any analyses.

from importlib import import_module
import json
import logging
from pathlib import Path

import ase.io
from ase.calculators.calculator import get_calculator_class
from ccu.workflows.calculation import run_calculation

from autojob.plugins import get_analysis
from autojob.tasks.calculation import CalculationInputs

logging.basicConfig(level=logging.DEBUG)

with Path("{{ settings.INPUTS_FILE }}").open(mode="r", encoding="utf-8") as file:
    inputs = json.load(file)

atoms = ase.io.read(inputs["task_inputs"]["atoms_filename"])

# Configure calculator
calc_inputs = CalculationInputs(**inputs["calculation_inputs"])
calculator = get_calculator_class(calc_inputs.calculator)
calc = calculator(**calc_inputs.calc_params)
atoms.calc = calc

# Configure optimizer
if calc_inputs.optimizer:
    *opt_mod_parts, opt_cls_name = calc_inputs.optimizer.split(".")
    opt_cls = getattr(import_module(".".join(opt_mod_parts)), opt_cls_name)
    opt = opt_cls(atoms, **calc_inputs.opt_params["init"])
    opt_run_params = calc_inputs.opt_params["run"]
else:
    opt = None
    opt_run_params = {}

run_calculation(atoms, opt=opt, **opt_run_params, output={{ settings.OUTPUT_ATOMS_FILE }})

# Run analyses
if calc_inputs.analyses:
    for name, (args, kwargs) in calc_inputs.analyses.items():
        analysis = get_analysis(name)
        _ = analysis(*args, **kwargs)

The run_calculation() function called on line 36 will result in the output trajectory being written.

vib.py.j2

vib.sh.j2 is the a calculation script template intended for use by Vibration tasks. This Python script template loads the input Atoms, configures the ASE Calculator and calls the ccu function run_vibration(), and then runs any analyses.

import json
import logging
from pathlib import Path

from ase.calculators.calculator import get_calculator_class
import ase.io
from ccu.workflows.vibration import run_vibration

from autojob.plugins import get_analysis
from autojob.tasks.calculation import CalculationInputs

logging.basicConfig(level=logging.DEBUG)

with Path("{{ settings.INPUTS_FILE }}").open(
    mode="r", encoding="utf-8"
) as file:
    inputs = json.load(file)

atoms = ase.io.read(inputs["task_inputs"]["atoms_filename"])

# Configure calculator
calc_inputs = CalculationInputs(**inputs["calculation_inputs"])
calculator = get_calculator_class(calc_inputs.calculator)
calc = calculator(**calc_inputs.calc_params)
atoms.calc = calc

run_vibration(atoms, **inputs["vibration_inputs"])

# Run analyses
if calc_inputs.analyses:
    for name, (args, kwargs) in calc_inputs.analyses.items():
        analysis = get_analysis(name)
        _ = analysis(*args, **kwargs)

The run_vibration() function called on line 27 will result in the output trajectory being written along with a JSON file containing the information needed to instantiate a VibrationsData object representing the results of the vibrational analysis.

See also

A Note on Templates