"""This is the script for `ray clusterbenchmark`."""

import time

import numpy as np

import ray
from ray.cluster_utils import Cluster


def main():
    cluster = Cluster(
        initialize_head=True,
        connect=True,
        head_node_args={"object_store_memory": 20 * 1024 * 1024 * 1024, "num_cpus": 16},
    )
    cluster.add_node(
        object_store_memory=20 * 1024 * 1024 * 1024, num_gpus=1, num_cpus=16
    )

    object_ref_list = []
    for i in range(0, 10):
        object_ref = ray.put(np.random.rand(1024 * 128, 1024))
        object_ref_list.append(object_ref)

    @ray.remote(num_gpus=1)
    def f(object_ref_list):
        diffs = []
        for object_ref in object_ref_list:
            before = time.time()
            ray.get(object_ref)
            after = time.time()
            diffs.append(after - before)
            time.sleep(1)
        return np.mean(diffs), np.std(diffs)

    time_diff, time_diff_std = ray.get(f.remote(object_ref_list))

    print(
        "latency to get an 1G object over network",
        round(time_diff, 2),
        "+-",
        round(time_diff_std, 2),
    )

    ray.shutdown()
    cluster.shutdown()


if __name__ == "__main__":
    main()