Lab 4.4 — Fix It: `ShuffleVertexManager` Zero-Partition & Auto-Parallelism Edge Cases

Lab type: Fix-It — investigate → reproduce the bug class → write a regression test → reason about the fix Estimated time: 120–150 min Tez component: tez-runtime-library → org.apache.tez.dag.library.vertexmanager.ShuffleVertexManager (+ ShuffleVertexManagerBase)

Background

ShuffleVertexManager decides, at runtime, how many reduce tasks a shuffle stage runs. It does this from statistics that map tasks push to it as VertexManagerEvents: each completed producer reports its output size and per-partition sizes. The manager accumulates those, and once it has seen enough data, computeRouting() divides expected total output by desired-task-input-size to pick a new (smaller) parallelism.

This lab targets a whole class of edge-case bug in that path: what happens at the boundaries — when a source vertex produces zero output (all records filtered), when parallelism math would divide by zero, and when large inputs overflow the arithmetic. A naive implementation crashes or picks nonsense parallelism in these cases. Your job is investigative: find out what current master actually does (it turns out the guards exist — you must locate them), reconstruct the historical bugs from git that motivated them, then write a regression test against the real harness and describe what a credible fix PR for this class of bug looks like.

This is not a "there is definitely a live NPE" lab. The honest finding on current master is that the zero-output and divide-by-zero cases are already guarded. The contributor skill being trained is: prove that with a test, understand why the guard is there by reading the commits that added it, and know how you would have written the fix if the guard were missing.

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

Why This Lab Matters for Contributors

Edge-case parallelism bugs in ShuffleVertexManager are high-impact: every Hive/Pig query that shuffles hits this code. A wrong reducer count wastes a cluster; a crash fails a production DAG. The class has a real history of exactly these bugs — integer overflow with large inputs (TEZ-3452, TEZ-3666) and the one-reducer special case (TEZ-1248). Reading how those were found and fixed, and adding a regression test that pins a boundary, is a textbook mergeable contribution — and reproducing before patching is the etiquette maintainers expect.

Prerequisites

Completed Lab 4.2; you know the slow-start / auto-parallelism split and the WAIT/SKIP/COMPUTE decision.
Tez checkout that builds; mvn -pl tez-runtime-library test-compile works.
These files open:
- tez-runtime-library/.../vertexmanager/ShuffleVertexManager.java
- tez-runtime-library/.../vertexmanager/ShuffleVertexManagerBase.java
- tez-runtime-library/src/test/.../vertexmanager/TestShuffleVertexManager.java
- tez-runtime-library/src/test/.../vertexmanager/TestShuffleVertexManagerUtils.java
A scratch branch: git switch -c lab-4.4-zeropart in your Tez checkout.

Step 1 — Locate the source and test files

find ~/src/oss-repos/tez -name "ShuffleVertexManager.java" | grep -v target
find ~/src/oss-repos/tez -name "ShuffleVertexManagerBase.java" | grep -v target
find ~/src/oss-repos/tez -name "TestShuffleVertexManager*.java" | grep -v target

You should find the manager and its base in tez-runtime-library/src/main/java/org/apache/tez/dag/library/vertexmanager/, and four test files: TestShuffleVertexManager.java, TestShuffleVertexManagerBase.java, TestShuffleVertexManagerUtils.java, TestFairShuffleVertexManager.java.

Step 2 — Trace the statistics path and find the existing guards

The event handler is in the base class. Read it:

grep -n "onVertexManagerEventReceived\|handleVertexManagerEvent\|getComputeRoutingAction\|numVertexManagerEventsReceived" \
  ~/src/oss-repos/tez/.../vertexmanager/ShuffleVertexManagerBase.java

Two guards handle the zero-output boundary. Guard 1 is in getComputeRoutingAction — quoted from module tez-runtime-library, class ShuffleVertexManagerBase:

} else if (numVertexManagerEventsReceived == 0 && totalNumBipartiteSourceTasks > 0) {
  // When source tasks don't have output data, there will be no VME.
  return ComputeRoutingAction.SKIP;
}

