<commit_before>import errno
import logging
import os
import Queue
import select
import signal
import socket
import sys
import time
from multiprocessing import queues

import redis
import setproctitle

from .config import Config
from .worker import Worker
from .job import Job, job_iterator
from .redis_extensions import extend_redis


class Master(object):

    signal_map = {
        "hup": "reload_config",
        "usr1": "relaunch",
        "ttin": "expand_workers",
        "ttou": "contract_workers",
        "int": "halt_current_jobs",
        "quit": "wind_down_gracefully",
        "term": "wind_down_immediately",
        "chld": "handle_worker_exit"
    }

    def __init__(self, config_file):
        self.pid = None
        self.pid_file_path = None
        self.reexec_pid = 0

        self.name = "master"

        self.logger = logging.getLogger(__name__)

        self.config_file = config_file
        self.load_config()

        self.workers = {}

        self.socket = None

        self.ready_queue = queues.Queue(maxsize=5)
        self.taken_queue = queues.Queue()
        self.results_queue = queues.Queue()

        self.redis = None

        self.wind_down_time = None

        args = sys.argv[:]
        args.insert(0, sys.executable)

        self.launch_context = {
            "executable": sys.executable,
            "args": args,
            "cwd": os.getcwd()
        }

    def load_config(self):
        self.config = Config(self.config_file)
        self.config.load()
        self.pid_file_path = self.config.master.pid_file
        self.number_of_workers = self.config.master.num_workers

    def setup_pid_file(self):
        if not os.path.isdir(os.path.dirname(self.pid_file_path)):
            raise RuntimeError(
                "%s does not exist! Can't create pid file" % self.pid_file_path
            )

        if os.path.exists(self.pid_file_path):
            raise RuntimeError(
                (
                    "pid file (%s) already exists! \n"
                    "If no master is running it could be stale."
                ) % self.pid_file_path
            )

        fd = open(self.pid_file_path, 'w')
        try:
            fd.write(u"%s\n" % self.pid)
        except:
            os.unlink(self.pid_file_path)
            raise
        else:
            fd.close()

    def setup_datastore(self):
        if ":" in self.config.master.redis:
            host, port = self.config.master.redis.split(":")
            self.redis = redis.Redis(host=host, port=port)
        else:
            self.redis = redis.Redis(host=self.config.master.redis)

        extend_redis(self.redis)

    def setup_socket(self):
        self.logger.info(
            "Listening on port %s", self.config.master.listen_port
        )

        self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.socket.setblocking(0)
        self.socket.bind(
            ('', int(self.config.master.listen_port))
        )
        self.socket.listen(5)

    def setup_signals(self):
        for signal_name, handler_name in self.signal_map.iteritems():
            signal.signal(
                getattr(signal, "SIG%s" % signal_name.upper()),
                getattr(self, handler_name)
            )

    def run(self):
        self.pid = os.getpid()
        self.logger.info("Starting %s (%d)", self.name, int(self.pid))

        self.setup_pid_file()

        self.setup_datastore()

        self.fill_ready_queue()

        self.setup_socket()
        self.setup_signals()

        setproctitle.setproctitle("distq: %s" % self.name)

        for i in range(self.number_of_workers):
            self.spawn_worker()

        while True:
            try:
                (input_sources, [], []) = select.select(
                    [
                        self.socket,
                        self.taken_queue._reader,
                        self.results_queue._reader
                    ],
                    [], [],
                    self.config.master.heartbeat_interval
                )

                if not input_sources:
                    self.heartbeat()
                    continue

                if input_sources[0] == self.socket:
                    self.handle_payload()

                if (
                    input_sources[0] == self.taken_queue._reader
                    or len(input_sources) > 1
                ):
                    self.handle_taken_job()

                if (
                    input_sources[0] == self.results_queue._reader
                    or len(input_sources) > 2
                ):
                    self.handle_worker_result()

            except select.error as e:
                if e.args[0] not in [errno.EAGAIN, errno.EINTR]:
                    raise
            except OSError as e:
                if e.errno not in [errno.EAGAIN, errno.EINTR]:
                    raise
            except IOError as e:
                if e.errno not in [errno.EBADF]:
                    raise
            except SystemExit:
                raise
            except KeyboardInterrupt:
                self.wind_down_gracefully()
            except Exception, e:
                self.logger.exception(e)
                sys.exit(-1)

    def expand_workers(self, *args):
        self.number_of_workers += 1
        self.logger.info(
            "Upping number of workers to %d",
            self.number_of_workers
        )
        self.spawn_worker()<commit_msg>Refactor main master loop's handling of input streams with data.<commit_after>import errno
import logging
import os
import Queue
import select
import signal
import socket
import sys
import time
from multiprocessing import queues

import redis
import setproctitle

