• Home
• com.linkedin.dagli
• core
• 15.0.0-beta7
• source code
• MultithreadedDAGExecutor.java

×

Project download

Many resources are needed to download a project. Please understand that we have to compensate our server costs. Thank you in advance.
Project price only 1 $

You can buy this project and download/modify it how often you want.

com.linkedin.dagli.dag.MultithreadedDAGExecutor Maven / Gradle / Ivy

package com.linkedin.dagli.dag;

import com.concurrentli.AtomicWriteOnceReference;
import com.concurrentli.ExclusiveIdempotentMethod;
import com.concurrentli.UncheckedInterruptedException;
import com.concurrentli.UnsafeCircularIntegerBuffer;
import com.concurrentli.UnsafeCircularReferenceBuffer;
import com.linkedin.dagli.generator.Constant;
import com.linkedin.dagli.generator.Generator;
import com.linkedin.dagli.objectio.ConcatenatedReader;
import com.linkedin.dagli.objectio.ObjectIterator;
import com.linkedin.dagli.objectio.ObjectReader;
import com.linkedin.dagli.objectio.ObjectWriter;
import com.linkedin.dagli.preparer.Preparer;
import com.linkedin.dagli.preparer.PreparerContext;
import com.linkedin.dagli.preparer.PreparerMode;
import com.linkedin.dagli.preparer.PreparerResult;
import com.linkedin.dagli.preparer.PreparerResultMixed;
import com.linkedin.dagli.producer.Producer;
import com.linkedin.dagli.transformer.PreparableTransformer;
import com.linkedin.dagli.transformer.PreparedTransformer;
import com.linkedin.dagli.transformer.Transformer;
import com.linkedin.dagli.tuple.Tuple2;
import com.linkedin.dagli.view.TransformerView;
import it.unimi.dsi.fastutil.ints.Int2IntOpenHashMap;
import java.util.ArrayDeque;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Semaphore;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicLong;
import java.util.concurrent.locks.ReentrantLock;
import java.util.function.Function;
import java.util.function.LongFunction;
import java.util.stream.Collectors;
import java.util.stream.Stream;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;


/**
 * The multithreaded DAG executor can be thought of as executing a DAG by doing the following:
 * (1) Split the input data into "batches" of constant size (except possibly the last batch, if you want to be pedantic)
 * (2) Multithreadedly process these batches across the DAG.  Parallelism includes both executing different nodes
 *     at once and executing the same node on multiple batches at once (where possible).
 * (3) {@link MultithreadedDAGExecutor} streams intermediate values between nodes rather than storing them
 *     (circular buffers are used for this purpose) where possible; when such values must be stored, the location is
 *     configurable, and may be heap (fastest, most memory intensive) or disk.
 *
 * {@link MultithreadedDAGExecutor} is generally used to train DAGs, and {@link FastPreparedDAGExecutor} used to do
 * inference, due to {@link FastPreparedDAGExecutor}'s lower overhead.  However, in rare cases where the DAG is large
 * and the nodes expensive to execute, it may be faster to execute the DAG via {@link MultithreadedDAGExecutor}.
 * Otherwise, use {@link LocalDAGExecutor}, which will perform training with {@link MultithreadedDAGExecutor} and
 * inference with {@link FastPreparedDAGExecutor} without having to manually specify.
 */
public class MultithreadedDAGExecutor extends AbstractDAGExecutor implements DAGExecutor {
  private static final long serialVersionUID = 1L;

  private static final int DEFAULT_THREAD_COUNT = 2 * Runtime.getRuntime().availableProcessors();

  /**
   * The default batch size used by this executor.  Larger batches reduce the per-batch overhead, but also reduce
   * opportunities for parallelization.  In general, larger batches are better when execution of the DAG is inexpensive.
   */
  public static final int DEFAULT_BATCH_SIZE = 5000;

  /**
   * During execution, each (non-root) node in the DAG has a circular buffer that caches its inputs.  This is the
   * default maximum number of batches that will be held in that buffer (in RAM) for each node.  Higher values can
   * improve parallelization at the expense of RAM.
   * machine.
   */
  public static final int DEFAULT_MAX_CONCURRENT_BATCHES = DEFAULT_THREAD_COUNT;

  private final int _batchSize;
  private final int _maxConcurrentBatches;
  private final int _maxThreadCount;
  private final LocalStorage _localStorage;

  /**
   * Sets the batch size.  Larger batches reduce the per-batch overhead, but also reduce opportunities for
   * parallelization.  In general, larger batches are better when execution of the DAG is inexpensive.
   *
   * @param batchSize the batch size to set
   * @return a copy of this executor with the specified batch size
   */
  public MultithreadedDAGExecutor withBatchSize(int batchSize) {
    return new MultithreadedDAGExecutor(batchSize, _maxConcurrentBatches, _maxThreadCount, _localStorage);
  }

