espsite.py.example.SLURMsherlock

# cluster-dependent definitions
# espresso.py will make use of all "self" variables in config class
import os

class config:
    def __init__(self):
        self.scratch = '/local-scratch/'+os.getenv('USER')
        if not os.path.exists(self.scratch):
            self.scratch = '/tmp'
        self.submitdir = os.getenv('SLURM_SUBMIT_DIR')
        self.batch = self.submitdir!=None
        #we typically run batch jobs without openmp threads here
        if self.batch:
            os.putenv('OMP_NUM_THREADS','1')
        #check for batch
        if self.batch:
            self.jobid = os.getenv('SLURM_JOBID')
            nodefile = self.submitdir+'/nodefile.'+self.jobid
            uniqnodefile = self.submitdir+'/uniqnodefile.'+self.jobid
            os.system('scontrol show hostnames $SLURM_JOB_NODELIST >'+uniqnodefile)
            taskspernode = os.getenv('SLURM_TASKS_PER_NODE')
            xtaskspernode = taskspernode.find('(')
            if xtaskspernode > -1:
        	taskspernode = taskspernode[:xtaskspernode]
            os.system("awk '{for(i=0;i<"+taskspernode+";i++)print}' "+uniqnodefile+" >"+nodefile)
            f = open(nodefile, 'r')
            procs = [x.strip() for x in f.readlines()]
            f.close()
            self.procs = procs

            nprocs = len(procs)
            self.nprocs = nprocs
            
            p = os.popen('wc -l <'+uniqnodefile, 'r')
            nnodes = p.readline().strip()
            p.close()
        
            self.perHostMpiExec = 'mpiexec -machinefile '+uniqnodefile+' -np '+nnodes
            self.perProcMpiExec = 'mpiexec -machinefile '+nodefile+' -np '+str(nprocs)+' -wdir %s %s'
            self.perSpecProcMpiExec = 'mpiexec -machinefile %s -np %d -wdir %s %s'


    def do_perProcMpiExec(self, workdir, program):
        return os.popen2(self.perProcMpiExec % (workdir, program))

    def do_perProcMpiExec_outputonly(self, workdir, program):
        return os.popen(self.perProcMpiExec % (workdir, program), 'r')

    def runonly_perProcMpiExec(self, workdir, program):
        os.system(self.perProcMpiExec % (workdir, program))

    def do_perSpecProcMpiExec(self, machinefile, nproc, workdir, program):
        return os.popen3(self.perSpecProcMpiExec % (machinefile, nproc, workdir, program))