<commit_before>"""
Submit jobs to slurm or torque, or with multiprocessing.

============================================================================

        AUTHOR: <NAME>, <EMAIL>
  ORGANIZATION: Stanford University
       LICENSE: MIT License, property of Stanford, use as you wish
       CREATED: 2016-44-20 23:03
 Last modified: 2016-03-23 00:34

   DESCRIPTION: Allows simple job submission with either torque, slurm, or
                with the multiprocessing module.
                To set the environement, set QUEUE to one of ['torque',
                'slurm', 'normal'], or run get_cluster_environment().
                To submit a job, run submit().

                All jobs write out a job file before submission, even though
                this is not necessary (or useful) with multiprocessing. In
                normal mode, this is a .cluster file, in slurm is is a
                .cluster.sbatch and a .cluster.script file, in torque it is a
                .cluster.qsub file.

                The name argument is required for submit, it is used to
                generate the STDOUT and STDERR files. Irrespective of mode
                the STDOUT file will be name.cluster.out and the STDERR file
                will be name.cluster.err.

                Note: `.cluster` is added to all names to make deletion less
                dangerous

                Dependency tracking is supported in torque or slurm mode,
                to use it pass a list of job ids to submit or submit_file with
                the `dependencies` keyword argument.

                To clean up cluster files, run clean(directory), if directory
                is not provided, the current directory is used.
                This will delete all files in that were generated by this
                script.

       CAUTION: The clean() function will delete **EVERY** file with
                extensions matching those in this file::
                    .cluster.err
                    .cluster.out
                    .cluster.sbatch & .cluster.script for slurm mode
                    .cluster.qsub for torque mode
                    .cluster for normal mode

============================================================================
"""
import os
import re
from time import sleep
from textwrap import dedent
from subprocess import check_output, CalledProcessError
from multiprocessing import Pool, pool

# Us
from ASEr import run
from ASEr import logme

#########################
#  Which system to use  #
#########################

# Default is normal, change to 'slurm' or 'torque' as needed.
QUEUE          = 'normal'
ALLOWED_QUEUES = ['torque', 'slurm', 'normal']

#########################################################
#  The multiprocessing pool, only used in 'local' mode  #
#########################################################

POOL = None


def get_cluster_environment():
    """Detect the local cluster environment and set QUEUE globally.

    Uses which to search for sbatch first, then qsub. If neither is found,
    QUEUE is set to local.

    :returns: QUEUE variable ('torque', 'slurm', or 'local')
    """
    global QUEUE
    if run.which('sbatch'):
        QUEUE = 'slurm'
    elif run.which('qsub'):
        QUEUE = 'torque'
    else:
        QUEUE = 'local'
    if QUEUE == 'slurm' or QUEUE == 'torque':
        logme.log('{} detected, using for cluster submissions'.format(QUEUE),
                  'debug')
    else:
        logme.log('No cluster environment detected, using multiprocessing',
                  'debug')
    return QUEUE


#####################################
#  Wait for cluster jobs to finish  #
#####################################


def wait(jobs):
    """Wait for jobs to finish.

    :jobs:    A single job or list of jobs to wait for. With torque or slurm,
              these should be job IDs, with normal mode, these are
              multiprocessing job objects (returned by submit())
    """
    check_queue()  # Make sure the QUEUE is usable

    # Sanitize argument
    if not isinstance(jobs, (list, tuple)):
        jobs = [jobs]
    for job in jobs:
        if not isinstance(job, (str, int, pool.ApplyResult)):
            raise ClusterError('job must be int, string, or ApplyResult, ' +
                                'is {}'.format(type(job)))

    if QUEUE == 'normal':
        for job in jobs:
            if not isinstance(job, pool.ApplyResult):
                raise ClusterError('jobs must be ApplyResult objects')
            job.wait()
    elif QUEUE == 'torque':
        # Wait for 5 seconds before checking, as jobs take a while to be queued
        # sometimes
        sleep(5)

        s = re.compile(r'  +')  # For splitting qstat output
        # Jobs must be strings for comparison operations
        jobs = [str(j) for j in jobs]
        while True:
            c = 0
            try:
                q = check_output(['qstat', '-a']).decode().rstrip().split('\n')
            except CalledProcessError:
                if c == 5:
                    raise
                c += 1
                sleep(2)
                continue
            # Check header
            if not re.split(r' {2,100}', q[3])[9]== 'S':
                raise ClusterError('Unrecognized torque qstat format')
            # Build a list of completed jobs
            complete = []
            for j in q[5:]:
                i = s.split(j)
                if i[9] == 'C':
                    complete.append(i[0].split('.')[0])
            # Build a list of all jobs
            all  = [s.split(i)[0].split('.')[0] for i in q[5:]]
            # Trim down job list
            jobs = [i for i in jobs if i in all]
            jobs = [i for i in jobs if i not in complete]
            if len(jobs) == 0:
                return
            sleep(2)
    elif QUEUE == 'slurm':
        # Wait for 5 seconds before checking, as jobs take a while to be queued
        # sometimes
        sleep(5)

        # Jobs must be strings for comparison operations<commit_msg>Change wait to 2 seconds for slurm<commit_after>"""
Submit jobs to slurm or torque, or with multiprocessing.

