Stage 8 — YARN Integration

What this stage teaches

Stage 8 lives at the Tez/YARN boundary — the layer where the AM talks to the ResourceManager, holds and reuses containers, and races its own shutdown. The bugs here are distributed-systems bugs: a container completes while the DAG is being torn down, a scheduler holds thousands of container objects for the lifetime of a long Hive session, an idle-timeout computation blows up on a boundary value, two priorities invert and the whole DAG hangs. You learn:

• The container lifecycle in the AM: YarnTaskSchedulerService / DagAwareYarnTaskScheduler request, hold, reuse, and release containers; the AMContainerMap / AMContainerImpl track their state machine.
• The shutdown/cleanup races: events about a container can arrive after the DAG that owned it has started tearing down. Every getAllContainers().get(id) can return null.
• Memory pressure in long-running session AMs: objects that accumulate per-DAG and are never pruned.
• Why almost every fix in this stage carries a MockClock / mock-AMRMClient test, because you cannot spin a real RM up in a unit test.

Patches are 2–150 lines of production code with a much larger test, because the only way to prove a race fix is to drive the race deterministically with a mock scheduler and a controllable clock.

What these bugs look like from the operator's chair: a Hive session that has been healthy for hours suddenly hangs a query, or the AM's heap creeps up over a day until it OOMs, or a scheduler NPEs in a stack trace that mentions container cleanup during what should have been a clean DAG completion. None of these reproduce on a single small DAG on your laptop — that is the tell that you are in Stage 8. The symptom always involves time (a long-running AM) or concurrency (shutdown overlapping normal work), never a simple deterministic input.

Prerequisite: Stage 7 plus the deep dives: YARN integration, scheduler, and container reuse. This stage assumes you know what a HeldContainer is and how container reuse is decided.

Reading order in the checkout

cd /Users/s0x/src/oss-repos/tez
grep -rln "class YarnTaskSchedulerService\|class DagAwareYarnTaskScheduler\|class TaskSchedulerManager\|class AMContainerMap" \
  tez-dag/src/main/java
ls tez-dag/src/main/java/org/apache/tez/dag/app/rm/
ls tez-dag/src/main/java/org/apache/tez/dag/app/rm/container/

Finding Stage 8 issues today

project = TEZ AND resolution = Unresolved
  AND component in ("tez-dag", "Scheduler")
  AND (summary ~ "container" OR summary ~ "scheduler" OR summary ~ "allocat"
       OR summary ~ "preempt" OR summary ~ "NullPointer" OR summary ~ "hang"
       OR summary ~ "reuse" OR summary ~ "AMRM")
ORDER BY priority DESC, updated DESC

Heuristic grep — the container-cleanup race has one shape, getAllContainers().get(id) followed by an immediate dereference:

grep -rn "getAllContainers().get(" tez-dag/src/main/java | grep -v "!= null\|== null" | head

Every hit is a candidate for the exact bug TEZ-3932 fixed below.

The subsystem knowledge you need

• The container state machine. An AMContainerImpl moves through ALLOCATED → (idle / running) → COMPLETED, plus error states. AMContainerMap owns the live set. Two components read that map constantly: the TaskSchedulerManager (routing attempt events) and the schedulers (reuse decisions). Any of them can observe a container mid-transition or already gone.

• Reuse and the two schedulers. Tez keeps containers alive between tasks to avoid JVM-launch cost. The legacy YarnTaskSchedulerService and the newer DagAwareYarnTaskScheduler implement the hold/match/release policy. A HeldContainer carries an idle-expiration timestamp; get its arithmetic wrong (TEZ-4081) and the scheduler crashes on a legal config. See container reuse.

• The session AM changes the rules. A single-DAG AM dies after its DAG, so leaks and stale references never bite. A session AM (the Hive default) runs for days across thousands of DAGs — so every per-DAG collection must be pruned at dagComplete(), every long-lived reference is a potential leak (TEZ-3643), and shutdown races (TEZ-3932) are routine because DAGs start and stop constantly.

• Two kinds of token, do not conflate them. The AMRMToken authenticates the AM to the ResourceManager on the allocate/heartbeat path; the RM rotates its master key on an interval, and a long session AM must tolerate that. Delegation tokens authenticate the AM and its tasks to HDFS and other services and have their own renewers and lifetimes. They live in different places, expire on different schedules, and a fix for one is not a fix for the other. When you read a token bug, first classify which token, then trace its renewal path.

Internalise this: the session AM is the adversary. Most Stage 8 bugs are single-DAG-invisible and session-fatal.

The container lifecycle and the race window, in one diagram:

flowchart TD
  RM[ResourceManager] -->|allocate| SCHED[TaskScheduler]
  SCHED -->|new AMContainer| MAP[AMContainerMap]
  MAP --> ALLOC[ALLOCATED]
  ALLOC --> RUN[running / idle-reuse]
  RUN --> COMP[COMPLETED]
  COMP -->|dagComplete prunes| GONE[removed from map]
  subgraph RACE[the race window]
    EV[late attempt-ended event] -.->|getAllContainers.get id| GONE
    GONE -.->|returns null → NPE| BOOM[TEZ-3932]
  end

The dotted path is the bug: an event that references a container arrives after dagComplete() pruned it. Every reader of AMContainerMap lives inside that window.

Case study A — TEZ-3932: NPE in a container-cleanup race

The archetypal Stage 8 bug. Read it:

git show 72c458a43   # TaskSchedulerManager NPE during DAGAppMaster container cleanup race

The symptom. During DAG shutdown, TaskSchedulerManager throws an NPE. The DAG was completing normally; a container-completion event arrived a beat too late, after the container had already been removed from the map.

The root cause. Three call sites did the same unguarded thing:

sendEvent(new AMNodeEventTaskAttemptEnded(appContext.getAllContainers()
    .get(attemptContainerId).getContainer().getNodeId(), ...));

getAllContainers().get(attemptContainerId) returns null once the container has been cleaned up, and .getContainer() NPEs. This is a textbook time-of-check race: the container existed when the attempt started, but not when the ended event was processed.

The fix — hoist the lookup, null-check, then use. Jonathan Eagles pulled the lookup into a local, checked it, and only sent the event if the container was still present:

AMContainer amContainer = appContext.getAllContainers().get(attemptContainerId);
// DAG can be shutting down so protect against container cleanup race
if (amContainer != null) {
  Container container = amContainer.getContainer();
  sendEvent(new AMNodeEventTaskAttemptEnded(container.getNodeId(),
      event.getSchedulerId(), attemptContainerId,
      attempt.getID(), event.getState() == TaskAttemptState.FAILED));
}

The same treatment was applied to the success path and the launch-request path, each with the same // DAG can be shutting down comment.

