Stage 8 — YARN Integration

What this stage teaches

Stage 8 lives at the Tez/YARN boundary. You learn:

  • How the Tez AM acquires and renews its AMRMToken, and the canonical bug: long-running session AMs (multi-day Hive sessions) whose AMRMToken expires while the AM is mid-RPC.
  • Log aggregation: how Tez's container exit hooks interact with the NM's LogAggregationService. The canonical symptom: missing container logs after AM crash because the AM never told the NM to flush.
  • The NM aux service: the Tez ShuffleHandler (or the MR ShuffleHandler when configured) lives in tez-plugins/tez-aux-services. Version mismatches between AM-side tez-runtime-library and NM-side aux service cause shuffle failures with cryptic error messages.
  • Kerberos delegation token renewal across DAG lifecycles, especially when multiple DAGs in a session use the same Credentials object.
  • TezClient AMRMToken handling: where the token lives in the submitter process versus the AM.

Patches in this stage are 50–400 lines but often require a Hadoop-version- specific code path, so the tez-plugins/tez-aux-services profile structure matters more than in other stages.

Reading order

  1. tez-api/src/main/java/org/apache/tez/client/TezClient.java
  2. tez-dag/src/main/java/org/apache/tez/dag/app/DAGAppMaster.java — focus on the AMRMToken handling and the credential propagation.
  3. tez-plugins/tez-aux-services/src/main/java/org/apache/tez/auxservices/ShuffleHandler.java
  4. The deep dive yarn-integration.
cd ~/tez-src
grep -rn "AMRMToken\|getCredentials\|TokenIdentifier" \
  tez-api/src/main/java tez-dag/src/main/java | head -30
ls tez-plugins/tez-aux-services/src/main/java/org/apache/tez/auxservices/

JIRA filter to find candidates

project = TEZ
  AND component in ("tez-dag", "tez-plugins")
  AND resolution = Unresolved
  AND (text ~ "AMRMToken" OR text ~ "kerberos" OR text ~ "delegation token"
       OR text ~ "log aggregation" OR text ~ "ShuffleHandler"
       OR text ~ "aux service" OR description ~ "TokenExpired")
ORDER BY updated DESC

A second filter focused on long-running session bugs:

project = TEZ AND text ~ "session" AND text ~ "expired"
  AND resolution = Unresolved

Walked example A — AMRMToken expiry on long DAGs

Symptom: a Hive session AM runs for 36 hours. On hour 24 it starts logging:

SecretManager$InvalidToken: AMRMToken for application appattempt_X has expired.

The AM crashes mid-DAG. The user loses the long-running session and resubmits all in-progress queries.

Step 1 — Trace token lifetime

cd ~/tez-src
grep -n "AMRMToken\|registerApplicationMaster" \
  tez-dag/src/main/java/org/apache/tez/dag/app/DAGAppMaster.java
grep -rn "renewMaxLifetime\|token-max-lifetime" tez-api tez-dag tez-common

YARN's yarn.resourcemanager.am-rm-tokens.master-key-rolling-interval-secs default is 24h. When the RM rotates the master key, the AM's cached AMRMToken becomes invalid. The fix is to detect a token-expired exception on the AMRM heartbeat path and re-acquire the token from the RM (which already exposes this via the heartbeat response in modern Hadoop versions).

Step 2 — Choose the right Hadoop version

tez-aux-services and tez-dag build against the configured Hadoop profile:

grep -rn "hadoop28\|hadoop29\|hadoop31" pom.xml | head

Token rollover handling differs across Hadoop minor versions. The patch must be a no-op on profiles where the Hadoop client already handles the rollover transparently. Confirm by:

grep -rn "AMRMToken" ~/hadoop-src/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-client | head

If AMRMClientAsyncImpl already loops on token expiry in Hadoop 3.x, your Tez patch is a Hadoop-2.x-only path guarded by an availability check.

Step 3 — Diff

--- a/tez-dag/src/main/java/org/apache/tez/dag/app/DAGAppMaster.java
+++ b/tez-dag/src/main/java/org/apache/tez/dag/app/DAGAppMaster.java
@@
   private void heartbeatLoop() {
     while (!shutdownRequested) {
       try {
         AllocateResponse resp = amRmClient.allocate(progress);
+        // Hadoop 2.x clients did not transparently refresh the AMRMToken
+        // on master-key rollover. Detect token-expired and re-acquire.
+        // See TEZ-XXXX.
+        if (resp.getAMRMToken() != null) {
+          UserGroupInformation.getCurrentUser().addToken(
+              ConverterUtils.convertFromYarn(resp.getAMRMToken(), null));
+        }
         processAllocations(resp);
       } catch (InvalidToken e) {
+        LOG.warn("AMRMToken invalid for {}, attempting re-register", appAttemptID);
+        try {
+          amRmClient.registerApplicationMaster(host, port, trackingUrl);
+          continue;
+        } catch (Exception reErr) {
+          LOG.error("Re-register failed; AM will exit", reErr);
+          throw new TezUncheckedException(reErr);
+        }
       } catch (Exception e) {
         ...
       }
     }
   }

Step 4 — Test

