From c03780d3eb094f20db8ce82ba7a12a83c8b4a569 Mon Sep 17 00:00:00 2001
From: Prabhjyot Singh <prabhjyot@acceldata.io>
Date: Wed, 20 Nov 2024 12:47:33 -0500
Subject: [PATCH] ODP-2642: TEZ-4580: Slow preemption of new containers when
 re-use is enabled (unmerged) #374 (#24)

(cherry picked from commit aea86ca87f1b9e038b3454b1ed483c179a689e07)
---
 .../dag/app/rm/YarnTaskSchedulerService.java  |  20 +++-
 .../tez/dag/app/rm/TestTaskScheduler.java     | 101 ++++++++++++++++++
 2 files changed, 118 insertions(+), 3 deletions(-)

diff --git a/tez-dag/src/main/java/org/apache/tez/dag/app/rm/YarnTaskSchedulerService.java b/tez-dag/src/main/java/org/apache/tez/dag/app/rm/YarnTaskSchedulerService.java
index 6f51566f75..75b6a662e0 100644
--- a/tez-dag/src/main/java/org/apache/tez/dag/app/rm/YarnTaskSchedulerService.java
+++ b/tez-dag/src/main/java/org/apache/tez/dag/app/rm/YarnTaskSchedulerService.java
@@ -1272,7 +1272,7 @@ boolean preemptIfNeeded() {
             + numHighestPriRequests + " pending requests at pri: "
             + highestPriRequest.getPriority());
       }
-
+      int newContainersReleased = 0;
       for (int i=0; i<numPendingRequestsToService; ++i) {
         // This request must have been considered for matching with all existing 
         // containers when request was made.
@@ -1308,7 +1308,7 @@ boolean preemptIfNeeded() {
               " with priority: " + lowestPriNewContainer.getPriority() + 
               " to free resource for request: " + highestPriRequest +
               " . Current free resources: " + freeResources);
-          numPendingRequestsToService--;
+          newContainersReleased++;
           releaseUnassignedContainers(Collections.singletonList(lowestPriNewContainer));
           // We are returning an unused resource back the RM. The RM thinks it 
           // has serviced our initial request and will not re-allocate this back
@@ -1321,7 +1321,7 @@ boolean preemptIfNeeded() {
           continue;
         }
       }
-      
+      numPendingRequestsToService -= newContainersReleased;
       if (numPendingRequestsToService < 1) {
         return true;
       }
