<commit_before>from __future__ import print_function

import zmq
import socket
import dill
import uuid
from collections import defaultdict
import itertools
from multiprocessing.pool import ThreadPool
import random
from datetime import datetime
from threading import Thread, Lock
from contextlib import contextmanager
import traceback
import sys
from time import sleep
from ..compatibility import Queue, unicode, Empty
try:
    import cPickle as pickle
except ImportError:
    import pickle

from ..core import get_dependencies, flatten
from .. import core
from ..async import finish_task, start_state_from_dask as dag_state_from_dask

with open('log.scheduler', 'w') as f:  # delete file
    pass


def log(*args):
    with open('log.scheduler', 'a') as f:
        print('\n', *args, file=f)


@contextmanager
def logerrors():
    try:
        yield
    except Exception as e:
        exc_type, exc_value, exc_traceback = sys.exc_info()
        tb = ''.join(traceback.format_tb(exc_traceback))
        log('Error!', str(e))
        log('Traceback', str(tb))
        raise

class Scheduler(object):
    """ Disitributed scheduler for dask computations

    Parameters
    ----------

    hostname: string
        hostname or IP address of this machine visible to outside world
    port_to_workers: int
        Port on which to listen to connections from workers
    port_to_clients: int
        Port on which to listen to connections from clients
    bind_to_workers: string
        Addresses from which we accept worker connections, defaults to *
    bind_to_clients: string
        Addresses from which we accept client connections, defaults to *
    block: bool
        Whether or not to block the process on creation

    State
    -----

    workers - dict
        Maps worker identities to information about that worker
    who_has - dict
        Maps data keys to sets of workers that own that data
    worker_has - dict
        Maps workers to data that they own
    data - dict
        Maps data keys to metadata about the computation that produced it
    to_workers - zmq.Socket (ROUTER)
        Socket to communicate to workers
    to_clients - zmq.Socket (ROUTER)
        Socket to communicate with users
    collections - dict
        Dict holding shared collections like bags and arrays
    """
    def __init__(self, port_to_workers=None, port_to_clients=None,
                 bind_to_workers='*', bind_to_clients='*',
                 hostname=None, block=False):
        self.context = zmq.Context()
        hostname = hostname or socket.gethostname()

        # Bind routers to addresses (and create addresses if necessary)
        self.to_workers = self.context.socket(zmq.ROUTER)
        if port_to_workers is None:
            port_to_workers = self.to_workers.bind_to_random_port('tcp://' + bind_to_workers)
        else:
            self.to_workers.bind('tcp://%s:%d' % (bind_to_workers, port_to_workers))
        self.address_to_workers = ('tcp://%s:%d' % (hostname, port_to_workers)).encode()
        self.worker_poller = zmq.Poller()
        self.worker_poller.register(self.to_workers, zmq.POLLIN)

        self.to_clients = self.context.socket(zmq.ROUTER)
        if port_to_clients is None:
            port_to_clients = self.to_clients.bind_to_random_port('tcp://' + bind_to_clients)
        else:
            self.to_clients.bind('tcp://%s:%d' % (bind_to_clients, port_to_clients))
        self.address_to_clients = ('tcp://%s:%d' % (hostname, port_to_clients)).encode()

        # State about my workers and computed data
        self.workers = dict()
        self.who_has = defaultdict(set)
        self.worker_has = defaultdict(set)
        self.available_workers = Queue()
        self.data = defaultdict(dict)
        self.collections = dict()

        self.send_to_workers_queue = Queue()
        self.send_to_workers_recv = self.context.socket(zmq.PAIR)
        self.send_to_workers_recv.bind('ipc://to-workers-signal')
        self.send_to_workers_send = self.context.socket(zmq.PAIR)
        self.send_to_workers_send.connect('ipc://to-workers-signal')
        self.worker_poller.register(self.send_to_workers_recv, zmq.POLLIN)

        self.pool = ThreadPool(100)
        self.lock = Lock()
        self.status = 'run'
        self.queues = dict()

        self._schedule_lock = Lock()

        # RPC functions that workers and clients can trigger
        self.worker_functions = {'register': self._worker_registration,
                                 'status': self._status_to_worker,
                                 'finished-task': self._worker_finished_task,
                                 'setitem-ack': self._setitem_ack,
                                 'getitem-ack': self._getitem_ack}
        self.client_functions = {'status': self._status_to_client,
                                 'schedule': self._schedule_from_client,
                                 'set-collection': self._set_collection,
                                 'get-collection': self._get_collection,
                                 'close': self._close}

        # Away we go!
        log(self.address_to_workers, 'Start')
        self._listen_to_workers_thread = Thread(target=self._listen_to_workers)
        self._listen_to_workers_thread.start()
        self._listen_to_clients_thread = Thread(target=self._listen_to_clients)
        self._listen_to_clients_thread.start()

        if block:
            self.block()

    def _listen_to_workers(self):
        """ Event loop: Listen to worker router """
        while self.status != 'closed':
            try:
                socks = dict(self.worker_poller.poll(100))
                if not socks:
                    continue
            except zmq.ZMQError:
                break
            if self.send_to_workers_recv in socks:
                self.send_to_workers_recv.recv()
                while not self.send_to_workers_queue.empty():
                    msg = self.send_to_workers_queue.get()
                    self.to_workers.send_multipart(msg)

            if self.to_workers in socks:
                address, header, payload = self.to_workers.recv_multipart()

                header = pickle.loads(header)
                if 'address' not in header:
                    header['address'] = address
                log(self.address_to_workers, 'Receive job from worker', header)

                try:
                    function = self.worker_functions[header['function']]
                except KeyError:
                    log(self.address_to_workers, 'Unknown function', header)
                else:
                    future = self.pool.apply_async(function, args=(header, payload))

    def _listen_to_clients(self):
        """ Event loop: Listen to client router """
        while self.status != 'closed':
            try:
                if not self.to_clients.poll(100):  # is this threadsafe?
                    continue
            except zmq.ZMQError:
                break
            with self.lock:
                address, header, payload = self.to_clients.recv_multipart()
            header = pickle.loads(header)
            if 'address' not in header:
                header['address'] = address
            log(self.address_to_clients, 'Receive job from client', header)

            try:
                function = self.client_functions[header['function']]
            except KeyError:
                log(self.address_to_clients, 'Unknown function', header)
            else:
                self.pool.apply_async(function, args=(header, payload))

    def block(self):
        """ Block until listener threads close

        Warning: If some other thread doesn't call `.close()` then, in the
        common case you can not easily escape from this.
        """
        self._listen_to_workers_thread.join()
        self._listen_to_clients_thread.join()

    def _worker_registration(self, header, payload):
        """ Worker came in, register them """
        payload = pickle.loads(payload)
        address = header['address']
        self.workers[address] = payload
        self.available_workers.put(address)

    def _worker_finished_task(self, header, payload):
        """ Worker reports back as having finished task, ready for more

        See also:
            Scheduler.trigger_task
            Scheduler.schedule<commit_msg>Add client registration to scheduler.<commit_after>from __future__ import print_function

import zmq
import socket
import dill
import uuid
from collections import defaultdict
import itertools
from multiprocessing.pool import ThreadPool
import random
from datetime import datetime
from threading import Thread, Lock
from contextlib import contextmanager
import traceback
import sys
from time import sleep
from ..compatibility import Queue, unicode, Empty
try:
    import cPickle as pickle
except ImportError:
    import pickle

from ..core import get_dependencies, flatten
from .. import core
from ..async import finish_task, start_state_from_dask as dag_state_from_dask

with open('log.scheduler', 'w') as f:  # delete file
    pass


def log(*args):
    with open('log.scheduler', 'a') as f:
        print('\n', *args, file=f)


@contextmanager
def logerrors():
    try:
        yield
    except Exception as e:
        exc_type, exc_value, exc_traceback = sys.exc_info()
        tb = ''.join(traceback.format_tb(exc_traceback))
        log('Error!', str(e))
        log('Traceback', str(tb))
        raise

class Scheduler(object):
    """ Disitributed scheduler for dask computations

    Parameters
    ----------

    hostname: string
        hostname or IP address of this machine visible to outside world
    port_to_workers: int
        Port on which to listen to connections from workers
    port_to_clients: int
        Port on which to listen to connections from clients
    bind_to_workers: string
        Addresses from which we accept worker connections, defaults to *
    bind_to_clients: string
        Addresses from which we accept client connections, defaults to *
    block: bool
        Whether or not to block the process on creation

    State
    -----

    workers - dict
        Maps worker identities to information about that worker
    who_has - dict
        Maps data keys to sets of workers that own that data
    worker_has - dict
        Maps workers to data that they own
    data - dict
        Maps data keys to metadata about the computation that produced it
    to_workers - zmq.Socket (ROUTER)
        Socket to communicate to workers
    to_clients - zmq.Socket (ROUTER)
        Socket to communicate with users
    collections - dict
        Dict holding shared collections like bags and arrays
    """
    def __init__(self, port_to_workers=None, port_to_clients=None,
                 bind_to_workers='*', bind_to_clients='*',
                 hostname=None, block=False):
        self.context = zmq.Context()
        hostname = hostname or socket.gethostname()

        # Bind routers to addresses (and create addresses if necessary)
        self.to_workers = self.context.socket(zmq.ROUTER)
        if port_to_workers is None:
            port_to_workers = self.to_workers.bind_to_random_port('tcp://' + bind_to_workers)
        else:
            self.to_workers.bind('tcp://%s:%d' % (bind_to_workers, port_to_workers))
        self.address_to_workers = ('tcp://%s:%d' % (hostname, port_to_workers)).encode()
        self.worker_poller = zmq.Poller()
        self.worker_poller.register(self.to_workers, zmq.POLLIN)

        self.to_clients = self.context.socket(zmq.ROUTER)
        if port_to_clients is None:
            port_to_clients = self.to_clients.bind_to_random_port('tcp://' + bind_to_clients)
        else:
            self.to_clients.bind('tcp://%s:%d' % (bind_to_clients, port_to_clients))
        self.address_to_clients = ('tcp://%s:%d' % (hostname, port_to_clients)).encode()

        # Client state
        self.clients = dict()

        # State about my workers and computed data
        self.workers = dict()
        self.who_has = defaultdict(set)
        self.worker_has = defaultdict(set)
        self.available_workers = Queue()
        self.data = defaultdict(dict)
        self.collections = dict()

        self.send_to_workers_queue = Queue()
        self.send_to_workers_recv = self.context.socket(zmq.PAIR)
        self.send_to_workers_recv.bind('ipc://to-workers-signal')
        self.send_to_workers_send = self.context.socket(zmq.PAIR)
        self.send_to_workers_send.connect('ipc://to-workers-signal')
        self.worker_poller.register(self.send_to_workers_recv, zmq.POLLIN)

        self.pool = ThreadPool(100)
        self.lock = Lock()
        self.status = 'run'
        self.queues = dict()

        self._schedule_lock = Lock()

        # RPC functions that workers and clients can trigger
        self.worker_functions = {'register': self._worker_registration,
                                 'status': self._status_to_worker,
                                 'finished-task': self._worker_finished_task,
                                 'setitem-ack': self._setitem_ack,
                                 'getitem-ack': self._getitem_ack}
        self.client_functions = {'status': self._status_to_client,
                                 'register': self._client_registration,
                                 'schedule': self._schedule_from_client,
                                 'set-collection': self._set_collection,
                                 'get-collection': self._get_collection,
                                 'close': self._close}

        # Away we go!
        log(self.address_to_workers, 'Start')
        self._listen_to_workers_thread = Thread(target=self._listen_to_workers)
        self._listen_to_workers_thread.start()
        self._listen_to_clients_thread = Thread(target=self._listen_to_clients)
        self._listen_to_clients_thread.start()

        if block:
            self.block()

    def _listen_to_workers(self):
        """ Event loop: Listen to worker router """
        while self.status != 'closed':
            try:
                socks = dict(self.worker_poller.poll(100))
                if not socks:
                    continue
            except zmq.ZMQError:
                break
            if self.send_to_workers_recv in socks:
                self.send_to_workers_recv.recv()
                while not self.send_to_workers_queue.empty():
                    msg = self.send_to_workers_queue.get()
                    self.to_workers.send_multipart(msg)

            if self.to_workers in socks:
                address, header, payload = self.to_workers.recv_multipart()

                header = pickle.loads(header)
                if 'address' not in header:
                    header['address'] = address
                log(self.address_to_workers, 'Receive job from worker', header)

                try:
                    function = self.worker_functions[header['function']]
                except KeyError:
                    log(self.address_to_workers, 'Unknown function', header)
                else:
                    future = self.pool.apply_async(function, args=(header, payload))

    def _listen_to_clients(self):
        """ Event loop: Listen to client router """
        while self.status != 'closed':
            try:
                if not self.to_clients.poll(100):  # is this threadsafe?
                    continue
            except zmq.ZMQError:
                break
            with self.lock:
                address, header, payload = self.to_clients.recv_multipart()
            header = pickle.loads(header)
            if 'address' not in header:
                header['address'] = address
            log(self.address_to_clients, 'Receive job from client', header)

            try:
                function = self.client_functions[header['function']]
            except KeyError:
                log(self.address_to_clients, 'Unknown function', header)
            else:
                self.pool.apply_async(function, args=(header, payload))

    def block(self):
        """ Block until listener threads close

        Warning: If some other thread doesn't call `.close()` then, in the
        common case you can not easily escape from this.
        """
        self._listen_to_workers_thread.join()
        self._listen_to_clients_thread.join()

    def _client_registration(self, header, payload):
        """ Client comes in, register it."""
        payload = pickle.loads(payload)
        address = header['address']
        self.clients[address] = payload

    def _worker_registration(self, header, payload):
        """ Worker came in, register them """
        payload = pickle.loads(payload)
        address = header['address']
        self.workers[address] = payload
        self.available_workers.put(address)

    def _worker_finished_task(self, header, payload):
        """ Worker reports back as having finished task, ready for more

        See also:
            Scheduler.trigger_task
            Scheduler.schedule
{'fraction_non_alphanumeric': 0.08711168314514714, 'fraction_numerical': 0.001053062657228105, 'mean_word_length': 2.1556211925419975, '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': 17, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4013303', 'n_tokens_mistral': 4748, 'n_tokens_neox': 4617, 'n_words': 1210}