Benchmark Script for AZFour

azfour_benchmark.py

import concurrent.futures
import numpy as np
from graphpipe import remote
import time


GP_SERVER = "http://127.0.0.1:9000"
WORKERS = 40


def do_work(num):
    count = 0
    for i in range(100//num):
        inp = np.zeros((num, 2, 6, 7)).astype(np.float32)
        remote.execute(GP_SERVER, inp)
        count += inp.shape[0]
    return count


def benchmark(batch_size):
    total = 0
    start = time.time()
    with concurrent.futures.ThreadPoolExecutor(max_workers=WORKERS) \
            as executor:
        futures = [executor.submit(do_work, batch_size)
                   for i in range(WORKERS)]
        for future in concurrent.futures.as_completed(futures):
            count = future.result()
            total += count

    duration = time.time() - start
    throughput = total/duration

    print("%s,%s" % (batch_size, throughput))


if __name__ == "__main__":
    for i in [1, 2, 3, 5, 7, 10, 15, 20, 30, 40, 50]:
        benchmark(i)