Lab 9.1 — Write Tests for Scheduler Behavior

Lab type: Build It — comprehensive test coverage Estimated time: 3–4 hours Tez module: tez-dag

Background

The Tez task scheduler is the AM-side machinery that turns "I have a task that wants to run on host X" into "I asked YARN for a container, YARN gave me one, and I assigned the task to it." It is a genuinely hard piece of concurrent code — it juggles pending requests, container reuse, locality fallback, preemption, blacklisting, and idle container expiry, all driven by asynchronous callbacks from the YARN AMRMClientAsync. It is also, historically, one of the areas where a subtle bug costs the most: a scheduling regression does not crash, it just quietly makes every DAG slower or wastes cluster capacity.

Tez ships two production schedulers, both in tez-dag/src/main/java/org/apache/tez/dag/app/rm/:

• DagAwareYarnTaskScheduler — the default since it was introduced (TezConfiguration.TEZ_AM_YARN_SCHEDULER_CLASS_DEFAULT points at it). It is DAG-topology-aware: it uses vertex-descendant information to decide which held container to reuse for which pending task.
• YarnTaskSchedulerService — the older scheduler, still supported and still tested, with its own delayed-container / preemption logic.

Both are routed to by a single TaskSchedulerManager, which receives AMSchedulerEvents from the rest of the AM and forwards them to the right scheduler instance. Read the Scheduler deep dive for the full two-layer picture before you continue.

Confirm the class inventory yourself — never trust a doc's list over the tree:

cd ~/src/oss-repos/tez
ls tez-dag/src/main/java/org/apache/tez/dag/app/rm/
# DagAwareYarnTaskScheduler.java  YarnTaskSchedulerService.java
# TaskSchedulerManager.java  LocalTaskSchedulerService.java  ...

ls tez-dag/src/test/java/org/apache/tez/dag/app/rm/
# TestDagAwareYarnTaskScheduler.java  TestTaskScheduler.java
# TestTaskSchedulerManager.java  TestTaskSchedulerHelpers.java  TestContainerReuse.java ...

Note on names: the test for YarnTaskSchedulerService is TestTaskScheduler.java (not TestYarnTaskSchedulerService), and there is no TestYarnTaskSchedulerService in the tree. Verify with the ls above — this is exactly the kind of stale assumption the house anti-staleness rule exists to kill.

Why This Lab Matters for Contributors

Scheduler tests are high-signal, high-trust contributions. They are hard to write well because the scheduler is asynchronous and stateful, so a naive test either races or asserts on the wrong thing. A test that drives a real scheduler through a precise event sequence and asserts on the exact YARN RM calls it produces is worth ten tests that call a method and check "no exception." Committers value them because scheduler behavior is where correctness and performance meet: a container that should be reused but is released is a performance bug; a task that should be re-requested after blacklisting but is dropped is a correctness bug. When you can extend this harness fluently, you can write regression tests for real scheduler JIRAs — which is precisely what got recent scheduler fixes like TEZ‑4580 (Slow preemption of new containers when re-use is enabled) merged with a 101-line test attached.

git show 9efa6f14d --stat        # TEZ-4580: fix + its regression test in TestTaskScheduler

Prerequisites

• You can build Tez and run a single test class: mvn test -pl tez-dag -Dtest=TestDagAwareYarnTaskScheduler.
• You have read the Scheduler deep dive.
• You understand Mockito basics: mock, spy, when(...).thenReturn(...), verify(mock, times(n)), and ArgumentCaptor.
• You understand why an asynchronous callback test needs to "drain" pending callbacks before asserting.

Step-by-Step Tasks

Step 1 — Dissect the mock harness

Every scheduler unit test is built on TestTaskSchedulerHelpers.java. Open it and read these pieces; they are the entire reason these tests are fast and deterministic instead of a mini-cluster.

rg -n "class MockAMRMClient|class AMRMClientAsyncForTest|setupMockTaskSchedulerContext|class TaskSchedulerContextDrainable|class CountingExecutorService" \
  tez-dag/src/test/java/org/apache/tez/dag/app/rm/TestTaskSchedulerHelpers.java

MockAMRMClient extends AMRMClientImpl<TaskRequest> — this is the fake YARN RM. It never talks to a cluster; it just records what the scheduler asked for and hands back a mocked registration response. Note it seeds cluster capacity directly:

// TestTaskSchedulerHelpers.MockAMRMClient
MockAMRMClient() {
  super();
  this.clusterAvailableResources = Resource.newInstance(4000, 4);
  this.clusterNodeCount = 5;
}

Because it is a real AMRMClientImpl subclass with serviceStart/serviceStop stubbed out, calls like addContainerRequest, removeContainerRequest, releaseAssignedContainer, and updateBlacklist are the real methods — which is why tests spy(...) it and verify(...) those exact calls.

