[WIP] Work stealing in Nemo (alternative implementation)

common/src/main/java/org/apache/nemo/common/ir/vertex/OperatorVertex.java

@@ -27,6 +27,7 @@
  */
 public class OperatorVertex extends IRVertex {
   private final Transform transform;
+  private final String transformFullName;
 
   /**
    * Constructor of OperatorVertex.
@@ -36,6 +37,12 @@ public class OperatorVertex extends IRVertex {
   public OperatorVertex(final Transform t) {
     super();
     this.transform = t;
+    this.transformFullName = "";
+  }
+
+  public OperatorVertex(final Transform t, final String transformFullName) {
+    this.transform = t;
+    this.transformFullName = transformFullName;
   }
 
   /**
@@ -46,6 +53,7 @@ public OperatorVertex(final Transform t) {
   private OperatorVertex(final OperatorVertex that) {
     super(that);
     this.transform = that.transform;
+    this.transformFullName = that.transformFullName;
   }
 
   @Override
@@ -60,6 +68,10 @@ public final Transform getTransform() {
     return transform;
   }
 
+  public final String getTransformFullName() {
+    return transformFullName;
+  }
+
   @Override
   public final ObjectNode getPropertiesAsJsonNode() {
     final ObjectNode node = getIRVertexPropertiesAsJsonNode();

common/src/main/java/org/apache/nemo/common/ir/vertex/executionproperty/WorkStealingStateProperty.java

@@ -0,0 +1,50 @@
+/*
+ * 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.vertex.executionproperty;
+
+import org.apache.nemo.common.ir.executionproperty.VertexExecutionProperty;
+
+/**
+ * Marks Work Stealing Strategy of the vertex.
+ *
+ * Currently, there are three types:
+ * SPLIT  : vertex which is the subject of work stealing
+ * MERGE  : vertex which merges the effect of work stealing
+ * DEFAULT : vertex which is not the subject of work stealing
+ */
+public class WorkStealingStateProperty extends VertexExecutionProperty<String> {
+  /**
+   * Default constructor.
+   *
+   * @param value value of the VertexExecutionProperty.
+   */
+  public WorkStealingStateProperty(final String value) {
+    super(value);
+  }
+
+  /**
+   * Static method exposing the constructor.
+   *
+   * @param value value of the new execution property.
+   * @return the newly created execution property.
+   */
+  public static WorkStealingStateProperty of(final String value) {
+    return new WorkStealingStateProperty(value);
+  }
+}

common/src/main/java/org/apache/nemo/common/ir/vertex/executionproperty/WorkStealingSubSplitProperty.java

@@ -0,0 +1,45 @@
+/*
+ * 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.vertex.executionproperty;
+
+import org.apache.nemo.common.ir.executionproperty.VertexExecutionProperty;
+
+/**
+ * Property to store the sub-split number of work stealing tasks.
+ */
+public class WorkStealingSubSplitProperty extends VertexExecutionProperty<Integer> {
+  /**
+   * Default constructor.
+   *
+   * @param value value of the VertexExecutionProperty.
+   */
+  public WorkStealingSubSplitProperty(final Integer value) {
+    super(value);
+  }
+
+  /**
+   * Static method exposing the constructor.
+   *
+   * @param value value of the new execution property.
+   * @return the newly created execution property.
+   */
+  public static WorkStealingSubSplitProperty of(final Integer value) {
+    return new WorkStealingSubSplitProperty(value);
+  }
+}

compiler/frontend/beam/src/main/java/org/apache/nemo/compiler/frontend/beam/PipelineTranslator.java

@@ -269,7 +269,8 @@ private static void parDoMultiOutputTranslator(final PipelineTranslationContext
   private static void groupByKeyTranslator(final PipelineTranslationContext ctx,
                                            final TransformHierarchy.Node beamNode,
                                            final GroupByKey<?, ?> transform) {
-    final IRVertex vertex = new OperatorVertex(createGBKTransform(ctx, beamNode));
+    final String fullName = beamNode.getFullName();
+    final IRVertex vertex = new OperatorVertex(createGBKTransform(ctx, beamNode), fullName);
     ctx.addVertex(vertex);
     beamNode.getInputs().values().forEach(input -> ctx.addEdgeTo(vertex, input));
     beamNode.getOutputs().values().forEach(output -> ctx.registerMainOutputFrom(beamNode, vertex, output));
@@ -324,7 +325,8 @@ private static void createPCollectionViewTranslator(final PipelineTranslationCon
   private static void flattenTranslator(final PipelineTranslationContext ctx,
                                         final TransformHierarchy.Node beamNode,
                                         final Flatten.PCollections<?> transform) {
-    final IRVertex vertex = new OperatorVertex(new FlattenTransform());
+    final String fullName = beamNode.getFullName();
+    final IRVertex vertex = new OperatorVertex(new FlattenTransform(), fullName);
     ctx.addVertex(vertex);
     beamNode.getInputs().values().forEach(input -> ctx.addEdgeTo(vertex, input));
     beamNode.getOutputs().values().forEach(output -> ctx.registerMainOutputFrom(beamNode, vertex, output));
@@ -350,6 +352,7 @@ private static Pipeline.PipelineVisitor.CompositeBehavior combinePerKeyTranslato
     final PTransform<?, ?> transform) {
 
     final Combine.PerKey perKey = (Combine.PerKey) transform;
+    final String fullName = beamNode.getFullName();
 
     // If there's any side inputs, translate each primitive transforms in this composite transform one by one.
     if (!perKey.getSideInputs().isEmpty()) {
@@ -382,8 +385,8 @@ private static Pipeline.PipelineVisitor.CompositeBehavior combinePerKeyTranslato
     // Choose between batch processing and stream processing based on window type and boundedness of data
     if (isMainInputBounded(beamNode, ctx.getPipeline()) && isGlobalWindow(beamNode, ctx.getPipeline())) {
       // Batch processing, using CombinePartialTransform and CombineFinalTransform
-      partialCombine = new OperatorVertex(new CombineFnPartialTransform<>(combineFn));
-      finalCombine = new OperatorVertex(new CombineFnFinalTransform<>(combineFn));
+      partialCombine = new OperatorVertex(new CombineFnPartialTransform<>(combineFn), fullName);
+      finalCombine = new OperatorVertex(new CombineFnFinalTransform<>(combineFn), fullName);
     } else {
       // Stream data processing, using GBKTransform
       final AppliedPTransform pTransform = beamNode.toAppliedPTransform(ctx.getPipeline());

compiler/optimizer/src/main/java/org/apache/nemo/compiler/optimizer/pass/compiletime/annotating/WorkStealingStatePass.java

@@ -0,0 +1,162 @@
+/*
+ * 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.compiler.optimizer.pass.compiletime.annotating;
+
+import org.apache.nemo.common.Pair;
+import org.apache.nemo.common.dag.Edge;
+import org.apache.nemo.common.ir.IRDAG;
+import org.apache.nemo.common.ir.edge.executionproperty.CommunicationPatternProperty;
+import org.apache.nemo.common.ir.vertex.IRVertex;
+import org.apache.nemo.common.ir.vertex.OperatorVertex;
+import org.apache.nemo.common.ir.vertex.executionproperty.WorkStealingStateProperty;
+import org.apache.nemo.common.ir.vertex.transform.Transform;
+import org.apache.nemo.compiler.optimizer.pass.compiletime.Requires;
+import org.apache.nemo.runtime.common.plan.StagePartitioner;
+
+import java.util.*;
+
+/**
+ * Optimization pass for annotating {@link WorkStealingStateProperty}.
+ */
+@Annotates(WorkStealingStateProperty.class)
+@Requires(CommunicationPatternProperty.class)
+public final class WorkStealingStatePass extends AnnotatingPass {
+  private static final String SPLIT_STRATEGY = "SPLIT";
+  private static final String MERGE_STRATEGY = "MERGE";
+  private static final String DEFAULT_STRATEGY = "DEFAULT";
+
+  private final StagePartitioner stagePartitioner = new StagePartitioner();
+
+  public WorkStealingStatePass() {
+    super(WorkStealingStatePass.class);
+  }
+
+  @Override
+  public IRDAG apply(final IRDAG irdag) {
+    irdag.topologicalDo(irVertex -> {
+      final boolean notConnectedToO2OEdge = irdag.getIncomingEdgesOf(irVertex).stream()
+        .map(edge -> edge.getPropertyValue(CommunicationPatternProperty.class).get())
+        .noneMatch(property -> property.equals(CommunicationPatternProperty.Value.ONE_TO_ONE));
+      if (irVertex instanceof OperatorVertex && notConnectedToO2OEdge) {
+        Transform transform = ((OperatorVertex) irVertex).getTransform();
+        String transformFullName = ((OperatorVertex) irVertex).getTransformFullName();
+        if (transform.toString().contains("work stealing") || transformFullName.contains("work stealing")) {
+          irVertex.setProperty(WorkStealingStateProperty.of(SPLIT_STRATEGY));
+        } else if (transform.toString().contains("merge") || transformFullName.contains("merge")) {
+          irVertex.setProperty(WorkStealingStateProperty.of(MERGE_STRATEGY));
+        } else {
+          irVertex.setProperty(WorkStealingStateProperty.of(DEFAULT_STRATEGY));
+        }
+      } else {
+        irVertex.setProperty(WorkStealingStateProperty.of(DEFAULT_STRATEGY));
+      }
+    });
+    return tidyWorkStealingAnnotation(irdag);
+  }
+
+
+  /**
+   * Tidy annotated dag.
+   * Cleanup conditions:
+   * -  The number of SPLIT annotations and MERGE annotations should be equal
+   * -  SPLIT and MERGE should not be together in one stage, but needs to be in adjacent stage.
+   * -  For now, nested work stealing optimizations are not provided. If detected, leave only the
+   *    innermost pair.
+   *
+   * @param irdag   irdag to cleanup.
+   * @return        cleaned irdag.
+   */
+  private IRDAG tidyWorkStealingAnnotation(final IRDAG irdag) {
+    String splitVertexId = null;
+
+    final List<Pair<String, String>> splitMergePairs = new ArrayList<>();
+    final Set<String> pairedVertices = new HashSet<>();
+    final Map<Integer, Set<IRVertex>> stageIdToStageVertices = new HashMap<>();
+
+    // Make SPLIT - MERGE vertex pair.
+    for (IRVertex vertex : irdag.getTopologicalSort()) {
+      if (vertex.getPropertyValue(WorkStealingStateProperty.class).get().equals(SPLIT_STRATEGY)) {
+        if (splitVertexId != null) {
+          // nested SPLIT vertex detected: delete the prior one.
+          irdag.getVertexById(splitVertexId).setProperty(WorkStealingStateProperty.of(DEFAULT_STRATEGY));
+        }
+        splitVertexId = vertex.getId();
+      } else if (vertex.getPropertyValue(WorkStealingStateProperty.class).get().equals(MERGE_STRATEGY)) {
+        if (splitVertexId != null) {
+          splitMergePairs.add(Pair.of(splitVertexId, vertex.getId()));
+          pairedVertices.add(splitVertexId);
+          pairedVertices.add(vertex.getId());
+          splitVertexId = null;
+        } else {
+          // no corresponding SPLIT vertex: delete
+          vertex.setProperty(WorkStealingStateProperty.of(DEFAULT_STRATEGY));
+        }
+      }
+    }
+
+    final Map<IRVertex, Integer> vertexToStageId = stagePartitioner.apply(irdag);
+
+    for (Pair<String, String> splitMergePair : splitMergePairs) {
+      IRVertex splitVertex = irdag.getVertexById(splitMergePair.left());
+      IRVertex mergeVertex = irdag.getVertexById(splitMergePair.right());
+
+      if (vertexToStageId.get(splitVertex) >= vertexToStageId.get(mergeVertex)
+        || irdag.getIncomingEdgesOf(mergeVertex).stream()
+        .map(Edge::getSrc)
+        .map(vertexToStageId::get)
+        .noneMatch(stageId -> stageId.equals(vertexToStageId.get(splitVertex)))) {
+        // split vertex is descendent of merge vertex or they are in the same stage,
+        // or they are not in adjacent stages
+        splitVertex.setProperty(WorkStealingStateProperty.of(DEFAULT_STRATEGY));
+        mergeVertex.setProperty(WorkStealingStateProperty.of(DEFAULT_STRATEGY));
+        pairedVertices.remove(splitVertex.getId());
+        pairedVertices.remove(mergeVertex.getId());
+      }
+    }
+
+    irdag.topologicalDo(vertex -> {
+      if (!vertex.getPropertyValue(WorkStealingStateProperty.class)
+        .orElse(DEFAULT_STRATEGY).equals(DEFAULT_STRATEGY)) {
+        if (!pairedVertices.contains(vertex.getId())) {
+          vertex.setProperty(WorkStealingStateProperty.of(DEFAULT_STRATEGY));
+        }
+      }
+    });
+
+    // update execution property of other vertices in same stage.
+    vertexToStageId.forEach((vertex, stageId) -> {
+      if (!stageIdToStageVertices.containsKey(stageId)) {
+        stageIdToStageVertices.put(stageId, new HashSet<>());
+      }
+      stageIdToStageVertices.get(stageId).add(vertex);
+    });
+
+    for (String vertexId : pairedVertices) {
+      IRVertex vertex = irdag.getVertexById(vertexId);
+      Set<IRVertex> stageVertices = stageIdToStageVertices.get(vertexToStageId.get(vertex));
+      String strategy = vertex.getPropertyValue(WorkStealingStateProperty.class)
+        .orElse(DEFAULT_STRATEGY);
+      for (IRVertex stageVertex : stageVertices) {
+        stageVertex.setProperty(WorkStealingStateProperty.of(strategy));
+      }
+    }
+
+    return irdag;
+  }
+}