From af9c5dc827c416b474ef65fde789bb99988be778 Mon Sep 17 00:00:00 2001 From: jaycedowell Date: Mon, 9 Dec 2024 13:46:14 -0700 Subject: [PATCH] Various cleanups to make the script a little nicer to use (#251). --- tools/thors_hammer.py | 51 ++++++++++++++++++++++++++++++++----------- 1 file changed, 38 insertions(+), 13 deletions(-) diff --git a/tools/thors_hammer.py b/tools/thors_hammer.py index 78fdabdf..96e9d0d3 100755 --- a/tools/thors_hammer.py +++ b/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,12 +103,12 @@ 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 @@ -91,27 +116,27 @@ def main(args): ## 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',