Prefetch shards for DiskDataset

import tempfile

import time

import shutil

from multiprocessing.dummy import Pool

__author__ = "Bharath Ramsundar"

__copyright__ = "Copyright 2016, Stanford University"

        shard_perm = np.random.permutation(num_shards)

      else:

        shard_perm = np.arange(num_shards)

      pool = Pool(1)

      next_shard = pool.apply_async(dataset.get_shard, (shard_perm[0],))

      for i in range(num_shards):

        X, y, w, ids = dataset.get_shard(shard_perm[i])

        X, y, w, ids = next_shard.get()

        if i < num_shards - 1:

          next_shard = pool.apply_async(dataset.get_shard, (shard_perm[i + 1],))

        n_samples = X.shape[0]

        # TODO(rbharath): This happens in tests sometimes, but don't understand why?

        # Handle edge case.