tests async repack

repack_trial.py

#!/usr/bin/env python

import os
import shutil
import re
import time
import gevent
import h5py
import numpy as np
from coverage_model.utils import prod
from coverage_model.hdf_utils import repack
from coverage_model.threads.sync import AsyncDispatcher


mtch = r'd_\d{2}.hdf5'
bp = 'test_data/repack_test'
oname = 'd.hdf5'
opth = os.path.join(bp, oname)
buname = '{0}_bu'.format(oname)
bupth = os.path.join(bp, buname)


def make_test_file():
    if not os.path.exists(bp):
        os.makedirs(bp)

    if not os.path.exists(opth):
        shp = (5000, 2500)
        with h5py.File(opth, 'a') as f:
            ds = f.require_dataset('test_ds', shape=shp, dtype='float32', chunks=None)
            ds[:] = np.arange(prod(shp)).reshape(*shp)


def make_copies(rep=5):
    make_test_file()

    shutil.copy(opth, bupth)

    for n in xrange(rep):
        shutil.copy(opth, os.path.join(bp, 'd_{0}.hdf5'.format('%02d' % n)))


def cleanup():
    for pth in [os.path.join(bp, p) for p in os.listdir(bp) if re.match(mtch, p) is not None]:
        os.remove(pth)

    os.remove(os.path.join(bp, buname))


def do_repack():
    make_copies()

    for pth in [os.path.join(bp, p) for p in os.listdir(bp) if re.match(mtch, p) is not None]:
        repack(pth)

    cleanup()


def try_repack_async():
    make_test_file()

    t1 = time.time()
    with AsyncDispatcher(do_repack) as dispatcher:
        v = dispatcher.wait(10)

    t1e = time.time() - t1

    t2 = time.time()
    with AsyncDispatcher(do_repack) as dispatcher:
        gevent.sleep(7)
        v = dispatcher.wait(10)

    t2e = time.time() - t2

    print 'Time for repack alone: %s', t1e
    print 'Time for repack w/ 7s sleep: %s', t2e