compute_node.py

#!/usr/bin/env python
import glob
import sys
import random
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 compute_node_thrift import *
from coordinator_node_thrift import *
from coordinator_node_thrift.ttypes import weights
from ML import *

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

'''
@brief Our compute node handler that receives tasks, load probability, and trains ML model off of given path
'''
class ComputeHandler:
    def __init__(self, load_probability):
        self.load_probability = load_probability
        self.is_laden = False

    def ping(self):
        return True

    '''
    @brief Initializes whether compute node is under load or not
    @return False if it is under load to reject task. True if it is not to accept task
    '''
    def receive_task(self):
        self.is_laden = random.random() < load_probability
        print(f'asked to do task, laden: {self.is_laden}')
        return not self.is_laden  # load prob is chance of rejecting, 80% (0.8) means 0.2 chance to accept
    '''
    @brief Trains local ML model using given file path and weights
    @param w Type weights which is our weights for V and W
    @param file_name Type string which is our file name for our data
    @return gradients Type weights which holds our weights for V and W after training
    '''
    def train(self, w, file_name, eta, epochs):
        print(f'training, file: {file_name}, laden: {self.is_laden}')
        # If it's under load, just sleep
        if self.is_laden:
            time.sleep(3)
            self.is_laden = False
        ml = mlp()

        if not ml.init_training_model(file_name, w.V, w.W):
            w.success = False
            return w

        if ml.train(eta, epochs) == -1:
            # print('Failed to train')
            w.success = False
            return w

        # Calculate and return gradients
        gradients = weights()
        gradients.V, gradients.W = ml.get_weights()
        gradients.V = calc_gradient(gradients.V, w.V)
        gradients.W = calc_gradient(gradients.W, w.W)
        gradients.success = True
        return gradients


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

    handler = ComputeHandler(load_probability)
    processor = compute.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.')