Source:
public void runMiniBatch(final Collection<NMFData> miniBatchTrainingData) { LOG.log(Level.INFO, "numTrainerTheads: {0}", numTrainerThreads); final CountDownLatch latch = new CountDownLatch(numTrainerThreads); final BlockingQueue<NMFData> instances = new ArrayBlockingQueue<>(miniBatchTrainingData.size()); instances.addAll(miniBatchTrainingData); final List<Integer> keys = getKeys(instances); final NMFModel model = pullModels(keys); final List<Future<Map<Integer, Vector>>> futures = new ArrayList<>(numTrainerThreads); try { final int drainSize = Math.max(instances.size() / numTrainerThreads / numTrainerThreads, 1); for (int threadIdx = 0; threadIdx < numTrainerThreads; threadIdx++) { final Future<Map<Integer, Vector>> future = executor.submit(() -> { final List<NMFData> drainedInstances = new ArrayList<>(drainSize); final Map<Integer, Vector> threadRGradient = new HashMap<>(); int count = 0; while (true) { final int numDrained = instances.drainTo(drainedInstances, drainSize); if (numDrained == 0) { break; } drainedInstances.forEach(instance -> updateGradient(instance, model, threadRGradient)); drainedInstances.clear(); count += numDrained; } latch.countDown(); LOG.log(Level.INFO, "{0} has computed {1} instances", new Object[] {Thread.currentThread().getName(), count}); return threadRGradient; }); futures.add(future); } latch.await(); } catch (final InterruptedException e) { LOG.log(Level.SEVERE, "Exception occurred.", e); throw new RuntimeException(e); } final List<Map<Integer, Vector>> totalRGradients = ThreadUtils.retrieveResults(futures); final Map<Integer, Vector> gradients = aggregateGradient(totalRGradients); pushAndResetGradients(gradients); }

Target:
public void runMiniBatch(final Collection<NMFData> miniBatchTrainingData) { final CountDownLatch latch = new CountDownLatch(numTrainerThreads); final BlockingQueue<NMFData> instances = new ArrayBlockingQueue<>(miniBatchTrainingData.size()); instances.addAll(miniBatchTrainingData); final List<Integer> keys = getKeys(instances); final NMFModel model = pullModels(keys); final List<Future<Map<Integer, Vector>>> futures = new ArrayList<>(numTrainerThreads); try { final int drainSize = Math.max(instances.size() / numTrainerThreads / numTrainerThreads, 1); for (int threadIdx = 0; threadIdx < numTrainerThreads; threadIdx++) { final Future<Map<Integer, Vector>> future = executor.submit(() -> { final List<NMFData> drainedInstances = new ArrayList<>(drainSize); final Map<Integer, Vector> threadRGradient = new HashMap<>(); int count = 0; while (true) { final int numDrained = instances.drainTo(drainedInstances, drainSize); if (numDrained == 0) { break; } drainedInstances.forEach(instance -> updateGradient(instance, model, threadRGradient)); drainedInstances.clear(); count += numDrained; } latch.countDown(); LOG.log(Level.INFO, "{0} has computed {1} instances", new Object[] {Thread.currentThread().getName(), count}); return threadRGradient; }); futures.add(future); } latch.await(); } catch (final InterruptedException e) { LOG.log(Level.SEVERE, "Exception occurred.", e); throw new RuntimeException(e); } final List<Map<Integer, Vector>> totalRGradients = ThreadUtils.retrieveResults(futures); final Map<Integer, Vector> gradients = aggregateGradient(totalRGradients); pushAndResetGradients(gradients); }