setupMockTaskSchedulerContext(...) builds the TaskSchedulerContext the scheduler calls back into (the AM's face to the scheduler). It is a Mockito mock with the essentials stubbed:

// TestTaskSchedulerHelpers.setupMockTaskSchedulerContext (abridged)
TaskSchedulerContext mockContext = mock(TaskSchedulerContext.class);
when(mockContext.getAppHostName()).thenReturn(appHost);
when(mockContext.getAMState()).thenReturn(TaskSchedulerContext.AMState.RUNNING_APP);
when(mockContext.getInitialUserPayload()).thenReturn(userPayload);   // your TezConfiguration
when(mockContext.getContainerSignatureMatcher())
    .thenReturn(new AlwaysMatchesContainerMatcher());

TaskSchedulerContextDrainable is the trick that makes async assertions deterministic. The scheduler delivers callbacks (taskAllocated, containerAllocated, …) through an executor; the drainable wrapper counts them so a test can call drain() and block until every queued callback has actually run before it asserts. This is why every real test interleaves drain() between injecting an event and verifying its effect.

Step 2 — Read one real test top-to-bottom

TestDagAwareYarnTaskScheduler.testNoReuse() is the canonical template. Read the whole method — it exercises registration, allocation, assignment, deallocation, container completion, blacklisting, and shutdown in one flow.

rg -n "public void testNoReuse" tez-dag/src/test/java/org/apache/tez/dag/app/rm/TestDagAwareYarnTaskScheduler.java

The setup boilerplate that opens every test in this class:

// TestDagAwareYarnTaskScheduler.testNoReuse (setup portion)
AMRMClientAsyncWrapperForTest mockRMClient = spy(new AMRMClientAsyncWrapperForTest());
Configuration conf = new Configuration();
conf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_ENABLED, false);
conf.setInt(TezConfiguration.TEZ_AM_RM_HEARTBEAT_INTERVAL_MS_MAX, 100);

DagInfo mockDagInfo = mock(DagInfo.class);
when(mockDagInfo.getTotalVertices()).thenReturn(10);
when(mockDagInfo.getVertexDescendants(anyInt())).thenReturn(new BitSet());
TaskSchedulerContext mockApp = setupMockTaskSchedulerContext(appHost, appPort, appUrl, conf);
when(mockApp.getCurrentDagInfo()).thenReturn(mockDagInfo);
TaskSchedulerContextDrainable drainableAppCallback = createDrainableContext(mockApp);

MockClock clock = new MockClock(1000);
NewTaskSchedulerForTest scheduler = new NewTaskSchedulerForTest(drainableAppCallback,
    mockRMClient, clock);
scheduler.initialize();
drainableAppCallback.drain();
scheduler.start();
drainableAppCallback.drain();

NewTaskSchedulerForTest is a thin subclass of the real DagAwareYarnTaskScheduler that swaps in a MockClock and a ControlledScheduledExecutorService so time and the delayed-container executor are deterministic. You are testing the real scheduler logic, not a stub — only its clock, executor, and RM client are faked.

Now trace the allocate → assign → verify rhythm. When the test injects an allocation and then hands the scheduler containers, it asserts on the RM calls the scheduler makes:

// allocate a task, then deliver containers from "YARN"
scheduler.allocateTask(mockTask1, mockCapability, hosts, racks, mockPriority, null, mockCookie1);
drainableAppCallback.drain();
verify(mockRMClient, times(1)).addContainerRequest(any());
...
scheduler.onContainersAllocated(containers);      // simulate YARN allocation callback
drainableAppCallback.drain();
verify(mockApp).taskAllocated(mockTask1, mockCookie1, mockContainer1);
verify(mockRMClient).releaseAssignedContainer(mockCId4);   // unwanted extra container released

Note the shape: inject an event on the scheduler, drain(), then verify on the mock RM and the mock app. That is the entire idiom.

Step 3 — Inventory scheduler behaviors that deserve tests

Before writing anything, list the behaviors and check which already have coverage. Run the method inventory:

rg -n "public void test" tez-dag/src/test/java/org/apache/tez/dag/app/rm/TestDagAwareYarnTaskScheduler.java
rg -n "public void test" tez-dag/src/test/java/org/apache/tez/dag/app/rm/TestContainerReuse.java
rg -n "public void test" tez-dag/src/test/java/org/apache/tez/dag/app/rm/TestTaskScheduler.java

The behaviors worth a dedicated, focused test — cross-referenced against the real production logic:

The blacklist re-request behavior is the good target for this lab: the production logic in informAppAboutAssignment() is nontrivial (it deallocates the bad container and re-submits the original request so the RM can pick a healthy node), but it is only tested tangled inside the 270-line testNoReuse. Splitting it into a focused, self-documenting test is exactly the kind of contribution committers welcome.

Read the production behavior before you test it:

rg -n "informAppAboutAssignment|blacklistedNodes|updateBlacklist" \
  tez-dag/src/main/java/org/apache/tez/dag/app/rm/DagAwareYarnTaskScheduler.java

You will find the invariant: when a container is allocated on a blacklisted node, the scheduler calls deallocateContainer(...) and then re-invokes allocateTask(...) with the original hints — it never drops the task, and the app is never told the task was allocated to the bad node.

Step 4 — Write a new focused behavior test end-to-end

Add this to TestDagAwareYarnTaskScheduler.java, following the harness exactly. Every symbol below is real in that file (spy, AMRMClientAsyncWrapperForTest, MockClock, NewTaskSchedulerForTest, setupMockTaskSchedulerContext, createDrainableContext, updateBlacklist).

@Test(timeout = 10000)
public void testBlacklistedNodeContainerReRequested() throws Exception {
  AMRMClientAsyncWrapperForTest mockRMClient = spy(new AMRMClientAsyncWrapperForTest());

  String appHost = "host";
  int appPort = 0;
  String appUrl = "url";

  Configuration conf = new Configuration();
  // Isolate the blacklist behavior: no reuse, so a bad container is released, not held.
  conf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_ENABLED, false);
  conf.setInt(TezConfiguration.TEZ_AM_RM_HEARTBEAT_INTERVAL_MS_MAX, 100);

  DagInfo mockDagInfo = mock(DagInfo.class);
  when(mockDagInfo.getTotalVertices()).thenReturn(10);
  when(mockDagInfo.getVertexDescendants(anyInt())).thenReturn(new BitSet());
  TaskSchedulerContext mockApp = setupMockTaskSchedulerContext(appHost, appPort, appUrl, conf);
  when(mockApp.getCurrentDagInfo()).thenReturn(mockDagInfo);
  TaskSchedulerContextDrainable drainableAppCallback = createDrainableContext(mockApp);

  MockClock clock = new MockClock(1000);
  NewTaskSchedulerForTest scheduler =
      new NewTaskSchedulerForTest(drainableAppCallback, mockRMClient, clock);

  scheduler.initialize();
  drainableAppCallback.drain();
  scheduler.start();
  drainableAppCallback.drain();

  // 1. Blacklist a node BEFORE any container from it arrives.
  String badHost = "badhost";
  NodeId badNodeId = NodeId.newInstance(badHost, 1);
  scheduler.blacklistNode(badNodeId);
  drainableAppCallback.drain();
  verify(mockRMClient, times(1))
      .updateBlacklist(eq(Collections.singletonList(badHost)), isNull());

  // 2. Request a container for a task (location hints intentionally null so any node matches).
  Object mockTask1 = new MockTask("task1");
  Object mockCookie1 = new Object();
  Resource mockCapability = Resources.createResource(1024, 1);
  Priority mockPriority = Priority.newInstance(1);
  scheduler.allocateTask(mockTask1, mockCapability, null, null, mockPriority, null, mockCookie1);
  drainableAppCallback.drain();
  verify(mockRMClient, times(1)).addContainerRequest(any());

  // 3. YARN hands us a container ON THE BLACKLISTED NODE.
  ApplicationAttemptId attemptId =
      ApplicationAttemptId.newInstance(ApplicationId.newInstance(1, 1), 1);
  ContainerId badCId = ContainerId.newContainerId(attemptId, 1);
  Container badContainer =
      Container.newInstance(badCId, badNodeId, null, mockCapability, mockPriority, null);
  List<Container> containers = new ArrayList<>();
  containers.add(badContainer);
  scheduler.onContainersAllocated(containers);
  drainableAppCallback.drain();

  // 4a. INVARIANT: the app is NEVER told the task ran on the bad node.
  verify(mockApp, times(0)).taskAllocated(any(), any(), any());
  // 4b. INVARIANT: the bad container is released back to YARN.
  verify(mockRMClient).releaseAssignedContainer(badCId);
  // 4c. INVARIANT: the task is re-requested (2nd addContainerRequest), not dropped.
  verify(mockRMClient, times(2)).addContainerRequest(any());

  // 5. Now a HEALTHY container arrives and the re-requested task lands on it.
  NodeId goodNodeId = NodeId.newInstance("goodhost", 2);
  ContainerId goodCId = ContainerId.newContainerId(attemptId, 2);
  Container goodContainer =
      Container.newInstance(goodCId, goodNodeId, null, mockCapability, mockPriority, null);
  containers.clear();
  containers.add(goodContainer);
  scheduler.onContainersAllocated(containers);
  drainableAppCallback.drain();
  verify(mockApp, times(1)).taskAllocated(mockTask1, mockCookie1, goodContainer);

  // Clean shutdown so the test doesn't leak the callback executor.
  AppFinalStatus finalStatus =
      new AppFinalStatus(FinalApplicationStatus.SUCCEEDED, "success", appUrl);
  when(mockApp.getFinalAppStatus()).thenReturn(finalStatus);
  scheduler.shutdown();
  drainableAppCallback.drain();
}

