<commit_before>import yaml
import tempfile
import os, re
import logging
import numpy as np
import subprocess
import time
import multiprocessing

from .batch import write_slurm_script, get_job_status

class PipelineStage(object):
    name = None
    id_str_fmt = None
    batch_compatible = True
    single_filter = False
    batch_type = 'slurm'

    _default_kwargs = {}

    def __init__(self, pipeline):
        self.pipeline = pipeline

    @property
    def _id_options(self):
        raise NotImplementedError("Need to define what id keys are valid.")

    @property
    def default_kwargs(self):
        return self._default_kwargs

    def id_str(self, filt=None):
        s = ''
        for key in self._id_options:
            try:
                if key=='visit':
                    s += 'visit={0} '.format(self.pipeline['visit'][filt])
                else:
                    fmt = '{0}={{0[{0}]}} '.format(key)
                    s += fmt.format(self.pipeline)
            except KeyError:
                continue

        if filt is not None:
            s += 'filter={0} '.format(filt)
        return s

    def jobname(self, filt=None):
        s = '{0[ticket]}-{0[field]}-{1}'.format(self.pipeline, self.name)
        if filt is not None:
            s += '-{}'.format(filt)
        return s

    def cmd_str(self, filt=None, test=False, **kwargs):
        cmd = '{0}.py {1[data_root]} --rerun {1.rerun} '.format(self.name, self.pipeline)

        if not isinstance(self, ManualBatchStage):
            cmd += '--job {0} '.format(self.jobname(filt))

        cmd += '--id {0} '.format(self.id_str(filt))

        if hasattr(self, 'selectId_str'):
            cmd += '--selectId {0} '.format(self.selectId_str(filt))

        kws = dict(self.default_kwargs)
        kws.update(self.pipeline['kwargs']['all'])
        kws.update(self.pipeline['kwargs'][self.name])
        kws.update(kwargs)
        for kw, val in kws.items():
            if kw == 'config':
                # val should be a dict
                for k, v in val.items():
                    cmd += '--config {0}={1} '.format(k,v)

            if val is False:
                continue

            # skip using "time" keyword if writing batch script manually
            if isinstance(self, ManualBatchStage) and kw == 'time':
                continue

            cmd += '--{0}'.format(kw)
            if val != '' and val is not True:
                # ensure quotation marks where necessary
                if isinstance(val, basestring):
                    val = '"{0}"'.format(val)

                cmd += '={0} '.format(val)
                if kw == 'nodes':
                    cmd += '--procs {0[cores_per_node]} '.format(self.pipeline)<commit_msg>fix config option<commit_after>import yaml
import tempfile
import os, re
import logging
import numpy as np
import subprocess
import time
import multiprocessing

from .batch import write_slurm_script, get_job_status

class PipelineStage(object):
    name = None
    id_str_fmt = None
    batch_compatible = True
    single_filter = False
    batch_type = 'slurm'

    _default_kwargs = {}

    def __init__(self, pipeline):
        self.pipeline = pipeline

    @property
    def _id_options(self):
        raise NotImplementedError("Need to define what id keys are valid.")

    @property
    def default_kwargs(self):
        return self._default_kwargs

    def id_str(self, filt=None):
        s = ''
        for key in self._id_options:
            try:
                if key=='visit':
                    s += 'visit={0} '.format(self.pipeline['visit'][filt])
                else:
                    fmt = '{0}={{0[{0}]}} '.format(key)
                    s += fmt.format(self.pipeline)
            except KeyError:
                continue

        if filt is not None:
            s += 'filter={0} '.format(filt)
        return s

    def jobname(self, filt=None):
        s = '{0[ticket]}-{0[field]}-{1}'.format(self.pipeline, self.name)
        if filt is not None:
            s += '-{}'.format(filt)
        return s

    def cmd_str(self, filt=None, test=False, **kwargs):
        cmd = '{0}.py {1[data_root]} --rerun {1.rerun} '.format(self.name, self.pipeline)

        if not isinstance(self, ManualBatchStage):
            cmd += '--job {0} '.format(self.jobname(filt))

        cmd += '--id {0} '.format(self.id_str(filt))

        if hasattr(self, 'selectId_str'):
            cmd += '--selectId {0} '.format(self.selectId_str(filt))

        kws = dict(self.default_kwargs)
        kws.update(self.pipeline['kwargs']['all'])
        kws.update(self.pipeline['kwargs'][self.name])
        kws.update(kwargs)
        for kw, val in kws.items():
            if kw == 'config':
                # val should be a dict
                for k, v in val.items():
                    cmd += '--config {0}={1} '.format(k,v)
                continue

            if val is False:
                continue

            # skip using "time" keyword if writing batch script manually
            if isinstance(self, ManualBatchStage) and kw == 'time':
                continue

            cmd += '--{0}'.format(kw)
            if val != '' and val is not True:
                # ensure quotation marks where necessary
                if isinstance(val, basestring):
                    val = '"{0}"'.format(val)

                cmd += '={0} '.format(val)
                if kw == 'nodes':
                    cmd += '--procs {0[cores_per_node]} '.format(self.pipeline)
{'fraction_non_alphanumeric': 0.1307734106140192, 'fraction_numerical': 0.007245064299945662, 'mean_word_length': 1.8116089613034623, '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': '3713456', 'n_tokens_mistral': 1673, 'n_tokens_neox': 1612, 'n_words': 432}