clusterutils.py

import tempfile
import os
import sys

batchscript_claix18 = '''#!/usr/bin/env zsh

# Autogenerated jobscript targeting the claix18 cluster.

#SBATCH --job-name={jobname}
#SBATCH --time={walltime}
#SBATCH --mem-per-cpu={mem_per_cpu}
#SBATCH --ntasks={num_cores}
#SBATCH --export=NONE
#SBATCH --output={jobname}.data/stdout.%j.log
#SBATCH --error={jobname}.data/stderr.%j.log

{custom_cmds}
{mpirun} $FLAGS_MPI_BATCH {mc_cmd}
{custom_post_cmds}
'''

batch_commands = {
    'claix18': 'sbatch {batchscript}'
}

valid_systems = ['local', 'claix18']

def generate_batchscript_claix18(cmd, jobname, jobconfig):
    template = batchscript_claix18

    custom_cmds = ''
    if 'project' in jobconfig:
        custom_cmds += '#SBATCH --account={}\n'.format(jobconfig['project'])
    custom_cmds += jobconfig.get('custom_cmds', '')
    custom_post_cmds = jobconfig.get('custom_post_cmds', '')

    try:
        return template.format(
            jobname=jobname,
            mpirun=jobconfig.get('mpirun','mpirun'),
            mem_per_cpu=jobconfig.get('mem_per_cpu','2G'),
            walltime=jobconfig['mc_walltime'],
            num_cores=jobconfig['num_cores'],
            custom_cmds=custom_cmds,
            custom_post_cmds=custom_post_cmds,
            mc_cmd=' '.join(cmd)
        )
    except KeyError as e:
        raise Exception('Error: required key "{}" missing in jobconfig'.format(e.args[0]))
        
def generate_batchscript(sysinfo, *args):
    if sysinfo == 'claix18':
        return generate_batchscript_claix18(*args)
    else:
        raise Exception('unknown system type "{}"'.format(sysinfo))

def determine_system():
    sysinfo = os.environ.get('MCLL_SYSTEM_INFO')

    if sysinfo == None:
        import socket
        hostname = socket.gethostname()
        if hostname.endswith('.hpc.itc.rwth-aachen.de'):
            sysinfo = 'claix18'

    if not sysinfo in valid_systems:
        raise Exception('unknown system type "{}"'.format(sysinfo))

    return sysinfo
    
def run(jobname, jobconfig, cmd):
    sysinfo = determine_system()

    print('running on system \'{}\''.format(sysinfo))

    if sysinfo == 'local':
        mpicmd = '{} -n {} {}'.format(jobconfig.get('mpirun', 'mpirun'), jobconfig['num_cores'], ' '.join(cmd))
        print('$ '+mpicmd)
        os.system(mpicmd)
    else:
        with tempfile.NamedTemporaryFile(mode='w',delete=False) as f:
            batchscript = generate_batchscript(sysinfo, cmd, jobname, jobconfig)
            print(batchscript)
            f.write(batchscript)
            bscriptname = f.name
        mpicmd = batch_commands[sysinfo].format(batchscript=bscriptname)
        print('$ '+mpicmd)
        os.system(mpicmd)
        os.unlink(f.name)