[NEMO-332] Refactor RunTimePass

client/src/main/java/org/apache/nemo/client/JobLauncher.java

@@ -21,12 +21,14 @@
 import com.google.common.annotations.VisibleForTesting;
 import com.google.protobuf.ByteString;
 import org.apache.nemo.common.ir.IRDAG;
+import org.apache.nemo.compiler.backend.nemo.NemoPlanRewriter;
 import org.apache.nemo.conf.JobConf;
 import org.apache.nemo.driver.NemoDriver;
 import org.apache.nemo.runtime.common.comm.ControlMessage;
 import org.apache.nemo.runtime.common.message.MessageEnvironment;
 import org.apache.nemo.runtime.common.message.MessageParameters;
 import org.apache.commons.lang3.SerializationUtils;
+import org.apache.nemo.runtime.common.plan.PlanRewriter;
 import org.apache.nemo.runtime.master.scheduler.Scheduler;
 import org.apache.reef.client.DriverConfiguration;
 import org.apache.reef.client.DriverLauncher;
@@ -320,6 +322,7 @@ private static Configuration getSchedulerConf(final Configuration jobConf)
     final JavaConfigurationBuilder jcb = Tang.Factory.getTang().newConfigurationBuilder();
     final Class schedulerImpl = ((Class<Scheduler>) Class.forName(classImplName));
     jcb.bindImplementation(Scheduler.class, schedulerImpl);
+    jcb.bindImplementation(PlanRewriter.class, NemoPlanRewriter.class);
     return jcb.build();
   }
 

common/src/main/java/org/apache/nemo/common/HashRange.java