============================================================================

        AUTHOR: <NAME>, <EMAIL>
  ORGANIZATION: Stanford University
       LICENSE: MIT License, property of Stanford, use as you wish
       CREATED: 2016-44-20 23:03
 Last modified: 2016-03-23 00:50

   DESCRIPTION: Allows simple job submission with either torque, slurm, or
                with the multiprocessing module.
                To set the environement, set QUEUE to one of ['torque',
                'slurm', 'normal'], or run get_cluster_environment().
                To submit a job, run submit().

                All jobs write out a job file before submission, even though
                this is not necessary (or useful) with multiprocessing. In
                normal mode, this is a .cluster file, in slurm is is a
                .cluster.sbatch and a .cluster.script file, in torque it is a
                .cluster.qsub file.

                The name argument is required for submit, it is used to
                generate the STDOUT and STDERR files. Irrespective of mode
                the STDOUT file will be name.cluster.out and the STDERR file
                will be name.cluster.err.

                Note: `.cluster` is added to all names to make deletion less
                dangerous

                Dependency tracking is supported in torque or slurm mode,
                to use it pass a list of job ids to submit or submit_file with
                the `dependencies` keyword argument.

                To clean up cluster files, run clean(directory), if directory
                is not provided, the current directory is used.
                This will delete all files in that were generated by this
                script.

       CAUTION: The clean() function will delete **EVERY** file with
                extensions matching those in this file::
                    .cluster.err
                    .cluster.out
                    .cluster.sbatch & .cluster.script for slurm mode
                    .cluster.qsub for torque mode
                    .cluster for normal mode

============================================================================
"""
import os
import re
from time import sleep
from textwrap import dedent
from subprocess import check_output, CalledProcessError
from multiprocessing import Pool, pool

# Us
from ASEr import run
from ASEr import logme

#########################
#  Which system to use  #
#########################

# Default is normal, change to 'slurm' or 'torque' as needed.
QUEUE          = 'normal'
ALLOWED_QUEUES = ['torque', 'slurm', 'normal']

#########################################################
#  The multiprocessing pool, only used in 'local' mode  #
#########################################################

POOL = None


def get_cluster_environment():
    """Detect the local cluster environment and set QUEUE globally.

    Uses which to search for sbatch first, then qsub. If neither is found,
    QUEUE is set to local.

    :returns: QUEUE variable ('torque', 'slurm', or 'local')
    """
    global QUEUE
    if run.which('sbatch'):
        QUEUE = 'slurm'
    elif run.which('qsub'):
        QUEUE = 'torque'
    else:
        QUEUE = 'local'
    if QUEUE == 'slurm' or QUEUE == 'torque':
        logme.log('{} detected, using for cluster submissions'.format(QUEUE),
                  'debug')
    else:
        logme.log('No cluster environment detected, using multiprocessing',
                  'debug')
    return QUEUE


#####################################
#  Wait for cluster jobs to finish  #
#####################################


def wait(jobs):
    """Wait for jobs to finish.

    :jobs:    A single job or list of jobs to wait for. With torque or slurm,
              these should be job IDs, with normal mode, these are
              multiprocessing job objects (returned by submit())
    """
    check_queue()  # Make sure the QUEUE is usable

    # Sanitize argument
    if not isinstance(jobs, (list, tuple)):
        jobs = [jobs]
    for job in jobs:
        if not isinstance(job, (str, int, pool.ApplyResult)):
            raise ClusterError('job must be int, string, or ApplyResult, ' +
                                'is {}'.format(type(job)))

    if QUEUE == 'normal':
        for job in jobs:
            if not isinstance(job, pool.ApplyResult):
                raise ClusterError('jobs must be ApplyResult objects')
            job.wait()
    elif QUEUE == 'torque':
        # Wait for 5 seconds before checking, as jobs take a while to be queued
        # sometimes
        sleep(5)

        s = re.compile(r'  +')  # For splitting qstat output
        # Jobs must be strings for comparison operations
        jobs = [str(j) for j in jobs]
        while True:
            c = 0
            try:
                q = check_output(['qstat', '-a']).decode().rstrip().split('\n')
            except CalledProcessError:
                if c == 5:
                    raise
                c += 1
                sleep(2)
                continue
            # Check header
            if not re.split(r' {2,100}', q[3])[9]== 'S':
                raise ClusterError('Unrecognized torque qstat format')
            # Build a list of completed jobs
            complete = []
            for j in q[5:]:
                i = s.split(j)
                if i[9] == 'C':
                    complete.append(i[0].split('.')[0])
            # Build a list of all jobs
            all  = [s.split(i)[0].split('.')[0] for i in q[5:]]
            # Trim down job list
            jobs = [i for i in jobs if i in all]
            jobs = [i for i in jobs if i not in complete]
            if len(jobs) == 0:
                return
            sleep(2)
    elif QUEUE == 'slurm':
        # Wait for 2 seconds before checking, as jobs take a while to be queued
        # sometimes
        sleep(2)

        # Jobs must be strings for comparison operations
{'fraction_non_alphanumeric': 0.1406288685318148, 'fraction_numerical': 0.007840224478006107, 'mean_word_length': 1.9348510535238557, 'pattern_counts': {'":': 0, '<': 7, '<?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': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3547920', 'n_tokens_mistral': 3279, 'n_tokens_neox': 3020, 'n_words': 1261}