When every source task produced zero output, no VertexManagerEvents are sent (numVertexManagerEventsReceived == 0), so the manager returns SKIP — it declines to compute a new routing and just schedules the pending tasks normally. This is the empty-partition guard.

Guard 2 is the parallelism floor in ShuffleVertexManager.computeRouting():

int desiredTaskParallelism = bigDesiredTaskParallelism.intValue();
if (desiredTaskParallelism < mgrConfig.getMinTaskParallelism()) {
  desiredTaskParallelism = mgrConfig.getMinTaskParallelism();   // floor, default 1
}
...
basePartitionRange = currentParallelism / desiredTaskParallelism;   // safe: divisor >= 1

Because desiredTaskParallelism is clamped to min-task-parallelism (default 1) before it is used as a divisor, the currentParallelism / desiredTaskParallelism division can never divide by zero. A third guard, in updatePendingTasks, bails when tasks <= 0.

A fourth, structural guard is updatePendingTasks, which refuses to (re)build the pending-task list when the vertex has no tasks:

void updatePendingTasks() {
  int tasks = getContext().getVertexNumTasks(getContext().getVertexName());
  if (tasks == pendingTasks.size() || tasks <= 0) {
    return;
  }
  ...
}

The tasks <= 0 clause means a managed vertex that itself has zero tasks never enters the scheduling loop at all. And canScheduleTasks() (Lab 4.2) blocks scheduling until every source vertex is vertexIsConfigured, precisely so that a source reporting numTasks == 0 is only trusted once it is known to be legitimate rather than merely not-yet-configured. Read all four guards together — they are a defense-in-depth around the same boundary, added over several JIRAs.

Questions

#	Question
1	What proto encodes the per-task statistics? (`VertexManagerEventPayloadProto`.)
2	Which counter distinguishes "sources ran but emitted nothing" from "sources haven't finished"? (`numVertexManagerEventsReceived` vs `numBipartiteSourceTasksCompleted`.)
3	Which action does `getComputeRoutingAction` return for the zero-output case, and what does the manager do next? (`SKIP`; it proceeds to schedule pending tasks without reconfiguring.)
4	Where is the divide-by-zero prevented? (The `min-task-parallelism` clamp in `computeRouting`.)
5	If the clamp were removed and `desiredTaskParallelism` became 0, what exception, on which line? (`ArithmeticException: / by zero` at `currentParallelism / desiredTaskParallelism`.)

Step 3 — Reconstruct the bug class from git history

The guards above are not obvious; they were added in response to real production failures. Read the history:

cd ~/src/oss-repos/tez
git log --oneline -- '*ShuffleVertexManager*.java' | head -20

Two commits define this bug class. Read both:

git show 8247a643f   # TEZ-3666: Integer overflow in ShuffleVertexManagerBase
git log -1 --format='%B' ed0361124   # TEZ-3452: Auto-reduce parallelism calculation can overflow with large inputs

TEZ-3452 — "Auto-reduce parallelism calculation can overflow with large inputs." With a large enough total output size, the old int/long arithmetic in the parallelism computation overflowed and produced a garbage (even negative) target parallelism. The fix moved the division to BigInteger — which is why computeRouting today uses BigInteger expectedTotalSourceTasksOutputSize and bigDesiredTaskParallelism.
TEZ-3666 — "Integer overflow in ShuffleVertexManagerBase." A follow-on overflow in the aggregation of per-source expected output; getExpectedTotalBipartiteSourceTasksOutputSize now accumulates into a BigInteger and guards against exceeding Integer.MAX_VALUE before calling intValue().

Also note TEZ-1248 — "Reduce slow-start should special case 1 reducer runs" (a boundary at the low end) and TEZ-3356 — "Fix initializing of stats when custom ShuffleVertexManager is used." Together these are the "parallelism edge case" family. Summarize in your log: the recurring failure mode is arithmetic on statistics at a boundary (zero, one, or Long.MAX_VALUE), and the recurring fix is a guard or a wider numeric type before the boundary operation.

How partition statistics are actually stored (why overflow was easy)

Read parsePartitionStats and the DATA_RANGE_IN_MB enum:

