coordinator_node.py

#!/usr/bin/env python
import glob
import sys
import threading
import csv
import queue
import enum
import time
sys.path.append('gen-py')
# Ignore thrift module files if they are not there. Thrift must be installed through pip if so.
try:
    sys.path.insert(0, glob.glob('../../thrift-0.19.0/lib/py/build/lib*')[0])
except:
    pass

from coordinator_node_thrift import *
from coordinator_node_thrift.ttypes import *
from compute_node_thrift import *
from ML import *

from thrift.transport import TSocket
from thrift.transport import TTransport
from thrift.protocol import TBinaryProtocol
from thrift.server import TServer

'''
Helper Class for ensuring thread-safe access to queues
'''
class CloseableQueue(queue.Queue):
    def __init__(self, closed=False, items_before_close=-1):
        super().__init__()
        self._closed = closed
        self._items_before_close = items_before_close
        self._is_closed = threading.Condition(self.mutex)

    def close(self):
        with self.mutex:
            self._close()

    def open(self):
        with self.mutex:
            self._closed = False

    def closed(self):
        with self.mutex:
            return self._closed

    def wait_for_close(self):
        with self._is_closed:
            self._is_closed.wait()

    def increment_items(self):
        with self.mutex:
            if self._items_before_close < 0:
                self._items_before_close = 1
            else:
                self._items_before_close += 1

    def decrement_items(self):
        with self.mutex:
            if self._items_before_close > 0:
                self._items_before_close -= 1
            if self._items_before_close == 0:
                self._close()

    def reset_items(self):
        with self.mutex:
            self._items_before_close = -1

    def _close(self):
        self._closed = True
        self._is_closed.notify()

    def debug(self):
        with self.mutex:
            return str(self.queue)


'''
Class for holding our compute nodes and testing connections
Stored in queue to allow access to various connections
'''
class ComputeNode:
    def __init__(self, active, transport, node, name):
        self.transport = transport
        self.node = node
        self.name = name
        self.active = active

    '''
    @brief Tests connection of the TSocket
    '''
    def test_connection(self):
        try:
            self.transport.open()
            ping_ret = self.node.ping()
            self.transport.close()
            if ping_ret:
                self.active = True
                return
        except:
            pass
        self.active = False

    def __str__(self):
        return f'<ComputeNode: {self.name}>'


