<commit_before>package org.deeplearning4j.parallelism;

import lombok.NonNull;
import org.deeplearning4j.api.storage.listener.RoutingIterationListener;
import org.deeplearning4j.datasets.iterator.AsyncDataSetIterator;
import org.deeplearning4j.datasets.iterator.AsyncMultiDataSetIterator;
import org.deeplearning4j.nn.api.Model;
import org.deeplearning4j.nn.api.Updater;
import org.deeplearning4j.nn.conf.MultiLayerConfiguration;
import org.deeplearning4j.nn.graph.ComputationGraph;
import org.deeplearning4j.nn.multilayer.MultiLayerNetwork;
import org.deeplearning4j.nn.updater.graph.ComputationGraphUpdater;
import org.deeplearning4j.optimize.api.IterationListener;
import org.nd4j.linalg.api.ndarray.INDArray;
import org.nd4j.linalg.api.ops.executioner.GridExecutioner;
import org.nd4j.linalg.dataset.api.DataSet;
import org.nd4j.linalg.dataset.api.MultiDataSet;
import org.nd4j.linalg.dataset.api.iterator.DataSetIterator;
import org.nd4j.linalg.dataset.api.iterator.MultiDataSetIterator;
import org.nd4j.linalg.factory.Nd4j;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.ArrayList;
import java.util.Collection;
import java.util.List;
import java.util.UUID;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicLong;

/**
 * This is simple data-parallel wrapper suitable for multi-cpu/multi-gpu environments.
 *
 * @author <EMAIL>
 */
public class ParallelWrapper implements AutoCloseable {
    private static Logger logger = LoggerFactory.getLogger(ParallelWrapper.class);
    private Model model;
    private int workers = 2;
    private int prefetchSize = 2;
    private int averagingFrequency = 1;
    private Trainer zoo[];
    private AtomicLong iterationsCounter = new AtomicLong(0);
    private boolean reportScore = false;
    private boolean averageUpdaters = true;
    private boolean legacyAveraging = false;
    private boolean wasAveraged = false;
    private AtomicBoolean stopFit = new AtomicBoolean(false);

    protected ParallelWrapper(Model model, int workers, int prefetchSize) {
        this.model = model;
        this.workers = workers;
        this.prefetchSize = prefetchSize;

        if (this.model instanceof MultiLayerNetwork) {
            ((MultiLayerNetwork) this.model).getUpdater();
        } else if (this.model instanceof ComputationGraph) {
            ((ComputationGraph) this.model).getUpdater();
        }

        zoo = new Trainer[workers];
        for (int cnt = 0; cnt < workers; cnt++) {
            zoo[cnt] = new Trainer(cnt, model);
            zoo[cnt].start();
        }
    }

    @Override
    public void close() throws Exception {
        if (zoo != null) {
            for (int i = 0; i < zoo.length; i++) {
                if (zoo[i] != null)
                    zoo[i].shutdown();
            }
            zoo = null;
        }
    }

