DAGAppMaster

DAGAppMaster is Tez's YARN ApplicationMaster: one JVM, launched by the ResourceManager, that owns one or more DAGs over its lifetime. It is the largest single class in the codebase (~2,500 lines) and the hub every other subsystem plugs into. This chapter maps its bring-up (serviceInit / serviceStart), the child services it composes, the AsyncDispatcher that is the only mutation path for DAG state, the session idle loop, the state machine, recovery on AM restart, and the plugin points that let you replace the scheduler, launcher, or task communicator.

After this chapter you should be able to map any AM log line in the first 60 seconds of operation to a method in DAGAppMaster.java, and explain what happens on attempt 2 after a crash.

Prerequisites: dag-model.md (the plan it executes), tez-client.md (who launches it), and event-routing.md (the events it dispatches). Hands-on: Lab 3.1.

Files to open

ls tez-dag/src/main/java/org/apache/tez/dag/app/

tez-dag/src/main/java/org/apache/tez/dag/app/
  DAGAppMaster.java              (the AM main class)
  DAGAppMasterState.java         (the AM-level state enum)
  TaskCommunicatorManager.java   (task umbilical multiplexer)
  TaskHeartbeatHandler.java      (task liveness)
  ContainerHeartbeatHandler.java (container liveness)
  RecoveryParser.java            (replay of the recovery log on restart)
  rm/
    TaskSchedulerManager.java    (routes AMSchedulerEvents to scheduler plugins)
    YarnTaskSchedulerService.java(default YARN AMRM scheduler)
    container/AMContainerImpl.java (container state machine)
  launcher/
    ContainerLauncherManager.java
    LocalContainerLauncher.java  (in-process, local mode)
  dag/impl/
    DAGImpl.java, VertexImpl.java, TaskImpl.java, TaskAttemptImpl.java
  history/
    HistoryEventHandler.java
    recovery/RecoveryService.java (writes the recovery event log)
tez-dag/src/main/java/org/apache/tez/dag/api/client/
  DAGClientServer.java           (the client-facing RPC server)
  DAGClientHandler.java          (dispatches client RPCs to the AM/DAG)

Bring-up: serviceInit then serviceStart

DAGAppMaster extends AbstractService. The static main() parses YARN environment (appAttemptId, submit time, client version, local/log dirs), installs the Tez classloader, constructs the AM, and runs init() then start(). serviceInit is where every child service is created and registered; serviceStart is where they are started and the first DAG begins.

grep -n "serviceInit\|initServices\|startServices\|serviceStart\|serviceStop" \
  tez-dag/src/main/java/org/apache/tez/dag/app/DAGAppMaster.java

The head of serviceInit establishes environment before any service exists:

// tez-dag: DAGAppMaster.serviceInit (trimmed, in order)
this.frameworkService = getFrameworkService(conf);          // YARN vs local plugin host
this.hadoopShim = new HadoopShimsLoader(this.amConf).getHadoopShim();
this.isLocal = conf.getBoolean(TezConfiguration.TEZ_LOCAL_MODE,
    TezConfiguration.TEZ_LOCAL_MODE_DEFAULT);
UserPayload defaultPayload = TezUtils.createUserPayloadFromConf(amConf);
PluginManager.PluginDescriptors pluginDescriptors =
    pluginManager.parseAllPlugins(isLocal, defaultPayload);   // scheduler/launcher/comm plugins
// ... client-AM version compatibility check (Simple2LevelVersionComparator) ...
dispatcher = createDispatcher();
context = new RunningAppContext(conf);
clientHandler = new DAGClientHandler(this);
addIfService(dispatcher, false);
recoveryEnabled = conf.getBoolean(TezConfiguration.DAG_RECOVERY_ENABLED,
    TezConfiguration.DAG_RECOVERY_ENABLED_DEFAULT);          // default true
initClientRpcServer();

Then the child services are created and added, each with a dispatcher registration. The real list, from serviceInit:

