<commit_before>import os
import string
import subprocess

import mvpipe.support
from mvpipe.runner import Runner, Job

def_options = {'env': True, 'wd': os.path.abspath(os.curdir), 'hold': False, 'nodes': 1}

'''
SLURM options:

global_hold - start all pipelines with a "placeholder" job to synchronize
              starts (and ensure that all jobs are submitted correctly
              before releasing the pipeline

account     - a default account to use

shell       - a shell to use for the script (default(s): /bin/bash, /usr/bin/bash, /usr/local/bin/bash, /bin/sh)

'''
class SlurmRunner(Runner):
    def __init__(self, dryrun, verbose, logger, global_hold=False, global_depends=None, account=None, interpreter=None):
        Runner.__init__(self, dryrun, verbose, logger)
        self.global_hold = global_hold
        self._holding_job = None
        self.global_depends = global_depends if global_depends else []
        self.account = account

        self.jobids = []

        self.testjobcount = 1

        if interpreter:
            self.interpreter = interpreter
        else:
            for intp in ['/bin/bash', '/usr/bin/bash', '/usr/local/bin/bash']:
                if os.path.exists(intp):
                    self.interpreter = intp
                    break
            if not self.interpreter:
                self.interpreter = '/bin/sh'


    def reset(self):
        pass

    def abort(self):
        if self._holding_job:
            self.cancel(self._holding_job.jobid)
        else:
            for jobid in self.jobids:
                self.cancel(jobid)

    def done(self):
        if not self.dryrun:
            if self._holding_job:
                self.release(self._holding_job.jobid)

    def _setup_holding_job(self):
        self._holding_job = Job('sleep 5', name="holding", hold=True, stdout='/dev/null', stderr='/dev/null', walltime="00:00:30")
        self.submit(self._holding_job)
        self.global_depends.append(self._holding_job.jobid)

    def check_jobid(self, jobid):
        with open('/dev/null', 'w') as devnull:
            proc = subprocess.Popen(["sacct", "-b", "-p", "-j", jobid], stdout=subprocess.PIPE, stderr=devnull)
            stdout, stderr = proc.communicate()
            for line in stdout.split('\n'):
                cols = line.strip().split('|')
                if cols[0] == jobid:
                    if cols[-1].split(':')[0] == '0':
                        return True
            return False

    def release(self, jobid):
        subprocess.call(["scontrol", "release", jobid])

    def cancel(self, jobid):
        subprocess.call(["scancel", jobid])

    def submit(self, job):
        if not job.src:
            return

        if self.global_hold and not self._holding_job:
            self._setup_holding_job()

        body = ''

        if job.pre:
            body = job.pre
            body += '\n'

        body += job.src

        if job.post:
            body += '\n'
            body += job.post

        jobopts = dict(def_options)
        for k in job.args:
            jobopts[k] = job.args[k]

        src = '#!%s\n' % self.interpreter
        src += '#SBATCH -J %s\n' % (job.name if job.name[0] in string.ascii_letters else 'mvp_%s' % job.name)

        if 'hold' in jobopts and jobopts['hold']:
            src += '#SBATCH -H\n'

        if 'env' in jobopts and jobopts['env']:
            src += '#SBATCH --export=ALL\n'

        if 'walltime' in jobopts:
            src += '#SBATCH -t %s\n' % mvpipe.support.calc_time(jobopts['walltime'])

        if 'procs' in jobopts and int(jobopts['procs']) > 1:
            src += '#SBATCH -n %s\n' % (jobopts['procs'])

        if 'nodes' in jobopts and int(jobopts['nodes']) > 1:
            src += '#SBATCH -N %s\n' % (jobopts['nodes'])

        if 'mem' in jobopts:
            if jobopts['mem'][-1] == 'M':
                mem = jobopts['mem'][:-1]
            elif jobopts['mem'][-1] == 'G':
                mem = int(jobopts['mem'][:-1]) * 1000
            else:
                mem = jobopts['mem']

            src += '#SBATCH --mem=%s\n' % mem

        # if 'stack' in jobopts:
        #     src += '#$ -l h_stack=%s\n' % jobopts['stack']

        if job.depids or self.global_depends:
            depids = job.depids
            if self.global_depends:
                depids.extend(self.global_depends)

            if depids:
                src += '#SBATCH -d afterok:%s\n' % ':'.join(depids)

        if 'qos' in jobopts:
            src += '#SBATCH --qos %s\n' % jobopts['qos']

        # if 'queue' in jobopts:
        #     src += '#$ -q %s\n' % jobopts['queue']

        if 'mail' in jobopts:
            src += '#SBATCH --mail-type %s\n' % jobopts['mail']

        if 'wd' in jobopts:
            src += '#SBATCH -D %s\n' % jobopts['wd']

        if 'account' in jobopts:
            src += '#SBATCH -A %s\n' % jobopts['account']
        elif self.account:
            src += '#SBATCH -A %s\n' % self.account

        if 'stdout' in jobopts:
            src += '#SBATCH -o %s\n' % jobopts['stdout']

        if 'stderr' in jobopts:
            src += '#SBATCH -e %s\n' % jobopts['stderr']

        src += 'set -o pipefail\nfunc () {\n  %s\n  return $?\n}\n' % body

        src += 'func\n'
        src += 'RETVAL=$?\n'
        src += 'exit $RETVAL\n'
<commit_msg>updates to slurm runner to make it call srun within the sbatch script. Also switched to use -c as opposed to -n to set the # of CPUs. We basically call srun $0, but only execute the job if a flag is set. If InactiveLimit is set to 0 in the slurm.conf file, then this doesn't matter. However, it doesn't hurt to call things this way.<commit_after>import os
import string
import subprocess