    /**
     * This method causes all threads used for parallel training to stop
     */
    public synchronized void shutdown() {
        try {
            close();
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    /**
    * Will stop a fit operation from continuing to iterate.
    */
    public void stopFit() { stopFit.set(true); }

    public synchronized void fit(@NonNull MultiDataSetIterator source) {
        stopFit.set(false);
        if (zoo == null) {
            zoo = new Trainer[workers];
            for (int cnt = 0; cnt < workers; cnt++) {
                // we pass true here, to tell Trainer to use MultiDataSet queue for training
                zoo[cnt] = new Trainer(cnt, model, true);
                zoo[cnt].start();
            }
        } else {
            for (int cnt = 0; cnt < workers; cnt++) {
                zoo[cnt].useMDS = true;
            }
        }
        source.reset();

        MultiDataSetIterator iterator;
        if (prefetchSize > 0 && source.asyncSupported()) {
            iterator = new AsyncMultiDataSetIterator(source, prefetchSize);
        } else iterator = source;

        AtomicInteger locker = new AtomicInteger(0);

        while (iterator.hasNext() && !stopFit.get()) {
            MultiDataSet dataSet = iterator.next();

            /*
             now dataSet should be dispatched to next free workers, until all workers are busy. And then we should block till all finished.
            */
            int pos = locker.getAndIncrement();
            zoo[pos].feedMultiDataSet(dataSet);

            /*
                if all workers are dispatched now, join till all are finished
            */
            if (pos + 1 == workers || !iterator.hasNext()) {
                iterationsCounter.incrementAndGet();

                for (int cnt = 0; cnt < workers && cnt < locker.get(); cnt ++) {
                    try {
                        zoo[cnt].waitTillRunning();
                    } catch (Exception e) {
                        throw new RuntimeException(e);
                    }
                }

                /*
                    average model, and propagate it to whole
                */
                if (iterationsCounter.get() % averagingFrequency == 0 && pos + 1 == workers) {
                    wasAveraged = true;
                    double score = 0.0;
                    if (!legacyAveraging || Nd4j.getAffinityManager().getNumberOfDevices() == 1) {
                        List<INDArray> params = new ArrayList<>();
                        for (int cnt = 0; cnt < workers && cnt < locker.get(); cnt++) {
                            params.add(zoo[cnt].getModel().params());
                            score += zoo[cnt].getModel().score();
                        }
                        Nd4j.averageAndPropagate(model.params(), params);
                    } else {
                        INDArray params = Nd4j.zeros(model.params().shape());
                        int cnt = 0;
                        for (; cnt < workers && cnt < locker.get(); cnt++) {
                            params.addi(zoo[cnt].getModel().params());
                            score += zoo[cnt].getModel().score();
                        }

                        params.divi(cnt);
                        model.setParams(params);
                    }

                    score /= Math.min(workers, locker.get());

                    // TODO: improve this
                    if (reportScore)
                        logger.info("Averaged score: " + score);

                    // averaging updaters state
                    if (model instanceof ComputationGraph) {
                        if (averageUpdaters) {
                            ComputationGraphUpdater updater = ((ComputationGraph) model).getUpdater();
                            int batchSize = 0;

                            if (updater != null && updater.getStateViewArray() != null) {
                                if (!legacyAveraging || Nd4j.getAffinityManager().getNumberOfDevices() == 1) {
                                    List<INDArray> updaters = new ArrayList<>();
                                    for (int cnt = 0; cnt < workers && cnt < locker.get(); cnt++) {
                                        ComputationGraph workerModel = (ComputationGraph) zoo[cnt].getModel();
                                        updaters.add(workerModel.getUpdater().getStateViewArray());
                                        batchSize += workerModel.batchSize();
                                    }
                                    Nd4j.averageAndPropagate(updater.getStateViewArray(), updaters);
                                } else {
                                    INDArray state = Nd4j.zeros(updater.getStateViewArray().shape());
                                    int cnt = 0;
                                    for (; cnt < workers && cnt < locker.get(); cnt++) {
                                        ComputationGraph workerModel = (ComputationGraph) zoo[cnt].getModel();
                                        state.addi(workerModel.getUpdater().getStateViewArray());
                                        batchSize += workerModel.batchSize();
                                    }
                                    state.divi(cnt);
                                    updater.setStateViewArray(state);
                                }
                            }
                        }

                        ((ComputationGraph) model).setScore(score);
                    } else throw new RuntimeException("MultiDataSet must only be used with ComputationGraph model");

                    if (legacyAveraging &&  Nd4j.getAffinityManager().getNumberOfDevices() > 1) {
                        for (int cnt = 0; cnt < workers; cnt++) {
                            zoo[cnt].updateModel(model);
                        }
                    }
                }
                locker.set(0);
            }
        }

        // sanity checks, or the dataset may never average
        if(!wasAveraged)
            logger.warn("Parameters were never averaged on current fit(). Ratios of batch size, num workers, and averaging frequency may be responsible.");
//            throw new IllegalStateException("Parameters were never averaged. Please check batch size ratios, number of workers, and your averaging frequency.");

        logger.debug("Iterations passed: {}", iterationsCounter.get());
//        iterationsCounter.set(0);
    }

    /**
     * This method takes DataSetIterator, and starts training over it by scheduling DataSets to different executors
     *
     * @param source
     */
    public synchronized void fit(@NonNull DataSetIterator source) {
        stopFit.set(false);
        if (zoo == null) {
            zoo = new Trainer[workers];
            for (int cnt = 0; cnt < workers; cnt++) {
                zoo[cnt] = new Trainer(cnt, model);
                zoo[cnt].start();
            }
        }<commit_msg>Thread UncaughtExceptionHandler<commit_after>package org.deeplearning4j.parallelism;

import lombok.NonNull;
import org.deeplearning4j.api.storage.listener.RoutingIterationListener;
import org.deeplearning4j.datasets.iterator.AsyncDataSetIterator;
import org.deeplearning4j.datasets.iterator.AsyncMultiDataSetIterator;
import org.deeplearning4j.nn.api.Model;
import org.deeplearning4j.nn.api.Updater;
import org.deeplearning4j.nn.conf.MultiLayerConfiguration;
import org.deeplearning4j.nn.graph.ComputationGraph;
import org.deeplearning4j.nn.multilayer.MultiLayerNetwork;
import org.deeplearning4j.nn.updater.graph.ComputationGraphUpdater;
import org.deeplearning4j.optimize.api.IterationListener;
import org.nd4j.linalg.api.ndarray.INDArray;
import org.nd4j.linalg.api.ops.executioner.GridExecutioner;
import org.nd4j.linalg.dataset.api.DataSet;
import org.nd4j.linalg.dataset.api.MultiDataSet;
import org.nd4j.linalg.dataset.api.iterator.DataSetIterator;
import org.nd4j.linalg.dataset.api.iterator.MultiDataSetIterator;
import org.nd4j.linalg.factory.Nd4j;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.ArrayList;
import java.util.Collection;
import java.util.List;
import java.util.UUID;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicLong;

/**
 * This is simple data-parallel wrapper suitable for multi-cpu/multi-gpu environments.
 *
 * @author <EMAIL>
 */
public class ParallelWrapper implements AutoCloseable {
    private static Logger logger = LoggerFactory.getLogger(ParallelWrapper.class);
    private Model model;
    private int workers = 2;
    private int prefetchSize = 2;
    private int averagingFrequency = 1;
    private Trainer zoo[];
    private AtomicLong iterationsCounter = new AtomicLong(0);
    private boolean reportScore = false;
    private boolean averageUpdaters = true;
    private boolean legacyAveraging = false;
    private boolean wasAveraged = false;
    private AtomicBoolean stopFit = new AtomicBoolean(false);

    // log uncaught exceptions
    Thread.UncaughtExceptionHandler handler = new Thread.UncaughtExceptionHandler() {
        public void uncaughtException(Thread th, Throwable ex) {
            logger.error("Uncaught exception: " + ex);
        }
    };

    protected ParallelWrapper(Model model, int workers, int prefetchSize) {
        this.model = model;
        this.workers = workers;
        this.prefetchSize = prefetchSize;

        if (this.model instanceof MultiLayerNetwork) {
            ((MultiLayerNetwork) this.model).getUpdater();
        } else if (this.model instanceof ComputationGraph) {
            ((ComputationGraph) this.model).getUpdater();
        }

        zoo = new Trainer[workers];
        for (int cnt = 0; cnt < workers; cnt++) {
            zoo[cnt] = new Trainer(cnt, model);
            zoo[cnt].setUncaughtExceptionHandler(handler);
            zoo[cnt].start();
        }
    }

    @Override
    public void close() throws Exception {
        if (zoo != null) {
            for (int i = 0; i < zoo.length; i++) {
                if (zoo[i] != null)
                    zoo[i].shutdown();
            }
            zoo = null;
        }
    }

    /**
     * This method causes all threads used for parallel training to stop
     */
    public synchronized void shutdown() {
        try {
            close();
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    /**
    * Will stop a fit operation from continuing to iterate.
    */
    public void stopFit() { stopFit.set(true); }

    public synchronized void fit(@NonNull MultiDataSetIterator source) {
        stopFit.set(false);
        if (zoo == null) {
            zoo = new Trainer[workers];
            for (int cnt = 0; cnt < workers; cnt++) {
                // we pass true here, to tell Trainer to use MultiDataSet queue for training
                zoo[cnt] = new Trainer(cnt, model, true);
                zoo[cnt].setUncaughtExceptionHandler(handler);
                zoo[cnt].start();
            }
        } else {
            for (int cnt = 0; cnt < workers; cnt++) {
                zoo[cnt].useMDS = true;
            }
        }
        source.reset();

        MultiDataSetIterator iterator;
        if (prefetchSize > 0 && source.asyncSupported()) {
            iterator = new AsyncMultiDataSetIterator(source, prefetchSize);
        } else iterator = source;

        AtomicInteger locker = new AtomicInteger(0);

        while (iterator.hasNext() && !stopFit.get()) {
            MultiDataSet dataSet = iterator.next();

            /*
             now dataSet should be dispatched to next free workers, until all workers are busy. And then we should block till all finished.
            */
            int pos = locker.getAndIncrement();
            zoo[pos].feedMultiDataSet(dataSet);

            /*
                if all workers are dispatched now, join till all are finished
            */
            if (pos + 1 == workers || !iterator.hasNext()) {
                iterationsCounter.incrementAndGet();

                for (int cnt = 0; cnt < workers && cnt < locker.get(); cnt ++) {
                    try {
                        zoo[cnt].waitTillRunning();
                    } catch (Exception e) {
                        throw new RuntimeException(e);
                    }
                }

                /*
                    average model, and propagate it to whole
                */
                if (iterationsCounter.get() % averagingFrequency == 0 && pos + 1 == workers) {
                    wasAveraged = true;
                    double score = 0.0;
                    if (!legacyAveraging || Nd4j.getAffinityManager().getNumberOfDevices() == 1) {
                        List<INDArray> params = new ArrayList<>();
                        for (int cnt = 0; cnt < workers && cnt < locker.get(); cnt++) {
                            params.add(zoo[cnt].getModel().params());
                            score += zoo[cnt].getModel().score();
                        }
                        Nd4j.averageAndPropagate(model.params(), params);
                    } else {
                        INDArray params = Nd4j.zeros(model.params().shape());
                        int cnt = 0;
                        for (; cnt < workers && cnt < locker.get(); cnt++) {
                            params.addi(zoo[cnt].getModel().params());
                            score += zoo[cnt].getModel().score();
                        }

                        params.divi(cnt);
                        model.setParams(params);
                    }

                    score /= Math.min(workers, locker.get());

                    // TODO: improve this
                    if (reportScore)
                        logger.info("Averaged score: " + score);

                    // averaging updaters state
                    if (model instanceof ComputationGraph) {
                        if (averageUpdaters) {
                            ComputationGraphUpdater updater = ((ComputationGraph) model).getUpdater();
                            int batchSize = 0;

                            if (updater != null && updater.getStateViewArray() != null) {
                                if (!legacyAveraging || Nd4j.getAffinityManager().getNumberOfDevices() == 1) {
                                    List<INDArray> updaters = new ArrayList<>();
                                    for (int cnt = 0; cnt < workers && cnt < locker.get(); cnt++) {
                                        ComputationGraph workerModel = (ComputationGraph) zoo[cnt].getModel();
                                        updaters.add(workerModel.getUpdater().getStateViewArray());
                                        batchSize += workerModel.batchSize();
                                    }
                                    Nd4j.averageAndPropagate(updater.getStateViewArray(), updaters);
                                } else {
                                    INDArray state = Nd4j.zeros(updater.getStateViewArray().shape());
                                    int cnt = 0;
                                    for (; cnt < workers && cnt < locker.get(); cnt++) {
                                        ComputationGraph workerModel = (ComputationGraph) zoo[cnt].getModel();
                                        state.addi(workerModel.getUpdater().getStateViewArray());
                                        batchSize += workerModel.batchSize();
                                    }
                                    state.divi(cnt);
                                    updater.setStateViewArray(state);
                                }
                            }
                        }

                        ((ComputationGraph) model).setScore(score);
                    } else throw new RuntimeException("MultiDataSet must only be used with ComputationGraph model");

                    if (legacyAveraging &&  Nd4j.getAffinityManager().getNumberOfDevices() > 1) {
                        for (int cnt = 0; cnt < workers; cnt++) {
                            zoo[cnt].updateModel(model);
                        }
                    }
                }
                locker.set(0);
            }
        }

        // sanity checks, or the dataset may never average
        if(!wasAveraged)
            logger.warn("Parameters were never averaged on current fit(). Ratios of batch size, num workers, and averaging frequency may be responsible.");
//            throw new IllegalStateException("Parameters were never averaged. Please check batch size ratios, number of workers, and your averaging frequency.");

        logger.debug("Iterations passed: {}", iterationsCounter.get());
//        iterationsCounter.set(0);
    }

    /**
     * This method takes DataSetIterator, and starts training over it by scheduling DataSets to different executors
     *
     * @param source
     */
    public synchronized void fit(@NonNull DataSetIterator source) {
        stopFit.set(false);
        if (zoo == null) {
            zoo = new Trainer[workers];
            for (int cnt = 0; cnt < workers; cnt++) {
                zoo[cnt] = new Trainer(cnt, model);
                zoo[cnt].setUncaughtExceptionHandler(handler);
                zoo[cnt].start();
            }
        }
{'fraction_non_alphanumeric': 0.09976719371423029, 'fraction_numerical': 0.005432146667960035, 'mean_word_length': 1.638050153531218, 'pattern_counts': {'":': 0, '<': 45, '<?xml version=': 0, '>': 17, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 38, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6083721', 'n_tokens_mistral': 5356, 'n_tokens_neox': 5145, 'n_words': 1248}