// tez-dag: DAGAppMaster.serviceInit (dispatcher registrations, trimmed)
dispatcher.register(AMContainerEventType.class, containers);
dispatcher.register(AMNodeEventType.class, nodes);
dispatcher.register(DAGAppMasterEventType.class, new DAGAppMasterEventHandler());
dispatcher.register(DAGEventType.class, dagEventDispatcher);
dispatcher.register(VertexEventType.class, vertexEventDispatcher);
if (!useConcurrentDispatcher) {
  dispatcher.register(TaskEventType.class, new TaskEventDispatcher());
  dispatcher.register(TaskAttemptEventType.class, new TaskAttemptEventDispatcher());
} else {
  AsyncDispatcherConcurrent shared = dispatcher.registerAndCreateDispatcher(
      TaskEventType.class, new TaskEventDispatcher(), "TaskAndAttemptEventThread", concurrency);
  dispatcher.registerWithExistingDispatcher(TaskAttemptEventType.class,
      new TaskAttemptEventDispatcher(), shared);
}
dispatcher.registerAndCreateDispatcher(SpeculatorEventType.class,
    new SpeculatorEventHandler(), "Speculator");
dispatcher.register(AMSchedulerEventType.class, taskSchedulerManager);
dispatcher.register(ContainerLauncherEventType.class, containerLauncherManager);

The useConcurrentDispatcher branch (tez.am.use.concurrent-dispatcher, off by default) is a scaling knob: for huge DAGs, task and attempt events are sharded across tez.am.concurrent-dispatcher.concurrency threads. It is off by default precisely because a single dispatch thread gives the simplest correctness story (see the single-thread invariant below).

initServices then calls init() on each added service, and startServices (from serviceStart) calls start() on each — order matters, which is why services are added with explicit dependencies (addIfServiceDependency).

The AsyncDispatcher — the one mutation path

Every state change to a DAG, vertex, task, or attempt happens by emitting an event onto the AsyncDispatcher, which delivers it to the registered handler on a single dispatch thread. The handler runs the entity's state-machine transition and may emit follow-on events. There is no other legal way to mutate state.

flowchart TB
    subgraph Sources
      TC["Task umbilical heartbeat"]
      SCH["Scheduler callback (RM)"]
      TL["Container launcher"]
      UC["Client: submitDAG / tryKillDAG"]
      RC["Recovery replay (restart)"]
    end
    TC --> D
    SCH --> D
    TL --> D
    UC --> D
    RC --> D
    D["AsyncDispatcher\n(single thread)"] --> DH["DagEventDispatcher"]
    D --> VH["VertexEventDispatcher"]
    D --> TH["TaskEventDispatcher"]
    D --> AH["TaskAttemptEventDispatcher"]
    D --> SH["TaskSchedulerManager"]
    D --> CL["ContainerLauncherManager"]
    D --> HE["HistoryEventHandler"]
    DH --> DI["DAGImpl"]
    VH --> VI["VertexImpl"]
    TH --> TI["TaskImpl"]
    AH --> TAI["TaskAttemptImpl"]

Handlers reach shared AM state through one read-only seam: AppContext (the RunningAppContext instance built in serviceInit). It exposes getCurrentDAG(), getCurrentDAGID(), getClock(), getAppAttemptID(), and the service handles, so a handler can look up the live DAGImpl without holding a reference to the AM. Treat AppContext as the AM's service locator — if you add a subsystem, you wire it in here so events can find it.

Why single-threaded matters: the state machines (state-machines.md) are not internally synchronized against concurrent transitions on the same entity; correctness relies on all events for a given entity being processed serially. A handler that blocks on I/O stalls the entire bus — hence the rule "handlers mutate state and emit events; they never do blocking I/O." When isLocal is false the AM calls dispatcher.enableExitOnDispatchException(), so an uncaught exception in a handler crashes the AM rather than silently corrupting state.

Session vs non-session at serviceStart

grep -n "public void serviceStart\|In Session mode\|In Non-Session mode\|DAGAppMasterState.IDLE\|recoverDAG" \
  tez-dag/src/main/java/org/apache/tez/dag/app/DAGAppMaster.java