'''
@brief Handler for our coordinator node which handles training the ML model
'''
class CoordinatorHandler:
    def __init__(self, scheduling_policy):
        self.ready_for_training = False
        self.scheduling_policy = scheduling_policy
        self.work_queue = None
        self.return_queue = None
        self.message_queue = None
        self.almighty = mlp()
        self.compute_nodes = dict()

    '''
    @brief Initializes coordinator node with training data and possible compute nodes
    @param nodes_file Type string which holds the possible compute nodes
    @return True if coordinator successfully initialized, False if it did not
    '''
    def init_coordinator(self, nodes_file):
        address_lines = list()
        # Attempt to read in the possible compute nodes
        try:
            with open(nodes_file, 'r') as file:
                for line in csv.reader(file):
                    address_lines.append(line)
        except Exception as e:
            print(f'Failed to read compute_nodes.txt: {str(e)}')
            self.ready_for_training = False
            return self.ready_for_training
        self.compute_nodes.clear()
        # For each possible compute node, append to compute_nodes queue
        for address_line in address_lines:
            active = False
            transport = TSocket.TSocket(address_line[0], int(address_line[1]))
            transport = TTransport.TBufferedTransport(transport)
            protocol = TBinaryProtocol.TBinaryProtocol(transport)
            c = compute.Client(protocol)
            name = f'{address_line[0]}:{address_line[1]}'
            try:
                transport.open()
                transport.close()
                active = True
            except Exception as e:
                print(f'Failed to add compute node <{name}>: {str(e)}')
            self.compute_nodes[name] = ComputeNode(active=active, transport=transport, node=c, name=name)

        # Initialize our work queue, return queue, message queue
        self._reset_queues()

        self.ready_for_training = True
        return self.ready_for_training

    '''
    @brief Initializes queues for use
    '''
    def _reset_queues(self):
        self.work_queue = CloseableQueue(closed=False, items_before_close=0)
        self.return_queue = queue.Queue()
        self.message_queue = queue.Queue()

    '''
    @brief Training which is called by our client
    @return validation Type double which is the validation error of our ML model
    '''
    def train(self, _dir, rounds, epochs, h, k, eta):
        # Ensuring that coordinator server is properly initialized and has compute nodes ready
        if not self.ready_for_training:
            print('Failed to train, run init_coordinator first')
            return -1
        if len(self.compute_nodes) == 0:
            print('Failed to train, no compute nodes reachable')
            return -1
        if not self.almighty.init_training_random(_dir + '/train_letters1.txt', k, h):
            print('Failed to train: init_training_random returned False')
            return -1

        # Attempt training
        training_tries = 3
        ret_val = self._training_loop(_dir, rounds, epochs, eta)
        while training_tries > 0:
            # If complete node failure, reattempt connection to all nodes
            if ret_val < 0:
                training_tries -= 1
                self._reset_queues()
                print('Attempting to restart training in 3...')
                time.sleep(3)
                print('Restarting')
                ret_val = self._training_loop(_dir, rounds - (abs(ret_val) - 1), epochs, eta)
            else:
                break  # training loop returned 0 (or greater), success
        # Failed to reattempt connections, abort
        if training_tries <= 0:
            print('Aborting training')
            return -1

        print('Training complete, validation:')
        validation = self.almighty.validate(_dir + '/validate_letters.txt')
        print(validation)

        return validation

    '''
    @brief Training loop that will initialize our work queue, create threads for compute node connection, and compute the gradients and weights
    @return 0 on success, -(_round + 1) to reattempt training on failure while maintaining same amount of rounds run
    '''
    def _training_loop(self, _dir, rounds, epochs, eta):
        compute_threads = list()
        ave_grads = weights()
        # Add training files into work_queue
        for _round in range(rounds):
            print(f'Round: {_round}')
            for file in glob.glob(_dir + "/train*"):
                self.work_queue.increment_items()
                self.work_queue.put(file)
            print(f'work queue: {self.work_queue.debug()}')

            active_nodes = 0
            # For each compute node, create a thread to contact
            for node_name in self.compute_nodes:
                compute_node = self.compute_nodes[node_name]
                # If inactive, reattempt connection to see if it is active now
                if not compute_node.active:
                    compute_node.test_connection()
                # Initialize thread's weights and compute node connection, store in compute_threads
                if compute_node.active:
                    active_nodes += 1
                    w = weights()
                    w.V, w.W = self.almighty.get_weights()
                    compute_threads.append(threading.Thread(target=self._compute_func, args=(compute_node, w, eta, epochs)))
            if active_nodes == 0:
                print('Failed to find an active training node')
                return -(_round + 1)
            # Begin training
            for compute_thread in compute_threads:
                compute_thread.start()

            # While there are threads working, check for any node failures
            while not self.work_queue.closed():
                try:
                    node_message = self.message_queue.get(timeout=1)
                except queue.Empty:
                    continue
                if not node_message[1]:
                    print(f'Node <{node_message[0]}> failed')
                    active_nodes -= 1
                if active_nodes <= 0:
                    print('All nodes failed')
                    for compute_thread in compute_threads:
                        compute_thread.join()
                    return -(_round + 1)

            for compute_thread in compute_threads:
                compute_thread.join()
            # Clear out our threads for new threads
            compute_threads.clear()

            # Reset our queue and open for thread access
            self.work_queue.reset_items()
            self.work_queue.open()

            # Compute the weights and gradients
            ave_grads.V = np.zeros(self.almighty.V.shape)
            ave_grads.W = np.zeros(self.almighty.W.shape)
            num_grads = 0
            while self.return_queue.qsize() != 0:
                single_grad = self.return_queue.get()
                num_grads += 1
                ave_grads.V = sum_matricies(ave_grads.V, single_grad.V)
                ave_grads.W = sum_matricies(ave_grads.W, single_grad.W)
            ave_grads.V = scale_matricies(ave_grads.V, 1 / num_grads)
            ave_grads.W = scale_matricies(ave_grads.W, 1 / num_grads)
            self.almighty.update_weights(ave_grads.V, ave_grads.W)
            print(self.almighty.validate(_dir + '/validate_letters.txt'))
        return 0

    '''
    @brief Our thread function which will communicate with the compute nodes
    @param compute_node Type ComputeNode which is holds our connection data
    @param input_weights Type weights which is our thrift type that holds our W and V matrices
    '''
    def _compute_func(self, compute_node, input_weights, eta, epochs):
        # Attempt a connection
        try:
            compute_node.transport.open()
        except Exception as e:
            print(f'Failed to connect to compute node {compute_node}: {str(e)}')
            return

        training_fails = 0
        # While the thread is open for thread access
        while not self.work_queue.closed():
            # Attempt to give compute node a task with load probability
            try:
                print(f'asking: {str(compute_node)}')
                answer = compute_node.node.receive_task()
            except Exception as e:
                print(f'Failed to receive_task in node {compute_node}: {str(e)}')
                compute_node.active = False
                break
            # Random load balance or node accepts work
            if (self.scheduling_policy == 1) or answer:
                # Attempt to retrieve file
                try:
                    file = self.work_queue.get(timeout=3)
                except:
                    # There was no file in the queue
                    print("Failed to get file, moving on")
                    continue
                # Train compute node
                try:
                    print(f'training: {str(compute_node)}')
                    return_gradients = compute_node.node.train(input_weights, file, eta, epochs)
                except Exception as e:
                    print(f'Failed to train in node {compute_node}: {str(e)}')
                    self.work_queue.put(file)  # return file to q to be processed elsewhere
                    compute_node.active = False
                    break
                if not return_gradients.success:
                    print("No return gradient, failed")
                    training_fails += 1
                    self.work_queue.put(file)
                # Add return gradient to our queue for later use
                else:
                    self.return_queue.put(return_gradients)
                    self.work_queue.decrement_items()
                if training_fails > 3:
                    print('Too many failed train attempts, aborting')
                    compute_node.active = False
                    break
            else:
                # Node was busy/under load, wait so another thread can execute
                time.sleep(3)
        compute_node.transport.close()
        self.message_queue.put((compute_node.name, compute_node.active))
        return


if __name__ == '__main__':
    if len(sys.argv) < 3:
        print("python3 coordinator_node.py <port> <scheduling_policy>")
        sys.exit(1)
    port = sys.argv[1]
    scheduling_policy = int(sys.argv[2])

    handler = CoordinatorHandler(scheduling_policy)
    processor = coordinator.Processor(handler)
    transport = TSocket.TServerSocket(port=port)
    tfactory = TTransport.TBufferedTransportFactory()
    pfactory = TBinaryProtocol.TBinaryProtocolFactory()

    server = TServer.TSimpleServer(processor, transport, tfactory, pfactory)

    print('Starting the server...')
    server.serve()
    print('done.')