pyFTS/pyFTS/distributed/dispy.py

import dispy as dispy, dispy.httpd, logging
from pyFTS.common import Util
import numpy as np


def start_dispy_cluster(method, nodes):
    """
    Start a new Dispy cluster on 'nodes' to execute the method 'method'

    :param method: function to be executed on each cluster node
    :param nodes: list of node names or IP's.
    :return: the dispy cluster instance and the http_server for monitoring
    """

    cluster = dispy.JobCluster(method, nodes=nodes, loglevel=logging.DEBUG, ping_interval=1000)

    http_server = dispy.httpd.DispyHTTPServer(cluster)

    return cluster, http_server


def stop_dispy_cluster(cluster, http_server):
    """
    Stop a dispy cluster and http_server

    :param cluster:
    :param http_server:
    :return:
    """
    #cluster.wait()  # wait for all jobs to finish

    cluster.print_status()

    http_server.shutdown()  # this waits until browser gets all updates
    cluster.close()


def get_number_of_cpus(cluster):
    cpus = 0
    for dispy_node in cluster.status().nodes:
        cpus += dispy_node.cpus

    return cpus


def simple_model_train(model, data, parameters):
    import time
    """
    Cluster function that receives a FTS instance 'model' and train using the 'data' and 'parameters'

    :param model: a FTS instance
    :param data: training dataset
    :param parameters: parameters for the training process
    :return: the trained model
    """
    _start = time.time()
    model.train(data, **parameters)
    _end = time.time()
    model.__dict__['training_time'] = _end - _start
    return model


def distributed_train(model, train_method, nodes, fts_method, data, num_batches=10,
                      train_parameters={}, **kwargs):
    import dispy, dispy.httpd, datetime

    batch_save = kwargs.get('batch_save', False)  # save model between batches

    batch_save_interval = kwargs.get('batch_save_interval', 1)

    file_path = kwargs.get('file_path', None)

    cluster, http_server = start_dispy_cluster(train_method, nodes)

    print("[{0: %H:%M:%S}] Distrituted Train Started with {1} CPU's"
          .format(datetime.datetime.now(), get_number_of_cpus(cluster)))

    jobs = []
    n = len(data)
    batch_size = int(n / num_batches)
    bcount = 1
    for ct in range(model.order, n, batch_size):
        if model.is_multivariate:
            ndata = data.iloc[ct - model.order:ct + batch_size]
        else:
            ndata = data[ct - model.order: ct + batch_size]

        tmp_model = fts_method()

        tmp_model.clone_parameters(model)

        job = cluster.submit(tmp_model, ndata, train_parameters)
        job.id = bcount  # associate an ID to identify jobs (if needed later)
        jobs.append(job)

        bcount += 1

    for job in jobs:
        print("[{0: %H:%M:%S}] Processing batch ".format(datetime.datetime.now()) + str(job.id))
        tmp = job()
        if job.status == dispy.DispyJob.Finished and tmp is not None:
            model.merge(tmp)
            if 'training_time' not in model.__dict__:
                model.__dict__['training_time'] = []
            model.__dict__['training_time'].append(tmp.__dict__['training_time'])

            if batch_save and (job.id % batch_save_interval) == 0:
                Util.persist_obj(model, file_path)

        else:
            print(job.exception)
            print(job.stdout)

        print("[{0: %H:%M:%S}] Finished batch ".format(datetime.datetime.now()) + str(job.id))

    print("[{0: %H:%M:%S}] Distrituted Train Finished".format(datetime.datetime.now()))

    stop_dispy_cluster(cluster, http_server)

    return model


def simple_model_predict(model, data, parameters):
    import time
    _start = time.time()
    forecasts = model.predict(data, **parameters)
    _stop = time.time()
    return forecasts, _stop - _start


def distributed_predict(model, parameters, nodes, data, num_batches, **kwargs):
    import dispy, dispy.httpd

    cluster, http_server = start_dispy_cluster(simple_model_predict, nodes)

    jobs = []
    n = len(data)
    batch_size = int(n / num_batches)
    bcount = 1
    for ct in range(model.order, n, batch_size):
        if model.is_multivariate:
            ndata = data.iloc[ct - model.order:ct + batch_size]
        else:
            ndata = data[ct - model.order: ct + batch_size]

        job = cluster.submit(model, ndata, parameters)
        job.id = bcount  # associate an ID to identify jobs (if needed later)
        jobs.append(job)

        bcount += 1

    ret = []

    for job in jobs:
        tmp = job()
        if job.status == dispy.DispyJob.Finished and tmp is not None:
            if job.id < batch_size:
                ret.extend(tmp[0][:-1])
            else:
                ret.extend(tmp[0])

            if 'forecasting_time' not in model.__dict__:
                model.__dict__['forecasting_time'] = []
            model.__dict__['forecasting_time'].append(tmp[1])

        else:
            print(job.exception)
            print(job.stdout)

    stop_dispy_cluster(cluster, http_server)

    return ret