// tez-dag: DAGAppMaster.serviceStart (trimmed)
startServices();
super.serviceStart();
if (isSession && !recoveryEnabled && appAttemptID.getAttemptId() > 1) {
  // A session AM on attempt 2 with recovery off cannot recover -> error out
  addDiagnostic(INVALID_SESSION_ERR_MSG);
  this.state = DAGAppMasterState.ERROR;
  this.taskSchedulerManager.setShouldUnregisterFlag();
  shutdownHandler.shutdown();
  return;
}
// emit AMStartedEvent to history
DAGRecoveryData recoveredDAGData = recoveryEnabled ? recoverDAG() : null;
if (!isSession) {
  LOG.info("In Non-Session mode.");
  dagPlan = readDAGPlanFile();                 // the staged tez-dag.pb
} else {
  LOG.info("In Session mode. Waiting for DAG over RPC");
  this.state = DAGAppMasterState.IDLE;
}

In non-session mode the AM immediately reads the plan it was launched with and starts it. In session mode it goes IDLE and waits for the client's submitDAG RPC, which lands in submitDAGToAppMaster:

// tez-dag: DAGAppMaster.submitDAGToAppMaster (trimmed)
appMasterReadinessService.waitToBeReady();
if (sessionStopped.get()) {
  throw new SessionNotRunning("AM unable to accept new DAG submissions. In the process of shutting down");
}
synchronized (this) {
  if (state.equals(DAGAppMasterState.ERROR) || sessionStopped.get()) {
    throw new SessionNotRunning("AM unable to accept new DAG submissions. ...");
  }
  if (currentDAG != null && !currentDAG.isComplete()) {
    throw new TezException("App master already running a DAG");   // the concurrent-DAG guard
  }
  startDAG(dagPlan, additionalResources);
  return currentDAG.getID().toString();
}

That "App master already running a DAG" is the server-side answer to the "can I submit two DAGs to one session?" question from tez-client.md.

The AM state machine

grep -n "enum DAGAppMasterState" -A 12 tez-dag/src/main/java/org/apache/tez/dag/app/DAGAppMasterState.java

// tez-dag: org.apache.tez.dag.app.DAGAppMasterState
public enum DAGAppMasterState {
  NEW, INITED, RECOVERING, IDLE, RUNNING, SUCCEEDED, FAILED, KILLED, ERROR
}

NEW → INITED during init; RUNNING while a DAG executes; IDLE between DAGs in a session; RECOVERING on attempt > 1 while replaying the log; the terminal four on shutdown. The session idle-timeout check (checkAndHandleSessionTimeout) only fires when the state is not RUNNING or RECOVERING — you cannot time out a session that is mid-DAG.

YARN-facing components

AMRM heartbeat (the resource conversation)

TaskSchedulerManager routes AMSchedulerEventType events to the active scheduler; the default is YarnTaskSchedulerService, which owns a YARN AMRMClient and heartbeats the RM. Requests (new containers), releases (freed containers), and progress go up; allocations and completed-container statuses come down.

grep -n "heartbeat\|AMRMClient\|allocate\|CompletedContainer" \
  tez-dag/src/main/java/org/apache/tez/dag/app/rm/YarnTaskSchedulerService.java | head

The scheduler is itself a plugin point (below), so a custom scheduler substitutes here without touching the AM. Details in scheduler.md and yarn-integration.md.

Container and task liveness

ContainerHeartbeatHandler and TaskHeartbeatHandler (both extend HeartbeatHandlerBase) track the wall-clock time of the last umbilical contact. When a container or task goes silent past its timeout the AM declares it lost and fails the attempt. The task umbilical itself is served by TaskCommunicatorManager, an in-AM RPC server tasks call for getTask, heartbeat/statusUpdate, done/fatalError, and to push runtime events. The protocol is TezTaskUmbilicalProtocol:

find . -name "TezTaskUmbilicalProtocol.java"
grep -n "class TaskCommunicatorManager" tez-dag/src/main/java/org/apache/tez/dag/app/TaskCommunicatorManager.java

The client RPC server

DAGClientServer (an AbstractService wrapping a Hadoop RPC Server) exposes DAGClientAMProtocol to clients; it runs tez.am.client.am.thread-count handler threads and binds within tez.am.client.am.port-range. Requests land in DAGClientHandler, which resolves the TezDAGID and forwards to DAGImpl — the same getDAGStatus/tryKillDAG path the reader in dag-client.md drives from the other end.

Recovery: surviving an AM restart

YARN restarts a failed AM up to tez.am.max.app.attempts (default 2), reusing the appId but incrementing the attempt id. When tez.dag.recovery.enabled (default true) and this is attempt > 1, the AM replays a durable event log instead of starting from scratch.

