[GitHub] [incubator-uniffle] zuston commented on a diff in pull request #871: [#870] feat(tez): Tez Shuffle Read supporting RemoteShuffle related app master and reader

["zuston (via GitHub)" <gi...@apache.org>]

zuston commented on code in PR #871:
URL: https://github.com/apache/incubator-uniffle/pull/871#discussion_r1190940365


##########
client-tez/src/main/java/org/apache/tez/dag/app/RssDAGAppMaster.java:
##########
@@ -0,0 +1,363 @@
+/*
+ * 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.dag.app;
+
+import java.lang.reflect.Field;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.ThreadFactory;
+import java.util.concurrent.TimeUnit;
+
+import org.apache.commons.cli.CommandLine;
+import org.apache.commons.cli.GnuParser;
+import org.apache.commons.cli.Options;
+import org.apache.hadoop.conf.Configuration;
+import org.apache.hadoop.fs.FSDataOutputStream;
+import org.apache.hadoop.fs.FileStatus;
+import org.apache.hadoop.fs.FileSystem;
+import org.apache.hadoop.fs.Path;
+import org.apache.hadoop.fs.permission.FsPermission;
+import org.apache.hadoop.mapreduce.JobSubmissionFiles;
+import org.apache.hadoop.security.Credentials;
+import org.apache.hadoop.security.UserGroupInformation;
+import org.apache.hadoop.util.ShutdownHookManager;
+import org.apache.hadoop.yarn.YarnUncaughtExceptionHandler;
+import org.apache.hadoop.yarn.api.ApplicationConstants;
+import org.apache.hadoop.yarn.api.records.ApplicationAttemptId;
+import org.apache.hadoop.yarn.api.records.ContainerId;
+import org.apache.hadoop.yarn.api.records.LocalResource;
+import org.apache.hadoop.yarn.conf.YarnConfiguration;
+import org.apache.hadoop.yarn.util.Clock;
+import org.apache.hadoop.yarn.util.ConverterUtils;
+import org.apache.hadoop.yarn.util.SystemClock;
+import org.apache.tez.common.RssTezConfig;
+import org.apache.tez.common.RssTezUtils;
+import org.apache.tez.common.TezClassLoader;
+import org.apache.tez.common.TezCommonUtils;
+import org.apache.tez.common.TezUtilsInternal;
+import org.apache.tez.common.VersionInfo;
+import org.apache.tez.dag.api.TezConfiguration;
+import org.apache.tez.dag.api.TezConstants;
+import org.apache.tez.dag.api.TezException;
+import org.apache.tez.dag.api.records.DAGProtos;
+import org.apache.tez.dag.api.records.DAGProtos.AMPluginDescriptorProto;
+import org.apache.tez.dag.library.vertexmanager.ShuffleVertexManager;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import org.apache.uniffle.client.api.ShuffleWriteClient;
+
+import static org.apache.tez.common.TezCommonUtils.TEZ_SYSTEM_SUB_DIR;
+
+public class RssDAGAppMaster extends DAGAppMaster {
+  private static final Logger LOG = LoggerFactory.getLogger(RssDAGAppMaster.class);
+  private ShuffleWriteClient shuffleWriteClient;
+  private TezRemoteShuffleManager tezRemoteShuffleManager;
+  private final String rssConfFileLocalResourceName = "rss_conf.xml";

Review Comment:
   Is this an immutable variable? If so, use the `static final` ?



##########
client-tez/src/main/java/org/apache/tez/dag/app/RssDAGAppMaster.java:
##########
@@ -0,0 +1,363 @@
+/*
+ * 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.dag.app;
+
+import java.lang.reflect.Field;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.ThreadFactory;
+import java.util.concurrent.TimeUnit;
+
+import org.apache.commons.cli.CommandLine;
+import org.apache.commons.cli.GnuParser;
+import org.apache.commons.cli.Options;
+import org.apache.hadoop.conf.Configuration;
+import org.apache.hadoop.fs.FSDataOutputStream;
+import org.apache.hadoop.fs.FileStatus;
+import org.apache.hadoop.fs.FileSystem;
+import org.apache.hadoop.fs.Path;
+import org.apache.hadoop.fs.permission.FsPermission;
+import org.apache.hadoop.mapreduce.JobSubmissionFiles;
+import org.apache.hadoop.security.Credentials;
+import org.apache.hadoop.security.UserGroupInformation;
+import org.apache.hadoop.util.ShutdownHookManager;
+import org.apache.hadoop.yarn.YarnUncaughtExceptionHandler;
+import org.apache.hadoop.yarn.api.ApplicationConstants;
+import org.apache.hadoop.yarn.api.records.ApplicationAttemptId;
+import org.apache.hadoop.yarn.api.records.ContainerId;
+import org.apache.hadoop.yarn.api.records.LocalResource;
+import org.apache.hadoop.yarn.conf.YarnConfiguration;
+import org.apache.hadoop.yarn.util.Clock;
+import org.apache.hadoop.yarn.util.ConverterUtils;
+import org.apache.hadoop.yarn.util.SystemClock;
+import org.apache.tez.common.RssTezConfig;
+import org.apache.tez.common.RssTezUtils;
+import org.apache.tez.common.TezClassLoader;
+import org.apache.tez.common.TezCommonUtils;
+import org.apache.tez.common.TezUtilsInternal;
+import org.apache.tez.common.VersionInfo;
+import org.apache.tez.dag.api.TezConfiguration;
+import org.apache.tez.dag.api.TezConstants;
+import org.apache.tez.dag.api.TezException;
+import org.apache.tez.dag.api.records.DAGProtos;
+import org.apache.tez.dag.api.records.DAGProtos.AMPluginDescriptorProto;
+import org.apache.tez.dag.library.vertexmanager.ShuffleVertexManager;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import org.apache.uniffle.client.api.ShuffleWriteClient;
+
+import static org.apache.tez.common.TezCommonUtils.TEZ_SYSTEM_SUB_DIR;
+
+public class RssDAGAppMaster extends DAGAppMaster {
+  private static final Logger LOG = LoggerFactory.getLogger(RssDAGAppMaster.class);
+  private ShuffleWriteClient shuffleWriteClient;
+  private TezRemoteShuffleManager tezRemoteShuffleManager;
+  private final String rssConfFileLocalResourceName = "rss_conf.xml";
+  private DAGProtos.PlanLocalResource rssConfFileLocalResource;
+  final ScheduledExecutorService scheduledExecutorService = Executors.newSingleThreadScheduledExecutor(
+    new ThreadFactory() {
+      @Override
+      public Thread newThread(Runnable r) {
+        Thread t = Executors.defaultThreadFactory().newThread(r);
+        t.setDaemon(true);
+        return t;
+      }
+    }
+  );
+
+  public RssDAGAppMaster(ApplicationAttemptId applicationAttemptId,
+                         ContainerId containerId, String nmHost, int nmPort, int nmHttpPort,
+                         Clock clock, long appSubmitTime, boolean isSession, String workingDirectory,
+                         String [] localDirs, String[] logDirs, String clientVersion,
+                         Credentials credentials, String jobUserName, AMPluginDescriptorProto pluginDescriptorProto) {
+    super(applicationAttemptId, containerId, nmHost, nmPort, nmHttpPort, clock, appSubmitTime, isSession,
+            workingDirectory, localDirs, logDirs, clientVersion, credentials, jobUserName, pluginDescriptorProto);
+  }
+
+  public ShuffleWriteClient getShuffleWriteClient() {
+    return shuffleWriteClient;
+  }
+
+  public void setShuffleWriteClient(ShuffleWriteClient shuffleWriteClient) {
+    this.shuffleWriteClient = shuffleWriteClient;
+  }
+
+  public TezRemoteShuffleManager getTezRemoteShuffleManager() {
+    return tezRemoteShuffleManager;
+  }
+
+  public void setTezRemoteShuffleManager(TezRemoteShuffleManager tezRemoteShuffleManager) {
+    this.tezRemoteShuffleManager = tezRemoteShuffleManager;
+  }
+
+  /**
+   * Init and Start Rss Client
+   * @param appMaster
+   * @param conf
+   * @param applicationAttemptId
+   * @throws Exception
+   */
+  public static void initAndStartRSSClient(final RssDAGAppMaster appMaster, Configuration conf,
+                                           ApplicationAttemptId  applicationAttemptId) throws Exception {
+
+    ShuffleWriteClient client = RssTezUtils.createShuffleClient(conf);
+    appMaster.setShuffleWriteClient(client);
+
+    String coordinators = conf.get(RssTezConfig.RSS_COORDINATOR_QUORUM);
+    LOG.info("Registering coordinators {}", coordinators);
+    client.registerCoordinators(coordinators);
+
+    String strAppAttemptId = applicationAttemptId.toString();
+    long heartbeatInterval = conf.getLong(RssTezConfig.RSS_HEARTBEAT_INTERVAL,
+            RssTezConfig.RSS_HEARTBEAT_INTERVAL_DEFAULT_VALUE);
+    long heartbeatTimeout = conf.getLong(RssTezConfig.RSS_HEARTBEAT_TIMEOUT, heartbeatInterval / 2);
+    client.registerApplicationInfo(strAppAttemptId, heartbeatTimeout, "user");
+
+    appMaster.scheduledExecutorService.scheduleAtFixedRate(
+            () -> {
+              try {
+                client.sendAppHeartbeat(strAppAttemptId, heartbeatTimeout);
+                LOG.info("Finish send heartbeat to coordinator and servers");

Review Comment:
   Should we need this log? It will print too much logs



##########
client-tez/src/main/java/org/apache/tez/dag/app/TezRemoteShuffleManager.java:
##########
@@ -0,0 +1,276 @@
+/*
+ * 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.dag.app;
+
+import java.io.IOException;
+import java.net.InetSocketAddress;
+import java.util.Arrays;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+
+import com.google.common.collect.Sets;
+import org.apache.commons.lang3.StringUtils;
+import org.apache.hadoop.conf.Configuration;
+import org.apache.hadoop.fs.CommonConfigurationKeysPublic;
+import org.apache.hadoop.ipc.ProtocolSignature;
+import org.apache.hadoop.ipc.RPC;
+import org.apache.hadoop.ipc.Server;
+import org.apache.hadoop.net.NetUtils;
+import org.apache.hadoop.security.authorize.PolicyProvider;
+import org.apache.hadoop.security.token.Token;
+import org.apache.tez.common.GetShuffleServerRequest;
+import org.apache.tez.common.GetShuffleServerResponse;
+import org.apache.tez.common.RssTezConfig;
+import org.apache.tez.common.RssTezUtils;
+import org.apache.tez.common.ServicePluginLifecycle;
+import org.apache.tez.common.ShuffleAssignmentsInfoWritable;
+import org.apache.tez.common.TezRemoteShuffleUmbilicalProtocol;
+import org.apache.tez.common.security.JobTokenIdentifier;
+import org.apache.tez.common.security.JobTokenSecretManager;
+import org.apache.tez.dag.api.TezConfiguration;
+import org.apache.tez.dag.api.TezException;
+import org.apache.tez.dag.api.TezUncheckedException;
+import org.apache.tez.dag.app.security.authorize.TezAMPolicyProvider;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import org.apache.uniffle.client.api.ShuffleWriteClient;
+import org.apache.uniffle.client.util.ClientUtils;
+import org.apache.uniffle.common.PartitionRange;
+import org.apache.uniffle.common.RemoteStorageInfo;
+import org.apache.uniffle.common.ShuffleAssignmentsInfo;
+import org.apache.uniffle.common.ShuffleDataDistributionType;
+import org.apache.uniffle.common.ShuffleServerInfo;
+import org.apache.uniffle.common.exception.RssException;
+import org.apache.uniffle.common.util.Constants;
+import org.apache.uniffle.common.util.RetryUtils;
+
+
+public class TezRemoteShuffleManager implements ServicePluginLifecycle {
+  private static final Logger LOG = LoggerFactory.getLogger(TezRemoteShuffleManager.class);
+
+  protected InetSocketAddress address;
+
+  protected volatile Server server;
+  private String tokenIdentifier;
+  private Token<JobTokenIdentifier> sessionToken;
+  private Configuration conf;
+  private TezRemoteShuffleUmbilicalProtocolImpl tezRemoteShuffleUmbilical;
+  private ShuffleWriteClient rssClient;
+  private String appId;
+
+  public TezRemoteShuffleManager(String tokenIdentifier, Token<JobTokenIdentifier> sessionToken,
+                                 Configuration conf, String appId, ShuffleWriteClient rssClient) {
+    this.tokenIdentifier = tokenIdentifier;
+    this.sessionToken = sessionToken;
+    this.conf = conf;
+    this.appId = appId;
+    this.rssClient = rssClient;
+    this.tezRemoteShuffleUmbilical = new TezRemoteShuffleUmbilicalProtocolImpl();
+  }
+
+  @Override
+  public void initialize() throws Exception {
+
+  }
+
+  @Override
+  public void start() throws Exception {
+    startRpcServer();
+  }
+
+  @Override
+  public void shutdown() throws Exception {
+    server.stop();
+  }
+
+  private class TezRemoteShuffleUmbilicalProtocolImpl implements TezRemoteShuffleUmbilicalProtocol {
+    private Map<Integer, ShuffleAssignmentsInfo> shuffleIdToShuffleAssignsInfo = new HashMap<>();
+
+    @Override
+    public long getProtocolVersion(String s, long l) throws IOException {
+      return versionID;
+    }
+
+    @Override
+    public ProtocolSignature getProtocolSignature(String protocol, long clientVersion,
+                                                  int clientMethodsHash) throws IOException {
+      return ProtocolSignature.getProtocolSignature(this, protocol,
+              clientVersion, clientMethodsHash);
+    }
+
+    @Override
+    public GetShuffleServerResponse getShuffleAssignments(GetShuffleServerRequest request)
+            throws IOException, TezException {
+      if (request != null) {
+        LOG.info("getShuffleAssignments with request = " + request);
+      } else {
+        LOG.info("getShuffleAssignments with request is null");
+      }
+      int shuffleId = request.getShuffleId();
+
+      GetShuffleServerResponse response = new GetShuffleServerResponse();
+      ShuffleAssignmentsInfo shuffleAssignmentsInfo;
+      try {
+        synchronized (TezRemoteShuffleUmbilicalProtocolImpl.class) {
+          if (shuffleIdToShuffleAssignsInfo.containsKey(shuffleId)) {
+            shuffleAssignmentsInfo = shuffleIdToShuffleAssignsInfo.get(shuffleId);
+          } else {
+            shuffleAssignmentsInfo = getShuffleWorks(request.getPartitionNum(), shuffleId);
+          }
+
+          if (shuffleAssignmentsInfo == null) {
+            response.setStatus(-1);
+            response.setRetMsg("shuffleAssignmentsInfo is null");
+          } else {
+            response.setStatus(0);
+            response.setRetMsg("");
+            response.setShuffleAssignmentsInfoWritable(new ShuffleAssignmentsInfoWritable(shuffleAssignmentsInfo));
+            shuffleIdToShuffleAssignsInfo.put(shuffleId, shuffleAssignmentsInfo);
+          }
+        }
+      } catch (Exception rssException) {
+        response.setStatus(-2);
+        response.setRetMsg(rssException.getMessage());
+      }
+
+      return response;
+    }
+  }
+
+  private ShuffleAssignmentsInfo getShuffleWorks(int partitionNum, int shuffleId) {
+    ShuffleAssignmentsInfo shuffleAssignmentsInfo;
+    int requiredAssignmentShuffleServersNum = RssTezUtils.getRequiredShuffleServerNumber(conf, 200, partitionNum);
+    // retryInterval must bigger than `rss.server.heartbeat.timeout`, or maybe it will return the same result
+    long retryInterval = conf.getLong(RssTezConfig.RSS_CLIENT_ASSIGNMENT_RETRY_INTERVAL,
+            RssTezConfig.RSS_CLIENT_ASSIGNMENT_RETRY_INTERVAL_DEFAULT_VALUE);
+    int retryTimes = conf.getInt(RssTezConfig.RSS_CLIENT_ASSIGNMENT_RETRY_TIMES,
+            RssTezConfig.RSS_CLIENT_ASSIGNMENT_RETRY_TIMES_DEFAULT_VALUE);
+
+    // Get the configured server assignment tags and it will also add default shuffle version tag.
+    Set<String> assignmentTags = new HashSet<>();

Review Comment:
   Put this part into a extra method.



##########
client-tez/src/main/java/org/apache/tez/dag/app/RssDAGAppMaster.java:
##########
@@ -0,0 +1,363 @@
+/*
+ * 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.dag.app;
+
+import java.lang.reflect.Field;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.ThreadFactory;
+import java.util.concurrent.TimeUnit;
+
+import org.apache.commons.cli.CommandLine;
+import org.apache.commons.cli.GnuParser;
+import org.apache.commons.cli.Options;
+import org.apache.hadoop.conf.Configuration;
+import org.apache.hadoop.fs.FSDataOutputStream;
+import org.apache.hadoop.fs.FileStatus;
+import org.apache.hadoop.fs.FileSystem;
+import org.apache.hadoop.fs.Path;
+import org.apache.hadoop.fs.permission.FsPermission;
+import org.apache.hadoop.mapreduce.JobSubmissionFiles;
+import org.apache.hadoop.security.Credentials;
+import org.apache.hadoop.security.UserGroupInformation;
+import org.apache.hadoop.util.ShutdownHookManager;
+import org.apache.hadoop.yarn.YarnUncaughtExceptionHandler;
+import org.apache.hadoop.yarn.api.ApplicationConstants;
+import org.apache.hadoop.yarn.api.records.ApplicationAttemptId;
+import org.apache.hadoop.yarn.api.records.ContainerId;
+import org.apache.hadoop.yarn.api.records.LocalResource;
+import org.apache.hadoop.yarn.conf.YarnConfiguration;
+import org.apache.hadoop.yarn.util.Clock;
+import org.apache.hadoop.yarn.util.ConverterUtils;
+import org.apache.hadoop.yarn.util.SystemClock;
+import org.apache.tez.common.RssTezConfig;
+import org.apache.tez.common.RssTezUtils;
+import org.apache.tez.common.TezClassLoader;
+import org.apache.tez.common.TezCommonUtils;
+import org.apache.tez.common.TezUtilsInternal;
+import org.apache.tez.common.VersionInfo;
+import org.apache.tez.dag.api.TezConfiguration;
+import org.apache.tez.dag.api.TezConstants;
+import org.apache.tez.dag.api.TezException;
+import org.apache.tez.dag.api.records.DAGProtos;
+import org.apache.tez.dag.api.records.DAGProtos.AMPluginDescriptorProto;
+import org.apache.tez.dag.library.vertexmanager.ShuffleVertexManager;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import org.apache.uniffle.client.api.ShuffleWriteClient;
+
+import static org.apache.tez.common.TezCommonUtils.TEZ_SYSTEM_SUB_DIR;
+
+public class RssDAGAppMaster extends DAGAppMaster {
+  private static final Logger LOG = LoggerFactory.getLogger(RssDAGAppMaster.class);
+  private ShuffleWriteClient shuffleWriteClient;
+  private TezRemoteShuffleManager tezRemoteShuffleManager;
+  private final String rssConfFileLocalResourceName = "rss_conf.xml";
+  private DAGProtos.PlanLocalResource rssConfFileLocalResource;
+  final ScheduledExecutorService scheduledExecutorService = Executors.newSingleThreadScheduledExecutor(
+    new ThreadFactory() {
+      @Override
+      public Thread newThread(Runnable r) {
+        Thread t = Executors.defaultThreadFactory().newThread(r);
+        t.setDaemon(true);
+        return t;
+      }
+    }
+  );
+
+  public RssDAGAppMaster(ApplicationAttemptId applicationAttemptId,
+                         ContainerId containerId, String nmHost, int nmPort, int nmHttpPort,
+                         Clock clock, long appSubmitTime, boolean isSession, String workingDirectory,
+                         String [] localDirs, String[] logDirs, String clientVersion,
+                         Credentials credentials, String jobUserName, AMPluginDescriptorProto pluginDescriptorProto) {
+    super(applicationAttemptId, containerId, nmHost, nmPort, nmHttpPort, clock, appSubmitTime, isSession,
+            workingDirectory, localDirs, logDirs, clientVersion, credentials, jobUserName, pluginDescriptorProto);
+  }
+
+  public ShuffleWriteClient getShuffleWriteClient() {
+    return shuffleWriteClient;
+  }
+
+  public void setShuffleWriteClient(ShuffleWriteClient shuffleWriteClient) {
+    this.shuffleWriteClient = shuffleWriteClient;
+  }
+
+  public TezRemoteShuffleManager getTezRemoteShuffleManager() {
+    return tezRemoteShuffleManager;
+  }
+
+  public void setTezRemoteShuffleManager(TezRemoteShuffleManager tezRemoteShuffleManager) {
+    this.tezRemoteShuffleManager = tezRemoteShuffleManager;
+  }
+
+  /**
+   * Init and Start Rss Client
+   * @param appMaster
+   * @param conf
+   * @param applicationAttemptId
+   * @throws Exception
+   */
+  public static void initAndStartRSSClient(final RssDAGAppMaster appMaster, Configuration conf,
+                                           ApplicationAttemptId  applicationAttemptId) throws Exception {
+
+    ShuffleWriteClient client = RssTezUtils.createShuffleClient(conf);
+    appMaster.setShuffleWriteClient(client);
+
+    String coordinators = conf.get(RssTezConfig.RSS_COORDINATOR_QUORUM);
+    LOG.info("Registering coordinators {}", coordinators);
+    client.registerCoordinators(coordinators);
+
+    String strAppAttemptId = applicationAttemptId.toString();
+    long heartbeatInterval = conf.getLong(RssTezConfig.RSS_HEARTBEAT_INTERVAL,
+            RssTezConfig.RSS_HEARTBEAT_INTERVAL_DEFAULT_VALUE);
+    long heartbeatTimeout = conf.getLong(RssTezConfig.RSS_HEARTBEAT_TIMEOUT, heartbeatInterval / 2);
+    client.registerApplicationInfo(strAppAttemptId, heartbeatTimeout, "user");
+
+    appMaster.scheduledExecutorService.scheduleAtFixedRate(
+            () -> {
+              try {
+                client.sendAppHeartbeat(strAppAttemptId, heartbeatTimeout);
+                LOG.info("Finish send heartbeat to coordinator and servers");
+              } catch (Exception e) {
+                LOG.warn("Fail to send heartbeat to coordinator and servers", e);
+              }
+            },
+            heartbeatInterval / 2,
+            heartbeatInterval,
+            TimeUnit.MILLISECONDS);
+
+    appMaster.setTezRemoteShuffleManager(new TezRemoteShuffleManager(strAppAttemptId, null, conf,
+            strAppAttemptId, client));
+    appMaster.getTezRemoteShuffleManager().initialize();
+    appMaster.getTezRemoteShuffleManager().start();
+
+    TezConfiguration extraConf = new TezConfiguration(false);
+    extraConf.clear();
+
+    String strAppId = applicationAttemptId.getApplicationId().toString();
+    extraConf.set(RssTezConfig.RSS_AM_SHUFFLE_MANAGER_ADDRESS,
+            appMaster.getTezRemoteShuffleManager().address.getHostName());
+    extraConf.setInt(RssTezConfig.RSS_AM_SHUFFLE_MANAGER_PORT,
+            appMaster.getTezRemoteShuffleManager().address.getPort());
+    writeExtraConf(appMaster, conf, extraConf, strAppId);
+
+    mayCloseTezSlowStart(conf);
+  }
+
+  @Override
+  public String submitDAGToAppMaster(DAGProtos.DAGPlan dagPlan,
+                                     Map<String, LocalResource> additionalResources) throws TezException {
+    List<DAGProtos.PlanLocalResource> planLocalResourceList = dagPlan.getLocalResourceList();
+
+    if (planLocalResourceList == null) {
+      LOG.warn("planLocalResourceList is null, add new list");
+      planLocalResourceList = new ArrayList<>();
+    } else {
+      planLocalResourceList = new ArrayList<>(planLocalResourceList);
+    }
+
+    try {
+      planLocalResourceList.add(rssConfFileLocalResource);
+      Field field = DAGProtos.DAGPlan.class.getDeclaredField("localResource_");
+      field.setAccessible(true);
+      field.set(dagPlan, planLocalResourceList);
+      field.setAccessible(false);
+    } catch (Exception e) {
+      LOG.error("submitDAGToAppMaster reflect error", e);
+      throw new TezException(e.getMessage());
+    }
+
+    if (LOG.isDebugEnabled()) {
+      for (DAGProtos.PlanLocalResource localResource : dagPlan.getLocalResourceList()) {
+        LOG.debug("localResource: {}", localResource.toString());
+      }
+    }
+
+    return super.submitDAGToAppMaster(dagPlan, additionalResources);
+  }
+
+  /**
+   * main method
+   * @param args
+   */
+  public static void main(String[] args) {
+    try {
+      // Install the tez class loader, which can be used add new resources
+      TezClassLoader.setupTezClassLoader();
+      Thread.setDefaultUncaughtExceptionHandler(new YarnUncaughtExceptionHandler());
+      final String pid = System.getenv().get("JVM_PID");
+      String containerIdStr =
+              System.getenv(ApplicationConstants.Environment.CONTAINER_ID.name());
+      String appSubmitTimeStr =
+              System.getenv(ApplicationConstants.APP_SUBMIT_TIME_ENV);
+      String clientVersion = System.getenv(TezConstants.TEZ_CLIENT_VERSION_ENV);
+      if (clientVersion == null) {
+        clientVersion = VersionInfo.UNKNOWN;
+      }
+
+      Objects.requireNonNull(appSubmitTimeStr,
+              ApplicationConstants.APP_SUBMIT_TIME_ENV + " is null");
+
+      ContainerId containerId = ConverterUtils.toContainerId(containerIdStr);
+      ApplicationAttemptId applicationAttemptId =
+              containerId.getApplicationAttemptId();
+
+
+      String jobUserName = System
+              .getenv(ApplicationConstants.Environment.USER.name());
+
+      // Command line options
+      Options opts = new Options();
+      opts.addOption(TezConstants.TEZ_SESSION_MODE_CLI_OPTION,
+              false, "Run Tez Application Master in Session mode");
+
+      CommandLine cliParser = new GnuParser().parse(opts, args);
+      boolean sessionModeCliOption = cliParser.hasOption(TezConstants.TEZ_SESSION_MODE_CLI_OPTION);
+
+      LOG.info("Creating DAGAppMaster for "
+              + "applicationId=" + applicationAttemptId.getApplicationId()
+              + ", attemptNum=" + applicationAttemptId.getAttemptId()
+              + ", AMContainerId=" + containerId
+              + ", jvmPid=" + pid
+              + ", userFromEnv=" + jobUserName
+              + ", cliSessionOption=" + sessionModeCliOption
+              + ", pwd=" + System.getenv(ApplicationConstants.Environment.PWD.name())
+              + ", localDirs=" + System.getenv(ApplicationConstants.Environment.LOCAL_DIRS.name())
+              + ", logDirs=" + System.getenv(ApplicationConstants.Environment.LOG_DIRS.name()));
+
+      Configuration conf = new Configuration(new YarnConfiguration());

Review Comment:
   +1. Loading the default yarn conf from classpath ? 



##########
client-tez/src/main/java/org/apache/tez/dag/app/RssDAGAppMaster.java:
##########
@@ -0,0 +1,363 @@
+/*
+ * 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.dag.app;
+
+import java.lang.reflect.Field;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.ThreadFactory;
+import java.util.concurrent.TimeUnit;
+
+import org.apache.commons.cli.CommandLine;
+import org.apache.commons.cli.GnuParser;
+import org.apache.commons.cli.Options;
+import org.apache.hadoop.conf.Configuration;
+import org.apache.hadoop.fs.FSDataOutputStream;
+import org.apache.hadoop.fs.FileStatus;
+import org.apache.hadoop.fs.FileSystem;
+import org.apache.hadoop.fs.Path;
+import org.apache.hadoop.fs.permission.FsPermission;
+import org.apache.hadoop.mapreduce.JobSubmissionFiles;
+import org.apache.hadoop.security.Credentials;
+import org.apache.hadoop.security.UserGroupInformation;
+import org.apache.hadoop.util.ShutdownHookManager;
+import org.apache.hadoop.yarn.YarnUncaughtExceptionHandler;
+import org.apache.hadoop.yarn.api.ApplicationConstants;
+import org.apache.hadoop.yarn.api.records.ApplicationAttemptId;
+import org.apache.hadoop.yarn.api.records.ContainerId;
+import org.apache.hadoop.yarn.api.records.LocalResource;
+import org.apache.hadoop.yarn.conf.YarnConfiguration;
+import org.apache.hadoop.yarn.util.Clock;
+import org.apache.hadoop.yarn.util.ConverterUtils;
+import org.apache.hadoop.yarn.util.SystemClock;
+import org.apache.tez.common.RssTezConfig;
+import org.apache.tez.common.RssTezUtils;
+import org.apache.tez.common.TezClassLoader;
+import org.apache.tez.common.TezCommonUtils;
+import org.apache.tez.common.TezUtilsInternal;
+import org.apache.tez.common.VersionInfo;
+import org.apache.tez.dag.api.TezConfiguration;
+import org.apache.tez.dag.api.TezConstants;
+import org.apache.tez.dag.api.TezException;
+import org.apache.tez.dag.api.records.DAGProtos;
+import org.apache.tez.dag.api.records.DAGProtos.AMPluginDescriptorProto;
+import org.apache.tez.dag.library.vertexmanager.ShuffleVertexManager;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import org.apache.uniffle.client.api.ShuffleWriteClient;
+
+import static org.apache.tez.common.TezCommonUtils.TEZ_SYSTEM_SUB_DIR;
+
+public class RssDAGAppMaster extends DAGAppMaster {
+  private static final Logger LOG = LoggerFactory.getLogger(RssDAGAppMaster.class);
+  private ShuffleWriteClient shuffleWriteClient;
+  private TezRemoteShuffleManager tezRemoteShuffleManager;
+  private final String rssConfFileLocalResourceName = "rss_conf.xml";
+  private DAGProtos.PlanLocalResource rssConfFileLocalResource;
+  final ScheduledExecutorService scheduledExecutorService = Executors.newSingleThreadScheduledExecutor(
+    new ThreadFactory() {
+      @Override
+      public Thread newThread(Runnable r) {
+        Thread t = Executors.defaultThreadFactory().newThread(r);
+        t.setDaemon(true);
+        return t;
+      }
+    }
+  );
+
+  public RssDAGAppMaster(ApplicationAttemptId applicationAttemptId,
+                         ContainerId containerId, String nmHost, int nmPort, int nmHttpPort,
+                         Clock clock, long appSubmitTime, boolean isSession, String workingDirectory,
+                         String [] localDirs, String[] logDirs, String clientVersion,
+                         Credentials credentials, String jobUserName, AMPluginDescriptorProto pluginDescriptorProto) {
+    super(applicationAttemptId, containerId, nmHost, nmPort, nmHttpPort, clock, appSubmitTime, isSession,
+            workingDirectory, localDirs, logDirs, clientVersion, credentials, jobUserName, pluginDescriptorProto);
+  }
+
+  public ShuffleWriteClient getShuffleWriteClient() {
+    return shuffleWriteClient;
+  }
+
+  public void setShuffleWriteClient(ShuffleWriteClient shuffleWriteClient) {
+    this.shuffleWriteClient = shuffleWriteClient;
+  }
+
+  public TezRemoteShuffleManager getTezRemoteShuffleManager() {
+    return tezRemoteShuffleManager;
+  }
+
+  public void setTezRemoteShuffleManager(TezRemoteShuffleManager tezRemoteShuffleManager) {
+    this.tezRemoteShuffleManager = tezRemoteShuffleManager;
+  }
+
+  /**
+   * Init and Start Rss Client
+   * @param appMaster
+   * @param conf
+   * @param applicationAttemptId
+   * @throws Exception
+   */
+  public static void initAndStartRSSClient(final RssDAGAppMaster appMaster, Configuration conf,
+                                           ApplicationAttemptId  applicationAttemptId) throws Exception {
+
+    ShuffleWriteClient client = RssTezUtils.createShuffleClient(conf);
+    appMaster.setShuffleWriteClient(client);
+
+    String coordinators = conf.get(RssTezConfig.RSS_COORDINATOR_QUORUM);
+    LOG.info("Registering coordinators {}", coordinators);
+    client.registerCoordinators(coordinators);
+
+    String strAppAttemptId = applicationAttemptId.toString();
+    long heartbeatInterval = conf.getLong(RssTezConfig.RSS_HEARTBEAT_INTERVAL,
+            RssTezConfig.RSS_HEARTBEAT_INTERVAL_DEFAULT_VALUE);
+    long heartbeatTimeout = conf.getLong(RssTezConfig.RSS_HEARTBEAT_TIMEOUT, heartbeatInterval / 2);
+    client.registerApplicationInfo(strAppAttemptId, heartbeatTimeout, "user");
+
+    appMaster.scheduledExecutorService.scheduleAtFixedRate(
+            () -> {
+              try {
+                client.sendAppHeartbeat(strAppAttemptId, heartbeatTimeout);
+                LOG.info("Finish send heartbeat to coordinator and servers");
+              } catch (Exception e) {
+                LOG.warn("Fail to send heartbeat to coordinator and servers", e);
+              }
+            },
+            heartbeatInterval / 2,
+            heartbeatInterval,
+            TimeUnit.MILLISECONDS);
+
+    appMaster.setTezRemoteShuffleManager(new TezRemoteShuffleManager(strAppAttemptId, null, conf,
+            strAppAttemptId, client));
+    appMaster.getTezRemoteShuffleManager().initialize();
+    appMaster.getTezRemoteShuffleManager().start();
+
+    TezConfiguration extraConf = new TezConfiguration(false);
+    extraConf.clear();
+
+    String strAppId = applicationAttemptId.getApplicationId().toString();
+    extraConf.set(RssTezConfig.RSS_AM_SHUFFLE_MANAGER_ADDRESS,
+            appMaster.getTezRemoteShuffleManager().address.getHostName());
+    extraConf.setInt(RssTezConfig.RSS_AM_SHUFFLE_MANAGER_PORT,
+            appMaster.getTezRemoteShuffleManager().address.getPort());
+    writeExtraConf(appMaster, conf, extraConf, strAppId);
+
+    mayCloseTezSlowStart(conf);
+  }
+
+  @Override
+  public String submitDAGToAppMaster(DAGProtos.DAGPlan dagPlan,
+                                     Map<String, LocalResource> additionalResources) throws TezException {
+    List<DAGProtos.PlanLocalResource> planLocalResourceList = dagPlan.getLocalResourceList();
+
+    if (planLocalResourceList == null) {
+      LOG.warn("planLocalResourceList is null, add new list");
+      planLocalResourceList = new ArrayList<>();
+    } else {
+      planLocalResourceList = new ArrayList<>(planLocalResourceList);
+    }
+
+    try {
+      planLocalResourceList.add(rssConfFileLocalResource);
+      Field field = DAGProtos.DAGPlan.class.getDeclaredField("localResource_");
+      field.setAccessible(true);
+      field.set(dagPlan, planLocalResourceList);
+      field.setAccessible(false);
+    } catch (Exception e) {
+      LOG.error("submitDAGToAppMaster reflect error", e);
+      throw new TezException(e.getMessage());
+    }
+
+    if (LOG.isDebugEnabled()) {
+      for (DAGProtos.PlanLocalResource localResource : dagPlan.getLocalResourceList()) {
+        LOG.debug("localResource: {}", localResource.toString());
+      }
+    }
+
+    return super.submitDAGToAppMaster(dagPlan, additionalResources);
+  }
+
+  /**
+   * main method
+   * @param args
+   */
+  public static void main(String[] args) {
+    try {
+      // Install the tez class loader, which can be used add new resources
+      TezClassLoader.setupTezClassLoader();
+      Thread.setDefaultUncaughtExceptionHandler(new YarnUncaughtExceptionHandler());
+      final String pid = System.getenv().get("JVM_PID");

Review Comment:
   JVM_PID should be as a `static final var`



##########
client-tez/src/main/java/org/apache/tez/dag/app/RssDAGAppMaster.java:
##########
@@ -0,0 +1,363 @@
+/*
+ * 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.dag.app;
+
+import java.lang.reflect.Field;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.ThreadFactory;
+import java.util.concurrent.TimeUnit;
+
+import org.apache.commons.cli.CommandLine;
+import org.apache.commons.cli.GnuParser;
+import org.apache.commons.cli.Options;
+import org.apache.hadoop.conf.Configuration;
+import org.apache.hadoop.fs.FSDataOutputStream;
+import org.apache.hadoop.fs.FileStatus;
+import org.apache.hadoop.fs.FileSystem;
+import org.apache.hadoop.fs.Path;
+import org.apache.hadoop.fs.permission.FsPermission;
+import org.apache.hadoop.mapreduce.JobSubmissionFiles;
+import org.apache.hadoop.security.Credentials;
+import org.apache.hadoop.security.UserGroupInformation;
+import org.apache.hadoop.util.ShutdownHookManager;
+import org.apache.hadoop.yarn.YarnUncaughtExceptionHandler;
+import org.apache.hadoop.yarn.api.ApplicationConstants;
+import org.apache.hadoop.yarn.api.records.ApplicationAttemptId;
+import org.apache.hadoop.yarn.api.records.ContainerId;
+import org.apache.hadoop.yarn.api.records.LocalResource;
+import org.apache.hadoop.yarn.conf.YarnConfiguration;
+import org.apache.hadoop.yarn.util.Clock;
+import org.apache.hadoop.yarn.util.ConverterUtils;
+import org.apache.hadoop.yarn.util.SystemClock;
+import org.apache.tez.common.RssTezConfig;
+import org.apache.tez.common.RssTezUtils;
+import org.apache.tez.common.TezClassLoader;
+import org.apache.tez.common.TezCommonUtils;
+import org.apache.tez.common.TezUtilsInternal;
+import org.apache.tez.common.VersionInfo;
+import org.apache.tez.dag.api.TezConfiguration;
+import org.apache.tez.dag.api.TezConstants;
+import org.apache.tez.dag.api.TezException;
+import org.apache.tez.dag.api.records.DAGProtos;
+import org.apache.tez.dag.api.records.DAGProtos.AMPluginDescriptorProto;
+import org.apache.tez.dag.library.vertexmanager.ShuffleVertexManager;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import org.apache.uniffle.client.api.ShuffleWriteClient;
+
+import static org.apache.tez.common.TezCommonUtils.TEZ_SYSTEM_SUB_DIR;
+
+public class RssDAGAppMaster extends DAGAppMaster {
+  private static final Logger LOG = LoggerFactory.getLogger(RssDAGAppMaster.class);
+  private ShuffleWriteClient shuffleWriteClient;
+  private TezRemoteShuffleManager tezRemoteShuffleManager;
+  private final String rssConfFileLocalResourceName = "rss_conf.xml";
+  private DAGProtos.PlanLocalResource rssConfFileLocalResource;
+  final ScheduledExecutorService scheduledExecutorService = Executors.newSingleThreadScheduledExecutor(
+    new ThreadFactory() {
+      @Override
+      public Thread newThread(Runnable r) {
+        Thread t = Executors.defaultThreadFactory().newThread(r);
+        t.setDaemon(true);
+        return t;
+      }
+    }
+  );
+
+  public RssDAGAppMaster(ApplicationAttemptId applicationAttemptId,
+                         ContainerId containerId, String nmHost, int nmPort, int nmHttpPort,
+                         Clock clock, long appSubmitTime, boolean isSession, String workingDirectory,
+                         String [] localDirs, String[] logDirs, String clientVersion,
+                         Credentials credentials, String jobUserName, AMPluginDescriptorProto pluginDescriptorProto) {
+    super(applicationAttemptId, containerId, nmHost, nmPort, nmHttpPort, clock, appSubmitTime, isSession,
+            workingDirectory, localDirs, logDirs, clientVersion, credentials, jobUserName, pluginDescriptorProto);
+  }
+
+  public ShuffleWriteClient getShuffleWriteClient() {
+    return shuffleWriteClient;
+  }
+
+  public void setShuffleWriteClient(ShuffleWriteClient shuffleWriteClient) {
+    this.shuffleWriteClient = shuffleWriteClient;
+  }
+
+  public TezRemoteShuffleManager getTezRemoteShuffleManager() {
+    return tezRemoteShuffleManager;
+  }
+
+  public void setTezRemoteShuffleManager(TezRemoteShuffleManager tezRemoteShuffleManager) {
+    this.tezRemoteShuffleManager = tezRemoteShuffleManager;
+  }
+
+  /**
+   * Init and Start Rss Client
+   * @param appMaster
+   * @param conf
+   * @param applicationAttemptId
+   * @throws Exception
+   */
+  public static void initAndStartRSSClient(final RssDAGAppMaster appMaster, Configuration conf,
+                                           ApplicationAttemptId  applicationAttemptId) throws Exception {
+
+    ShuffleWriteClient client = RssTezUtils.createShuffleClient(conf);
+    appMaster.setShuffleWriteClient(client);
+
+    String coordinators = conf.get(RssTezConfig.RSS_COORDINATOR_QUORUM);
+    LOG.info("Registering coordinators {}", coordinators);
+    client.registerCoordinators(coordinators);
+
+    String strAppAttemptId = applicationAttemptId.toString();
+    long heartbeatInterval = conf.getLong(RssTezConfig.RSS_HEARTBEAT_INTERVAL,
+            RssTezConfig.RSS_HEARTBEAT_INTERVAL_DEFAULT_VALUE);
+    long heartbeatTimeout = conf.getLong(RssTezConfig.RSS_HEARTBEAT_TIMEOUT, heartbeatInterval / 2);
+    client.registerApplicationInfo(strAppAttemptId, heartbeatTimeout, "user");
+
+    appMaster.scheduledExecutorService.scheduleAtFixedRate(
+            () -> {
+              try {
+                client.sendAppHeartbeat(strAppAttemptId, heartbeatTimeout);
+                LOG.info("Finish send heartbeat to coordinator and servers");
+              } catch (Exception e) {
+                LOG.warn("Fail to send heartbeat to coordinator and servers", e);
+              }
+            },
+            heartbeatInterval / 2,
+            heartbeatInterval,
+            TimeUnit.MILLISECONDS);
+
+    appMaster.setTezRemoteShuffleManager(new TezRemoteShuffleManager(strAppAttemptId, null, conf,
+            strAppAttemptId, client));
+    appMaster.getTezRemoteShuffleManager().initialize();
+    appMaster.getTezRemoteShuffleManager().start();
+
+    TezConfiguration extraConf = new TezConfiguration(false);
+    extraConf.clear();
+
+    String strAppId = applicationAttemptId.getApplicationId().toString();
+    extraConf.set(RssTezConfig.RSS_AM_SHUFFLE_MANAGER_ADDRESS,
+            appMaster.getTezRemoteShuffleManager().address.getHostName());
+    extraConf.setInt(RssTezConfig.RSS_AM_SHUFFLE_MANAGER_PORT,
+            appMaster.getTezRemoteShuffleManager().address.getPort());
+    writeExtraConf(appMaster, conf, extraConf, strAppId);
+
+    mayCloseTezSlowStart(conf);
+  }
+
+  @Override
+  public String submitDAGToAppMaster(DAGProtos.DAGPlan dagPlan,
+                                     Map<String, LocalResource> additionalResources) throws TezException {
+    List<DAGProtos.PlanLocalResource> planLocalResourceList = dagPlan.getLocalResourceList();
+
+    if (planLocalResourceList == null) {
+      LOG.warn("planLocalResourceList is null, add new list");
+      planLocalResourceList = new ArrayList<>();
+    } else {
+      planLocalResourceList = new ArrayList<>(planLocalResourceList);
+    }
+
+    try {
+      planLocalResourceList.add(rssConfFileLocalResource);
+      Field field = DAGProtos.DAGPlan.class.getDeclaredField("localResource_");
+      field.setAccessible(true);
+      field.set(dagPlan, planLocalResourceList);
+      field.setAccessible(false);
+    } catch (Exception e) {
+      LOG.error("submitDAGToAppMaster reflect error", e);
+      throw new TezException(e.getMessage());
+    }
+
+    if (LOG.isDebugEnabled()) {
+      for (DAGProtos.PlanLocalResource localResource : dagPlan.getLocalResourceList()) {
+        LOG.debug("localResource: {}", localResource.toString());
+      }
+    }
+
+    return super.submitDAGToAppMaster(dagPlan, additionalResources);
+  }
+
+  /**
+   * main method
+   * @param args
+   */
+  public static void main(String[] args) {
+    try {
+      // Install the tez class loader, which can be used add new resources
+      TezClassLoader.setupTezClassLoader();
+      Thread.setDefaultUncaughtExceptionHandler(new YarnUncaughtExceptionHandler());
+      final String pid = System.getenv().get("JVM_PID");
+      String containerIdStr =
+              System.getenv(ApplicationConstants.Environment.CONTAINER_ID.name());
+      String appSubmitTimeStr =
+              System.getenv(ApplicationConstants.APP_SUBMIT_TIME_ENV);
+      String clientVersion = System.getenv(TezConstants.TEZ_CLIENT_VERSION_ENV);
+      if (clientVersion == null) {
+        clientVersion = VersionInfo.UNKNOWN;
+      }
+
+      Objects.requireNonNull(appSubmitTimeStr,
+              ApplicationConstants.APP_SUBMIT_TIME_ENV + " is null");
+
+      ContainerId containerId = ConverterUtils.toContainerId(containerIdStr);
+      ApplicationAttemptId applicationAttemptId =
+              containerId.getApplicationAttemptId();
+
+
+      String jobUserName = System
+              .getenv(ApplicationConstants.Environment.USER.name());
+
+      // Command line options
+      Options opts = new Options();
+      opts.addOption(TezConstants.TEZ_SESSION_MODE_CLI_OPTION,
+              false, "Run Tez Application Master in Session mode");
+
+      CommandLine cliParser = new GnuParser().parse(opts, args);
+      boolean sessionModeCliOption = cliParser.hasOption(TezConstants.TEZ_SESSION_MODE_CLI_OPTION);
+
+      LOG.info("Creating DAGAppMaster for "
+              + "applicationId=" + applicationAttemptId.getApplicationId()
+              + ", attemptNum=" + applicationAttemptId.getAttemptId()
+              + ", AMContainerId=" + containerId
+              + ", jvmPid=" + pid
+              + ", userFromEnv=" + jobUserName
+              + ", cliSessionOption=" + sessionModeCliOption
+              + ", pwd=" + System.getenv(ApplicationConstants.Environment.PWD.name())
+              + ", localDirs=" + System.getenv(ApplicationConstants.Environment.LOCAL_DIRS.name())
+              + ", logDirs=" + System.getenv(ApplicationConstants.Environment.LOG_DIRS.name()));
+
+      Configuration conf = new Configuration(new YarnConfiguration());
+
+      DAGProtos.ConfigurationProto confProto = TezUtilsInternal.readUserSpecifiedTezConfiguration(
+                      System.getenv(ApplicationConstants.Environment.PWD.name()));
+      TezUtilsInternal.addUserSpecifiedTezConfiguration(conf, confProto.getConfKeyValuesList());
+
+      AMPluginDescriptorProto amPluginDescriptorProto = null;
+      if (confProto.hasAmPluginDescriptor()) {
+        amPluginDescriptorProto = confProto.getAmPluginDescriptor();
+      }
+
+      UserGroupInformation.setConfiguration(conf);
+      Credentials credentials = UserGroupInformation.getCurrentUser().getCredentials();
+
+      TezUtilsInternal.setSecurityUtilConfigration(LOG, conf);
+
+      String nodeHostString = System.getenv(ApplicationConstants.Environment.NM_HOST.name());
+      String nodePortString = System.getenv(ApplicationConstants.Environment.NM_PORT.name());
+      String nodeHttpPortString =
+              System.getenv(ApplicationConstants.Environment.NM_HTTP_PORT.name());
+      long appSubmitTime = Long.parseLong(appSubmitTimeStr);
+      RssDAGAppMaster appMaster =
+              new RssDAGAppMaster(applicationAttemptId, containerId, nodeHostString,
+                      Integer.parseInt(nodePortString),
+                      Integer.parseInt(nodeHttpPortString), new SystemClock(), appSubmitTime,
+                      sessionModeCliOption,
+                      System.getenv(ApplicationConstants.Environment.PWD.name()),
+                      TezCommonUtils.getTrimmedStrings(
+                              System.getenv(ApplicationConstants.Environment.LOCAL_DIRS.name())),
+                      TezCommonUtils.getTrimmedStrings(System.getenv(ApplicationConstants.Environment.LOG_DIRS.name())),
+                      clientVersion, credentials, jobUserName, amPluginDescriptorProto);
+      ShutdownHookManager.get().addShutdownHook(
+              new DAGAppMasterShutdownHook(appMaster), SHUTDOWN_HOOK_PRIORITY);
+      ShutdownHookManager.get().addShutdownHook(new RssDAGAppMasterShutdownHook(appMaster), SHUTDOWN_HOOK_PRIORITY);
+
+
+      // log the system properties
+      if (LOG.isInfoEnabled()) {
+        String systemPropsToLog = TezCommonUtils.getSystemPropertiesToLog(conf);
+        if (systemPropsToLog != null) {
+          LOG.info(systemPropsToLog);
+        }
+      }
+
+      initAndStartRSSClient(appMaster, conf, applicationAttemptId);

Review Comment:
   `startClient` looks weird. How about `initRssClient`



##########
client-tez/src/main/java/org/apache/tez/dag/app/RssDAGAppMaster.java:
##########
@@ -0,0 +1,363 @@
+/*
+ * 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.dag.app;
+
+import java.lang.reflect.Field;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.ThreadFactory;
+import java.util.concurrent.TimeUnit;
+
+import org.apache.commons.cli.CommandLine;
+import org.apache.commons.cli.GnuParser;
+import org.apache.commons.cli.Options;
+import org.apache.hadoop.conf.Configuration;
+import org.apache.hadoop.fs.FSDataOutputStream;
+import org.apache.hadoop.fs.FileStatus;
+import org.apache.hadoop.fs.FileSystem;
+import org.apache.hadoop.fs.Path;
+import org.apache.hadoop.fs.permission.FsPermission;
+import org.apache.hadoop.mapreduce.JobSubmissionFiles;
+import org.apache.hadoop.security.Credentials;
+import org.apache.hadoop.security.UserGroupInformation;
+import org.apache.hadoop.util.ShutdownHookManager;
+import org.apache.hadoop.yarn.YarnUncaughtExceptionHandler;
+import org.apache.hadoop.yarn.api.ApplicationConstants;
+import org.apache.hadoop.yarn.api.records.ApplicationAttemptId;
+import org.apache.hadoop.yarn.api.records.ContainerId;
+import org.apache.hadoop.yarn.api.records.LocalResource;
+import org.apache.hadoop.yarn.conf.YarnConfiguration;
+import org.apache.hadoop.yarn.util.Clock;
+import org.apache.hadoop.yarn.util.ConverterUtils;
+import org.apache.hadoop.yarn.util.SystemClock;
+import org.apache.tez.common.RssTezConfig;
+import org.apache.tez.common.RssTezUtils;
+import org.apache.tez.common.TezClassLoader;
+import org.apache.tez.common.TezCommonUtils;
+import org.apache.tez.common.TezUtilsInternal;
+import org.apache.tez.common.VersionInfo;
+import org.apache.tez.dag.api.TezConfiguration;
+import org.apache.tez.dag.api.TezConstants;
+import org.apache.tez.dag.api.TezException;
+import org.apache.tez.dag.api.records.DAGProtos;
+import org.apache.tez.dag.api.records.DAGProtos.AMPluginDescriptorProto;
+import org.apache.tez.dag.library.vertexmanager.ShuffleVertexManager;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import org.apache.uniffle.client.api.ShuffleWriteClient;
+
+import static org.apache.tez.common.TezCommonUtils.TEZ_SYSTEM_SUB_DIR;
+
+public class RssDAGAppMaster extends DAGAppMaster {
+  private static final Logger LOG = LoggerFactory.getLogger(RssDAGAppMaster.class);
+  private ShuffleWriteClient shuffleWriteClient;
+  private TezRemoteShuffleManager tezRemoteShuffleManager;
+  private final String rssConfFileLocalResourceName = "rss_conf.xml";
+  private DAGProtos.PlanLocalResource rssConfFileLocalResource;
+  final ScheduledExecutorService scheduledExecutorService = Executors.newSingleThreadScheduledExecutor(
+    new ThreadFactory() {
+      @Override
+      public Thread newThread(Runnable r) {
+        Thread t = Executors.defaultThreadFactory().newThread(r);
+        t.setDaemon(true);
+        return t;
+      }
+    }
+  );
+
+  public RssDAGAppMaster(ApplicationAttemptId applicationAttemptId,
+                         ContainerId containerId, String nmHost, int nmPort, int nmHttpPort,
+                         Clock clock, long appSubmitTime, boolean isSession, String workingDirectory,
+                         String [] localDirs, String[] logDirs, String clientVersion,
+                         Credentials credentials, String jobUserName, AMPluginDescriptorProto pluginDescriptorProto) {
+    super(applicationAttemptId, containerId, nmHost, nmPort, nmHttpPort, clock, appSubmitTime, isSession,
+            workingDirectory, localDirs, logDirs, clientVersion, credentials, jobUserName, pluginDescriptorProto);
+  }
+
+  public ShuffleWriteClient getShuffleWriteClient() {
+    return shuffleWriteClient;
+  }
+
+  public void setShuffleWriteClient(ShuffleWriteClient shuffleWriteClient) {
+    this.shuffleWriteClient = shuffleWriteClient;
+  }
+
+  public TezRemoteShuffleManager getTezRemoteShuffleManager() {
+    return tezRemoteShuffleManager;
+  }
+
+  public void setTezRemoteShuffleManager(TezRemoteShuffleManager tezRemoteShuffleManager) {
+    this.tezRemoteShuffleManager = tezRemoteShuffleManager;
+  }
+
+  /**
+   * Init and Start Rss Client
+   * @param appMaster
+   * @param conf
+   * @param applicationAttemptId
+   * @throws Exception
+   */
+  public static void initAndStartRSSClient(final RssDAGAppMaster appMaster, Configuration conf,
+                                           ApplicationAttemptId  applicationAttemptId) throws Exception {
+
+    ShuffleWriteClient client = RssTezUtils.createShuffleClient(conf);
+    appMaster.setShuffleWriteClient(client);
+
+    String coordinators = conf.get(RssTezConfig.RSS_COORDINATOR_QUORUM);
+    LOG.info("Registering coordinators {}", coordinators);
+    client.registerCoordinators(coordinators);
+
+    String strAppAttemptId = applicationAttemptId.toString();
+    long heartbeatInterval = conf.getLong(RssTezConfig.RSS_HEARTBEAT_INTERVAL,
+            RssTezConfig.RSS_HEARTBEAT_INTERVAL_DEFAULT_VALUE);
+    long heartbeatTimeout = conf.getLong(RssTezConfig.RSS_HEARTBEAT_TIMEOUT, heartbeatInterval / 2);
+    client.registerApplicationInfo(strAppAttemptId, heartbeatTimeout, "user");
+
+    appMaster.scheduledExecutorService.scheduleAtFixedRate(
+            () -> {
+              try {
+                client.sendAppHeartbeat(strAppAttemptId, heartbeatTimeout);
+                LOG.info("Finish send heartbeat to coordinator and servers");
+              } catch (Exception e) {
+                LOG.warn("Fail to send heartbeat to coordinator and servers", e);
+              }
+            },
+            heartbeatInterval / 2,
+            heartbeatInterval,
+            TimeUnit.MILLISECONDS);
+
+    appMaster.setTezRemoteShuffleManager(new TezRemoteShuffleManager(strAppAttemptId, null, conf,
+            strAppAttemptId, client));
+    appMaster.getTezRemoteShuffleManager().initialize();
+    appMaster.getTezRemoteShuffleManager().start();
+
+    TezConfiguration extraConf = new TezConfiguration(false);
+    extraConf.clear();
+
+    String strAppId = applicationAttemptId.getApplicationId().toString();
+    extraConf.set(RssTezConfig.RSS_AM_SHUFFLE_MANAGER_ADDRESS,
+            appMaster.getTezRemoteShuffleManager().address.getHostName());
+    extraConf.setInt(RssTezConfig.RSS_AM_SHUFFLE_MANAGER_PORT,
+            appMaster.getTezRemoteShuffleManager().address.getPort());
+    writeExtraConf(appMaster, conf, extraConf, strAppId);
+
+    mayCloseTezSlowStart(conf);
+  }
+
+  @Override
+  public String submitDAGToAppMaster(DAGProtos.DAGPlan dagPlan,
+                                     Map<String, LocalResource> additionalResources) throws TezException {
+    List<DAGProtos.PlanLocalResource> planLocalResourceList = dagPlan.getLocalResourceList();
+
+    if (planLocalResourceList == null) {
+      LOG.warn("planLocalResourceList is null, add new list");
+      planLocalResourceList = new ArrayList<>();
+    } else {
+      planLocalResourceList = new ArrayList<>(planLocalResourceList);
+    }
+
+    try {
+      planLocalResourceList.add(rssConfFileLocalResource);
+      Field field = DAGProtos.DAGPlan.class.getDeclaredField("localResource_");
+      field.setAccessible(true);
+      field.set(dagPlan, planLocalResourceList);
+      field.setAccessible(false);
+    } catch (Exception e) {
+      LOG.error("submitDAGToAppMaster reflect error", e);
+      throw new TezException(e.getMessage());
+    }
+
+    if (LOG.isDebugEnabled()) {
+      for (DAGProtos.PlanLocalResource localResource : dagPlan.getLocalResourceList()) {
+        LOG.debug("localResource: {}", localResource.toString());
+      }
+    }
+
+    return super.submitDAGToAppMaster(dagPlan, additionalResources);
+  }
+
+  /**
+   * main method
+   * @param args
+   */
+  public static void main(String[] args) {
+    try {
+      // Install the tez class loader, which can be used add new resources
+      TezClassLoader.setupTezClassLoader();
+      Thread.setDefaultUncaughtExceptionHandler(new YarnUncaughtExceptionHandler());
+      final String pid = System.getenv().get("JVM_PID");
+      String containerIdStr =
+              System.getenv(ApplicationConstants.Environment.CONTAINER_ID.name());
+      String appSubmitTimeStr =
+              System.getenv(ApplicationConstants.APP_SUBMIT_TIME_ENV);
+      String clientVersion = System.getenv(TezConstants.TEZ_CLIENT_VERSION_ENV);
+      if (clientVersion == null) {
+        clientVersion = VersionInfo.UNKNOWN;
+      }
+
+      Objects.requireNonNull(appSubmitTimeStr,
+              ApplicationConstants.APP_SUBMIT_TIME_ENV + " is null");
+
+      ContainerId containerId = ConverterUtils.toContainerId(containerIdStr);
+      ApplicationAttemptId applicationAttemptId =
+              containerId.getApplicationAttemptId();
+
+
+      String jobUserName = System
+              .getenv(ApplicationConstants.Environment.USER.name());
+
+      // Command line options
+      Options opts = new Options();
+      opts.addOption(TezConstants.TEZ_SESSION_MODE_CLI_OPTION,
+              false, "Run Tez Application Master in Session mode");
+
+      CommandLine cliParser = new GnuParser().parse(opts, args);
+      boolean sessionModeCliOption = cliParser.hasOption(TezConstants.TEZ_SESSION_MODE_CLI_OPTION);
+
+      LOG.info("Creating DAGAppMaster for "
+              + "applicationId=" + applicationAttemptId.getApplicationId()
+              + ", attemptNum=" + applicationAttemptId.getAttemptId()
+              + ", AMContainerId=" + containerId
+              + ", jvmPid=" + pid
+              + ", userFromEnv=" + jobUserName
+              + ", cliSessionOption=" + sessionModeCliOption
+              + ", pwd=" + System.getenv(ApplicationConstants.Environment.PWD.name())
+              + ", localDirs=" + System.getenv(ApplicationConstants.Environment.LOCAL_DIRS.name())
+              + ", logDirs=" + System.getenv(ApplicationConstants.Environment.LOG_DIRS.name()));
+
+      Configuration conf = new Configuration(new YarnConfiguration());
+
+      DAGProtos.ConfigurationProto confProto = TezUtilsInternal.readUserSpecifiedTezConfiguration(
+                      System.getenv(ApplicationConstants.Environment.PWD.name()));
+      TezUtilsInternal.addUserSpecifiedTezConfiguration(conf, confProto.getConfKeyValuesList());
+
+      AMPluginDescriptorProto amPluginDescriptorProto = null;
+      if (confProto.hasAmPluginDescriptor()) {
+        amPluginDescriptorProto = confProto.getAmPluginDescriptor();
+      }
+
+      UserGroupInformation.setConfiguration(conf);
+      Credentials credentials = UserGroupInformation.getCurrentUser().getCredentials();
+
+      TezUtilsInternal.setSecurityUtilConfigration(LOG, conf);
+
+      String nodeHostString = System.getenv(ApplicationConstants.Environment.NM_HOST.name());
+      String nodePortString = System.getenv(ApplicationConstants.Environment.NM_PORT.name());
+      String nodeHttpPortString =
+              System.getenv(ApplicationConstants.Environment.NM_HTTP_PORT.name());
+      long appSubmitTime = Long.parseLong(appSubmitTimeStr);
+      RssDAGAppMaster appMaster =
+              new RssDAGAppMaster(applicationAttemptId, containerId, nodeHostString,
+                      Integer.parseInt(nodePortString),
+                      Integer.parseInt(nodeHttpPortString), new SystemClock(), appSubmitTime,
+                      sessionModeCliOption,
+                      System.getenv(ApplicationConstants.Environment.PWD.name()),
+                      TezCommonUtils.getTrimmedStrings(
+                              System.getenv(ApplicationConstants.Environment.LOCAL_DIRS.name())),
+                      TezCommonUtils.getTrimmedStrings(System.getenv(ApplicationConstants.Environment.LOG_DIRS.name())),
+                      clientVersion, credentials, jobUserName, amPluginDescriptorProto);
+      ShutdownHookManager.get().addShutdownHook(
+              new DAGAppMasterShutdownHook(appMaster), SHUTDOWN_HOOK_PRIORITY);
+      ShutdownHookManager.get().addShutdownHook(new RssDAGAppMasterShutdownHook(appMaster), SHUTDOWN_HOOK_PRIORITY);
+
+
+      // log the system properties
+      if (LOG.isInfoEnabled()) {
+        String systemPropsToLog = TezCommonUtils.getSystemPropertiesToLog(conf);
+        if (systemPropsToLog != null) {
+          LOG.info(systemPropsToLog);
+        }
+      }
+
+      initAndStartRSSClient(appMaster, conf, applicationAttemptId);
+      initAndStartAppMaster(appMaster, conf);
+    } catch (Throwable t) {
+      LOG.error("Error starting DAGAppMaster", t);

Review Comment:
   Errors on starting DAGAppMaster. How about this? 



##########
client-tez/src/main/java/org/apache/tez/dag/app/TezRemoteShuffleManager.java:
##########
@@ -0,0 +1,276 @@
+/*
+ * 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.dag.app;
+
+import java.io.IOException;
+import java.net.InetSocketAddress;
+import java.util.Arrays;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+
+import com.google.common.collect.Sets;
+import org.apache.commons.lang3.StringUtils;
+import org.apache.hadoop.conf.Configuration;
+import org.apache.hadoop.fs.CommonConfigurationKeysPublic;
+import org.apache.hadoop.ipc.ProtocolSignature;
+import org.apache.hadoop.ipc.RPC;
+import org.apache.hadoop.ipc.Server;
+import org.apache.hadoop.net.NetUtils;
+import org.apache.hadoop.security.authorize.PolicyProvider;
+import org.apache.hadoop.security.token.Token;
+import org.apache.tez.common.GetShuffleServerRequest;
+import org.apache.tez.common.GetShuffleServerResponse;
+import org.apache.tez.common.RssTezConfig;
+import org.apache.tez.common.RssTezUtils;
+import org.apache.tez.common.ServicePluginLifecycle;
+import org.apache.tez.common.ShuffleAssignmentsInfoWritable;
+import org.apache.tez.common.TezRemoteShuffleUmbilicalProtocol;
+import org.apache.tez.common.security.JobTokenIdentifier;
+import org.apache.tez.common.security.JobTokenSecretManager;
+import org.apache.tez.dag.api.TezConfiguration;
+import org.apache.tez.dag.api.TezException;
+import org.apache.tez.dag.api.TezUncheckedException;
+import org.apache.tez.dag.app.security.authorize.TezAMPolicyProvider;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import org.apache.uniffle.client.api.ShuffleWriteClient;
+import org.apache.uniffle.client.util.ClientUtils;
+import org.apache.uniffle.common.PartitionRange;
+import org.apache.uniffle.common.RemoteStorageInfo;
+import org.apache.uniffle.common.ShuffleAssignmentsInfo;
+import org.apache.uniffle.common.ShuffleDataDistributionType;
+import org.apache.uniffle.common.ShuffleServerInfo;
+import org.apache.uniffle.common.exception.RssException;
+import org.apache.uniffle.common.util.Constants;
+import org.apache.uniffle.common.util.RetryUtils;
+
+
+public class TezRemoteShuffleManager implements ServicePluginLifecycle {
+  private static final Logger LOG = LoggerFactory.getLogger(TezRemoteShuffleManager.class);
+
+  protected InetSocketAddress address;
+
+  protected volatile Server server;
+  private String tokenIdentifier;
+  private Token<JobTokenIdentifier> sessionToken;
+  private Configuration conf;
+  private TezRemoteShuffleUmbilicalProtocolImpl tezRemoteShuffleUmbilical;
+  private ShuffleWriteClient rssClient;
+  private String appId;
+
+  public TezRemoteShuffleManager(String tokenIdentifier, Token<JobTokenIdentifier> sessionToken,
+                                 Configuration conf, String appId, ShuffleWriteClient rssClient) {
+    this.tokenIdentifier = tokenIdentifier;
+    this.sessionToken = sessionToken;
+    this.conf = conf;
+    this.appId = appId;
+    this.rssClient = rssClient;
+    this.tezRemoteShuffleUmbilical = new TezRemoteShuffleUmbilicalProtocolImpl();
+  }
+
+  @Override
+  public void initialize() throws Exception {
+
+  }
+
+  @Override
+  public void start() throws Exception {
+    startRpcServer();
+  }
+
+  @Override
+  public void shutdown() throws Exception {
+    server.stop();
+  }
+
+  private class TezRemoteShuffleUmbilicalProtocolImpl implements TezRemoteShuffleUmbilicalProtocol {
+    private Map<Integer, ShuffleAssignmentsInfo> shuffleIdToShuffleAssignsInfo = new HashMap<>();
+
+    @Override
+    public long getProtocolVersion(String s, long l) throws IOException {
+      return versionID;
+    }
+
+    @Override
+    public ProtocolSignature getProtocolSignature(String protocol, long clientVersion,
+                                                  int clientMethodsHash) throws IOException {
+      return ProtocolSignature.getProtocolSignature(this, protocol,
+              clientVersion, clientMethodsHash);
+    }
+
+    @Override
+    public GetShuffleServerResponse getShuffleAssignments(GetShuffleServerRequest request)
+            throws IOException, TezException {
+      if (request != null) {
+        LOG.info("getShuffleAssignments with request = " + request);
+      } else {
+        LOG.info("getShuffleAssignments with request is null");
+      }
+      int shuffleId = request.getShuffleId();
+
+      GetShuffleServerResponse response = new GetShuffleServerResponse();
+      ShuffleAssignmentsInfo shuffleAssignmentsInfo;
+      try {
+        synchronized (TezRemoteShuffleUmbilicalProtocolImpl.class) {
+          if (shuffleIdToShuffleAssignsInfo.containsKey(shuffleId)) {
+            shuffleAssignmentsInfo = shuffleIdToShuffleAssignsInfo.get(shuffleId);
+          } else {
+            shuffleAssignmentsInfo = getShuffleWorks(request.getPartitionNum(), shuffleId);
+          }
+
+          if (shuffleAssignmentsInfo == null) {
+            response.setStatus(-1);
+            response.setRetMsg("shuffleAssignmentsInfo is null");
+          } else {
+            response.setStatus(0);
+            response.setRetMsg("");
+            response.setShuffleAssignmentsInfoWritable(new ShuffleAssignmentsInfoWritable(shuffleAssignmentsInfo));
+            shuffleIdToShuffleAssignsInfo.put(shuffleId, shuffleAssignmentsInfo);
+          }
+        }
+      } catch (Exception rssException) {
+        response.setStatus(-2);
+        response.setRetMsg(rssException.getMessage());
+      }
+
+      return response;
+    }
+  }
+
+  private ShuffleAssignmentsInfo getShuffleWorks(int partitionNum, int shuffleId) {
+    ShuffleAssignmentsInfo shuffleAssignmentsInfo;
+    int requiredAssignmentShuffleServersNum = RssTezUtils.getRequiredShuffleServerNumber(conf, 200, partitionNum);

Review Comment:
   `RssTezUtils` exists?



##########
client-tez/src/main/java/org/apache/tez/dag/app/RssDAGAppMaster.java:
##########
@@ -0,0 +1,363 @@
+/*
+ * 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.dag.app;
+
+import java.lang.reflect.Field;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.ThreadFactory;
+import java.util.concurrent.TimeUnit;
+
+import org.apache.commons.cli.CommandLine;
+import org.apache.commons.cli.GnuParser;
+import org.apache.commons.cli.Options;
+import org.apache.hadoop.conf.Configuration;
+import org.apache.hadoop.fs.FSDataOutputStream;
+import org.apache.hadoop.fs.FileStatus;
+import org.apache.hadoop.fs.FileSystem;
+import org.apache.hadoop.fs.Path;
+import org.apache.hadoop.fs.permission.FsPermission;
+import org.apache.hadoop.mapreduce.JobSubmissionFiles;
+import org.apache.hadoop.security.Credentials;
+import org.apache.hadoop.security.UserGroupInformation;
+import org.apache.hadoop.util.ShutdownHookManager;
+import org.apache.hadoop.yarn.YarnUncaughtExceptionHandler;
+import org.apache.hadoop.yarn.api.ApplicationConstants;
+import org.apache.hadoop.yarn.api.records.ApplicationAttemptId;
+import org.apache.hadoop.yarn.api.records.ContainerId;
+import org.apache.hadoop.yarn.api.records.LocalResource;
+import org.apache.hadoop.yarn.conf.YarnConfiguration;
+import org.apache.hadoop.yarn.util.Clock;
+import org.apache.hadoop.yarn.util.ConverterUtils;
+import org.apache.hadoop.yarn.util.SystemClock;
+import org.apache.tez.common.RssTezConfig;
+import org.apache.tez.common.RssTezUtils;
+import org.apache.tez.common.TezClassLoader;
+import org.apache.tez.common.TezCommonUtils;
+import org.apache.tez.common.TezUtilsInternal;
+import org.apache.tez.common.VersionInfo;
+import org.apache.tez.dag.api.TezConfiguration;
+import org.apache.tez.dag.api.TezConstants;
+import org.apache.tez.dag.api.TezException;
+import org.apache.tez.dag.api.records.DAGProtos;
+import org.apache.tez.dag.api.records.DAGProtos.AMPluginDescriptorProto;
+import org.apache.tez.dag.library.vertexmanager.ShuffleVertexManager;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import org.apache.uniffle.client.api.ShuffleWriteClient;
+
+import static org.apache.tez.common.TezCommonUtils.TEZ_SYSTEM_SUB_DIR;
+
+public class RssDAGAppMaster extends DAGAppMaster {
+  private static final Logger LOG = LoggerFactory.getLogger(RssDAGAppMaster.class);
+  private ShuffleWriteClient shuffleWriteClient;
+  private TezRemoteShuffleManager tezRemoteShuffleManager;
+  private final String rssConfFileLocalResourceName = "rss_conf.xml";
+  private DAGProtos.PlanLocalResource rssConfFileLocalResource;
+  final ScheduledExecutorService scheduledExecutorService = Executors.newSingleThreadScheduledExecutor(
+    new ThreadFactory() {
+      @Override
+      public Thread newThread(Runnable r) {
+        Thread t = Executors.defaultThreadFactory().newThread(r);
+        t.setDaemon(true);
+        return t;
+      }
+    }
+  );
+
+  public RssDAGAppMaster(ApplicationAttemptId applicationAttemptId,
+                         ContainerId containerId, String nmHost, int nmPort, int nmHttpPort,
+                         Clock clock, long appSubmitTime, boolean isSession, String workingDirectory,
+                         String [] localDirs, String[] logDirs, String clientVersion,
+                         Credentials credentials, String jobUserName, AMPluginDescriptorProto pluginDescriptorProto) {
+    super(applicationAttemptId, containerId, nmHost, nmPort, nmHttpPort, clock, appSubmitTime, isSession,
+            workingDirectory, localDirs, logDirs, clientVersion, credentials, jobUserName, pluginDescriptorProto);
+  }
+
+  public ShuffleWriteClient getShuffleWriteClient() {
+    return shuffleWriteClient;
+  }
+
+  public void setShuffleWriteClient(ShuffleWriteClient shuffleWriteClient) {
+    this.shuffleWriteClient = shuffleWriteClient;
+  }
+
+  public TezRemoteShuffleManager getTezRemoteShuffleManager() {
+    return tezRemoteShuffleManager;
+  }
+
+  public void setTezRemoteShuffleManager(TezRemoteShuffleManager tezRemoteShuffleManager) {
+    this.tezRemoteShuffleManager = tezRemoteShuffleManager;
+  }
+
+  /**
+   * Init and Start Rss Client
+   * @param appMaster
+   * @param conf
+   * @param applicationAttemptId
+   * @throws Exception
+   */
+  public static void initAndStartRSSClient(final RssDAGAppMaster appMaster, Configuration conf,
+                                           ApplicationAttemptId  applicationAttemptId) throws Exception {
+
+    ShuffleWriteClient client = RssTezUtils.createShuffleClient(conf);
+    appMaster.setShuffleWriteClient(client);
+
+    String coordinators = conf.get(RssTezConfig.RSS_COORDINATOR_QUORUM);
+    LOG.info("Registering coordinators {}", coordinators);
+    client.registerCoordinators(coordinators);
+
+    String strAppAttemptId = applicationAttemptId.toString();
+    long heartbeatInterval = conf.getLong(RssTezConfig.RSS_HEARTBEAT_INTERVAL,
+            RssTezConfig.RSS_HEARTBEAT_INTERVAL_DEFAULT_VALUE);
+    long heartbeatTimeout = conf.getLong(RssTezConfig.RSS_HEARTBEAT_TIMEOUT, heartbeatInterval / 2);
+    client.registerApplicationInfo(strAppAttemptId, heartbeatTimeout, "user");
+
+    appMaster.scheduledExecutorService.scheduleAtFixedRate(
+            () -> {
+              try {
+                client.sendAppHeartbeat(strAppAttemptId, heartbeatTimeout);
+                LOG.info("Finish send heartbeat to coordinator and servers");
+              } catch (Exception e) {
+                LOG.warn("Fail to send heartbeat to coordinator and servers", e);
+              }
+            },
+            heartbeatInterval / 2,
+            heartbeatInterval,
+            TimeUnit.MILLISECONDS);
+
+    appMaster.setTezRemoteShuffleManager(new TezRemoteShuffleManager(strAppAttemptId, null, conf,
+            strAppAttemptId, client));
+    appMaster.getTezRemoteShuffleManager().initialize();
+    appMaster.getTezRemoteShuffleManager().start();
+
+    TezConfiguration extraConf = new TezConfiguration(false);
+    extraConf.clear();
+
+    String strAppId = applicationAttemptId.getApplicationId().toString();
+    extraConf.set(RssTezConfig.RSS_AM_SHUFFLE_MANAGER_ADDRESS,
+            appMaster.getTezRemoteShuffleManager().address.getHostName());
+    extraConf.setInt(RssTezConfig.RSS_AM_SHUFFLE_MANAGER_PORT,
+            appMaster.getTezRemoteShuffleManager().address.getPort());
+    writeExtraConf(appMaster, conf, extraConf, strAppId);
+
+    mayCloseTezSlowStart(conf);
+  }
+
+  @Override
+  public String submitDAGToAppMaster(DAGProtos.DAGPlan dagPlan,
+                                     Map<String, LocalResource> additionalResources) throws TezException {
+    List<DAGProtos.PlanLocalResource> planLocalResourceList = dagPlan.getLocalResourceList();
+
+    if (planLocalResourceList == null) {
+      LOG.warn("planLocalResourceList is null, add new list");
+      planLocalResourceList = new ArrayList<>();
+    } else {
+      planLocalResourceList = new ArrayList<>(planLocalResourceList);
+    }
+
+    try {
+      planLocalResourceList.add(rssConfFileLocalResource);
+      Field field = DAGProtos.DAGPlan.class.getDeclaredField("localResource_");
+      field.setAccessible(true);
+      field.set(dagPlan, planLocalResourceList);
+      field.setAccessible(false);
+    } catch (Exception e) {
+      LOG.error("submitDAGToAppMaster reflect error", e);
+      throw new TezException(e.getMessage());
+    }
+
+    if (LOG.isDebugEnabled()) {
+      for (DAGProtos.PlanLocalResource localResource : dagPlan.getLocalResourceList()) {
+        LOG.debug("localResource: {}", localResource.toString());
+      }
+    }
+
+    return super.submitDAGToAppMaster(dagPlan, additionalResources);
+  }
+
+  /**
+   * main method
+   * @param args
+   */
+  public static void main(String[] args) {
+    try {
+      // Install the tez class loader, which can be used add new resources
+      TezClassLoader.setupTezClassLoader();
+      Thread.setDefaultUncaughtExceptionHandler(new YarnUncaughtExceptionHandler());
+      final String pid = System.getenv().get("JVM_PID");
+      String containerIdStr =
+              System.getenv(ApplicationConstants.Environment.CONTAINER_ID.name());
+      String appSubmitTimeStr =
+              System.getenv(ApplicationConstants.APP_SUBMIT_TIME_ENV);
+      String clientVersion = System.getenv(TezConstants.TEZ_CLIENT_VERSION_ENV);
+      if (clientVersion == null) {
+        clientVersion = VersionInfo.UNKNOWN;
+      }
+
+      Objects.requireNonNull(appSubmitTimeStr,
+              ApplicationConstants.APP_SUBMIT_TIME_ENV + " is null");
+
+      ContainerId containerId = ConverterUtils.toContainerId(containerIdStr);
+      ApplicationAttemptId applicationAttemptId =
+              containerId.getApplicationAttemptId();
+
+
+      String jobUserName = System
+              .getenv(ApplicationConstants.Environment.USER.name());
+
+      // Command line options
+      Options opts = new Options();
+      opts.addOption(TezConstants.TEZ_SESSION_MODE_CLI_OPTION,
+              false, "Run Tez Application Master in Session mode");
+
+      CommandLine cliParser = new GnuParser().parse(opts, args);
+      boolean sessionModeCliOption = cliParser.hasOption(TezConstants.TEZ_SESSION_MODE_CLI_OPTION);
+
+      LOG.info("Creating DAGAppMaster for "

Review Comment:
   Do we need blank line? And I prefer using the {} to construct log



##########
client-tez/src/main/java/org/apache/tez/dag/app/TezRemoteShuffleManager.java:
##########
@@ -0,0 +1,276 @@
+/*
+ * 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.dag.app;
+
+import java.io.IOException;
+import java.net.InetSocketAddress;
+import java.util.Arrays;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+
+import com.google.common.collect.Sets;
+import org.apache.commons.lang3.StringUtils;
+import org.apache.hadoop.conf.Configuration;
+import org.apache.hadoop.fs.CommonConfigurationKeysPublic;
+import org.apache.hadoop.ipc.ProtocolSignature;
+import org.apache.hadoop.ipc.RPC;
+import org.apache.hadoop.ipc.Server;
+import org.apache.hadoop.net.NetUtils;
+import org.apache.hadoop.security.authorize.PolicyProvider;
+import org.apache.hadoop.security.token.Token;
+import org.apache.tez.common.GetShuffleServerRequest;
+import org.apache.tez.common.GetShuffleServerResponse;
+import org.apache.tez.common.RssTezConfig;
+import org.apache.tez.common.RssTezUtils;
+import org.apache.tez.common.ServicePluginLifecycle;
+import org.apache.tez.common.ShuffleAssignmentsInfoWritable;
+import org.apache.tez.common.TezRemoteShuffleUmbilicalProtocol;
+import org.apache.tez.common.security.JobTokenIdentifier;
+import org.apache.tez.common.security.JobTokenSecretManager;
+import org.apache.tez.dag.api.TezConfiguration;
+import org.apache.tez.dag.api.TezException;
+import org.apache.tez.dag.api.TezUncheckedException;
+import org.apache.tez.dag.app.security.authorize.TezAMPolicyProvider;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import org.apache.uniffle.client.api.ShuffleWriteClient;
+import org.apache.uniffle.client.util.ClientUtils;
+import org.apache.uniffle.common.PartitionRange;
+import org.apache.uniffle.common.RemoteStorageInfo;
+import org.apache.uniffle.common.ShuffleAssignmentsInfo;
+import org.apache.uniffle.common.ShuffleDataDistributionType;
+import org.apache.uniffle.common.ShuffleServerInfo;
+import org.apache.uniffle.common.exception.RssException;
+import org.apache.uniffle.common.util.Constants;
+import org.apache.uniffle.common.util.RetryUtils;
+
+
+public class TezRemoteShuffleManager implements ServicePluginLifecycle {
+  private static final Logger LOG = LoggerFactory.getLogger(TezRemoteShuffleManager.class);
+
+  protected InetSocketAddress address;
+
+  protected volatile Server server;
+  private String tokenIdentifier;
+  private Token<JobTokenIdentifier> sessionToken;
+  private Configuration conf;
+  private TezRemoteShuffleUmbilicalProtocolImpl tezRemoteShuffleUmbilical;
+  private ShuffleWriteClient rssClient;
+  private String appId;
+
+  public TezRemoteShuffleManager(String tokenIdentifier, Token<JobTokenIdentifier> sessionToken,
+                                 Configuration conf, String appId, ShuffleWriteClient rssClient) {
+    this.tokenIdentifier = tokenIdentifier;
+    this.sessionToken = sessionToken;
+    this.conf = conf;
+    this.appId = appId;
+    this.rssClient = rssClient;
+    this.tezRemoteShuffleUmbilical = new TezRemoteShuffleUmbilicalProtocolImpl();
+  }
+
+  @Override
+  public void initialize() throws Exception {
+
+  }
+
+  @Override
+  public void start() throws Exception {
+    startRpcServer();
+  }
+
+  @Override
+  public void shutdown() throws Exception {
+    server.stop();
+  }
+
+  private class TezRemoteShuffleUmbilicalProtocolImpl implements TezRemoteShuffleUmbilicalProtocol {
+    private Map<Integer, ShuffleAssignmentsInfo> shuffleIdToShuffleAssignsInfo = new HashMap<>();
+
+    @Override
+    public long getProtocolVersion(String s, long l) throws IOException {
+      return versionID;
+    }
+
+    @Override
+    public ProtocolSignature getProtocolSignature(String protocol, long clientVersion,
+                                                  int clientMethodsHash) throws IOException {
+      return ProtocolSignature.getProtocolSignature(this, protocol,
+              clientVersion, clientMethodsHash);
+    }
+
+    @Override
+    public GetShuffleServerResponse getShuffleAssignments(GetShuffleServerRequest request)
+            throws IOException, TezException {
+      if (request != null) {
+        LOG.info("getShuffleAssignments with request = " + request);
+      } else {
+        LOG.info("getShuffleAssignments with request is null");
+      }
+      int shuffleId = request.getShuffleId();
+
+      GetShuffleServerResponse response = new GetShuffleServerResponse();
+      ShuffleAssignmentsInfo shuffleAssignmentsInfo;
+      try {
+        synchronized (TezRemoteShuffleUmbilicalProtocolImpl.class) {
+          if (shuffleIdToShuffleAssignsInfo.containsKey(shuffleId)) {
+            shuffleAssignmentsInfo = shuffleIdToShuffleAssignsInfo.get(shuffleId);
+          } else {
+            shuffleAssignmentsInfo = getShuffleWorks(request.getPartitionNum(), shuffleId);
+          }
+
+          if (shuffleAssignmentsInfo == null) {
+            response.setStatus(-1);
+            response.setRetMsg("shuffleAssignmentsInfo is null");
+          } else {
+            response.setStatus(0);
+            response.setRetMsg("");
+            response.setShuffleAssignmentsInfoWritable(new ShuffleAssignmentsInfoWritable(shuffleAssignmentsInfo));
+            shuffleIdToShuffleAssignsInfo.put(shuffleId, shuffleAssignmentsInfo);
+          }
+        }
+      } catch (Exception rssException) {
+        response.setStatus(-2);
+        response.setRetMsg(rssException.getMessage());
+      }
+
+      return response;
+    }
+  }
+
+  private ShuffleAssignmentsInfo getShuffleWorks(int partitionNum, int shuffleId) {

Review Comment:
   `getShuffleWorks` -> `getAssignment` ?



##########
client-tez/src/main/java/org/apache/tez/dag/app/TezRemoteShuffleManager.java:
##########
@@ -0,0 +1,276 @@
+/*
+ * 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.dag.app;
+
+import java.io.IOException;
+import java.net.InetSocketAddress;
+import java.util.Arrays;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+
+import com.google.common.collect.Sets;
+import org.apache.commons.lang3.StringUtils;
+import org.apache.hadoop.conf.Configuration;
+import org.apache.hadoop.fs.CommonConfigurationKeysPublic;
+import org.apache.hadoop.ipc.ProtocolSignature;
+import org.apache.hadoop.ipc.RPC;
+import org.apache.hadoop.ipc.Server;
+import org.apache.hadoop.net.NetUtils;
+import org.apache.hadoop.security.authorize.PolicyProvider;
+import org.apache.hadoop.security.token.Token;
+import org.apache.tez.common.GetShuffleServerRequest;
+import org.apache.tez.common.GetShuffleServerResponse;
+import org.apache.tez.common.RssTezConfig;
+import org.apache.tez.common.RssTezUtils;
+import org.apache.tez.common.ServicePluginLifecycle;
+import org.apache.tez.common.ShuffleAssignmentsInfoWritable;
+import org.apache.tez.common.TezRemoteShuffleUmbilicalProtocol;
+import org.apache.tez.common.security.JobTokenIdentifier;
+import org.apache.tez.common.security.JobTokenSecretManager;
+import org.apache.tez.dag.api.TezConfiguration;
+import org.apache.tez.dag.api.TezException;
+import org.apache.tez.dag.api.TezUncheckedException;
+import org.apache.tez.dag.app.security.authorize.TezAMPolicyProvider;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import org.apache.uniffle.client.api.ShuffleWriteClient;
+import org.apache.uniffle.client.util.ClientUtils;
+import org.apache.uniffle.common.PartitionRange;
+import org.apache.uniffle.common.RemoteStorageInfo;
+import org.apache.uniffle.common.ShuffleAssignmentsInfo;
+import org.apache.uniffle.common.ShuffleDataDistributionType;
+import org.apache.uniffle.common.ShuffleServerInfo;
+import org.apache.uniffle.common.exception.RssException;
+import org.apache.uniffle.common.util.Constants;
+import org.apache.uniffle.common.util.RetryUtils;
+
+
+public class TezRemoteShuffleManager implements ServicePluginLifecycle {
+  private static final Logger LOG = LoggerFactory.getLogger(TezRemoteShuffleManager.class);
+
+  protected InetSocketAddress address;
+
+  protected volatile Server server;
+  private String tokenIdentifier;
+  private Token<JobTokenIdentifier> sessionToken;
+  private Configuration conf;
+  private TezRemoteShuffleUmbilicalProtocolImpl tezRemoteShuffleUmbilical;
+  private ShuffleWriteClient rssClient;
+  private String appId;
+
+  public TezRemoteShuffleManager(String tokenIdentifier, Token<JobTokenIdentifier> sessionToken,
+                                 Configuration conf, String appId, ShuffleWriteClient rssClient) {
+    this.tokenIdentifier = tokenIdentifier;
+    this.sessionToken = sessionToken;
+    this.conf = conf;
+    this.appId = appId;
+    this.rssClient = rssClient;
+    this.tezRemoteShuffleUmbilical = new TezRemoteShuffleUmbilicalProtocolImpl();
+  }
+
+  @Override
+  public void initialize() throws Exception {
+
+  }
+
+  @Override
+  public void start() throws Exception {
+    startRpcServer();
+  }
+
+  @Override
+  public void shutdown() throws Exception {
+    server.stop();
+  }
+
+  private class TezRemoteShuffleUmbilicalProtocolImpl implements TezRemoteShuffleUmbilicalProtocol {
+    private Map<Integer, ShuffleAssignmentsInfo> shuffleIdToShuffleAssignsInfo = new HashMap<>();
+
+    @Override
+    public long getProtocolVersion(String s, long l) throws IOException {
+      return versionID;
+    }
+
+    @Override
+    public ProtocolSignature getProtocolSignature(String protocol, long clientVersion,
+                                                  int clientMethodsHash) throws IOException {
+      return ProtocolSignature.getProtocolSignature(this, protocol,
+              clientVersion, clientMethodsHash);
+    }
+
+    @Override
+    public GetShuffleServerResponse getShuffleAssignments(GetShuffleServerRequest request)
+            throws IOException, TezException {
+      if (request != null) {
+        LOG.info("getShuffleAssignments with request = " + request);
+      } else {
+        LOG.info("getShuffleAssignments with request is null");
+      }
+      int shuffleId = request.getShuffleId();
+
+      GetShuffleServerResponse response = new GetShuffleServerResponse();
+      ShuffleAssignmentsInfo shuffleAssignmentsInfo;
+      try {
+        synchronized (TezRemoteShuffleUmbilicalProtocolImpl.class) {

Review Comment:
   Why this is locked? And why lock the `class` rather than `instance`



##########
client-tez/src/main/java/org/apache/tez/dag/app/RssDAGAppMaster.java:
##########
@@ -0,0 +1,363 @@
+/*
+ * 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.dag.app;
+
+import java.lang.reflect.Field;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.ThreadFactory;
+import java.util.concurrent.TimeUnit;
+
+import org.apache.commons.cli.CommandLine;
+import org.apache.commons.cli.GnuParser;
+import org.apache.commons.cli.Options;
+import org.apache.hadoop.conf.Configuration;
+import org.apache.hadoop.fs.FSDataOutputStream;
+import org.apache.hadoop.fs.FileStatus;
+import org.apache.hadoop.fs.FileSystem;
+import org.apache.hadoop.fs.Path;
+import org.apache.hadoop.fs.permission.FsPermission;
+import org.apache.hadoop.mapreduce.JobSubmissionFiles;
+import org.apache.hadoop.security.Credentials;
+import org.apache.hadoop.security.UserGroupInformation;
+import org.apache.hadoop.util.ShutdownHookManager;
+import org.apache.hadoop.yarn.YarnUncaughtExceptionHandler;
+import org.apache.hadoop.yarn.api.ApplicationConstants;
+import org.apache.hadoop.yarn.api.records.ApplicationAttemptId;
+import org.apache.hadoop.yarn.api.records.ContainerId;
+import org.apache.hadoop.yarn.api.records.LocalResource;
+import org.apache.hadoop.yarn.conf.YarnConfiguration;
+import org.apache.hadoop.yarn.util.Clock;
+import org.apache.hadoop.yarn.util.ConverterUtils;
+import org.apache.hadoop.yarn.util.SystemClock;
+import org.apache.tez.common.RssTezConfig;
+import org.apache.tez.common.RssTezUtils;
+import org.apache.tez.common.TezClassLoader;
+import org.apache.tez.common.TezCommonUtils;
+import org.apache.tez.common.TezUtilsInternal;
+import org.apache.tez.common.VersionInfo;
+import org.apache.tez.dag.api.TezConfiguration;
+import org.apache.tez.dag.api.TezConstants;
+import org.apache.tez.dag.api.TezException;
+import org.apache.tez.dag.api.records.DAGProtos;
+import org.apache.tez.dag.api.records.DAGProtos.AMPluginDescriptorProto;
+import org.apache.tez.dag.library.vertexmanager.ShuffleVertexManager;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import org.apache.uniffle.client.api.ShuffleWriteClient;
+
+import static org.apache.tez.common.TezCommonUtils.TEZ_SYSTEM_SUB_DIR;
+
+public class RssDAGAppMaster extends DAGAppMaster {
+  private static final Logger LOG = LoggerFactory.getLogger(RssDAGAppMaster.class);
+  private ShuffleWriteClient shuffleWriteClient;
+  private TezRemoteShuffleManager tezRemoteShuffleManager;
+  private final String rssConfFileLocalResourceName = "rss_conf.xml";
+  private DAGProtos.PlanLocalResource rssConfFileLocalResource;
+  final ScheduledExecutorService scheduledExecutorService = Executors.newSingleThreadScheduledExecutor(
+    new ThreadFactory() {
+      @Override
+      public Thread newThread(Runnable r) {
+        Thread t = Executors.defaultThreadFactory().newThread(r);
+        t.setDaemon(true);
+        return t;
+      }
+    }
+  );
+
+  public RssDAGAppMaster(ApplicationAttemptId applicationAttemptId,
+                         ContainerId containerId, String nmHost, int nmPort, int nmHttpPort,
+                         Clock clock, long appSubmitTime, boolean isSession, String workingDirectory,
+                         String [] localDirs, String[] logDirs, String clientVersion,
+                         Credentials credentials, String jobUserName, AMPluginDescriptorProto pluginDescriptorProto) {
+    super(applicationAttemptId, containerId, nmHost, nmPort, nmHttpPort, clock, appSubmitTime, isSession,
+            workingDirectory, localDirs, logDirs, clientVersion, credentials, jobUserName, pluginDescriptorProto);
+  }
+
+  public ShuffleWriteClient getShuffleWriteClient() {
+    return shuffleWriteClient;
+  }
+
+  public void setShuffleWriteClient(ShuffleWriteClient shuffleWriteClient) {
+    this.shuffleWriteClient = shuffleWriteClient;
+  }
+
+  public TezRemoteShuffleManager getTezRemoteShuffleManager() {
+    return tezRemoteShuffleManager;
+  }
+
+  public void setTezRemoteShuffleManager(TezRemoteShuffleManager tezRemoteShuffleManager) {
+    this.tezRemoteShuffleManager = tezRemoteShuffleManager;
+  }
+
+  /**
+   * Init and Start Rss Client
+   * @param appMaster
+   * @param conf
+   * @param applicationAttemptId
+   * @throws Exception
+   */
+  public static void initAndStartRSSClient(final RssDAGAppMaster appMaster, Configuration conf,
+                                           ApplicationAttemptId  applicationAttemptId) throws Exception {

Review Comment:
   I think you should load the checkstyle file into IDEA .



##########
client-tez/src/main/java/org/apache/tez/dag/app/TezRemoteShuffleManager.java:
##########
@@ -0,0 +1,276 @@
+/*
+ * 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.dag.app;
+
+import java.io.IOException;
+import java.net.InetSocketAddress;
+import java.util.Arrays;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+
+import com.google.common.collect.Sets;
+import org.apache.commons.lang3.StringUtils;
+import org.apache.hadoop.conf.Configuration;
+import org.apache.hadoop.fs.CommonConfigurationKeysPublic;
+import org.apache.hadoop.ipc.ProtocolSignature;
+import org.apache.hadoop.ipc.RPC;
+import org.apache.hadoop.ipc.Server;
+import org.apache.hadoop.net.NetUtils;
+import org.apache.hadoop.security.authorize.PolicyProvider;
+import org.apache.hadoop.security.token.Token;
+import org.apache.tez.common.GetShuffleServerRequest;
+import org.apache.tez.common.GetShuffleServerResponse;
+import org.apache.tez.common.RssTezConfig;
+import org.apache.tez.common.RssTezUtils;
+import org.apache.tez.common.ServicePluginLifecycle;
+import org.apache.tez.common.ShuffleAssignmentsInfoWritable;
+import org.apache.tez.common.TezRemoteShuffleUmbilicalProtocol;
+import org.apache.tez.common.security.JobTokenIdentifier;
+import org.apache.tez.common.security.JobTokenSecretManager;
+import org.apache.tez.dag.api.TezConfiguration;
+import org.apache.tez.dag.api.TezException;
+import org.apache.tez.dag.api.TezUncheckedException;
+import org.apache.tez.dag.app.security.authorize.TezAMPolicyProvider;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import org.apache.uniffle.client.api.ShuffleWriteClient;
+import org.apache.uniffle.client.util.ClientUtils;
+import org.apache.uniffle.common.PartitionRange;
+import org.apache.uniffle.common.RemoteStorageInfo;
+import org.apache.uniffle.common.ShuffleAssignmentsInfo;
+import org.apache.uniffle.common.ShuffleDataDistributionType;
+import org.apache.uniffle.common.ShuffleServerInfo;
+import org.apache.uniffle.common.exception.RssException;
+import org.apache.uniffle.common.util.Constants;
+import org.apache.uniffle.common.util.RetryUtils;
+
+
+public class TezRemoteShuffleManager implements ServicePluginLifecycle {
+  private static final Logger LOG = LoggerFactory.getLogger(TezRemoteShuffleManager.class);
+
+  protected InetSocketAddress address;
+
+  protected volatile Server server;
+  private String tokenIdentifier;
+  private Token<JobTokenIdentifier> sessionToken;
+  private Configuration conf;
+  private TezRemoteShuffleUmbilicalProtocolImpl tezRemoteShuffleUmbilical;
+  private ShuffleWriteClient rssClient;
+  private String appId;
+
+  public TezRemoteShuffleManager(String tokenIdentifier, Token<JobTokenIdentifier> sessionToken,
+                                 Configuration conf, String appId, ShuffleWriteClient rssClient) {
+    this.tokenIdentifier = tokenIdentifier;
+    this.sessionToken = sessionToken;
+    this.conf = conf;
+    this.appId = appId;
+    this.rssClient = rssClient;
+    this.tezRemoteShuffleUmbilical = new TezRemoteShuffleUmbilicalProtocolImpl();
+  }
+
+  @Override
+  public void initialize() throws Exception {
+
+  }
+
+  @Override
+  public void start() throws Exception {
+    startRpcServer();
+  }
+
+  @Override
+  public void shutdown() throws Exception {
+    server.stop();
+  }
+
+  private class TezRemoteShuffleUmbilicalProtocolImpl implements TezRemoteShuffleUmbilicalProtocol {
+    private Map<Integer, ShuffleAssignmentsInfo> shuffleIdToShuffleAssignsInfo = new HashMap<>();
+
+    @Override
+    public long getProtocolVersion(String s, long l) throws IOException {
+      return versionID;
+    }
+
+    @Override
+    public ProtocolSignature getProtocolSignature(String protocol, long clientVersion,
+                                                  int clientMethodsHash) throws IOException {
+      return ProtocolSignature.getProtocolSignature(this, protocol,
+              clientVersion, clientMethodsHash);
+    }
+
+    @Override
+    public GetShuffleServerResponse getShuffleAssignments(GetShuffleServerRequest request)
+            throws IOException, TezException {
+      if (request != null) {
+        LOG.info("getShuffleAssignments with request = " + request);
+      } else {
+        LOG.info("getShuffleAssignments with request is null");
+      }
+      int shuffleId = request.getShuffleId();
+
+      GetShuffleServerResponse response = new GetShuffleServerResponse();
+      ShuffleAssignmentsInfo shuffleAssignmentsInfo;
+      try {
+        synchronized (TezRemoteShuffleUmbilicalProtocolImpl.class) {
+          if (shuffleIdToShuffleAssignsInfo.containsKey(shuffleId)) {
+            shuffleAssignmentsInfo = shuffleIdToShuffleAssignsInfo.get(shuffleId);
+          } else {
+            shuffleAssignmentsInfo = getShuffleWorks(request.getPartitionNum(), shuffleId);
+          }
+
+          if (shuffleAssignmentsInfo == null) {
+            response.setStatus(-1);
+            response.setRetMsg("shuffleAssignmentsInfo is null");
+          } else {
+            response.setStatus(0);
+            response.setRetMsg("");
+            response.setShuffleAssignmentsInfoWritable(new ShuffleAssignmentsInfoWritable(shuffleAssignmentsInfo));
+            shuffleIdToShuffleAssignsInfo.put(shuffleId, shuffleAssignmentsInfo);
+          }
+        }
+      } catch (Exception rssException) {
+        response.setStatus(-2);
+        response.setRetMsg(rssException.getMessage());
+      }
+
+      return response;
+    }
+  }
+
+  private ShuffleAssignmentsInfo getShuffleWorks(int partitionNum, int shuffleId) {
+    ShuffleAssignmentsInfo shuffleAssignmentsInfo;
+    int requiredAssignmentShuffleServersNum = RssTezUtils.getRequiredShuffleServerNumber(conf, 200, partitionNum);
+    // retryInterval must bigger than `rss.server.heartbeat.timeout`, or maybe it will return the same result
+    long retryInterval = conf.getLong(RssTezConfig.RSS_CLIENT_ASSIGNMENT_RETRY_INTERVAL,

Review Comment:
   Where is `RssTezConfig`



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@uniffle.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@uniffle.apache.org
For additional commands, e-mail: issues-help@uniffle.apache.org