JobScript.py

Go to the documentation of this file.

#!/bin/env python3
 
import os
from pathlib import Path
import subprocess
 
class JobScriptBase():
    '''
    Describes an HPC job script including
    + reading a configuration
    + generating the script
    + submitting the script
 
    Each HPC job submission system (e.g., PBSPro on Cheyenne)
    will have its own derived class that defines
    the job header and submission command
    generic config elements:
        required config parameter(s):
        basescript (required) - either a list of str's containing individual lines of the script
                                or a str giving the location of the script
 
        optional config parameter(s):
        env - linux environment of the script (e.g., csh, bash, sh, tcsh)
        name - job name
        nppernode - processors per node
        nnode - number of nodes
        walltime - walltime
        olog - output log name
        elog - error log name
    '''
    def __init__(self, conf):
        ## job descriptors
        self.namename = conf.get('name','PyJobScript')
        self.nppernodenppernode = conf.get('nppernode',1)
        self.nnodennode = conf.get('nnode',1)
        self.walltimewalltime = conf.get('walltime','01:00:00')
        self.ologolog = conf.get('olog','log.job.out')
        self.elogelog = conf.get('elog','log.job.err')
 
        ## submission descriptors
        self.envenv = conf.get('env','csh')
        self.basescriptbasescript = conf['script']
        assert (isinstance(self.basescriptbasescript,list) or isinstance(self.basescriptbasescript,str)), \
            "JobScriptBase : basescript must either be a list or a string"
        self.jobpathjobpath = Path(conf.get('path','./'))
        self.scriptscript = self.namename+'.job.'+self.envenv
 
        self.commandcommand = './'
        self.headerheader = []
 
    def create(self):
        ## initialize env
        joblines = ['#!/bin/'+self.envenv+'\n']
 
        ## concatenate header
        for line in self.headerheader: joblines.append(line+'\n')
 
        ## concatenate script
        if isinstance(self.basescriptbasescript,list):
            # assume self.basescript is a list of lines
            for line in self.basescriptbasescript: joblines.append(line)
 
        elif isinstance(self.basescriptbasescript,str):
            # assume self.basescript is a file
            bs = open(self.basescriptbasescript,'r')
            for line in bs: joblines.append(line)
            bs.close()
 
        ## create the job path
        self.jobpathjobpath.mkdir(parents=True, exist_ok=True)
 
        ## write job script
        script = str(self.jobpathjobpath/self.scriptscript)
        if os.path.exists(script):
            os.remove(script)
        js = open(script,'w')
        js.writelines(joblines)
        js.close()
        os.system('chmod 744 '+script)
 
    def submit(self):
        ## submit job
        command = self.commandcommand+self.scriptscript
        CWD = os.getcwd()
        os.chdir(str(self.jobpathjobpath))
        print(command+" in "+os.getcwd())
        os.system(command)
        os.chdir(CWD)
 
 
class PBSProCheyenne(JobScriptBase):
    '''
    PBSPro job script on Cheyenne
    unique config elements compared to base class:
        account - cheyenne account for charging
        queue   - name of job submission queue (see qavail)
        memory  - amount of memory requested per node (see mavail)
 
    NOTE: Cheyenne has a maximum of 36 processors available per node
    '''
    qavail = ['economy', 'regular', 'premium']
    mavail = [45, 109]
    maxnppernode = 36
    def __init__(self, conf):
        # Initialize derived config settings
        super().__init__(conf)
 
        # Initialize config settings that are specific to PBSProCheyenne
        self.accountaccount = conf.get('account','NMMM0043')
        self.queuequeue = conf.get('queue','regular')
        assert self.queuequeue in self.qavailqavail, ("ERROR: PBSProCheyenne requires queue to be any of ",self.qavailqavail)
        self.memorymemory = conf.get('memory',109)
        assert self.memorymemory in self.mavailmavail, ("ERROR: PBSProCheyenne requires memory (in GB) to be any  of", self.mavailmavail)
        assert self.nppernodenppernode <= self.maxnppernodemaxnppernode, ("ERROR: PBSProCheyenne requires nppernode <= ", self.maxnppernodemaxnppernode)
 
        self.headerheaderheader = [
            '#PBS -N '+self.namename,
            '#PBS -A '+self.accountaccount,
            '#PBS -q '+self.queuequeue,
            '#PBS -l select='+str(self.nnodennode)+':ncpus='+str(self.nppernodenppernode)+':mpiprocs='+str(self.nppernodenppernode)+':mem='+str(self.memorymemory)+'GB',
            '#PBS -l walltime='+self.walltimewalltime,
            '#PBS -m ae',
            '#PBS -k eod',
            '#PBS -o '+self.ologolog,
            '#PBS -e '+self.elogelog,
        ]
 
        self.commandcommandcommand = 'qsub '
 
 
class SLURMCasper(JobScriptBase):
    '''
    SLURM job script on Casper
    unique config elements compared to base class:
        account - casper account for charging
        partition - name of casper partition (see pavail)
        memory  - amount of memory requested per node (see maxmemory)
 
    NOTE: Casper has a maximum of 36 processors available per node
    '''
    pavail = ['dav']
    maxnppernode = 36
    maxmemory = 1100
    def __init__(self, conf):
        # Initialize derived config settings
        super().__init__(conf)
 
        # Initialize config settings that are specific to SLURMCasper
        self.accountaccount = conf.get('account','NMMM0015')
        self.partitionpartition = conf.get('partition','dav')
        assert self.partitionpartition in self.pavailpavail, ("ERROR: SLURMCasper requires partition to be any of ",self.pavailpavail)
        self.memorymemory = conf.get('memory',300)
        assert self.memorymemory <= self.maxmemorymaxmemory, ("ERROR: SLURMCasper requires memory (in GB) to be <= ", self.maxmemorymaxmemory)
 
        assert self.nppernodenppernode <= self.maxnppernodemaxnppernode, ("ERROR: SLURMCasper requires nppernode <= ", self.maxnppernodemaxnppernode)
 
        self.headerheaderheader = [
            '#SBATCH --job-name='+self.namename,
            '#SBATCH --account='+self.accountaccount,
            '#SBATCH --ntasks='+str(self.nnodennode),
            '#SBATCH --cpus-per-task='+str(self.nppernodenppernode),
            '#SBATCH --mem='+str(self.memorymemory)+'G',
            '#SBATCH --time='+self.walltimewalltime,
            '#SBATCH --partition='+self.partitionpartition,
            '#SBATCH --output='+self.ologolog,
        ]
 
        self.commandcommandcommand = 'sbatch '
 
JobScriptDict = {
    ## cheyenne
    # login nodes
    'cheyenne': PBSProCheyenne,
    # cron jobs
    'chadmin': PBSProCheyenne,
    ## casper
    'casper': SLURMCasper
}
 
 
def JobScriptFactory(conf):
    ## get system name
    conf['sysname'] = subprocess.run(['uname','-n'],
                      stdout=subprocess.PIPE).stdout.decode('utf-8')
 
    ## match system name with JobScriptDict or return base class object by default
    for key, jobclass in JobScriptDict.items():
        if key in conf['sysname']:
            return jobclass(conf)
    return JobScriptBase(conf)