memverge_benchmark.py

import argparse
import time

import ray
import numpy as np

GiB = 1024 * 1024 * 1024
 
def make_np(size_gib: int):
    total_size = GiB * size_gib
    record_dim = 1024
    record_size = record_dim * 8
    num_records = int(total_size / record_size)
    dataset = np.random.randint(1000, size=(num_records, record_dim))
    return dataset

@ray.remote
def echo(a):
    print(f'running on {ray.get_runtime_context().get_node_id()}')
    return 0

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--loop-count", type=int, default=10)
    parser.add_argument("--remote", type=int, default=1)
    args = parser.parse_args()

    if args.remote > 0:
        print("Executing remotely")
    else:
        print("Executing locally")

    ds = make_np(1)
    print (f"{ds.size * ds.itemsize}")

    ds_ref = ray.put(ds)

    times = []
    for x in range(args.loop_count):
        ds_ref = ray.put(ds)
        start = time.time()
        if args.remote > 0:
            ray.get(echo.options(resources={"worker": 0.01}).remote(ds_ref))
        else:
            ray.get(echo.options(resources={"head": 0.01}).remote(ds_ref))
        time_diff = time.time() - start
        del ds_ref
        times.append(time_diff)

    print("Time spent {0}".format(sum(times) / len(times)))