<commit_before>import luigi
import parameter
import time

# ================================================================================

# A few 'constants'
RUNMODE_LOCAL = 'runmode_local'
RUNMODE_HPC = 'runmode_hpc'
RUNMODE_MPI = 'runmode_mpi'

# ================================================================================

class SlurmInfo():
    runmode = None # One of RUNMODE_LOCAL|RUNMODE_HPC|RUNMODE_MPI
    project = None
    partition = None
    cores = None
    time = None
    jobname = None
    threads = None

    def __init__(self, runmode, project, partition, cores, time, jobname, threads):
        self.runmode = runmode
        self.project = project
        self.partition = partition
        self.cores = cores
        self.time = time
        self.jobname = jobname
        self.threads = threads

    def get_argstr_hpc(self):
        '''
        Return a formatted string with arguments and option flags to SLURM
        commands such as salloc and sbatch, for non-MPI, HPC jobs.
        '''
        #FIXME: Double-check the format of this one!
        argstr = ' -A {pr} -p {pt} -n {c} -t {t} -J {j} '.format(
                pr = self.project,
                pt = self.partition,
                c = self.cores,
                t = self.time,
                j = self.jobname,
                thr = self.threads)
        return argstr

    def get_argstr_mpi(self):
        '''
        Return a formatted string with arguments and option flags to SLURM
        commands such as salloc and sbatch, for MPI jobs.
        '''
        argstr = ' -A {pr} -p {pt} -n {c} -t {t} -J {j} srun -n 1 -c {thr} '.format(
                pr = self.project,
                pt = self.partition,
                c = self.cores,
                t = self.time,
                j = self.jobname,
                thr = self.threads)
        return argstr


# ================================================================================

class SlurmInfoParameter(parameter.Parameter):
    def parse(self, x):
        # TODO: Possibly do something more fancy here
        return x

# ================================================================================

class SlurmHelpers():
    '''
    Mixin with various convenience methods for executing jobs via SLURM
    '''
    # Other class-fields
    slurminfo = SlurmInfoParameter(default=None) # Class: SlurmInfo

    # Main Execution methods
    def ex(self, command):
        '''
        Execute either locally or via SLURM, depending on config
        '''
        if self.slurminfo.runmode == RUNMODE_LOCAL:
            self.ex_local(command) # Defined in task.py
        elif self.slurminfo.runmode == RUNMODE_HPC:
            self.ex_hpc(command)
        elif self.slurminfo.runmode == RUNMODE_MPI:
            self.ex_mpi(command)


    def ex_hpc(self, command):
        if isinstance(command, list):
            command = ' '.join(command)

        (status, output) = self.ex_local(slurm_part + command)

        # TODO: Do this only if audit logging is activated!
        self.log_slurm_info(output)
        return (status, output)


    def ex_mpi(self, command):
        if isinstance(command, list):
            command = ' '.join(command)

        fullcommand = 'salloc %s %s' % (self.slurminfo.get_argstr_mpi(), command)

        (status, output) = self.ex_local(fullcommand)

        # TODO: Do this only if audit logging is activated!
        self.log_slurm_info(output)
        return (status, output)


    # Various convenience methods

    def assert_matches_character_class(self, char_class, a_string):
        if not bool(re.match('^{c}+$'.format(c=char_class), a_string)):
            raise Exception('String {s} does not match character class {cc}'.format(s=a_string, cc=char_class))
<commit_msg>Bugfixing HPC and MPI salloc strings<commit_after>import luigi
import parameter
import time

# ================================================================================

# A few 'constants'
RUNMODE_LOCAL = 'runmode_local'
RUNMODE_HPC = 'runmode_hpc'
RUNMODE_MPI = 'runmode_mpi'

# ================================================================================

class SlurmInfo():
    runmode = None # One of RUNMODE_LOCAL|RUNMODE_HPC|RUNMODE_MPI
    project = None
    partition = None
    cores = None
    time = None
    jobname = None
    threads = None

    def __init__(self, runmode, project, partition, cores, time, jobname, threads):
        self.runmode = runmode
        self.project = project
        self.partition = partition
        self.cores = cores
        self.time = time
        self.jobname = jobname
        self.threads = threads

    def get_argstr_hpc(self):
        '''
        Return a formatted string with arguments and option flags to SLURM
        commands such as salloc and sbatch, for non-MPI, HPC jobs.
        '''
        argstr = ' -A {pr} -p {pt} -n {c} -t {t} -J {j} srun -n 1 -c {thr} '.format(
                pr = self.project,
                pt = self.partition,
                c = self.cores,
                t = self.time,
                j = self.jobname,
                thr = self.threads)
        return argstr

    def get_argstr_mpi(self):
        '''
        Return a formatted string with arguments and option flags to SLURM
        commands such as salloc and sbatch, for MPI jobs.
        '''
        argstr = ' -A {pr} -p {pt} -n {c} -t {t} -J {j} mpirun -v -np {c} '.format(
                pr = self.project,
                pt = self.partition,
                c = self.cores,
                t = self.time,
                j = self.jobname)
        return argstr


# ================================================================================

class SlurmInfoParameter(parameter.Parameter):
    def parse(self, x):
        # TODO: Possibly do something more fancy here
        return x

# ================================================================================

class SlurmHelpers():
    '''
    Mixin with various convenience methods for executing jobs via SLURM
    '''
    # Other class-fields
    slurminfo = SlurmInfoParameter(default=None) # Class: SlurmInfo

    # Main Execution methods
    def ex(self, command):
        '''
        Execute either locally or via SLURM, depending on config
        '''
        if self.slurminfo.runmode == RUNMODE_LOCAL:
            self.ex_local(command) # Defined in task.py
        elif self.slurminfo.runmode == RUNMODE_HPC:
            self.ex_hpc(command)
        elif self.slurminfo.runmode == RUNMODE_MPI:
            self.ex_mpi(command)


    def ex_hpc(self, command):
        if isinstance(command, list):
            command = ' '.join(command)

        fullcommand = 'salloc %s %s' % (self.slurminfo.get_argstr_hpc(), command)
        (status, output) = self.ex_local(fullcommand)

        # TODO: Do this only if audit logging is activated!
        self.log_slurm_info(output)
        return (status, output)


    def ex_mpi(self, command):
        if isinstance(command, list):
            command = ' '.join(command)

        fullcommand = 'salloc %s %s' % (self.slurminfo.get_argstr_mpi(), command)
        (status, output) = self.ex_local(fullcommand)

        # TODO: Do this only if audit logging is activated!
        self.log_slurm_info(output)
        return (status, output)


    # Various convenience methods

    def assert_matches_character_class(self, char_class, a_string):
        if not bool(re.match('^{c}+$'.format(c=char_class), a_string)):
            raise Exception('String {s} does not match character class {cc}'.format(s=a_string, cc=char_class))

{'fraction_non_alphanumeric': 0.17490144546649145, 'fraction_numerical': 0.00026281208935611036, 'mean_word_length': 2.3721754541426674, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, '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': '303594', 'n_tokens_mistral': 2247, 'n_tokens_neox': 2126, 'n_words': 613}