<commit_before>'''
Created on Aug 2, 2013

@author: chadcumba

Parallel workflow execution with SLURM
'''
from __future__ import print_function, division, unicode_literals, absolute_import
from builtins import open

import os
import re
from time import sleep

from ...interfaces.base import CommandLine
from .base import (SGELikeBatchManagerBase, logger, iflogger, logging)

import subprocess


class SLURMPlugin(SGELikeBatchManagerBase):
    '''
    Execute using SLURM

    The plugin_args input to run can be used to control the SLURM execution.
    Currently supported options are:

    - template : template to use for batch job submission

    - sbatch_args: arguments to pass prepend to the sbatch call


    '''

    def __init__(self, **kwargs):

        template = "#!/bin/bash"

        self._retry_timeout = 2
        self._max_tries = 2
        self._template = template
        self._sbatch_args = None
        self._jobid_re = "Submitted batch job ([0-9]*)"

        if 'plugin_args' in kwargs and kwargs['plugin_args']:
            if 'retry_timeout' in kwargs['plugin_args']:
                self._retry_timeout = kwargs['plugin_args']['retry_timeout']
            if 'max_tries' in kwargs['plugin_args']:
                self._max_tries = kwargs['plugin_args']['max_tries']
            if 'jobid_re' in kwargs['plugin_args']:
                self._jobid_re = kwargs['plugin_args']['jobid_re']
            if 'template' in kwargs['plugin_args']:
                self._template = kwargs['plugin_args']['template']
                if os.path.isfile(self._template):
                    with open(self._template) as f:
                        self._template = f.read()
            if 'sbatch_args' in kwargs['plugin_args']:
                self._sbatch_args = kwargs['plugin_args']['sbatch_args']
        self._pending = {}
        super(SLURMPlugin, self).__init__(self._template, **kwargs)

    def _is_pending(self, taskid):
        #  subprocess.Popen requires taskid to be a string
        proc = subprocess.Popen(["squeue", '-j', '%s' % taskid],
                                stdout=subprocess.PIPE,
                                stderr=subprocess.PIPE)
        o, _ = proc.communicate()
        return o.find(str(taskid)) > -1

    def _submit_batchtask(self, scriptfile, node):
        """
        This is more or less the _submit_batchtask from sge.py with flipped
        variable names, different command line switches, and different output
        formatting/processing<commit_msg>fix: remove temp slurm fix<commit_after>'''
Created on Aug 2, 2013

@author: chadcumba

Parallel workflow execution with SLURM
'''
from __future__ import print_function, division, unicode_literals, absolute_import
from builtins import open

import os
import re
from time import sleep

from ...interfaces.base import CommandLine
from .base import (SGELikeBatchManagerBase, logger, iflogger, logging)



class SLURMPlugin(SGELikeBatchManagerBase):
    '''
    Execute using SLURM

    The plugin_args input to run can be used to control the SLURM execution.
    Currently supported options are:

    - template : template to use for batch job submission

    - sbatch_args: arguments to pass prepend to the sbatch call


    '''

    def __init__(self, **kwargs):

        template = "#!/bin/bash"

        self._retry_timeout = 2
        self._max_tries = 2
        self._template = template
        self._sbatch_args = None
        self._jobid_re = "Submitted batch job ([0-9]*)"

        if 'plugin_args' in kwargs and kwargs['plugin_args']:
            if 'retry_timeout' in kwargs['plugin_args']:
                self._retry_timeout = kwargs['plugin_args']['retry_timeout']
            if 'max_tries' in kwargs['plugin_args']:
                self._max_tries = kwargs['plugin_args']['max_tries']
            if 'jobid_re' in kwargs['plugin_args']:
                self._jobid_re = kwargs['plugin_args']['jobid_re']
            if 'template' in kwargs['plugin_args']:
                self._template = kwargs['plugin_args']['template']
                if os.path.isfile(self._template):
                    with open(self._template) as f:
                        self._template = f.read()
            if 'sbatch_args' in kwargs['plugin_args']:
                self._sbatch_args = kwargs['plugin_args']['sbatch_args']
        self._pending = {}
        super(SLURMPlugin, self).__init__(self._template, **kwargs)

    def _is_pending(self, taskid):
        #  subprocess.Popen requires taskid to be a string
        res = CommandLine('squeue',
                          args=' '.join(['-j', '%s' % taskid]),
                          terminal_output='allatonce').run()
        return res.runtime.stdout.find(str(taskid)) > -1

    def _submit_batchtask(self, scriptfile, node):
        """
        This is more or less the _submit_batchtask from sge.py with flipped
        variable names, different command line switches, and different output
        formatting/processing
{'fraction_non_alphanumeric': 0.09911717495987159, 'fraction_numerical': 0.0040128410914927765, 'mean_word_length': 2.5031623330990866, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 5, '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': '5264571', 'n_tokens_mistral': 1445, 'n_tokens_neox': 1398, 'n_words': 393}