TEZ-4137: Input/Output/Processor should merge payload to local conf

tez-api/src/main/java/org/apache/tez/common/TezUtils.java

@@ -30,6 +30,7 @@
 import com.google.protobuf.ByteString;
 
 import com.google.protobuf.CodedInputStream;
+import org.apache.tez.runtime.api.TaskContext;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.apache.hadoop.classification.InterfaceAudience;
@@ -120,6 +121,27 @@ public static Configuration createConfFromByteString(ByteString byteString) thro
     }
   }
 
+  public static Configuration createConfFromBaseConfAndPayload(TaskContext context)
+      throws IOException {
+    Configuration baseConf = context.getContainerConfiguration();
+    Configuration configuration = new Configuration(baseConf);
+    UserPayload payload = context.getUserPayload();
+    ByteString byteString = ByteString.copyFrom(payload.getPayload());
+    try(SnappyInputStream uncompressIs = new SnappyInputStream(byteString.newInput())) {
+      DAGProtos.ConfigurationProto confProto = DAGProtos.ConfigurationProto.parseFrom(uncompressIs);
+      readConfFromPB(confProto, configuration);
+      return configuration;
+    }
+  }
+
+  public static void addToConfFromByteString(Configuration configuration, ByteString byteString)
+      throws IOException {
+    try(SnappyInputStream uncompressIs = new SnappyInputStream(byteString.newInput())) {
+      DAGProtos.ConfigurationProto confProto = DAGProtos.ConfigurationProto.parseFrom(uncompressIs);
+      readConfFromPB(confProto, configuration);
+    }
+  }
+
   /**
    * Convert an instance of {@link org.apache.tez.dag.api.UserPayload} to {@link
    * org.apache.hadoop.conf.Configuration}

tez-api/src/main/java/org/apache/tez/runtime/api/InputInitializerContext.java

@@ -24,6 +24,7 @@
 
 import org.apache.hadoop.classification.InterfaceAudience.Public;
 import org.apache.hadoop.classification.InterfaceStability.Unstable;
+import org.apache.hadoop.conf.Configuration;
 import org.apache.hadoop.yarn.api.records.ApplicationId;
 import org.apache.hadoop.yarn.api.records.Resource;
 import org.apache.tez.common.counters.TezCounters;
@@ -49,7 +50,13 @@ public interface InputInitializerContext {
    * @return DAG name
    */
   String getDAGName();
-
+
+  /**
+   * Get vertex configuration
+   * @return Vertex configuration
+   */
+  Configuration getVertexConfiguration();
+
   /**
    * Get the name of the input
    * @return Input name

tez-api/src/main/java/org/apache/tez/runtime/api/TaskContext.java

@@ -27,6 +27,7 @@
 import org.apache.hadoop.classification.InterfaceAudience.Private;
 import org.apache.hadoop.classification.InterfaceAudience.Public;
 import org.apache.hadoop.classification.InterfaceStability.Unstable;
+import org.apache.hadoop.conf.Configuration;
 import org.apache.hadoop.yarn.api.records.ApplicationId;
 import org.apache.tez.common.counters.TezCounters;
 import org.apache.tez.dag.api.UserPayload;
@@ -62,6 +63,12 @@ public interface TaskContext {
    */
   public int getTaskAttemptNumber();
 