  /**
   * Sets the maximum number of concurrent batches to queue per node.  During execution, each (non-root) node in the DAG
   * has a circular buffer that caches its inputs (in RAM).  Higher values can improve parallelization at the expense of
   * memory utilization.  By default, this is twice the number of logical processors on the machine.
   *
   * @param maxConcurrentBatches sets the maximum number of concurrent batches to be used by the executor
   * @return a copy of this executor that will use the specified number of concurrent batches
   */
  public MultithreadedDAGExecutor withConcurrentBatches(int maxConcurrentBatches) {
    return new MultithreadedDAGExecutor(_batchSize, maxConcurrentBatches, _maxThreadCount, _localStorage);
  }

  /**
   * Sets the maximum number of threads used by the executor to process nodes in the DAG.  Note that, because values are
   * processed in batches, a single node will often be executed in parallel with itself (on different batches of
   * inputs).
   *
   * The default value is twice the number of logical processors on the machine.  This is because some nodes may, e.g.
   * block for IO or locks and thus a surplus of threads can help improve parallelism.
   *
   * Note that this is the number of threads used by the executor.  Individual nodes are free to (and often do)
   * implement their own multithreading, which may result in a greater number of total concurrent threads in your
   * program.
   *
   * @param maxThreadCount the maximum number of threads to use
   * @return a copy of this executor that will use the specified maximum number of threads
   */
  public MultithreadedDAGExecutor withMaxThreads(int maxThreadCount) {
    return new MultithreadedDAGExecutor(_batchSize, _maxConcurrentBatches, maxThreadCount, _localStorage);
  }

  /**
   * Sets the storage used by the executor to store intermediate values and results.
   *
   * Although {@link MultithreadedDAGExecutor} avoids storing values wherever possible (preferring to stream data
   * between nodes), in some cases--for example, when a preparable node must be prepared (trained) and then applied
   * (inference) to get the result--a node must see its inputs a second time, which requires that the values be stored
   * so they can be read back.
   *
   * Storing on the heap (in RAM) is of course the fastest, but may exhaust machine replaceholders when the amount of
   * intermediate values and outputs being stored is large.  For such problems, disk-backed storage (e.g. via Kryo) is
   * a more scalable solution.
   *
   * @param storage the storage method to use for intermediate and outputs that must be cached by the executor
   * @return a copy of this executor that will use the specified storage method
   */
  public MultithreadedDAGExecutor withStorage(LocalStorage storage) {
    return new MultithreadedDAGExecutor(_batchSize, _maxConcurrentBatches, _maxThreadCount, storage);
  }

  private LongFunction> getAppendableGenerator() {
    return _localStorage == null ? LocalStorage.MEMORY_HEAP._objectWriterGenerator
        : _localStorage._objectWriterGenerator;
  }

  /**
   * Creates a new executor with the default batch size, maximum concurrent batches to queue for each node, threads and
   * storage mechanism for intermediate and output values.
   */
  public MultithreadedDAGExecutor() {
    this(DEFAULT_BATCH_SIZE, DEFAULT_MAX_CONCURRENT_BATCHES, DEFAULT_THREAD_COUNT, LocalStorage.MEMORY_HEAP);
  }

  @Override
  public String toString() {
    return "MultithreadedDAGExecutor (batch size: " + _batchSize + "; concurrent batches: " + _maxConcurrentBatches
        + "; max threads: " + _maxThreadCount + ")";
  }

  private MultithreadedDAGExecutor(int batchSize, int maxConcurrentBatches, int maxThreadCount,
      LocalStorage storage) {
    _batchSize = batchSize;
    _maxConcurrentBatches = maxConcurrentBatches;
    _maxThreadCount = maxThreadCount;
    _localStorage = storage;
  }

  /**
   * How Scheduling Works
   *
   * 0. Nodes track the state of a given producer (input and output buffers) and schedule tasks
   * 1. Supplier tasks read values into buffers
   * 2. Generator tasks generate values into buffers
   * 3. Preparable tasks use inputs to prepare the transformer
   * 4. Prepared tasks transform their inputs to produce results
   *
   * If the result of a producer is an output or is fed to a later phase, it must be written to
   * a ObjectWriter so that it may be re-read in the future.
   */
  private static class Scheduler {
    private static final Logger LOGGER = LogManager.getLogger();

    public final ExecutorService _threadPool;
    public final MultithreadedDAGExecutor _executor;
    public final DAGStructure _dag;
    public final long _count;
    public final int _batchSize;
    public final long _batchCount;

    private Object _threadExceptionMutex = new Object();
    private Exception _threadException = null;

    private static class ReducableSemaphore extends Semaphore {
      public ReducableSemaphore(int permits) {
        super(permits);
      }

      @Override
      public void reducePermits(int reduction) {
        super.reducePermits(reduction);
      }
    }

    public final ReducableSemaphore _pendingTaskSemaphore = new ReducableSemaphore(0);

    private final ObjectReader