Source code for polaris.job

import importlib.resources as imp_res
import os

import numpy as np
from jinja2 import Template


[docs]def write_job_script(config, machine, target_cores, min_cores, work_dir, suite=''): """ Parameters ---------- config : polaris.config.PolarisConfigParser Configuration options for this test case, a combination of user configs and the defaults for the machine and component machine : {str, None} The name of the machine target_cores : int The target number of cores for the job to use min_cores : int The minimum number of cores for the job to use work_dir : str The work directory where the job script should be written suite : str, optional The name of the suite """ if config.has_option('parallel', 'account'): account = config.get('parallel', 'account') else: account = '' cores_per_node = config.getint('parallel', 'cores_per_node') # as a rule of thumb, let's do the geometric mean between min and target cores = np.sqrt(target_cores * min_cores) nodes = int(np.ceil(cores / cores_per_node)) partition, qos, constraint, wall_time = get_slurm_options( config, machine, nodes) job_name = config.get('job', 'job_name') if job_name == '<<<default>>>': if suite == '': job_name = 'polaris' else: job_name = f'polaris_{suite}' template = Template(imp_res.files('polaris.job').joinpath( 'job_script.template').read_text()) text = template.render(job_name=job_name, account=account, nodes=f'{nodes}', wall_time=wall_time, qos=qos, partition=partition, constraint=constraint, suite=suite) text = _clean_up_whitespace(text) if suite == '': script_filename = 'job_script.sh' else: script_filename = f'job_script.{suite}.sh' script_filename = os.path.join(work_dir, script_filename) with open(script_filename, 'w') as handle: handle.write(text)
[docs]def get_slurm_options(config, machine, nodes): """ Get Slurm options Parameters ---------- config : polaris.config.PolarisConfigParser Config options machine : str Name of the machine nodes : int Number of nodes Returns ------- partition : str Slurm partition qos : str Slurm quality of service constraint : str Slurm constraint wall_time : str Slurm wall time """ partition = config.get('job', 'partition') if partition == '<<<default>>>': if machine == 'anvil': # choose the partition based on the number of nodes if nodes <= 5: partition = 'acme-small' elif nodes <= 60: partition = 'acme-medium' else: partition = 'acme-large' elif config.has_option('parallel', 'partitions'): # get the first, which is the default partition = config.getlist('parallel', 'partitions')[0] else: partition = '' qos = config.get('job', 'qos') if qos == '<<<default>>>': if config.has_option('parallel', 'qos'): # get the first, which is the default qos = config.getlist('parallel', 'qos')[0] else: qos = '' constraint = config.get('job', 'constraint') if constraint == '<<<default>>>': if config.has_option('parallel', 'constraints'): # get the first, which is the default constraint = config.getlist('parallel', 'constraints')[0] else: constraint = '' wall_time = config.get('job', 'wall_time') return partition, qos, constraint, wall_time
def _clean_up_whitespace(text): prev_line = None lines = text.split('\n') trimmed = list() # remove extra blank lines for line in lines: if line != '' or prev_line != '': trimmed.append(line) prev_line = line line = '' lines = list() # remove blank lines between comments for next_line in trimmed: if line != '' or not next_line.startswith('#'): lines.append(line) line = next_line # add the last line that we missed and an extra blank line lines.extend([trimmed[-1], '']) text = '\n'.join(lines) return text