A unit test stubs the AMRMClient to return an InvalidToken once then a healthy response, and asserts that registerApplicationMaster was called once and the loop continued. Pattern:

@Test(timeout = 10000)
public void testAmrmTokenReacquiredOnInvalidToken() throws Exception {
  AMRMClient mockRm = mock(AMRMClient.class);
  when(mockRm.allocate(anyFloat()))
      .thenThrow(new InvalidToken("expired"))
      .thenReturn(emptyAllocateResponse());
  DAGAppMaster am = createTestAM(mockRm);
  am.runOneHeartbeatIteration();
  verify(mockRm).registerApplicationMaster(anyString(), anyInt(), anyString());
  am.runOneHeartbeatIteration();
  // second iteration must succeed
}

A MiniYARNCluster test that triggers an actual key rollover is possible but slow; the unit test above is sufficient for review.

Walked example B — log aggregation race on AM crash

Symptom: an AM crashes (OutOfMemoryError). The cluster operator runs yarn logs -applicationId ... and gets nothing. The NodeManager's LogAggregationService reports the logs as never finalised.

Root cause: the JVM crashed before Tez's DAGAppMaster.shutdown() could flag the logs as aggregation-ready. NM's default is "wait for the AM to mark finalisation" rather than aggregating on container exit.

The fix

Tez registers a JVM shutdown hook (Runtime.getRuntime().addShutdownHook) that calls into the YARN LogAggregationContext to force-finalise. The hook must run before the JVM's normal exit handlers.

cd ~/tez-src
grep -n "addShutdownHook\|LogAggregationContext" \
  tez-dag/src/main/java/org/apache/tez/dag/app/DAGAppMaster.java

If a shutdown hook is registered but does not handle OutOfMemoryError, add a defensive try/catch (Throwable) and ensure the hook is the first shutdown hook registered (so it runs last and after other hooks have cleaned up).

The diff is small; the test is hard. The accepted pattern is a logged-evidence test: spin up a MiniYARNCluster, submit a DAG, kill -9 the AM process, and assert that the NM log aggregation finalised the logs within a bounded time. This test belongs in tez-tests and is slow (~30s).

Walked example C — NM aux service version mismatch

Symptom: a cluster operator deploys Tez 0.10.3 but the NMs still run the Tez 0.10.1 aux service. Shuffle fails with:

IOException: Unknown shuffle handler version: 2; expected 1

The fix is in tez-aux-services plus a docs note: the aux service on every NM must match the AM-side tez-runtime-library minor version. The Tez patch is twofold:

  1. The aux service must report its version in the protocol handshake.
  2. The client side must produce a self-describing error message that names the NM, the version it reported, and the version the AM expected.
--- a/tez-runtime-library/src/main/java/org/apache/tez/runtime/library/common/shuffle/orderedgrouped/Fetcher.java
+++ b/tez-runtime-library/src/main/java/org/apache/tez/runtime/library/common/shuffle/orderedgrouped/Fetcher.java
@@
-    if (serverVersion != EXPECTED_SHUFFLE_VERSION) {
-      throw new IOException("Unknown shuffle handler version: " + serverVersion);
+    if (serverVersion != EXPECTED_SHUFFLE_VERSION) {
+      throw new IOException(String.format(
+          "Tez shuffle handler version mismatch on %s:%d: server=%d, expected=%d. "
+              + "Likely cause: NodeManager aux-service jar is older than the AM. "
+              + "Ensure tez-aux-services-%s.jar is deployed to every NM.",
+          host, port, serverVersion, EXPECTED_SHUFFLE_VERSION,
+          TezVersionInfo.getVersion()));
     }

The patch is one improved error message and one documentation update in docs/src/site/markdown/install.md.

Pitfalls

  • Don't add a new JVM shutdown hook without considering ordering. Java does not guarantee shutdown hook order; if two hooks rely on each other, you must serialise them explicitly.
  • Don't catch Throwable outside a shutdown path. Catching Throwable in the heartbeat loop will swallow OutOfMemoryError and leave the AM in an undefined state.
  • Don't conflate AMRMToken with delegation tokens. AMRMToken authenticates the AM to the RM; delegation tokens authenticate the AM/tasks to HDFS or other services. Renewal paths and lifetimes are different.
  • Don't deploy a fix that requires the operator to redeploy tez-aux-services without saying so in the release notes. Aux service upgrades require an NM restart; that is operationally expensive.
  • Don't assume the Hadoop version on disk is the Hadoop version in production. Test against the minimum Hadoop version supported by your Tez release line (see pom.xml profile defs).
  • Don't hard-code token renewal intervals. Use the YARN-side configuration keys directly (yarn.resourcemanager.am-rm-tokens.master-key-rolling-interval-secs).

Exit criteria — when you're ready for the next stage

Move to Stage 9 when:

  • You have shipped one YARN-integration patch with evidence (in the JIRA description) of which Hadoop minor versions you tested against.
  • You can describe the AMRMToken lifecycle in five sentences including the master-key rollover.
  • You have read the LogAggregationContext API in the Hadoop source and understand the logIncludePattern / logExcludePattern interplay.
  • You have a tez-plugins/tez-aux-services build that runs locally and you understand which NMs need it.

Stage 9 returns to the in-repo skill set with a focus on test stability.