glm-bench.py

from distributed import Worker, Scheduler, Client, Nanny
from distributed.client import _wait

from tornado import gen
from tornado.ioloop import IOLoop

loop = IOLoop.current()

s = Scheduler(loop=loop)
s.start(0)
workers = [Worker(s.ip, s.port, loop=loop) for i in range(4)]


from dask import persist
import dask.array as da
import numpy as np


def loglike(Xbeta, y):
    return da.log1p(da.exp(Xbeta)).sum() - y.dot(Xbeta)

N = 1000000
M = 5
chunks = 1000

beta = np.ones(M)
step = np.ones(M)
Xbeta = da.ones(N, chunks=chunks)
Xstep = da.ones(N, chunks=chunks)
y = da.ones(N, chunks=chunks)
curr_val = 1
stepSize = 1

obeta, oXbeta = beta, Xbeta



@gen.coroutine
def f():
    yield [w._start(0) for w in workers]
    c = Client(s.address, loop=loop, start=False)
    yield c._start()

    oXbeta, y2, Xstep2 = persist(Xbeta, y, Xstep)
    yield _wait([oXbeta, y2, Xstep2])

    for i in range(5):
        Xbeta2 = oXbeta - stepSize * Xstep2
        func = loglike(Xbeta2, y2)
        Xbeta2, func = persist(Xbeta2, func)
        yield _wait(func)

loop.run_sync(f)