Logic to kill orphan DAGs left behind by crashed HS2

Author: tanishq-chugh

ql/src/java/org/apache/hadoop/hive/ql/exec/tez/TezExternalSessionState.java

@@ -29,8 +29,12 @@
 import org.apache.hadoop.security.Credentials;
 import org.apache.hadoop.yarn.api.records.ApplicationId;
 import org.apache.tez.client.TezClient;
+import org.apache.tez.dag.api.DAG;
 import org.apache.tez.dag.api.TezConfiguration;
 import org.apache.tez.dag.api.TezException;
+import org.apache.tez.dag.api.client.DAGClient;
+import org.apache.tez.dag.api.client.rpc.DAGClientAMProtocolBlockingPB;
+import org.apache.tez.dag.api.client.rpc.DAGClientAMProtocolRPC;
 import org.apache.tez.serviceplugins.api.ServicePluginsDescriptor;
 
 /**
@@ -187,4 +191,35 @@ public boolean killQuery(String reason) throws HiveException {
     killQuery.killQuery(queryId, reason, conf, false);
     return true;
   }
+
+  @Override
+  public DAGClient submitDAG(DAG dag) throws TezException, IOException {
+    try {
+      return getTezClient().submitDAG(dag);
+    } catch (TezException e) {
+      if (e.getMessage() == null || !e.getMessage().contains("App master already running a DAG")) {
+        throw e;
+      }
+      tryKillRunningDAGs(getTezClient());
+      return getTezClient().submitDAG(dag);
+    }
+  }
+
+  private void tryKillRunningDAGs(TezClient session) throws TezException {
+    LOG.info("External session has an AM which is already running a DAG on app ID {}", externalAppId);
+    DAGClientAMProtocolBlockingPB proxy = session.sendAMHeartbeat(null);
+    if (proxy == null) {
+      throw new TezException("Error while trying to connect to AM for app ID " + externalAppId);
+    }
+    try {
+      DAGClientAMProtocolRPC.GetAllDAGsResponseProto allDAGSResponse =
+          proxy.getAllDAGs(null, DAGClientAMProtocolRPC.GetAllDAGsRequestProto.newBuilder().build());
+      for (String dagId : allDAGSResponse.getDagIdList()) {
+        LOG.info("External session: attempting to kill dagId {} on app ID {}", dagId, externalAppId);
+        proxy.tryKillDAG(null, DAGClientAMProtocolRPC.TryKillDAGRequestProto.newBuilder().setDagId(dagId).build());
+      }
+    } catch (Exception e) {
+      throw new TezException("Error while trying to kill existing DAG running on app ID " + externalAppId, e);
+    }
+  }
 }

ql/src/java/org/apache/hadoop/hive/ql/exec/tez/TezSession.java

@@ -34,7 +34,9 @@
 import org.apache.hadoop.hive.ql.wm.WmContext;
 import org.apache.hadoop.yarn.api.records.LocalResource;
 import org.apache.tez.client.TezClient;
+import org.apache.tez.dag.api.DAG;
 import org.apache.tez.dag.api.TezException;
+import org.apache.tez.dag.api.client.DAGClient;
 import org.apache.tez.dag.api.client.DAGStatus;
 
 /**
@@ -86,6 +88,7 @@ public String toString() {
 
   HiveConf getConf();
   TezClient getTezClient();
+  DAGClient submitDAG(DAG dag) throws TezException, IOException;
   boolean isOpen();
   boolean isOpening();
   boolean getDoAsEnabled();

ql/src/java/org/apache/hadoop/hive/ql/exec/tez/TezSessionPoolSession.java

@@ -34,7 +34,9 @@
 import org.apache.hadoop.hive.registry.impl.TezAmInstance;
 import org.apache.hadoop.yarn.api.records.LocalResource;
 import org.apache.tez.client.TezClient;
+import org.apache.tez.dag.api.DAG;
 import org.apache.tez.dag.api.TezException;
+import org.apache.tez.dag.api.client.DAGClient;
 import org.apache.tez.dag.api.client.DAGStatus;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -337,6 +339,11 @@ public TezClient getTezClient() {
     return baseSession.getTezClient();
   }
 
+  @Override
+  public DAGClient submitDAG(DAG dag) throws TezException, IOException {
+    return baseSession.submitDAG(dag);
+  }
+
   @Override
   public boolean isOpening() {
     return baseSession.isOpening();

ql/src/java/org/apache/hadoop/hive/ql/exec/tez/TezSessionState.java

@@ -76,11 +76,13 @@
 import org.apache.hadoop.yarn.conf.YarnConfiguration;
 import org.apache.tez.client.TezClient;
 import org.apache.tez.common.TezUtils;
+import org.apache.tez.dag.api.DAG;
 import org.apache.tez.dag.api.PreWarmVertex;
 import org.apache.tez.dag.api.SessionNotRunning;
 import org.apache.tez.dag.api.TezConfiguration;
 import org.apache.tez.dag.api.TezException;
 import org.apache.tez.dag.api.UserPayload;
+import org.apache.tez.dag.api.client.DAGClient;
 import org.apache.tez.dag.api.client.DAGStatus;
 import org.apache.tez.dag.api.client.Progress;
 import org.apache.tez.mapreduce.hadoop.DeprecatedKeys;
@@ -820,6 +822,11 @@ public TezClient getTezClient() {
     return session;
   }
 
+  @Override
+  public DAGClient submitDAG(DAG dag) throws TezException, IOException {
+    return getTezClient().submitDAG(dag);
+  }
+
   @Override
   public LocalResource getAppJarLr() {
     return appJarLr;

ql/src/java/org/apache/hadoop/hive/ql/exec/tez/TezTask.java

@@ -692,7 +692,7 @@ DAGClient submit(DAG dag, Ref<TezSession> sessionStateRef) throws Exception {
 
   private DAGClient submitInternal(DAG dag, TezSession sessionState) throws TezException, IOException {
     runtimeContext.init(sessionState);
-    return sessionState.getTezClient().submitDAG(dag);
+    return sessionState.submitDAG(dag);
   }
 
   private void sessionDestroyOrReturnToPool(Ref<TezSession> sessionStateRef,