+  /**
+   * Get container configuration
+   * @return Container configuration
+   */
+  public Configuration getContainerConfiguration();
+
   /**
    * Get the name of the DAG
    * @return the DAG name

tez-dag/src/main/java/org/apache/tez/dag/app/dag/impl/TezRootInputInitializerContextImpl.java

@@ -23,6 +23,7 @@
 import java.util.Set;
 import java.util.Objects;
 
+import org.apache.hadoop.conf.Configuration;
 import org.apache.hadoop.yarn.api.records.ApplicationId;
 import org.apache.hadoop.yarn.api.records.Resource;
 import org.apache.tez.common.counters.TezCounters;
@@ -85,7 +86,12 @@ public UserPayload getInputUserPayload() {
   public UserPayload getUserPayload() {
     return this.input.getControllerDescriptor().getUserPayload();
   }
-
+
+  @Override
+  public Configuration getVertexConfiguration() {
+    return vertex.getConf();
+  }
+
   @Override 
   public int getNumTasks() {
     return vertex.getTotalTasks();

tez-mapreduce/src/main/java/org/apache/tez/mapreduce/common/MRInputAMSplitGenerator.java

@@ -30,7 +30,6 @@
 import org.apache.hadoop.classification.InterfaceStability.Evolving;
 import org.apache.hadoop.conf.Configuration;
 import org.apache.hadoop.mapred.JobConf;
-import org.apache.hadoop.mapreduce.split.TezMapReduceSplitsGrouper;
 import org.apache.hadoop.security.UserGroupInformation;
 import org.apache.tez.common.TezUtils;
 import org.apache.tez.dag.api.VertexLocationHint;
@@ -80,8 +79,8 @@ public List<Event> initialize() throws Exception {
           + sw.now(TimeUnit.MILLISECONDS));
     }
     sw.reset().start();
-    Configuration conf = TezUtils.createConfFromByteString(userPayloadProto
-        .getConfigurationBytes());
+    Configuration conf = new JobConf(getContext().getVertexConfiguration());
+    TezUtils.addToConfFromByteString(conf, userPayloadProto.getConfigurationBytes());
 
     sendSerializedEvents = conf.getBoolean(
         MRJobConfig.MR_TEZ_INPUT_INITIALIZER_SERIALIZE_EVENT_PAYLOAD,

tez-mapreduce/src/main/java/org/apache/tez/mapreduce/input/base/MRInputBase.java

@@ -72,8 +72,9 @@ public List<Event> initialize() throws IOException {
     boolean isGrouped = mrUserPayload.getGroupingEnabled();
     Preconditions.checkArgument(mrUserPayload.hasSplits() == false,
         "Split information not expected in " + this.getClass().getName());
-    Configuration conf = TezUtils
-        .createConfFromByteString(mrUserPayload.getConfigurationBytes());
+
+    Configuration conf = new JobConf(getContext().getContainerConfiguration());
+    TezUtils.addToConfFromByteString(conf, mrUserPayload.getConfigurationBytes());
     this.jobConf = new JobConf(conf);
     useNewApi = this.jobConf.getUseNewMapper();
     if (isGrouped) {

tez-mapreduce/src/main/java/org/apache/tez/mapreduce/output/MROutput.java

@@ -29,6 +29,7 @@
 import java.util.concurrent.atomic.AtomicBoolean;
 
 import com.google.common.annotations.VisibleForTesting;
+import com.google.protobuf.ByteString;
 import org.apache.tez.common.Preconditions;
 import org.apache.commons.lang.StringUtils;
 import org.apache.hadoop.mapreduce.lib.output.LazyOutputFormat;
@@ -398,8 +399,9 @@ protected List<Event> initializeBase() throws IOException, InterruptedException
     taskNumberFormat.setGroupingUsed(false);
     nonTaskNumberFormat.setMinimumIntegerDigits(3);
     nonTaskNumberFormat.setGroupingUsed(false);
-    Configuration conf = TezUtils.createConfFromUserPayload(getContext().getUserPayload());
-    this.jobConf = new JobConf(conf);
+    UserPayload userPayload = getContext().getUserPayload();
+    this.jobConf = new JobConf(getContext().getContainerConfiguration());
+    TezUtils.addToConfFromByteString(this.jobConf, ByteString.copyFrom(userPayload.getPayload()));
     // Add tokens to the jobConf - in case they are accessed within the RW / OF
     jobConf.getCredentials().mergeAll(UserGroupInformation.getCurrentUser().getCredentials());
     this.isMapperOutput = jobConf.getBoolean(MRConfig.IS_MAP_PROCESSOR,

tez-mapreduce/src/test/java/org/apache/tez/mapreduce/TezTestUtils.java

@@ -17,6 +17,7 @@
  */
 package org.apache.tez.mapreduce;
 