What the test did. TestTaskSchedulerManager was extended with a scenario that drives a task-attempt-ended event for a container that is no longer in the map, and asserts no NPE and that the downstream event is (or isn't) sent. It uses the drainable-context test harness so the events are processed synchronously.

The lesson. In the AM, shutdown is concurrent with everything. Any lookup into a shared map (containers, nodes, attempts) can miss because another thread is tearing down. The fix pattern is always the same: hoist to a local, null-check, proceed. Learn to see the map.get(x).method() shape as a bug on sight.

Case study B — TEZ-3491: a DAG hangs on container priority inversion

git show a93dbf0b2   # Tez job can hang due to container priority inversion

The symptom. A DAG hangs — never completes, never fails. It happens when the scheduler holds a new container at one priority while tasks that need that priority are waiting, and nothing ever rematches them.

The root cause. YarnTaskSchedulerService had rescheduling logic inline in two different places. When a delayed container completed, the code path just logged "Ignoring unknown container" and did nothing to re-drive the waiting requests at that priority. The task requests sat forever.

The fix — extract the rematch, call it from both paths. Jason Lowe factored the "reschedule every task request at this priority" loop into one method:

private void maybeRescheduleContainerAtPriority(Priority priority) {
  for (Map.Entry<Object, CookieContainerRequest> entry : taskRequests.entrySet()) {
    Object task = entry.getKey();
    CookieContainerRequest request = entry.getValue();
    if (request.getPriority().equals(priority)) {
      LOG.info("Resending request for task again: " + task);
      deallocateTask(task, true, null, null);
      allocateTask(task, request.getCapability(), ...);
      break;
    }
  }
}

and called it from the container-completion path that previously did nothing:

if (delayedContainer != null) {
  LOG.info("Delayed container {} completed", containerStatus.getContainerId());
  maybeRescheduleContainerAtPriority(delayedContainer.getContainer().getPriority());
} else {
  LOG.info("Ignoring unknown container: " + containerStatus.getContainerId());
}

What the test did. TestTaskScheduler.testContainerExpired builds a mock TezAMRMClientAsync, configures reuse-locality delay and idle-release timeouts to zero so matching happens in one pass, allocates two tasks at two priorities, expires a container, and asserts the waiting request is re-driven. It uses the TaskSchedulerContextDrainable wrapper so callbacks are deterministic — no sleeps.

The lesson. A "hang" bug is a missing state transition, not a crash. When you see logic duplicated in two event handlers and one of them is a no-op ("Ignoring unknown container"), suspect that the no-op branch is missing a wake-up. The scheduler must always re-drive waiting requests when the resource situation changes.

Case study C — TEZ-3643: a long session AM runs out of memory

git show 605154203   # Long running AMs can go out of memory due to retained AMContainer instances

The symptom. A multi-day Hive session AM slowly grows its heap and eventually OOMs. Each DAG in the session allocates containers; the AMContainer objects were never released after the DAG finished.

The root cause. AMContainerMap accumulated every AMContainer ever created and only removed them on AM shutdown. For a session running thousands of DAGs, that is thousands of retained container objects, each holding references to signatures, credentials, and node info.

The fix — prune completed containers on DAG completion. Siddharth Seth added an isInErrorState() accessor to the AMContainer interface and a cleanup pass that dagComplete() now invokes:

private void cleanupCompletedContainers() {
  Iterator<Map.Entry<ContainerId, AMContainer>> iterator = containerMap.entrySet().iterator();
  int count = 0;
  while (iterator.hasNext()) {
    Map.Entry<ContainerId, AMContainer> entry = iterator.next();
    AMContainer amContainer = entry.getValue();
    if (AMContainerState.COMPLETED.equals(amContainer.getState())
        || amContainer.isInErrorState()) {
      iterator.remove();
      count++;
    }
  }
  LOG.info("Cleaned up completed containers on dagComplete. Removed={}, Remaining={}",
      count, containerMap.size());
}

To make the map testable, the field was exposed @VisibleForTesting and the AMContainerImpl construction was extracted into an overridable createAmContainer(...) factory so TestAMContainerMap can inject mock containers in COMPLETED/error states and assert they are pruned while running ones survive.

The lesson. Long-running AMs turn any per-DAG collection into a leak. When you add a map keyed by container/attempt/vertex, ask: what removes entries, and when? "On AM shutdown" is the wrong answer for a session AM. The right boundary is dagComplete(). And notice the fix's testability changes — @VisibleForTesting and a factory method — are part of the patch, not an afterthought.

Boundary-value sibling: TEZ-4081 (git show b078e3a25) is a two-line fix in DagAwareYarnTaskScheduler: when the idle-release min and max timeouts are set equal, random.nextLong(min, max) throws IllegalArgumentException (empty range). The guard picks min directly when min == max. Read it as a reminder that a config the operator is allowed to set (min == max) must not crash the scheduler. Its test, testMinMaxContainerIdleMillisAreEqual, uses a MockClock.

The contribution playbook for this class

1. Grep the race shape first. getAllContainers().get(id).method() and friends are TEZ-3932 waiting to happen.
2. Reproduce with the mock scheduler harness, never a real cluster. The TaskSchedulerContextDrainable / MockClock / mock-AMRMClient pattern (see TestTaskScheduler, TestDagAwareYarnTaskScheduler, TestTaskSchedulerManager) is how every fix above was proven. Copy the closest existing test.
3. For a leak, name the removal boundary. Add the prune where the owning scope ends (dagComplete), and assert both "removed the dead" and "kept the live."
4. For a hang, find the missing transition. Look for a no-op branch in an event handler that should be re-driving requests.
5. Make it testable in the same patch. @VisibleForTesting, a factory method, or an injectable clock — reviewers expect them.
6. Run the scheduler suites.

mvn -pl tez-dag test -Dtest=TestTaskScheduler,TestTaskSchedulerManager,TestDagAwareYarnTaskScheduler -q
mvn -pl tez-dag test -Dtest=TestAMContainerMap,TestAMContainer -q

A note on evidence, because YARN fixes are hard to prove: a reviewer will not take "I tested it on a cluster" — clusters are nondeterministic. The currency here is a unit test that forces the exact interleaving (drive the late event, expire the container, set min == max) and asserts the outcome. If you cannot express the bug as a deterministic mock-scheduler scenario, you do not yet understand it well enough to fix it. That is not a bureaucratic hurdle; it is the difference between a patch that fixes the symptom you saw and one that fixes the race.

Common mistakes

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

• You have shipped one YARN-integration patch with a deterministic mock-scheduler test (drainable context + MockClock), no Thread.sleep.
• You can spot the map.get(id).method() race on sight and explain why TEZ-3932's null-check-and-hoist fixes it.
• You can explain the difference between a hang (TEZ-3491, missing transition) and a crash (TEZ-3932, NPE), and how you'd reproduce each.
• You can name the removal boundary for a per-DAG collection and know why "on AM shutdown" is wrong for a session AM.
• You have read YarnTaskSchedulerService and AMContainerMap around the fixed sites without feeling lost.

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