test_storage.py

# Storage performance test with reading/writing a stack of tiff files

# Necessary python packages: tiffile, threading, numpy

# Syntax:
# python test_storage.py <data size in each dimension> <number of threads> <path to store data>

# Example:
# python test_storage.py 1024 8 /data/tmp/tiff_stack/

# Output for tomodata1:
'''
init random data
write [1024,1024,1024] float64 dataset to /data/tmp/tiff_stack/ in 8 nthreads. Test 0.
write time 0.544s, write speed 14.69514213427615GB/s
write [1024,1024,1024] float64 dataset to /data/tmp/tiff_stack/ in 8 nthreads. Test 1.
write time 0.505s, write speed 15.849991497401984GB/s
write [1024,1024,1024] float64 dataset to /data/tmp/tiff_stack/ in 8 nthreads. Test 2.
write time 0.499s, write speed 16.03298871340673GB/s

read [1024,1024,1024] float64 dataset from /data/tmp/tiff_stack/ in 8 nthreads. Test 0.
read time 0.49s, read speed 16.312251580936973GB/s
read [1024,1024,1024] float64 dataset from /data/tmp/tiff_stack/ in 8 nthreads. Test 1.
read time 0.587s, read speed 13.633091166314339GB/s
read [1024,1024,1024] float64 dataset from /data/tmp/tiff_stack/ in 8 nthreads. Test 2.
read time 0.471s, read speed 16.987633320119663GB/s
'''
# Output for tomo1:
'''
init random data
write [1024,1024,1024] float64 dataset to /data/tmp/tiff_stack in 8 nthreads. Test 0.
write time 5.24s, write speed 1.5254433066478903GB/s
write [1024,1024,1024] float64 dataset to /data/tmp/tiff_stack in 8 nthreads. Test 1.
write time 5.19s, write speed 1.5403242176852137GB/s
write [1024,1024,1024] float64 dataset to /data/tmp/tiff_stack in 8 nthreads. Test 2.
write time 5.17s, write speed 1.5483579214395435GB/s

read [1024,1024,1024] float64 dataset from /data/tmp/tiff_stack in 8 nthreads. Test 0.
read time 0.955s, read speed 8.378170893609287GB/s
read [1024,1024,1024] float64 dataset from /data/tmp/tiff_stack in 8 nthreads. Test 1.
read time 0.844s, read speed 9.481811032544455GB/s
read [1024,1024,1024] float64 dataset from /data/tmp/tiff_stack in 8 nthreads. Test 2.
read time 0.721s, read speed 11.09047508558865GB/s
'''

import tifffile
import time
import threading
import os
import sys
import numpy as np

def write_tiff_chunk(data, fname, k, lchunk):
    st = k*lchunk
    end = min((k+1)*lchunk, data.shape[0])
    for k in range(st, end):
        tifffile.imwrite(f'{fname}/d_{k:05}.tiff', data[k])


def write_tiff(data, fname, nthread=8):
    lchunk = int(np.ceil(data.shape[0]/nthread))
    threads = []
    for k in range(nthread):
        write_thread = threading.Thread(
            target=write_tiff_chunk, args=(data, fname, k, lchunk))
        threads.append(write_thread)
        write_thread.start()
    for thread in threads:
        thread.join()


def read_tiff_chunk(data, fname, k, lchunk):
    st = k*lchunk
    end = min((k+1)*lchunk, data.shape[0])
    for k in range(st, end):
        data[k] = tifffile.imread(f'{fname}/d_{k:05}.tiff')


def read_tiff(data, fname, nthread=8):
    lchunk = int(np.ceil(data.shape[0]/nthread))
    threads = []
    for k in range(nthread):
        read_thread = threading.Thread(
            target=read_tiff_chunk, args=(data, fname, k, lchunk))
        threads.append(read_thread)
        read_thread.start()
    for thread in threads:
        thread.join()
    return data


n = int(sys.argv[1])
nthreads = int(sys.argv[2])
path = sys.argv[3]

print('init random data')
data = np.random.random([n, n, n])
os.system(f'mkdir -p {path}')

for k in range(3):
    os.system(f'rm -rf {path}/*')
    print(
        f'write [{n},{n},{n}] {data.dtype} dataset to {path} in {nthreads} nthreads. Test {k}.')
    t = time.time()
    write_tiff(data, path,nthreads)
    tt = time.time()
    print(
        f'write time {(tt-t):.3}s, write speed {data.nbytes/(tt-t)/1024/1024/1024}GB/s')
print("")

for k in range(3):
    print(
        f'read [{n},{n},{n}] {data.dtype} dataset from {path} in {nthreads} nthreads. Test {k}.')
    t = time.time()
    data = read_tiff(data, path, nthreads)
    tt = time.time()
    print(
        f'read time {(tt-t):.3}s, read speed {data.nbytes/(tt-t)/1024/1024/1024}GB/s')