Read what each block proves. Steps 4a–4c are the three separate invariants the omnibus testNoReuse only checks in passing; here they are named, isolated, and would each fail with a distinct message if the blacklist logic regressed. That is what makes this a better test, not just another one.

Step 5 — Run and verify

cd ~/src/oss-repos/tez
mvn test -pl tez-dag -Dtest=TestDagAwareYarnTaskScheduler#testBlacklistedNodeContainerReRequested -q 2>&1 | tail -20

Expected: Tests run: 1, Failures: 0, Errors: 0, Skipped: 0, BUILD SUCCESS.

Then run the whole class to confirm you did not disturb the shared harness:

mvn test -pl tez-dag -Dtest=TestDagAwareYarnTaskScheduler -q 2>&1 | tail -20

Step 6 — Prove the test actually tests something (planted-bug drill)

A test you have never seen fail is not yet a test. Temporarily break the production invariant and confirm your test catches it. In DagAwareYarnTaskScheduler.informAppAboutAssignment(...), the re-request path calls allocateTask(...) after deallocating the bad container. Comment out that allocateTask(...) call, rebuild, and re-run:

mvn test -pl tez-dag -Dtest=TestDagAwareYarnTaskScheduler#testBlacklistedNodeContainerReRequested -q 2>&1 | tail -20

Your assertion verify(mockRMClient, times(2)).addContainerRequest(any()) should now fail (only one request was ever made — the task was dropped). Restore the line, re-run, confirm green. You have now watched your test detect the exact regression it exists to prevent — the difference between a real test and decoration.

Deliverables

• A new @Test method in TestDagAwareYarnTaskScheduler.java that drives the real scheduler through a precise event sequence against the mocked AMRMClient.
• The test asserts on RM interactions (addContainerRequest, releaseAssignedContainer, updateBlacklist) and app callbacks (taskAllocated), not on private scheduler state.
• Every injected event is followed by drainableAppCallback.drain() before its assertions.
• The test passes in isolation and as part of the full class.
• You completed the planted-bug drill and watched the test go red, then green.
• A short written note: which uncovered behavior you chose, why it deserved a focused test, and which production method carries the invariant.

Troubleshooting

Stretch Goals

• Write the same behavior test against YarnTaskSchedulerService using the TestTaskScheduler.java harness (TaskSchedulerWithDrainableContext, AMRMClientAsyncForTest). Note how the two harnesses differ and why.
• Add a locality-fallback test: request a task with a host hint, deliver a container on a different host in the same rack, and assert reuse happens only when TEZ_AM_CONTAINER_REUSE_RACK_FALLBACK_ENABLED is true. Cross-check against HeldContainer.moveToNextMatchingLevel().
• Read the TEZ‑4580 regression test in TestTaskScheduler (git show 9efa6f14d), reproduce the pre-fix behavior by reverting the production change on a scratch branch, and watch that test fail. This is how a committer confirms a regression test guards the fix.
• Write a TestTaskSchedulerManager-level test that drives an AMSchedulerEventTALaunchRequest through the manager and asserts the resulting AMContainerEventAssignTA is dispatched — testing the routing layer, not the scheduler. Model it on testSimpleAllocate.

Validation / Self-check

Answer all of these before marking the lab complete:

1. Why does every real scheduler test spy(...) the AMRMClient and verify(...) its methods, rather than asserting on the scheduler's internal collections?
2. What does TaskSchedulerContextDrainable.drain() actually do, and what class of bug appears in tests that forget to call it?
3. MockAMRMClient seeds clusterAvailableResources and clusterNodeCount in its constructor. Why does the scheduler need those before it can service any request?
4. What is the difference between DagAwareYarnTaskScheduler and YarnTaskSchedulerService, and which one does the AM instantiate by default?
5. In the blacklist behavior, name the three separate invariants a focused test should assert, and the production method that enforces them.
6. Your test passes but you never saw it fail. Describe the planted-bug drill you would run to prove it detects the regression it targets.
7. What would a MiniTezCluster integration test of this behavior make harder to set up, and what would it let you assert that the unit test cannot?