+import org.apache.hadoop.conf.Configuration;
 import org.apache.hadoop.yarn.api.records.ApplicationId;
 import org.apache.hadoop.yarn.api.records.Resource;
 import org.apache.tez.common.counters.TezCounters;
@@ -59,10 +60,12 @@ public static class TezRootInputInitializerContextForTest implements
 
     private final ApplicationId appId;
     private final UserPayload payload;
+    private final Configuration vertexConfig;
 
-    public TezRootInputInitializerContextForTest(UserPayload payload) throws IOException {
+    public TezRootInputInitializerContextForTest(UserPayload payload, Configuration vertexConfig) throws IOException {
       appId = ApplicationId.newInstance(1000, 200);
       this.payload = payload == null ? UserPayload.create(null) : payload;
+      this.vertexConfig = vertexConfig;
     }
 
     @Override
@@ -75,6 +78,11 @@ public String getDAGName() {
       return "FakeDAG";
     }
 
+    @Override
+    public Configuration getVertexConfiguration() {
+      return vertexConfig;
+    }
+
     @Override
     public String getInputName() {
       return "MRInput";

tez-mapreduce/src/test/java/org/apache/tez/mapreduce/common/TestMRInputAMSplitGenerator.java

@@ -96,7 +96,7 @@ private void testGroupSplitsAndSortSplits(boolean groupSplitsEnabled,
     UserPayload userPayload = dataSource.getInputDescriptor().getUserPayload();
 
     InputInitializerContext context =
-        new TezTestUtils.TezRootInputInitializerContextForTest(userPayload);
+        new TezTestUtils.TezRootInputInitializerContextForTest(userPayload, new Configuration(false));
     MRInputAMSplitGenerator splitGenerator =
         new MRInputAMSplitGenerator(context);
 

tez-mapreduce/src/test/java/org/apache/tez/mapreduce/common/TestMRInputSplitDistributor.java

@@ -70,7 +70,8 @@ public void testSerializedPayload() throws IOException {
     UserPayload userPayload =
         UserPayload.create(payloadProto.build().toByteString().asReadOnlyByteBuffer());
 
-    InputInitializerContext context = new TezTestUtils.TezRootInputInitializerContextForTest(userPayload);
+    InputInitializerContext context = new TezTestUtils.TezRootInputInitializerContextForTest(userPayload,
+        new Configuration(false));
     MRInputSplitDistributor splitDist = new MRInputSplitDistributor(context);
 
     List<Event> events = splitDist.initialize();
@@ -119,7 +120,8 @@ public void testDeserializedPayload() throws IOException {
     UserPayload userPayload =
         UserPayload.create(payloadProto.build().toByteString().asReadOnlyByteBuffer());
 
-    InputInitializerContext context = new TezTestUtils.TezRootInputInitializerContextForTest(userPayload);
+    InputInitializerContext context = new TezTestUtils.TezRootInputInitializerContextForTest(userPayload,
+        new Configuration(false));
     MRInputSplitDistributor splitDist = new MRInputSplitDistributor(context);
 
     List<Event> events = splitDist.initialize();

tez-mapreduce/src/test/java/org/apache/tez/mapreduce/input/MRInputForTest.java

@@ -0,0 +1,35 @@
+/**
+ * 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.tez.mapreduce.input;
+
+import org.apache.hadoop.conf.Configuration;
+import org.apache.tez.runtime.api.InputContext;
+
+/**
+ * This is used for inspecting jobConf in test.
+ */
+public class MRInputForTest extends MRInput {
+  public MRInputForTest(InputContext inputContext, int numPhysicalInputs) {
+    super(inputContext, numPhysicalInputs);
+  }
+
+  public Configuration getConfiguration() {
+    return jobConf;
+  }
+}

tez-mapreduce/src/test/java/org/apache/tez/mapreduce/input/MultiMRInputForTest.java

@@ -0,0 +1,35 @@
+/**
+ * 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.tez.mapreduce.input;
+
+import org.apache.hadoop.conf.Configuration;
+import org.apache.tez.runtime.api.InputContext;
+
+/**
+ * This is used for inspecting jobConf in test.
+ */
+public class MultiMRInputForTest extends MultiMRInput {
+  public MultiMRInputForTest(InputContext inputContext, int numPhysicalInputs) {
+    super(inputContext, numPhysicalInputs);
+  }
+
+  public Configuration getConfiguration() {
+    return jobConf;
+  }
+}

tez-mapreduce/src/test/java/org/apache/tez/mapreduce/input/TestMRInput.java

@@ -69,6 +69,7 @@ public void test0PhysicalInputs() throws IOException {
     doReturn(1).when(inputContext).getTaskIndex();
     doReturn(1).when(inputContext).getTaskAttemptNumber();
     doReturn(new TezCounters()).when(inputContext).getCounters();
+    doReturn(new JobConf(false)).when(inputContext).getContainerConfiguration();
 
 
     MRInput mrInput = new MRInput(inputContext, 0);
@@ -120,6 +121,7 @@ public void testAttributesInJobConf() throws Exception {
     doReturn(TEST_ATTRIBUTES_INPUT_NAME).when(inputContext).getSourceVertexName();
     doReturn(TEST_ATTRIBUTES_APPLICATION_ID).when(inputContext).getApplicationId();
     doReturn(TEST_ATTRIBUTES_UNIQUE_IDENTIFIER).when(inputContext).getUniqueIdentifier();
+    doReturn(new Configuration(false)).when(inputContext).getContainerConfiguration();
 
 
     DataSourceDescriptor dsd = MRInput.createConfigBuilder(new Configuration(false),
@@ -147,6 +149,43 @@ public void testAttributesInJobConf() throws Exception {
     assertTrue(TestInputFormat.invoked.get());
   }
 
+  @Test(timeout = 5000)
+  public void testConfigMerge() throws Exception {
+    JobConf jobConf = new JobConf(false);
+    jobConf.set("payload-key", "payload-value");
+
+    Configuration localConfig = new Configuration(false);
+    localConfig.set("local-key", "local-value");
+
+    InputContext inputContext = mock(InputContext.class);
+
+    DataSourceDescriptor dsd = MRInput.createConfigBuilder(jobConf,
+        TestInputFormat.class).groupSplits(false).build();
+
+    doReturn(dsd.getInputDescriptor().getUserPayload()).when(inputContext).getUserPayload();
+    doReturn(TEST_ATTRIBUTES_DAG_INDEX).when(inputContext).getDagIdentifier();
+    doReturn(TEST_ATTRIBUTES_VERTEX_INDEX).when(inputContext).getTaskVertexIndex();
+    doReturn(TEST_ATTRIBUTES_TASK_INDEX).when(inputContext).getTaskIndex();
+    doReturn(TEST_ATTRIBUTES_TASK_ATTEMPT_INDEX).when(inputContext).getTaskAttemptNumber();
+    doReturn(TEST_ATTRIBUTES_INPUT_INDEX).when(inputContext).getInputIndex();
+    doReturn(TEST_ATTRIBUTES_DAG_ATTEMPT_NUMBER).when(inputContext).getDAGAttemptNumber();
+    doReturn(TEST_ATTRIBUTES_DAG_NAME).when(inputContext).getDAGName();
+    doReturn(TEST_ATTRIBUTES_VERTEX_NAME).when(inputContext).getTaskVertexName();
+    doReturn(TEST_ATTRIBUTES_INPUT_NAME).when(inputContext).getSourceVertexName();
+    doReturn(TEST_ATTRIBUTES_APPLICATION_ID).when(inputContext).getApplicationId();
+    doReturn(TEST_ATTRIBUTES_UNIQUE_IDENTIFIER).when(inputContext).getUniqueIdentifier();
+    doReturn(localConfig).when(inputContext).getContainerConfiguration();
+    doReturn(new TezCounters()).when(inputContext).getCounters();
+
+    MRInputForTest input = new MRInputForTest(inputContext, 1);
+    input.initialize();
+
+    Configuration mergedConfig = input.getConfiguration();
+
+    assertEquals("local-value", mergedConfig.get("local-key"));
+    assertEquals("payload-value", mergedConfig.get("payload-key"));
+  }
+
   /**
    * Test class to verify
    */