@@ -24,24 +24,21 @@
  * Descriptor for hash range.
  */
 public final class HashRange implements KeyRange<Integer> {
-  private static final HashRange ALL = new HashRange(0, Integer.MAX_VALUE, false);
+  private static final HashRange ALL = new HashRange(0, Integer.MAX_VALUE);
   private final int rangeBeginInclusive;
   private final int rangeEndExclusive;
-  private boolean isSkewed;
 
   /**
    * Private constructor.
    * @param rangeBeginInclusive point at which the hash range starts (inclusive).
    * @param rangeEndExclusive point at which the hash range ends (exclusive).
-   * @param isSkewed whether or not the range is skewed
    */
-  private HashRange(final int rangeBeginInclusive, final int rangeEndExclusive, final boolean isSkewed) {
+  private HashRange(final int rangeBeginInclusive, final int rangeEndExclusive) {
     if (rangeBeginInclusive < 0 || rangeEndExclusive < 0) {
       throw new RuntimeException("Each boundary value of the range have to be non-negative.");
     }
     this.rangeBeginInclusive = rangeBeginInclusive;
     this.rangeEndExclusive = rangeEndExclusive;
-    this.isSkewed = isSkewed;
   }
 
   /**
@@ -54,11 +51,10 @@ public static HashRange all() {
   /**
    * @param rangeStartInclusive the start of the range (inclusive)
    * @param rangeEndExclusive   the end of the range (exclusive)
-   * @param isSkewed  whether or not the range is skewed
    * @return A hash range descriptor representing [{@code rangeBeginInclusive}, {@code rangeEndExclusive})
    */
-  public static HashRange of(final int rangeStartInclusive, final int rangeEndExclusive, final boolean isSkewed) {
-    return new HashRange(rangeStartInclusive, rangeEndExclusive, isSkewed);
+  public static HashRange of(final int rangeStartInclusive, final int rangeEndExclusive) {
+    return new HashRange(rangeStartInclusive, rangeEndExclusive);
   }
 
   /**
@@ -114,8 +110,7 @@ public boolean equals(final Object o) {
     }
     final HashRange hashRange = (HashRange) o;
     if (rangeBeginInclusive != hashRange.rangeBeginInclusive
-        || rangeEndExclusive != hashRange.rangeEndExclusive
-        || isSkewed != hashRange.isSkewed) {
+      || rangeEndExclusive != hashRange.rangeEndExclusive) {
       return false;
     }
     return true;
@@ -127,16 +122,8 @@ public boolean equals(final Object o) {
   @Override
   public int hashCode() {
     return Arrays.hashCode(new Object[] {
-        rangeBeginInclusive,
-        rangeEndExclusive,
-        isSkewed,
+      rangeBeginInclusive,
+      rangeEndExclusive,
     });
   }
-
-  /**
-   * @return whether or not the range is skewed.
-   */
-  public boolean isSkewed() {
-    return isSkewed;
-  }
 }

common/src/main/java/org/apache/nemo/common/KeyRange.java

@@ -26,7 +26,7 @@
  */
 public interface KeyRange<K extends Serializable> extends Serializable {
 
- /**
+  /**
    * @return whether this instance represents the entire range or not.
    */
   boolean isAll();

common/src/main/java/org/apache/nemo/common/dag/DAGBuilder.java

@@ -21,7 +21,7 @@
 import org.apache.nemo.common.exception.CompileTimeOptimizationException;
 import org.apache.nemo.common.ir.edge.IREdge;
 import org.apache.nemo.common.ir.edge.executionproperty.DataFlowProperty;
-import org.apache.nemo.common.ir.edge.executionproperty.MetricCollectionProperty;
+import org.apache.nemo.common.ir.edge.executionproperty.MessageIdProperty;
 import org.apache.nemo.common.ir.vertex.*;
 import org.apache.nemo.common.exception.IllegalVertexOperationException;
 import org.apache.nemo.common.ir.vertex.system.MessageAggregatorVertex;
@@ -260,7 +260,7 @@ private void sinkCheck() {
   private void executionPropertyCheck() {
     // DataSizeMetricCollection is not compatible with Push (All data have to be stored before the data collection)
     vertices.forEach(v -> incomingEdges.get(v).stream().filter(e -> e instanceof IREdge).map(e -> (IREdge) e)
-        .filter(e -> e.getPropertyValue(MetricCollectionProperty.class).isPresent())
+        .filter(e -> e.getPropertyValue(MessageIdProperty.class).isPresent())
         .filter(e -> !(e.getDst() instanceof OperatorVertex
           && e.getDst() instanceof MessageAggregatorVertex))
         .filter(e -> DataFlowProperty.Value.Push.equals(e.getPropertyValue(DataFlowProperty.class).get()))

common/src/main/java/org/apache/nemo/common/ir/IRDAG.java

@@ -35,6 +35,7 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import javax.annotation.concurrent.NotThreadSafe;
 import java.util.List;
 import java.util.Set;
 import java.util.concurrent.atomic.AtomicInteger;
@@ -54,6 +55,7 @@
  * - Annotation: setProperty(), getPropertyValue() on each IRVertex/IREdge
  * - Reshaping: insert(), delete() on the IRDAG
  */
+@NotThreadSafe
 public final class IRDAG implements DAGInterface<IRVertex, IREdge> {
   private static final Logger LOG = LoggerFactory.getLogger(IRDAG.class.getName());
 
@@ -207,7 +209,7 @@ public void insert(final MessageBarrierVertex messageBarrierVertex,
           final IREdge edgeToOriginalDst =
             new IREdge(edge.getPropertyValue(CommunicationPatternProperty.class).get(), edge.getSrc(), v);
           edge.copyExecutionPropertiesTo(edgeToOriginalDst);
-          edgeToOriginalDst.setPropertyPermanently(MetricCollectionProperty.of(currentMetricCollectionId));
+          edgeToOriginalDst.setPropertyPermanently(MessageIdProperty.of(currentMetricCollectionId));
           builder.connectVertices(edgeToOriginalDst);
         } else {
           // NO MATCH, so simply connect vertices as before.
@@ -247,7 +249,7 @@ private IREdge edgeBetweenMessageVertices(final MessageBarrierVertex mbv,
     newEdge.setProperty(DataStoreProperty.of(DataStoreProperty.Value.LocalFileStore));
     newEdge.setProperty(DataPersistenceProperty.of(DataPersistenceProperty.Value.Keep));
     newEdge.setProperty(DataFlowProperty.of(DataFlowProperty.Value.Push));
-    newEdge.setPropertyPermanently(MetricCollectionProperty.of(currentMetricCollectionId));
+    newEdge.setPropertyPermanently(MessageIdProperty.of(currentMetricCollectionId));
     final KeyExtractor pairKeyExtractor = (element) -> {
       if (element instanceof Pair) {
         return ((Pair) element).left();

common/src/main/java/org/apache/nemo/common/ir/edge/executionproperty/MetricCollectionProperty.java → common/src/main/java/org/apache/nemo/common/ir/edge/executionproperty/MessageIdProperty.java

@@ -21,14 +21,14 @@
 import org.apache.nemo.common.ir.executionproperty.EdgeExecutionProperty;
 
 /**
- * MetricCollection ExecutionProperty that indicates the edge of which data metric will be collected.
+ * Edges with the same MessageId are subject to the same run-time optimization.
  */
-public final class MetricCollectionProperty extends EdgeExecutionProperty<Integer> {
+public final class MessageIdProperty extends EdgeExecutionProperty<Integer> {
   /**
    * Constructor.
    * @param value value of the execution property.
    */
-  private MetricCollectionProperty(final Integer value) {
+  private MessageIdProperty(final Integer value) {
     super(value);
   }
 
@@ -37,7 +37,7 @@ private MetricCollectionProperty(final Integer value) {
    * @param value value of the new execution property.
    * @return the newly created execution property.
    */
-  public static MetricCollectionProperty of(final Integer value) {
-    return new MetricCollectionProperty(value);
+  public static MessageIdProperty of(final Integer value) {
+    return new MessageIdProperty(value);
   }
 }

common/src/main/java/org/apache/nemo/common/ir/edge/executionproperty/PartitionSetProperty.java

@@ -0,0 +1,60 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.nemo.common.ir.edge.executionproperty;
+
+import org.apache.nemo.common.KeyRange;
+import org.apache.nemo.common.ir.executionproperty.EdgeExecutionProperty;
+
+import java.util.ArrayList;
+
+/**
+ * This property decides which partitions the tasks of the destination IRVertex should fetch.
+ * The position of a KeyRange in the list corresponds to the offset of the destination task.
+ *
+ * For example, in the following setup:
+ * Source IRVertex (Parallelism=2) - IREdge (Partitioner.Num=4) - Destination IRVertex (Parallelism=2)
+ *
+ * Setting PartitionSetProperty([0, 3), [3, 3)) on the IREdge with will enforce the following behaviors.
+ * - The first destination task fetches the first 3 partitions from each of the 2 data blocks
+ * - The second destination task fetches the last partitions from each of the 2 data blocks
+ *
+ * This property is useful for handling data skews.
+ * For example, if the size ratio of the 4 partitions in the above setup are (17%, 16%, 17%, 50%),
+ * then each of the destination task will evenly handle 50% of the load.
+ */
+public final class PartitionSetProperty extends EdgeExecutionProperty<ArrayList<KeyRange>> {
+  /**
+   * Constructor.
+   *
+   * @param value value of the execution property.
+   */
+  private PartitionSetProperty(final ArrayList<KeyRange> value) {
+    super(value);
+  }
+
+  /**
+   * Static method exposing the constructor.
+   *
+   * @param value value of the new execution property.
+   * @return the newly created execution property.
+   */
+  public static PartitionSetProperty of(final ArrayList<KeyRange> value) {
+    return new PartitionSetProperty(value);
+  }
+}

common/src/main/java/org/apache/nemo/common/ir/edge/executionproperty/PartitionerProperty.java

@@ -18,38 +18,63 @@
  */
 package org.apache.nemo.common.ir.edge.executionproperty;
 
+import org.apache.nemo.common.Pair;
 import org.apache.nemo.common.ir.executionproperty.EdgeExecutionProperty;
 
 /**
  * Partitioner ExecutionProperty.
  */
-public final class PartitionerProperty extends EdgeExecutionProperty<PartitionerProperty.Value> {
+public final class PartitionerProperty
+  extends EdgeExecutionProperty<Pair<PartitionerProperty.Type, Integer>> {
+
+  // Lazily use the number of destination parallelism to minimize the number of partitions.
+  public static final int NUM_EQUAL_TO_DST_PARALLELISM = 0;
+
   /**
    * Constructor.
    *
    * @param value value of the execution property.
    */
-  private PartitionerProperty(final Value value) {
+  private PartitionerProperty(final Pair<Type, Integer> value) {
     super(value);
   }
 
   /**
-   * Static method exposing the constructor.
-   *
-   * @param value value of the new execution property.
-   * @return the newly created execution property.
+   * @param type of the partitioner.
+   * @return the property.
+   */
+  public static PartitionerProperty of(final Type type) {
+    return PartitionerProperty.of(type, NUM_EQUAL_TO_DST_PARALLELISM, true);
+  }
+
+  /**
+   * @param type of the partitioner.
+   * @param numOfPartitions to create.
+   * @return the property.
+   */
+  public static PartitionerProperty of(final Type type, final int numOfPartitions) {
+    return PartitionerProperty.of(type, numOfPartitions, false);
+  }
+
+  /**
+   * @param type of the partitioner.
+   * @param numOfPartitions to create.
+   * @param auto if the number of partitions is auto.
+   * @return the property.
    */
-  public static PartitionerProperty of(final Value value) {
-    return new PartitionerProperty(value);
+  private static PartitionerProperty of(final Type type, final int numOfPartitions, final boolean auto) {
+    if (!auto && numOfPartitions <= 0) {
+      throw new IllegalArgumentException(String.valueOf(numOfPartitions));
+    }
+    return new PartitionerProperty(Pair.of(type, numOfPartitions));
   }
 
   /**
-   * Possible values of Partitioner ExecutionProperty.
+   * Partitioning types.
    */
-  public enum Value {
-    DataSkewHashPartitioner,
-    HashPartitioner,
-    IntactPartitioner,
-    DedicatedKeyPerElementPartitioner
+  public enum Type {
+    Hash,
+    Intact,
+    DedicatedKeyPerElement
   }
 }

common/src/main/java/org/apache/nemo/common/ir/executionproperty/ExecutionProperty.java

@@ -58,4 +58,14 @@ public final boolean equals(final Object o) {
   public final int hashCode() {
     return value != null ? value.hashCode() : 0;
   }
+
+  @Override
+  public final String toString() {
+    final StringBuilder sb = new StringBuilder();
+    sb.append(this.getClass().getSimpleName());
+    sb.append("(");
+    sb.append(value.toString());
+    sb.append(")");
+    return sb.toString();
+  }
 }

common/src/main/java/org/apache/nemo/common/ir/executionproperty/ExecutionPropertyMap.java

@@ -27,7 +27,7 @@
 import org.apache.nemo.common.ir.edge.executionproperty.*;
 import org.apache.nemo.common.ir.vertex.IRVertex;
 import org.apache.nemo.common.ir.vertex.executionproperty.ResourcePriorityProperty;
-import org.apache.nemo.common.ir.vertex.executionproperty.ParallelismProperty;
+import org.apache.nemo.common.ir.vertex.executionproperty.MinParallelismProperty;
 
 import com.google.common.annotations.VisibleForTesting;
 import org.apache.commons.lang3.builder.HashCodeBuilder;
@@ -74,20 +74,19 @@ public static ExecutionPropertyMap<EdgeExecutionProperty> of(
     map.put(DecoderProperty.of(DecoderFactory.DUMMY_DECODER_FACTORY));
     switch (commPattern) {
       case Shuffle:
-        map.put(PartitionerProperty.of(PartitionerProperty.Value.HashPartitioner));
+        map.put(PartitionerProperty.of(PartitionerProperty.Type.Hash));
         map.put(DataStoreProperty.of(DataStoreProperty.Value.LocalFileStore));
         break;
       case BroadCast:
-        map.put(PartitionerProperty.of(PartitionerProperty.Value.IntactPartitioner));
+        map.put(PartitionerProperty.of(PartitionerProperty.Type.Intact));
         map.put(DataStoreProperty.of(DataStoreProperty.Value.LocalFileStore));
         break;
       case OneToOne:
-        map.put(PartitionerProperty.of(PartitionerProperty.Value.IntactPartitioner));
+        map.put(PartitionerProperty.of(PartitionerProperty.Type.Intact));
         map.put(DataStoreProperty.of(DataStoreProperty.Value.MemoryStore));
         break;
       default:
-        map.put(PartitionerProperty.of(PartitionerProperty.Value.HashPartitioner));
-        map.put(DataStoreProperty.of(DataStoreProperty.Value.LocalFileStore));
+        throw new IllegalStateException(commPattern.toString());
     }
     return map;
   }
@@ -99,7 +98,7 @@ public static ExecutionPropertyMap<EdgeExecutionProperty> of(
    */
   public static ExecutionPropertyMap<VertexExecutionProperty> of(final IRVertex irVertex) {
     final ExecutionPropertyMap<VertexExecutionProperty> map = new ExecutionPropertyMap<>(irVertex.getId());
-    map.put(ParallelismProperty.of(1));
+    map.put(MinParallelismProperty.of(1));
     map.put(ResourcePriorityProperty.of(ResourcePriorityProperty.NONE));
     return map;
   }