// tez-dag: DAGAppMaster.recoverDAG (trimmed)
private DAGRecoveryData recoverDAG() throws IOException, TezException {
  if (recoveryEnabled) {
    if (this.appAttemptID.getAttemptId() > 1) {
      LOG.info("Recovering data from previous attempts, currentAttemptId="
          + this.appAttemptID.getAttemptId());
      this.state = DAGAppMasterState.RECOVERING;
      return parseDAGFromRecoveryData();     // RecoveryParser reads the prior attempt's log
    }
  }
  return null;
}

The write side is RecoveryService (tez-dag, extends AbstractService), a history handler that persists a subset of DAGHistoryEvents to ${tez.staging-dir}/<appId>/recovery/<attemptId>/ via per-DAG RecoveryStreams and a summaryStream. On restart, RecoveryParser reconstructs a DAGRecoveryData, and serviceStart fans out on it:

• Prior DAG completed or non-recoverable → the AM re-emits a DAGEventRecoverEvent with the final DAGState (or FAILED if non-recoverable), records a DAGRecoveredEvent to history, and goes RUNNING just long enough to report the outcome.
• Prior DAG recoverable and mid-flight → the AM emits a DAGEventRecoverEvent with the recovered data and calls currentDAG.onStart() to re-initialize vertex services (including speculators), then resumes.
• Prior attempt crashed while shutting down (isSessionStopped) → continue the shutdown and report SUCCEEDED.

Note: Recovery is per-DAG, not per-task. Completed tasks stay completed; tasks that were in flight get fresh attempts; a vertex that was RUNNING becomes RUNNING again. A DAG containing a CONCURRENT edge is currently deemed unrecoverable and is re-run from scratch (the source calls out TEZ-4017 for proper failover). Full failure semantics are in failure-handling.md.

sequenceDiagram
    participant RM as YARN RM
    participant AM2 as DAGAppMaster (attempt 2)
    participant RP as RecoveryParser
    participant DISP as AsyncDispatcher
    participant DAG as DAGImpl

    RM-->>AM2: relaunch (same appId, attemptId=2)
    AM2->>AM2: serviceStart; recoveryEnabled && attempt>1
    AM2->>RP: parseRecoveryData()
    RP-->>AM2: DAGRecoveryData (last durable state)
    AM2->>DISP: DAGEventRecoverEvent
    DISP->>DAG: replay -> rebuild in-memory state
    AM2->>DAG: onStart() (re-init vertex services)
    Note over AM2,DAG: completed tasks stay done;\nin-flight tasks re-attempted

Shutdown: unregister before you exit

An AM does not just System.exit. It must first tell the RM it is done — otherwise YARN treats the disappearance as a failure and (up to the attempt limit) relaunches it. Two paths matter:

• DAGAppMasterShutdownHandler.shutdown() — an inner class that runs the ordered teardown on its own thread, honoring tez.am.sleep.time.before.exit.millis (TEZ_AM_SLEEP_TIME_BEFORE_EXIT_MILLIS) so late history/ATS events can drain before the JVM dies. Every terminal path — version mismatch, invalid session, DAG completion in non-session mode, session timeout, an explicit shutdownSession RPC — funnels through it.
• serviceStop() — the AbstractService teardown:

// tez-dag: DAGAppMaster.serviceStop (trimmed)
if (isSession) { sessionStopped.set(true); }
if (this.dagSubmissionTimer != null) { this.dagSubmissionTimer.cancel(); }
initiateStop();          // release held containers BEFORE stopping services (TEZ-2687)
stopServices();          // stop child services in reverse dependency order
// delete the staging scratch dir only if the scheduler successfully unregistered
if (deleteTezScratchData && taskSchedulerManager != null
    && taskSchedulerManager.hasUnregistered()) {
  // appMasterUgi.doAs(... delete tezSystemStagingDir ...)
}

The load-bearing details: containers are released before services stop so the RM is not left holding allocations; and the scratch dir is deleted only if taskSchedulerManager.hasUnregistered() is true — under pre-emption an unclean exit deliberately leaves scratch data so a relaunched attempt can recover. Any terminal path that wants YARN to not retry sets taskSchedulerManager.setShouldUnregisterFlag() first, which is why you see that call before every shutdownHandler.shutdown().

