bmk: add benchmark for writer

benchmark/benchmark.py

@@ -16,15 +16,21 @@ def print_underline(msg):
     print("-" * len(msg))
 
 
-def run_benchmark(bm_path):
+def run_benchmark(bm_path, repeats=5):
     print_underline(f"Running {bm_path}")
     bm_path = pathlib.Path(bm_path).resolve()
     os.chdir(f"{bm_path.parent}")
     bm_mod = importlib.import_module(f"{bm_path.stem}")
-    t = timeit.Timer(bm_mod.main)
-    res = t.repeat(repeat=5, number=10)
-    print(f"best={min(res):.3g}, mean={np.mean(res):.3g}")
-    return res
+
+    reps = []
+    print("Running...", end="\r")
+    for ii in range(repeats):
+        t = timeit.timeit(bm_mod.main, setup=bm_mod.setup, number=1)
+        reps.append(t)
+        print(f"Running {ii + 1}/{repeats}", end="\r")
+    print(reps)
+    print(f"best={min(reps):.3g}, mean={np.mean(reps):.3g}")
+    return reps
 
 
 if __name__ == "__main__":

benchmark/bm_write_deque_writer_thread.py

@@ -0,0 +1,43 @@
+from collections import deque
+import pathlib
+import tempfile
+
+import multiprocessing as mp
+
+import numpy as np
+
+from dcnum import write
+
+
+mp_spawn = mp.get_context('spawn')
+
+
+def setup():
+    global path_out
+    global writer_dq
+    total_frames = 3000
+    batch_size = 500
+    num_batches = 6
+    assert batch_size * num_batches == total_frames
+
+    writer_dq = deque()
+    # Create 1000 events with at most two repetitions in a frame
+    np.random.seed(42)
+    rng = np.random.default_rng()
+
+    # create a sample event
+    for ii in range(num_batches):
+        writer_dq.append(("mask", rng.random((batch_size, 80, 320)) > .5))
+        writer_dq.append(("temp", rng.normal(23, size=batch_size)))
+
+    path_out = pathlib.Path(
+        tempfile.mkdtemp(prefix=pathlib.Path(__file__).name)) / "out.rtdc"
+
+
+def main():
+    thr_drw = write.DequeWriterThread(
+        path_out=path_out,
+        dq=writer_dq,
+    )
+    thr_drw.may_stop_loop = True
+    thr_drw.run()

benchmark/bm_write_queue_collector_thread.py

@@ -1,5 +1,6 @@
 from collections import deque
 import multiprocessing as mp
+
 import numpy as np
 
 from dcnum import write
@@ -8,9 +9,12 @@
 mp_spawn = mp.get_context('spawn')
 
 
-def main():
+def setup():
+    global event_queue
+    global writer_dq
+    global feat_nevents
     batch_size = 1000
-    num_batches = 5
+    num_batches = 3
     num_events = batch_size * num_batches
     event_queue = mp.Queue()
     writer_dq = deque()
@@ -30,6 +34,8 @@ def main():
         for idx in number_order:
             event_queue.put((ii*batch_size + idx, event))
 
+
+def main():
     thr_coll = write.QueueCollectorThread(
         event_queue=event_queue,
         writer_dq=writer_dq,