Various cleanups to make the script a little nicer to use (#251).

tools/thors_hammer.py

@@ -36,6 +36,30 @@
 from bifrost.ndarray import copy_array
 from bifrost.packet_writer import HeaderInfo, DiskWriter
 
+def data_size(string):
+    """
+    Convert a string like '100GiB' into a size in bytes.
+    """
+
+    if string.endswith('T') or string.endswith('TiB'):
+        value, _ = string.split('T', 1)
+        value = int(value) * 1024**4
+    elif string.endswith('G') or string.endswith('GiB'):
+        value, _ = string.split('G', 1)
+        value = int(value) * 1024**3
+    elif string.endswith('M') or string.endswith('MiB'):
+        value, _ = string.split('M', 1)
+        value = int(value) * 1024**2
+    elif string.endswith('k') or string.endswith('kiB'):
+        value, _ = string.split('k', 1)
+        value = int(value) * 1024
+    else:
+        try:
+            value = int(string)
+        except ValueError:
+            raise ArgumentTypeError(f"Cannot parse '{string}' as a data size")
+    return value
+
 class DummyFileHandle:
     """
     Wrapper around a Python file descriptor as returned by os.open() to make
@@ -53,7 +77,7 @@ def close(self):
 
 def main(args):
     # Create an empty data array
-    data = np.empty(shape=(args.repeats, args.ntime, 1, args.bytes_per_time), dtype=np.uint8)
+    data = np.empty(shape=(args.repeats, args.npacket, 1, args.packet_size), dtype=np.uint8)
     data = BFArray(data, space='system')
 
     # Setup the file flags, enabling O_DIRECT | O_SYNC as needed
@@ -62,8 +86,9 @@ def main(args):
         file_flags |= os.O_DIRECT | os.O_SYNC
 
         ## Make sure that the request write width matches what O_DIRECT expects
-        if args.ntime*args.bytes_per_time % 512 != 0:
-            raise RuntimeError(f"ntime x bytes_per_time ({args.ntime*args.bytes_per_time}) % 512 != 0")
+        if args.npacket*args.packet_size % 512 != 0:
+            write_width = args.npacket*args.packet_size
+            raise RuntimeError(f"npacket x packet_size ({write_width}) % 512 != 0")
 
     # Header descriptor (not populated in the case of 'generic'
     desc = HeaderInfo()
@@ -78,40 +103,40 @@ def main(args):
         ## Open the i-th file
         fd = os.open(f"{args.filename}_{i+1}", file_flags, mode=0o664)
         fh = DummyFileHandle(fd)
-        udt = DiskWriter(f"generic_{args.bytes_per_time}", fh)
+        udt = DiskWriter(f"generic_{args.packet_size}", fh)
 
         ## Write a collection of repeats to the file
         for j in range(args.repeats):
             udt.send(desc, 0, 1, 0, 1, data[j,...])
-            nbyteF += args.ntime*args.bytes_per_time
+            nbyteF += args.npacket*args.packet_size
 
         ## Close it out
         del udt
         fh.close()
 
         ## Report the throughput for this file
         t1F = time.time()
-        print(f"  #{i+1} in {t1F-t0F:.3f}s -> {nbyteF/1e9/(t1F-t0F):.1f} GB/s")
+        print(f"  #{i+1} in {t1F-t0F:.3f}s -> {nbyteF/1024**3/(t1F-t0F):.1f} GiB/s")
 
         nbyte += nbyteF
 
     # Final throughput report across all file_count files
     t1 = time.time()
-    print(f"Finished in {t1-t0:.3f}s -> {nbyte/1e9/(t1-t0):.1f} GB/s")
+    print(f"Finished in {t1-t0:.3f}s -> {nbyte/1024**3/(t1-t0):.1f} GiB/s")
 
 if __name__ == '__main__':
     parser = argparse.ArgumentParser(
         description='test disk throughput as a function of different Bifrost parameters',
         formatter_class=argparse.ArgumentDefaultsHelpFormatter
         )
-    parser.add_argument('-b', '--bytes-per-time', type=int, default=67584,
-                        help='number of bytes to write per time sample')
-    parser.add_argument('-n', '--ntime', type=int, default=1920,
-                        help='number of time samples to write per Bifrost DiskWriter() call')
+    parser.add_argument('-p', '--packet-size', type=data_size, default=67584,
+                        help='number of bytes to write per packet')
+    parser.add_argument('-n', '--npacket', type=int, default=1920,
+                        help='number of packets to write per Bifrost DiskWriter() call')
     parser.add_argument('-r', '--repeats', type=int, default=1,
-                        help='number of times to repeat the ntime x bytes-per-time writes')
+                        help='number of times to repeat the npacket x packet-size sized writes')
     parser.add_argument('-c', '--file-count', type=int, default=1,
-                        help='number of files containing repeats x ntime x bytes-per-time to write')
+                        help='number of files containing repeats x npacket x packet-size to write')
     parser.add_argument('-d', '--direct-io', action='store_true',
                         help='enable O_DIRECT | O_SYNC when opening the file(s)')
     parser.add_argument('-f', '--filename', type=str, default='hammer.test',