grep -n "parsePartitionStats\|DATA_RANGE_IN_MB\|statsInMB\|getExpectedTotalBipartiteSourceTasksOutputSize" \
  ~/src/oss-repos/tez/.../vertexmanager/ShuffleVertexManagerBase.java

Each producer's per-partition sizes arrive as a compressed RoaringBitmap that indexes into a set of size buckets (DATA_RANGE_IN_MB); parsePartitionStats accumulates each partition's bucketed size into srcInfo.statsInMB[index]. The manager then projects a total from the sampled tasks in getExpectedTotalBipartiteSourceTasksOutputSize:

BigInteger expectedSrcOutputSize = srcOutputSize.multiply(srcNumTasks).divide(srcNumVMEventsReceived);

That outputSize * numTasks multiply is precisely where the old long arithmetic overflowed with large inputs — projecting from a few completed tasks to thousands of tasks multiplies a large number by a large number. TEZ-3452/TEZ-3666 rewrote it in BigInteger, and getExpectedStatsAtIndex even logs and clamps to Long.MAX_VALUE before narrowing. Note the guard if (srcInfo.numTasks > 0 && srcInfo.numVMEventsReceived > 0) — this is what stops a divide-by-zero in the projection when a source produced no events, complementing the SKIP guard from Step 2.

Step 4 — Read the test harness

grep -n "createVertexManagerContext\|getVertexManagerEvent\|createManager\|scheduleTasks\|reconfigureVertex" \
  ~/src/oss-repos/tez/.../vertexmanager/TestShuffleVertexManagerUtils.java

The harness is TestShuffleVertexManagerUtils. Its key helpers (quoted from module tez-runtime-library, class TestShuffleVertexManagerUtils):

final VertexManagerPluginContext mockContext = mock(VertexManagerPluginContext.class);
when(mockContext.getInputVertexEdgeProperties()).thenReturn(mockInputVertices);
when(mockContext.getVertexName()).thenReturn(mockManagedVertexId);
when(mockContext.getVertexNumTasks(mockSrcVertexId1)).thenReturn(numTasksSrcVertexId1);
...
doAnswer(new ScheduledTasksAnswer(scheduledTasks)).when(mockContext).scheduleTasks(anyList());
doAnswer(new reconfigVertexAnswer(mockContext, mockManagedVertexId, newEdgeManagers))
    .when(mockContext).reconfigureVertex(anyInt(), any(), anyMap());

and the VM-event builder, which lets you construct a producer event with a chosen partition-size array:

VertexManagerEvent getVertexManagerEvent(long[] partitionSizes,
    long uncompressedTotalSize, String vertexName, boolean reportDetailedStats)

To simulate a zero-output producer you pass a payload whose total size is 0. To simulate "the source ran but never sent a VME" you simply drive onSourceTaskCompleted(...) without ever calling onVertexManagerEventReceived(...). Read testAutoParallelismConfig and testSchedulingWithPartitionStats in TestShuffleVertexManagerBase for the exact call order: createManager(...) → onVertexStarted(...) → per-source onSourceTaskCompleted(...) / onVertexManagerEventReceived(...) → verify(mockContext).reconfigureVertex(eq(N), any(), anyMap()).

Step 5 — Write the regression test

Add a test to TestShuffleVertexManager.java (or TestShuffleVertexManagerBase) that pins the zero-output boundary: with auto-parallelism on, drive all source tasks to completion without any output statistics, and assert the manager does not crash and does not reconfigure parallelism (it should SKIP and schedule pending tasks). Adapt the setup from the nearest existing test; the shape is:

@Test(timeout = 5000)
public void testZeroOutputSourceSkipsReconfigure() throws Exception {
  Configuration conf = new Configuration();
  final VertexManagerPluginContext mockContext =
      createVertexManagerContext(/* small src + managed vertex task counts */);
  ShuffleVertexManager manager = (ShuffleVertexManager) createManager(
      ShuffleVertexManager.class, conf, mockContext,
      /* enableAutoParallelism */ true, /* desiredTaskInputSize */ 1000L,
      /* min */ 0.01f, /* max */ 0.75f);

  manager.onVertexStarted(emptyCompletions());

  // Complete every source task, but never deliver a VertexManagerEvent
  // (i.e. every source produced zero output -> no VME).
  for (int i = 0; i < totalSourceTasks; i++) {
    manager.onSourceTaskCompleted(createTaskAttemptIdentifier(srcVertexName, i));
  }

  // Must NOT crash and must NOT reduce parallelism (SKIP path).
  verify(mockContext, times(0)).reconfigureVertex(anyInt(), any(), anyMap());
  // Pending tasks are still scheduled normally.
  verify(mockContext, atLeastOnce()).scheduleTasks(anyList());
}

Run it:

cd ~/src/oss-repos/tez
mvn test -pl tez-runtime-library \
  -Dtest=TestShuffleVertexManager#testZeroOutputSourceSkipsReconfigure -q 2>&1 | tail -30

On current master this test should pass — proving the SKIP guard works. Record that result: a green regression test that pins an existing guard is a legitimate, mergeable contribution (it stops a future refactor from silently removing the guard).

Step 6 — Demonstrate the bug the guard prevents (then revert)

To see the failure the guard suppresses, temporarily defeat it on your scratch branch. Comment out the min-task-parallelism clamp in computeRouting:

// if (desiredTaskParallelism < mgrConfig.getMinTaskParallelism()) {
//   desiredTaskParallelism = mgrConfig.getMinTaskParallelism();
// }

then construct a case where desiredTaskParallelism computes to 0 (tiny expected output, large desired-task-input-size) and assert the resulting ArithmeticException at currentParallelism / desiredTaskParallelism. This shows why the clamp exists. Revert the edit immediately — the committed code must stay guarded.

git checkout -- tez-runtime-library/.../vertexmanager/ShuffleVertexManager.java

Step 7 — What a credible fix PR for this class contains

Even though master is guarded today, write up the anatomy of the fix as if you were submitting it — this is the deliverable a maintainer would merge:

A minimal guard at the boundary. For a hypothetical missing case, the fix is a few lines: either return ComputeRoutingAction.SKIP; when there is no data, or clamp the divisor to a floor (>= 1). Choose the semantically correct one: for scheduling, skipping the reconfigure and running the pre-planned parallelism is safer than forcing parallelism to 1.
No signature changes, no reformatting. Touch only the guarded method and the test.
A regression test like Step 5 that fails before the guard and passes after.
A CHANGELOG/JIRA reference. Real Tez PRs cite a TEZ-xxxx and carry a DCO Signed-off-by line (git commit -s).

Green suite + Checkstyle:

mvn test -pl tez-runtime-library -q 2>&1 | tail -20
mvn checkstyle:check -pl tez-runtime-library -q 2>&1 | grep -iE "error|violation" | head

Draft the JIRA text:

Summary: ShuffleVertexManager auto-parallelism: guard <boundary> in computeRouting

Description:
  When <boundary condition, e.g. every source vertex produces zero output>,
  ShuffleVertexManager's auto-parallelism path <SKIPs / must clamp the divisor>.
  Without the guard, computeRouting would <ArithmeticException / overflow / choose
  parallelism 0>. Reproduced by TestShuffleVertexManager#testZeroOutputSourceSkipsReconfigure.

  Fix: <SKIP the reconfigure when numVertexManagerEventsReceived == 0> /
       <clamp desiredTaskParallelism to min-task-parallelism before the division>.

Component: tez-dag / tez-runtime-library
Affects Version: <your build>

Deliverables

The two current guards located and quoted (the SKIP action and the min-task-parallelism clamp), plus the updatePendingTasks tasks <= 0 guard.
A written reconstruction of the bug class from TEZ-3452 and TEZ-3666 (overflow → BigInteger), with TEZ-1248 and TEZ-3356 noted as siblings.
A passing regression test that drives the zero-output path and asserts no crash / no reconfigure.
Step 6 done and reverted: you triggered and then re-suppressed the divide-by-zero.
A fix-PR write-up (guard + test + JIRA text) matching the class of bug.

Troubleshooting