Plugin points: ServicePluginsDescriptor

The scheduler, container launcher, and task communicator are all replaceable without forking the AM. The client expresses this with ServicePluginsDescriptor (org.apache.tez.serviceplugins.api), which is serialized into the plan as AMPluginDescriptorProto and parsed by PluginManager.parseAllPlugins during serviceInit:

grep -n "public static ServicePluginsDescriptor create" \
  tez-api/src/main/java/org/apache/tez/serviceplugins/api/ServicePluginsDescriptor.java
ls tez-api/src/main/java/org/apache/tez/serviceplugins/api/

// tez-api: ServicePluginsDescriptor.create (one overload)
public static ServicePluginsDescriptor create(
    TaskSchedulerDescriptor[] taskSchedulerDescriptor,
    ContainerLauncherDescriptor[] containerLauncherDescriptor,
    TaskCommunicatorDescriptor[] taskCommunicatorDescriptor) { ... }
// also: create(boolean enableUber), create(enableContainers, enableUber, ...), etc.

Each descriptor names a class implementing the corresponding SPI: TaskScheduler, ContainerLauncher, or the task-communicator interface. The AM builds a named list per SPI (the YARN default plus any custom ones), and a vertex's VertexExecutionContext (from dag-model.md) selects which named plugin handles it. This is exactly how the LLAP/external-service and uber (in-AM task) execution modes are wired — different scheduler + launcher + communicator, same DAGAppMaster. The enableUber flags map onto the in-process execution path shared with local-mode.md.

Reading exercise

# Bring-up boundaries
grep -n "serviceInit\|initServices\|startServices\|serviceStart" \
  tez-dag/src/main/java/org/apache/tez/dag/app/DAGAppMaster.java
sed -n '422,625p' tez-dag/src/main/java/org/apache/tez/dag/app/DAGAppMaster.java

# Every dispatcher registration
grep -n "dispatcher.register\|registerAndCreateDispatcher\|addIfService" \
  tez-dag/src/main/java/org/apache/tez/dag/app/DAGAppMaster.java

# Session control + recovery hookup
grep -n "isSession\|IDLE\|recoverDAG\|checkAndHandleSessionTimeout" \
  tez-dag/src/main/java/org/apache/tez/dag/app/DAGAppMaster.java

# Plugins
grep -n "parseAllPlugins\|ServicePluginsDescriptor\|VertexExecutionContext" \
  tez-dag/src/main/java/org/apache/tez/dag/app/DAGAppMaster.java

Answer, with citations:

1. In what order are the child services created in serviceInit, and which two dependencies are declared explicitly via addIfServiceDependency? Why does TaskSchedulerManager depend on DAGClientServer?
2. List the first three events that flow through the dispatcher when a non-session AM starts a DAG.
3. What thread does DAGImpl.handle run on relative to VertexImpl.handle when tez.am.use.concurrent-dispatcher is off vs on?
4. Where is the appAttemptId > 1 decision made, and what special case does a session AM with recovery disabled hit on attempt 2?
5. Distinguish checkAndHandleSessionTimeout (idle shutdown) from the client-heartbeat timeout (dead-client shutdown). Which config keys drive each?
6. Trace an umbilical heartbeat reporting task completion all the way to a TaskAttemptImpl terminal transition.

Common bugs and symptoms

Validation: prove you understand this

1. From memory, list ten EventType → handler registrations from serviceInit and name the subsystem each drives.
2. Draw the event path from TezTaskUmbilicalProtocol.heartbeat (task reports done) to TaskAttemptImpl.handle(TA_DONE), naming every dispatcher hop.
3. Bring up a single-DAG, non-session AM on a MiniCluster and identify, in the AM log, the line emitted by each child-service start().
4. Read RecoveryService and classify which DAGHistoryEvent types are persisted vs which are in-memory-only, then predict what a mid-vertex-init crash recovers to.
5. Explain, with the state machines from state-machines.md, why the dispatcher must be single-threaded per entity and exactly what breaks if two events for the same VertexImpl were processed concurrently.
6. Write a ServicePluginsDescriptor that keeps the YARN scheduler but adds a custom ContainerLauncher, and identify where in serviceInit your class name is turned into a live launcher.