@@ -1570,6 +1570,9 @@ private void releaseContainer(ContainerId containerId) {
     if (delayedContainer != null) {
       Resources.subtractFrom(allocatedResources,
           delayedContainer.getContainer().getResource());
+      if (shouldReuseContainers) {
+        delayedContainerManager.removeDelayedContainer(delayedContainer);
+      }
     }
     if (delayedContainer != null || !shouldReuseContainers) {
       amRmClient.releaseAssignedContainer(containerId);
@@ -2158,6 +2161,17 @@ void addDelayedContainer(Container container,
       }
     }
 
+    void removeDelayedContainer(HeldContainer container) {
+      if (container != null) {
+        synchronized(this) {
+          if (delayedContainers.remove(container)) {
+            LOG.info("Removed {} from delayed containers", container.getContainer().getId());
+          } else {
+            LOG.warn("Unknown container {} sent for removal. Ignoring.", container.getContainer().getId());
+          }
+        }
+      }
+    }
   }
   
   synchronized void determineMinHeldContainers() {
diff --git a/tez-dag/src/test/java/org/apache/tez/dag/app/rm/TestTaskScheduler.java b/tez-dag/src/test/java/org/apache/tez/dag/app/rm/TestTaskScheduler.java
index ff34084d37..df555ffcb0 100644
--- a/tez-dag/src/test/java/org/apache/tez/dag/app/rm/TestTaskScheduler.java
+++ b/tez-dag/src/test/java/org/apache/tez/dag/app/rm/TestTaskScheduler.java
@@ -989,6 +989,107 @@ public void testTaskSchedulerRandomReuseExpireTime() throws Exception {
     scheduler2.shutdown();
   }
 
+  @Test(timeout = 5000)
+  public void testTaskSchedulerPreemptionWithLowAndHighPriorityRequests() throws Exception {
+    TezAMRMClientAsync<CookieContainerRequest> mockRMClient = spy(
+            new AMRMClientAsyncForTest(new AMRMClientForTest(), 100));
+
+    Configuration conf = new Configuration();
+    conf.setInt(TezConfiguration.TEZ_AM_PREEMPTION_PERCENTAGE, 50);
+
+    TaskSchedulerContext mockApp = setupMockTaskSchedulerContext(DEFAULT_APP_HOST, DEFAULT_APP_PORT, DEFAULT_APP_URL,
+            false, null, null, new PreemptionMatcher(), conf);
+    final TaskSchedulerContextDrainable drainableAppCallback = createDrainableContext(mockApp);
+
+    final TaskSchedulerWithDrainableContext scheduler =
+            new TaskSchedulerWithDrainableContext(drainableAppCallback, mockRMClient);
+    scheduler.initialize();
+    scheduler.start();
+    int initialRmCapacity = 4;
+    int lowPriorityTasks = 5;
+    int highPriorityTasks = 6;
+    Resource taskAsk = Resource.newInstance(1000, 1);
+
+    Resource totalResource = Resource.newInstance(4000, 4);
+    when(mockRMClient.getAvailableResources()).thenReturn(totalResource);
+
+    // Add lower priority tasks
+    Priority lowPriority = Priority.newInstance(74);
+    for (int i = 0; i < lowPriorityTasks; i++) {
+      Object low = new Object();
+      TaskAttempt ta = mock(TaskAttempt.class);
+      scheduler.allocateTask(ta, taskAsk, null, null, lowPriority, low, null);
+    }
+
+    scheduler.getProgress();  // Will update the highest priority
+    drainableAppCallback.drain();
+    // 5 containers requested for lower priority tasks
+    verify(mockRMClient, times(5)).addContainerRequest(any(CookieContainerRequest.class));
+
+    // Allocate requested containers
+    List<Container> lowPriorityContainers = new ArrayList<>();
+    for (int i = 0; i < initialRmCapacity; i++) {
+      ContainerId containerId = ContainerId.newContainerId(
+              ApplicationAttemptId.newInstance(ApplicationId.newInstance(1L, 1), 1), i);
+      NodeId nodeId = NodeId.newInstance("host-" + i, 8041);
+      Container container = Container.newInstance(containerId, nodeId, "host-" + i, taskAsk, lowPriority, null);
+      lowPriorityContainers.add(container);
+    }
+
+    totalResource = Resource.newInstance(0, 0);
+    when(mockRMClient.getAvailableResources()).thenReturn(totalResource);
+
+    // We don't want containers to be assigned to a task by delayedContainerManager as it invokes another preemption flow
+    // Delayed thread first takes lock on delayedContainerManager instance to check if there are any containers
+    // We block the thread, ensure all delayed containers have schedule time beyond test's runtime to avoid assignment.
+    synchronized (scheduler.delayedContainerManager) {
+      scheduler.onContainersAllocated(lowPriorityContainers);
+      drainableAppCallback.drain();
+      for (HeldContainer container : scheduler.delayedContainerManager.delayedContainers) {
+        // Set next schedule beyond this test's time to avoid any assignment
+        container.setNextScheduleTime(System.currentTimeMillis() + 10000);
+        // No preemption if assignment attempt of new container < 3
+        container.incrementAssignmentAttempts();
+        container.incrementAssignmentAttempts();
+        container.incrementAssignmentAttempts();
+      }
+    }
+
+    // No releases so far
+    verify(mockRMClient, times(0)).releaseAssignedContainer(any());
+
+    // Add higher priority task
+    Priority highPriority = Priority.newInstance(71);
+    for (int i = 0; i < highPriorityTasks; i++) {
+      Object high = new Object();
+      TaskAttempt ta = mock(TaskAttempt.class);
+      scheduler.allocateTask(ta, taskAsk, null, null, highPriority, high, null);
+    }
+
+    drainableAppCallback.drain();
+    // low priority tasks + high priority tasks
+    verify(mockRMClient, times(11)).addContainerRequest(any(CookieContainerRequest.class));
+
+    // Trigger preemption to release containers as 50% of pending high priority requests
+    scheduler.getProgress();
+    drainableAppCallback.drain();
+
+    // 50% of 6 high priority requests = 3, 4 containers were held - hence 3 will be released
+    verify(mockRMClient, times(3)).releaseAssignedContainer(any());
+
+    // Trigger another preemption cycle
+    scheduler.getProgress();
+    drainableAppCallback.drain();
+    // 50% of 6 high priority requests = 3, but only 1 container is held - which will be released,
+    // incrementing total to 4
+    verify(mockRMClient, times(4)).releaseAssignedContainer(any());
+    AppFinalStatus finalStatus =
+            new AppFinalStatus(FinalApplicationStatus.SUCCEEDED, "", DEFAULT_APP_URL);
+    when(mockApp.getFinalAppStatus()).thenReturn(finalStatus);
+    scheduler.shutdown();
+    drainableAppCallback.drain();
+  }
+
   @SuppressWarnings({ "unchecked", "rawtypes" })
   @Test (timeout=5000)
   public void testTaskSchedulerPreemption() throws Exception {