from .config import Config
from .worker import Worker
from .job import Job, job_iterator
from .redis_extensions import extend_redis


class Master(object):

    signal_map = {
        "hup": "reload_config",
        "usr1": "relaunch",
        "ttin": "expand_workers",
        "ttou": "contract_workers",
        "int": "halt_current_jobs",
        "quit": "wind_down_gracefully",
        "term": "wind_down_immediately",
        "chld": "handle_worker_exit"
    }

    def __init__(self, config_file):
        self.pid = None
        self.pid_file_path = None
        self.reexec_pid = 0

        self.name = "master"

        self.logger = logging.getLogger(__name__)

        self.config_file = config_file
        self.load_config()

        self.workers = {}

        self.socket = None

        self.ready_queue = queues.Queue(maxsize=5)
        self.taken_queue = queues.Queue()
        self.results_queue = queues.Queue()

        self.redis = None

        self.wind_down_time = None

        args = sys.argv[:]
        args.insert(0, sys.executable)

        self.launch_context = {
            "executable": sys.executable,
            "args": args,
            "cwd": os.getcwd()
        }

    def load_config(self):
        self.config = Config(self.config_file)
        self.config.load()
        self.pid_file_path = self.config.master.pid_file
        self.number_of_workers = self.config.master.num_workers

    def setup_pid_file(self):
        if not os.path.isdir(os.path.dirname(self.pid_file_path)):
            raise RuntimeError(
                "%s does not exist! Can't create pid file" % self.pid_file_path
            )

        if os.path.exists(self.pid_file_path):
            raise RuntimeError(
                (
                    "pid file (%s) already exists! \n"
                    "If no master is running it could be stale."
                ) % self.pid_file_path
            )

        fd = open(self.pid_file_path, 'w')
        try:
            fd.write(u"%s\n" % self.pid)
        except:
            os.unlink(self.pid_file_path)
            raise
        else:
            fd.close()

    def setup_datastore(self):
        if ":" in self.config.master.redis:
            host, port = self.config.master.redis.split(":")
            self.redis = redis.Redis(host=host, port=port)
        else:
            self.redis = redis.Redis(host=self.config.master.redis)

        extend_redis(self.redis)

    def setup_socket(self):
        self.logger.info(
            "Listening on port %s", self.config.master.listen_port
        )

        self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.socket.setblocking(0)
        self.socket.bind(
            ('', int(self.config.master.listen_port))
        )
        self.socket.listen(5)

    def setup_signals(self):
        for signal_name, handler_name in self.signal_map.iteritems():
            signal.signal(
                getattr(signal, "SIG%s" % signal_name.upper()),
                getattr(self, handler_name)
            )

    def run(self):
        self.pid = os.getpid()
        self.logger.info("Starting %s (%d)", self.name, int(self.pid))

        self.setup_pid_file()

        self.setup_datastore()

        self.fill_ready_queue()

        self.setup_socket()
        self.setup_signals()

        setproctitle.setproctitle("distq: %s" % self.name)

        for i in range(self.number_of_workers):
            self.spawn_worker()

        while True:
            try:
                (input_sources, [], []) = select.select(
                    [
                        self.socket,
                        self.taken_queue._reader,
                        self.results_queue._reader
                    ],
                    [], [],
                    self.config.master.heartbeat_interval
                )

                if not input_sources:
                    self.heartbeat()
                    continue

                while len(input_sources) > 0:
                    source_with_data = input_sources.pop(0)

                    if source_with_data == self.socket:
                        self.handle_payload()
                    elif source_with_data == self.taken_queue._reader:
                        self.handle_taken_job()
                    elif source_with_data == self.results_queue._reader:
                        self.handle_worker_result()

            except select.error as e:
                if e.args[0] not in [errno.EAGAIN, errno.EINTR]:
                    raise
            except OSError as e:
                if e.errno not in [errno.EAGAIN, errno.EINTR]:
                    raise
            except IOError as e:
                if e.errno not in [errno.EBADF]:
                    raise
            except SystemExit:
                raise
            except KeyboardInterrupt:
                self.wind_down_gracefully()
            except Exception, e:
                self.logger.exception(e)
                sys.exit(-1)

    def expand_workers(self, *args):
        self.number_of_workers += 1
        self.logger.info(
            "Upping number of workers to %d",
            self.number_of_workers
        )
        self.spawn_worker()
{'fraction_non_alphanumeric': 0.09699179580674568, 'fraction_numerical': 0.0022789425706472195, 'mean_word_length': 1.5466573816155988, 'pattern_counts': {'":': 26, '<': 3, '<?xml version=': 0, '>': 6, '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': '4462509', 'n_tokens_mistral': 3268, 'n_tokens_neox': 3076, 'n_words': 631}