TEZ-4559: Fix Retry logic in case of Recovery

apache · May 6, 2024 · 4a970a6 · 4a970a6
1 parent 66a6ca6
commit 4a970a6
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 2 deletions.
diff --git a/tez-api/src/main/java/org/apache/tez/dag/api/client/DAGClientImpl.java b/tez-api/src/main/java/org/apache/tez/dag/api/client/DAGClientImpl.java
@@ -410,8 +410,14 @@ private DAGStatus getDAGStatusViaAM(@Nullable Set<StatusGetOpts> statusOptions,
       LOG.info("DAG is no longer running - application not found by YARN", e);
       dagCompleted = true;
     } catch (NoCurrentDAGException e) {
-      LOG.info("Got NoCurrentDAGException from AM, returning a failed DAG", e);
-      return dagLost();
+      if (conf.getBoolean(TezConfiguration.DAG_RECOVERY_ENABLED,
+          TezConfiguration.DAG_RECOVERY_ENABLED_DEFAULT)) {
+        LOG.info("Got NoCurrentDAGException from AM, going on as recovery is enabled", e);
+      } else {
+        // if recovery is disabled, we're not expecting the DAG to be finished any time in the future
+        LOG.info("Got NoCurrentDAGException from AM, returning a failed DAG as recovery is disabled", e);
+        return dagLost();
+      }
     } catch (TezException e) {
       // can be either due to a n/w issue or due to AM completed.
       LOG.info("Cannot retrieve DAG Status due to TezException: {}", e.getMessage());

diff --git a/tez-api/src/test/java/org/apache/tez/dag/api/client/rpc/TestDAGClient.java b/tez-api/src/test/java/org/apache/tez/dag/api/client/rpc/TestDAGClient.java
@@ -715,6 +715,7 @@ public void testGetDagStatusWithCachedStatusExpiration() throws Exception {
   @Test
   public void testDagClientReturnsFailedDAGOnNoCurrentDAGException() throws Exception {
     TezConfiguration tezConf = new TezConfiguration();
+    tezConf.setBoolean(TezConfiguration.DAG_RECOVERY_ENABLED, false);
 
     try (DAGClientImplForTest dagClientImpl = new DAGClientImplForTest(mockAppId, dagIdStr, tezConf, null)) {