Symptom	Likely cause	Fix
`grep` for `parseStatsHeader` / `numPartitions` finds nothing	Those names don't exist in current Tez	The real path is `handleVertexManagerEvent` → `parsePartitionStats`; the count that matters is `numVertexManagerEventsReceived`
Your zero-partition test throws instead of skipping	You delivered a VME with 0 partitions instead of no VME	The `SKIP` guard keys on `numVertexManagerEventsReceived == 0`; don't send any event to simulate zero output
`reconfigureVertex` is never called even with data	Not enough data seen vs `desired-task-input-size`, or reduction < half	Lower `desiredTaskInputSize` in the test, or raise the source output sizes
`createManager` returns null	Wrong manager class passed	Pass `ShuffleVertexManager.class` (or `FairShuffleVertexManager.class`) to the factory
Mockito `cannot mock` in a standalone run	JDK/Mockito mismatch outside Maven	Run via `mvn -pl tez-runtime-library test`, which uses the project's managed versions
Overflow reproduces even on master	You are on a pre-TEZ-3666 tag	`git log` the file; the `BigInteger` fix must be present

Issues to Practice On

Parallelism and slow-start behavior is a steady source of real, tractable issues. Reproduce first, add a test, then propose the guard — the etiquette maintainers expect. Target apache/tez.

# Component and keyword searches (labels move; confirm on the tracker):
gh issue list --repo apache/tez --search "ShuffleVertexManager parallelism" --state open
gh issue list --repo apache/tez --search "auto reduce parallelism" --state all
gh issue list --repo apache/tez --search "slow start reducer" --state open

Representative patterns worth a PR:

A boundary that produces a silly reducer count. "Query with a highly selective filter spawns hundreds of empty reducers" — reproduce with a source that emits near-zero output, confirm the SKIP path, and add the regression test from Step 5 if one is missing.
A flaky TestShuffleVertexManager* assertion. Timing/ordering assumptions in the mocked scheduling order. Stabilize and pin the intended order.

Etiquette: claim the issue, reproduce before patching, keep the diff minimal, and every PR needs a test + a TEZ-xxxx reference + a DCO Signed-off-by (git commit -s). See event-routing.md for how the VM events reach this code and the Level 4 overview for the surrounding class map.

Stretch Goals

Pin the overflow fix. Write a test that feeds a source output size near Long.MAX_VALUE and asserts computeRouting does not overflow (it either declines to reduce or returns a sane parallelism). This is the regression test for TEZ-3452/TEZ-3666.
The one-reducer boundary (TEZ-1248). Construct a case where auto-parallelism would reduce to a single reducer and confirm the "reduce by at least half" rule (basePartitionRange <= 1 returns null) prevents a pointless reconfigure.
FairShuffleVertexManager. It overrides computeRouting. Does it share the same floor/BigInteger protections? If any boundary is unguarded there, that is a real, fileable issue — write the reproducing test.

Validation / Self-check

Answer in your own words, citing the class:

When a source vertex produces zero output, why are there no VertexManagerEvents, and which ComputeRoutingAction does the manager take as a result?
Where exactly is the divide-by-zero in the parallelism computation prevented, and what is the default value of the clamp?
What was the historical bug class in TEZ-3452 / TEZ-3666, and what change in numeric type fixed it?
How do you simulate a zero-output producer in the test harness — and how is that different from sending a VertexManagerEvent that reports zero bytes?
Why is a passing regression test (pinning an existing guard) still a worthwhile contribution?
What are the five parts of a credible fix PR for a boundary bug in this class, per Step 7?
If you removed the min-task-parallelism clamp, what exception would you get and on which operation?
Name the four defense-in-depth guards around the zero/empty boundary (SKIP in getComputeRoutingAction, the min-task-parallelism clamp in computeRouting, the numTasks > 0 && numVMEventsReceived > 0 guard in the projection, and tasks <= 0 in updatePendingTasks) and say which JIRA family motivated the BigInteger rewrite.

When you can locate both guards, explain the overflow history, and produce a green regression test that pins the zero-output boundary, you have completed Lab 4.4 — and Level 4's labs. Return to the Level 4 overview to review the deep-dives before Level 5.

Open-Source Engineer & Contributor