Lab 4.2: VertexManager Deep Dive

Background

The VertexManager is the hook that makes Tez more than a static DAG scheduler. By plugging in a VertexManagerPlugin, an application controls, at runtime, when a vertex's tasks are scheduled and how many of them run — slow start, auto-parallelism, skew handling, custom routing — without patching the Application Master. When StartTransition fires in the state machine you read in Lab 4.1, the very next thing VertexImpl.startVertex() does is call vertexManager.onVertexStarted(...). That single call is the seam between the core AM and pluggable scheduling policy.

This lab reads the real contract in tez-api, dissects the two production plugins (ImmediateStartVertexManager and ShuffleVertexManager), and pins down the exact configuration keys and algorithm that Tez uses to auto-reduce reducer parallelism. It is a code-reading lab; you write no code here (you build a plugin in Lab 4.3 and investigate a ShuffleVertexManager edge case in Lab 4.4).

Deep-dive companions: vertex-lifecycle.md, event-routing.md, and the Level 4 overview.

Verify, never assume. VertexManager method names have changed across releases (scheduleVertexTasks → scheduleTasks, setVertexParallelism → reconfigureVertex). Every API in this lab is quoted from the source on current master. When you read your own checkout, re-grep; if a signature differs, trust your source over this page and note the drift.

Why This Lab Matters for Contributors

ShuffleVertexManager is one of the most consequential classes in Tez: it decides how many reducers every shuffle stage runs, on every Hive-on-Tez and Pig-on-Tez query in production. Its bugs show up as "too many tiny reducers", "auto-parallelism didn't kick in", "reducers started too early and starved mappers of slots", and — the subject of Lab 4.4 — edge cases when a source produces no output. To triage or fix any of these you must be able to read the slow-start fraction math and the auto-parallelism reduction algorithm and connect a config key to the line that reads it. This lab builds exactly that fluency.

Prerequisites

• Completed Lab 4.1; you know StartTransition calls the VertexManager.
• Tez checkout that builds.
• These files open:
  • tez-api/src/main/java/org/apache/tez/dag/api/VertexManagerPlugin.java
  • tez-api/src/main/java/org/apache/tez/dag/api/VertexManagerPluginContext.java
  • tez-dag/src/main/java/org/apache/tez/dag/app/dag/impl/ImmediateStartVertexManager.java
  • tez-runtime-library/src/main/java/org/apache/tez/dag/library/vertexmanager/ShuffleVertexManager.java
  • tez-runtime-library/src/main/java/org/apache/tez/dag/library/vertexmanager/ShuffleVertexManagerBase.java
• A reading log: mkdir -p ~/tez-notes && : > ~/tez-notes/reading-log-4.2.md

Step-by-Step Tasks

Step 1 — Read the VertexManagerPlugin contract

cat tez-api/src/main/java/org/apache/tez/dag/api/VertexManagerPlugin.java

It is an abstract class, not an interface, and its constructor takes the context (quoted from module tez-api, class VertexManagerPlugin):

public abstract class VertexManagerPlugin {
  private final VertexManagerPluginContext context;

  public VertexManagerPlugin(VertexManagerPluginContext context) {
    this.context = context;
  }

  public abstract void initialize() throws Exception;

  public void onVertexStarted(List<TaskAttemptIdentifier> completions) throws Exception { ... }

  public void onSourceTaskCompleted(TaskAttemptIdentifier attempt) throws Exception { ... }

  public abstract void onVertexManagerEventReceived(VertexManagerEvent vmEvent) throws Exception;

  public abstract void onRootVertexInitialized(String inputName,
      InputDescriptor inputDescriptor, List<Event> events) throws Exception;

  public final VertexManagerPluginContext getContext() { return this.context; }
}

Note the details you must get right:

• The context arrives via the constructor, then initialize() is called — a subclass must provide a (VertexManagerPluginContext) constructor so Tez can reflectively instantiate it.
• onVertexStarted and onSourceTaskCompleted each have a modern overload taking TaskAttemptIdentifier and a @Deprecated overload taking Map<String,List<Integer>> / (String, Integer). The modern overloads default to adapting the call down to the deprecated ones, so a plugin may override either. Record which each production plugin overrides.
• onVertexManagerEventReceived and onRootVertexInitialized are abstract — every plugin must implement them, even if with an empty body.
• onVertexStateUpdated is a non-abstract callback for VertexStateUpdate notifications (the public VertexState enum from Lab 4.1) that fires only after registerForVertexStateUpdates.

Step 2 — Read VertexManagerPluginContext — the callbacks into the AM

cat tez-api/src/main/java/org/apache/tez/dag/api/VertexManagerPluginContext.java

The methods a plugin actually uses (all quoted/verified from module tez-api, interface VertexManagerPluginContext):

ScheduleTaskRequest is a static nested class with a factory:

public static ScheduleTaskRequest create(int taskIndex, @Nullable TaskLocationHint locationHint) { ... }

Log it: the current scheduling method is scheduleTasks(List<ScheduleTaskRequest>); the current reconfigure method is reconfigureVertex(...). The scheduleVertexTasks and setVertexParallelism variants are @Deprecated — do not use them in new code.

How the callbacks are driven. The plugin never runs on the AM's main thread directly. The AM wraps every plugin in an internal VertexManager dispatcher (tez-dag, org.apache.tez.dag.app.dag.impl.VertexManager) that receives VertexManagerEvents and lifecycle signals off the event queue and calls the plugin's methods, catching any exception and turning it into the V_MANAGER_USER_CODE_ERROR event you met in Lab 4.1 (VertexManagerUserCodeErrorTransition). That is why a throw from your plugin fails the vertex rather than the AM. Grep it: grep -n "class VertexManager" tez-dag/src/main/java/org/apache/tez/dag/app/dag/impl/VertexManager.java.

Step 3 — Read ImmediateStartVertexManager (the baseline)

cat tez-dag/src/main/java/org/apache/tez/dag/app/dag/impl/ImmediateStartVertexManager.java

This is the default manager for vertices that need no slow start. Answer from the code:

