benchmark_just_do_work.py


1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31

# A script that does some work. Useful for external benchmarking tools, e.g. hyperfine.
# Compare two commits like this:
#
# hyperfine --warmup 1 \
#   --command-name no_sync --prepare 'git checkout c8d45e8 && (rm -r data || true)' 'python3 benchmark_just_do_work.py data' \
#   --command-name sync --prepare 'git checkout b6c1a44 && (rm -r data || true)' 'python3 benchmark_just_do_work.py data' \
#   --cleanup 'rm -r data'

import argparse
from femtoqueue import FemtoQueue

def benchmark_femtoqueue(data_dir: str, num_tasks: int = 1000):
    queue = FemtoQueue(data_dir=data_dir, node_id="node1")

    data = b"x" * 100  # 100-byte payload

    for _ in range(num_tasks):
        queue.push(data)

    while True:
        task = queue.pop()
        if not task:
            break
        queue.done(task)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="FemtoQueue: Just run some tasks for external benchmarking")
    parser.add_argument("data_dir", type=str, help="Data directory. Won't be cleaned up!")
    args = parser.parse_args()

    benchmark_femtoqueue(args.data_dir)