import mvpipe.support
from mvpipe.runner import Runner, Job

def_options = {'env': True, 'wd': os.path.abspath(os.curdir), 'hold': False, 'nodes': 1}

'''
SLURM options:

global_hold - start all pipelines with a "placeholder" job to synchronize
              starts (and ensure that all jobs are submitted correctly
              before releasing the pipeline

account     - a default account to use

shell       - a shell to use for the script (default(s): /bin/bash, /usr/bin/bash, /usr/local/bin/bash, /bin/sh)

'''
class SlurmRunner(Runner):
    def __init__(self, dryrun, verbose, logger, global_hold=False, global_depends=None, account=None, interpreter=None):
        Runner.__init__(self, dryrun, verbose, logger)
        self.global_hold = global_hold
        self._holding_job = None
        self.global_depends = global_depends if global_depends else []
        self.account = account

        self.jobids = []

        self.testjobcount = 1

        if interpreter:
            self.interpreter = interpreter
        else:
            for intp in ['/bin/bash', '/usr/bin/bash', '/usr/local/bin/bash']:
                if os.path.exists(intp):
                    self.interpreter = intp
                    break
            if not self.interpreter:
                self.interpreter = '/bin/sh'


    def reset(self):
        pass

    def abort(self):
        if self._holding_job:
            self.cancel(self._holding_job.jobid)
        else:
            for jobid in self.jobids:
                self.cancel(jobid)

    def done(self):
        if not self.dryrun:
            if self._holding_job:
                self.release(self._holding_job.jobid)

    def _setup_holding_job(self):
        self._holding_job = Job('sleep 5', name="holding", hold=True, stdout='/dev/null', stderr='/dev/null', walltime="00:00:30")
        self.submit(self._holding_job)
        self.global_depends.append(self._holding_job.jobid)

    def check_jobid(self, jobid):
        with open('/dev/null', 'w') as devnull:
            proc = subprocess.Popen(["sacct", "-b", "-p", "-j", jobid], stdout=subprocess.PIPE, stderr=devnull)
            stdout, stderr = proc.communicate()
            for line in stdout.split('\n'):
                cols = line.strip().split('|')
                if cols[0] == jobid:
                    if cols[-1].split(':')[0] == '0':
                        return True
            return False

    def release(self, jobid):
        subprocess.call(["scontrol", "release", jobid])

    def cancel(self, jobid):
        subprocess.call(["scancel", jobid])

    def submit(self, job):
        if not job.src:
            return

        if self.global_hold and not self._holding_job:
            self._setup_holding_job()

        body = ''

        if job.pre:
            body = job.pre
            body += '\n'

        body += job.src

        if job.post:
            body += '\n'
            body += job.post

        jobopts = dict(def_options)
        for k in job.args:
            jobopts[k] = job.args[k]

        src = '#!%s\n' % self.interpreter
        src += '#SBATCH -J %s\n' % (job.name if job.name[0] in string.ascii_letters else 'mvp_%s' % job.name)

        if 'hold' in jobopts and jobopts['hold']:
            src += '#SBATCH -H\n'

        if 'env' in jobopts and jobopts['env']:
            src += '#SBATCH --export=ALL\n'

        if 'walltime' in jobopts:
            src += '#SBATCH -t %s\n' % mvpipe.support.calc_time(jobopts['walltime'])

        if 'procs' in jobopts and int(jobopts['procs']) > 1:
            src += '#SBATCH -c %s\n' % (jobopts['procs'])

        if 'nodes' in jobopts and int(jobopts['nodes']) > 1:
            src += '#SBATCH -N %s\n' % (jobopts['nodes'])

        if 'mem' in jobopts:
            if jobopts['mem'][-1] == 'M':
                mem = jobopts['mem'][:-1]
            elif jobopts['mem'][-1] == 'G':
                mem = int(jobopts['mem'][:-1]) * 1000
            else:
                mem = jobopts['mem']

            src += '#SBATCH --mem=%s\n' % mem

        # if 'stack' in jobopts:
        #     src += '#$ -l h_stack=%s\n' % jobopts['stack']

        if job.depids or self.global_depends:
            depids = job.depids
            if self.global_depends:
                depids.extend(self.global_depends)

            if depids:
                src += '#SBATCH -d afterok:%s\n' % ':'.join(depids)

        if 'qos' in jobopts:
            src += '#SBATCH --qos %s\n' % jobopts['qos']

        # if 'queue' in jobopts:
        #     src += '#$ -q %s\n' % jobopts['queue']

        if 'mail' in jobopts:
            src += '#SBATCH --mail-type %s\n' % jobopts['mail']

        if 'wd' in jobopts:
            src += '#SBATCH -D %s\n' % jobopts['wd']

        if 'account' in jobopts:
            src += '#SBATCH -A %s\n' % jobopts['account']
        elif self.account:
            src += '#SBATCH -A %s\n' % self.account

        if 'stdout' in jobopts:
            src += '#SBATCH -o %s\n' % jobopts['stdout']

        if 'stderr' in jobopts:
            src += '#SBATCH -e %s\n' % jobopts['stderr']

        src += 'set -o pipefail\nfunc () {\n  %s\n  return $?\n}\n' % body

        src += 'if [ "$0" == "" ]; then\n'
        src += '    srun $0 run\n'
        src += 'else\n'
        src += '    func\n'
        src += '    RETVAL=$?\n'
        src += '    exit $RETVAL\n'
        src += 'fi\n'

{'fraction_non_alphanumeric': 0.14464333903031393, 'fraction_numerical': 0.004677520913915625, 'mean_word_length': 1.8047426841574168, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5794419', 'n_tokens_mistral': 3698, 'n_tokens_neox': 3518, 'n_words': 934}