1. What does initialize() do? (Very little — it inspects source-edge configuration and sets a flag; the interesting behavior is in onVertexStarted.)
2. In onVertexStarted(List<TaskAttemptIdentifier>), it sets onVertexStartedDone and calls the private scheduleTasks(), which builds a list of ScheduleTaskRequest for all tasks and calls getContext().scheduleTasks(tasksToStart) once. Why one call and not one per task? (Batching avoids N separate scheduling round-trips through the AM's event queue.)
3. Does it override onSourceTaskCompleted? What does it use it for? (It tracks source completion so it can start once inputs are ready.)

This is the "schedule everything immediately" reference point from which ShuffleVertexManager diverges by adding slow start and auto-parallelism.

Step 4 — ShuffleVertexManager: the configuration keys

grep -n "TEZ_SHUFFLE_VERTEX_MANAGER" \
  tez-runtime-library/src/main/java/org/apache/tez/dag/library/vertexmanager/ShuffleVertexManager.java | head -30

The real keys and defaults (quoted from module tez-runtime-library, class ShuffleVertexManager):

Note the two orthogonal features these split into:

• Slow start is governed by min-src-fraction / max-src-fraction and is always on.
• Auto-parallelism (reducing reducer count) is governed by enable.auto-parallel (off by default), desired-task-input-size, and min-task-parallelism (the floor — never reduce below this; default 1). Record that the floor is 1 — it matters in Lab 4.4.

Step 5 — Slow start: the scheduling fraction formula

Most of the shared logic lives in the base class:

grep -n "getNumOfTasksToSchedule\|percentRange\|tasksFractionToSchedule" \
  tez-runtime-library/src/main/java/org/apache/tez/dag/library/vertexmanager/ShuffleVertexManagerBase.java | head

Read getNumOfTasksToSchedule(float). The core (quoted from module tez-runtime-library, class ShuffleVertexManagerBase):

if (numBipartiteSourceTasksCompleted == totalNumBipartiteSourceTasks) {
  LOG.info("All source tasks completed. Ramping up {} remaining tasks ...");
  return numPendingTasks;
}
// linearly increase the number of scheduled tasks such that all tasks are
// scheduled when source tasks completed fraction reaches max
float tasksFractionToSchedule = 1;
float percentRange = config.getMaxFraction() - config.getMinFraction();
if (percentRange > 0) {
  tasksFractionToSchedule =
      (minSourceVertexCompletedTaskFraction - config.getMinFraction()) / percentRange;
} else {
  // min and max are equal. schedule 100% on reaching min
  if (minSourceVertexCompletedTaskFraction < config.getMinFraction()) {
    tasksFractionToSchedule = 0;
  }
}

Answer:

1. Which field tracks how many source tasks have completed? (numBipartiteSourceTasksCompleted, against totalNumBipartiteSourceTasks.)
2. Below min-src-fraction, how many downstream tasks are scheduled? (Zero — the fraction goes negative and is clamped.) At max-src-fraction and above? (All of them.)
3. State the linear rule in one sentence: between min and max, the fraction of downstream tasks scheduled rises linearly with the fraction of completed source tasks.

There is one more subtlety in the constructor: the effective max fraction is clamped so it is never below min. Grep the config setup:

grep -n "Math.max(slowStartMinFraction\|Math.max(1, conf" \
  tez-runtime-library/src/main/java/org/apache/tez/dag/library/vertexmanager/ShuffleVertexManager.java

You will find the max fraction is built as Math.max(slowStartMinFraction, <configured max>) — a mis-set max-src-fraction below min-src-fraction is silently corrected up, so the ramp is always well-formed. Also note canScheduleTasks() in the base: it returns false until every source vertex reports vertexIsConfigured, because until a source is configured Tez cannot know whether numTasks == 0 is legitimate — a detail that matters directly in Lab 4.4.

Step 6 — Auto-parallelism: what a VertexManagerEvent carries

When auto-parallelism is on, each completed upstream task sends a VertexManagerEvent whose payload is a protobuf carrying its output size and per-partition statistics. Read the handler:

grep -n "onVertexManagerEventReceived\|handleVertexManagerEvent\|VertexManagerEventPayloadProto\|getOutputSize\|hasPartitionStats" \
  tez-runtime-library/src/main/java/org/apache/tez/dag/library/vertexmanager/ShuffleVertexManagerBase.java | head

The decode (quoted from ShuffleVertexManagerBase.handleVertexManagerEvent):

VertexManagerEventPayloadProto proto;
try {
  proto = VertexManagerEventPayloadProto.parseFrom(
      ByteString.copyFrom(vmEvent.getUserPayload()));
} catch (InvalidProtocolBufferException e) {
  throw new TezUncheckedException(e);
}
sourceTaskOutputSize = proto.getOutputSize();

if (proto.hasPartitionStats()) {
  ... RoaringBitmap partitionStats ...
  parsePartitionStats(srcInfo, partitionStats);
} else if (proto.hasDetailedPartitionStats()) {
  parseDetailedPartitionStats(srcInfo, proto.getDetailedPartitionStats().getSizeInMbList());
}
srcInfo.numVMEventsReceived++;
srcInfo.outputSize += sourceTaskOutputSize;
completedSourceTasksOutputSize += sourceTaskOutputSize;

Answer:

1. What protobuf message is decoded? (VertexManagerEventPayloadProto, from ShuffleUserPayloads.)
2. What is accumulated across all events? (completedSourceTasksOutputSize — total bytes — plus a per-partition statsInMB array via parsePartitionStats, which decodes a compressed RoaringBitmap of per-partition size buckets.)
3. Why is handleVertexManagerEvent guarded by taskWithVmEvents.add(producerTask)? (Multiple attempts of the same task produce identical output, so duplicates are ignored.)

Step 7 — The reduction algorithm: getComputeRoutingAction + computeRouting

Two methods decide whether and how to reduce parallelism.

First, whether — getComputeRoutingAction in the base returns one of three actions (enum ComputeRoutingAction { WAIT, SKIP, COMPUTE }):

if (getNumOfTasksToSchedule(minSourceVertexCompletedTaskFraction) <= 0 &&
    numBipartiteSourceTasksCompleted != totalNumBipartiteSourceTasks) {
  return ComputeRoutingAction.WAIT;          // not enough completed source tasks yet
} else if (numVertexManagerEventsReceived == 0 && totalNumBipartiteSourceTasks > 0) {
  return ComputeRoutingAction.SKIP;          // sources produced no output -> no VMEs
} else if (completedSourceTasksOutputSize < config.getDesiredTaskInputDataSize()
    && (minSourceVertexCompletedTaskFraction < config.getMaxFraction())) {
  return ComputeRoutingAction.WAIT;          // too little data seen; wait for more up to max
} else {
  return ComputeRoutingAction.COMPUTE;
}

Then, how — ShuffleVertexManager.computeRouting(). The essential arithmetic (quoted from module tez-runtime-library, class ShuffleVertexManager, method computeRouting):

int currentParallelism = pendingTasks.size();
BigInteger expectedTotalSourceTasksOutputSize =
    getExpectedTotalBipartiteSourceTasksOutputSize();
BigInteger desiredTaskInputDataSize = BigInteger.valueOf(config.getDesiredTaskInputDataSize());
// ceil(expected / desired)
BigInteger bigDesiredTaskParallelism =
    expectedTotalSourceTasksOutputSize.add(desiredTaskInputDataSizeMinusOne)
        .divide(desiredTaskInputDataSize);
...
int desiredTaskParallelism = bigDesiredTaskParallelism.intValue();
if (desiredTaskParallelism < mgrConfig.getMinTaskParallelism()) {
  desiredTaskParallelism = mgrConfig.getMinTaskParallelism();   // floor, default 1
}
if (desiredTaskParallelism >= currentParallelism) {
  return null;   // never increase; only reduce
}
basePartitionRange = currentParallelism / desiredTaskParallelism;
if (basePartitionRange <= 1) {
  return null;   // reducing by less than half isn't worth breaking the desired size
}

Answer, faithfully to the code:

1. How is the target parallelism derived? (Ceiling of expected total output size ÷ desired-task-input-size; the expectation projects from completed tasks via getExpectedTotalBipartiteSourceTasksOutputSize, using BigInteger to avoid the integer overflow that TEZ-3452 / TEZ-3666 once caused.)
2. What is the floor? (min-task-parallelism, default 1 — this clamp is what prevents a divide-by-zero at currentParallelism / desiredTaskParallelism.)
3. When does computeRouting decline to reduce? (When desiredTaskParallelism >= currentParallelism, or when basePartitionRange <= 1, i.e. it would only combine fewer than two partitions per task.)
4. When the reduction is applied, what does it call? Trace to getContext().reconfigureVertex(finalTaskParallelism, null, edgeProperties) and note that a CustomShuffleEdgeManager is installed so upstream events re-route to the coalesced partitions.

Step 7b — Work a reduction by hand

Pin the arithmetic with real numbers. Suppose a shuffle vertex is planned with currentParallelism = 100 reducers, desired-task-input-size = 100 MB, min-task-parallelism = 1, and the accumulated statistics project expectedTotalSourceTasksOutputSize = 900 MB. Follow computeRouting:

• bigDesiredTaskParallelism = ceil(900 MB / 100 MB) = 9.
• 9 >= min-task-parallelism (1), so no clamp.
• 9 < currentParallelism (100), so a reduction is warranted.
• basePartitionRange = 100 / 9 = 11 (> 1, so it proceeds).
• numShufflersWithBaseRange = 100 / 11 = 9; remainderRangeForLastShuffler = 100 % 11 = 1.
• finalTaskParallelism = 9 + 1 = 10 (remainder is non-zero, so one extra shuffler mops up the leftover partition).

So 100 reducers collapse to 10, each consuming ~10 of the original 100 partitions, and reconfigureVertex(10, null, edgeProperties) is called. Now redo it for a tiny input: expectedTotalSourceTasksOutputSize = 5 MB. Then bigDesiredTaskParallelism = 1, which equals the floor; basePartitionRange = 100 / 1 = 100; the manager reduces 100 → 1. Confirm from the code that this is legal (the floor is 1, not 0) and that the basePartitionRange <= 1 early-return does not fire here because the range is 100. This second case is exactly the boundary Lab 4.4 probes.

Step 8 — Read the test harness

find tez-runtime-library/src/test -name "TestShuffleVertexManager*.java"

TestShuffleVertexManagerUtils.createVertexManagerContext(...) builds a Mockito mock of VertexManagerPluginContext, stubbing getInputVertexEdgeProperties, getVertexName, getVertexNumTasks, and installing doAnswer(...) handlers that record scheduleTasks and reconfigureVertex calls. getVertexManagerEvent(long[] partitionSizes, long total, String vertex) builds a real VertexManagerEventPayloadProto. Read one test — testAutoParallelismConfig or testSchedulingWithPartitionStats — and note the pattern: build manager → feed source completions and VM events → verify(mockContext).reconfigureVertex(eq(N), any(), anyMap()). You will reuse this exact harness in Lab 4.4.

Deliverables

• The VertexManagerPlugin lifecycle methods listed, with which are abstract and which have deprecated overloads.
• The context method table with the current scheduling (scheduleTasks) and reconfigure (reconfigureVertex) methods identified, deprecated ones flagged.
• ImmediateStartVertexManager behavior summarized (one batched scheduleTasks call for all tasks).
• The five ShuffleVertexManager config keys with their property strings and defaults.
• The slow-start linear formula, in your own words, tied to getNumOfTasksToSchedule.
• The auto-parallelism decision (WAIT/SKIP/COMPUTE) and reduction algorithm, including the min-task-parallelism floor and the "never increase / reduce by at least half" rules.

Troubleshooting

Stretch Goals

1. CartesianProductVertexManager. Read find tez-runtime-library/src/main/java -name "CartesianProductVertexManager.java". What cross-product scheduling does it coordinate, and how does it use reconfigureVertex differently from ShuffleVertexManager?
2. FairShuffleVertexManager. It shares ShuffleVertexManagerBase but overrides computeRouting. Diff the two computeRouting implementations and describe what "fair" routing changes about partition-to-task assignment.
3. A real JIRA. Find a resolved ShuffleVertexManager issue (git log --oneline -- '*ShuffleVertexManager*'), read its commit, and state the invariant it restored. TEZ-3666 (integer overflow) and TEZ-1248 (special-case 1 reducer) are good starting points and directly relevant to Lab 4.4.

Validation / Self-check

Answer in your own words, citing the class for each:

1. Which state-machine transition invokes the VertexManager, and which plugin method does it call?

2. How does a plugin schedule tasks and reconfigure parallelism on current master, and which older methods are deprecated?

3. Describe ImmediateStartVertexManager in one sentence, including why it batches its scheduling into a single scheduleTasks call.

4. Which two config keys govern slow start, and what happens to the number of scheduled tasks below min-src-fraction, between the fractions, and above max-src-fraction?

5. Which config key enables auto-parallelism, which sets the target input size per task, and what is the parallelism floor (and its default)?

6. Walk the auto-parallelism decision: what makes getComputeRoutingAction return WAIT, SKIP, and COMPUTE, and how does computeRouting turn accumulated output size into a new parallelism?

7. What does a VertexManagerEvent payload carry, and why are duplicate events from multiple attempts of the same task ignored?

8. In the worked example (Step 7b), 100 reducers with 900 MB of expected output reduce to 10. Show the intermediate values (basePartitionRange, remainderRangeForLastShuffler, finalTaskParallelism) and explain the "+1" for the remainder shuffler.

When you can trace a config key to the line that reads it and explain the reduction as arithmetic on real fields, you have completed Lab 4.2. Continue to Lab 4.3: Build It — WavingVertexManager.