[FLINK-33390][runtime] Support slot balancing at TM level for Adaptive Scheduler

Author: RocMarshal

flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/adaptive/AdaptiveSchedulerFactory.java

@@ -22,6 +22,7 @@
 import org.apache.flink.configuration.DeploymentOptions;
 import org.apache.flink.configuration.JobManagerOptions;
 import org.apache.flink.configuration.StateRecoveryOptions;
+import org.apache.flink.configuration.TaskManagerOptions;
 import org.apache.flink.core.failure.FailureEnricher;
 import org.apache.flink.runtime.blob.BlobWriter;
 import org.apache.flink.runtime.blocklist.BlocklistOperations;
@@ -123,7 +124,9 @@ public SchedulerNG createInstance(
                         jobMasterConfiguration.get(StateRecoveryOptions.LOCAL_RECOVERY),
                         jobMasterConfiguration.get(DeploymentOptions.TARGET),
                         jobMasterConfiguration.get(
-                                JobManagerOptions.SCHEDULER_PREFER_MINIMAL_TASKMANAGERS_ENABLED));
+                                JobManagerOptions.SCHEDULER_PREFER_MINIMAL_TASKMANAGERS_ENABLED),
+                        jobMasterConfiguration.get(
+                                TaskManagerOptions.TASK_MANAGER_LOAD_BALANCE_MODE));
 
         final ExecutionGraphFactory executionGraphFactory =
                 new DefaultExecutionGraphFactory(
@@ -169,13 +172,15 @@ public static SlotSharingSlotAllocator createSlotSharingSlotAllocator(
             DeclarativeSlotPool declarativeSlotPool,
             boolean localRecoveryEnabled,
             @Nullable String executionTarget,
-            boolean minimalTaskManagerPreferred) {
+            boolean minimalTaskManagerPreferred,
+            TaskManagerOptions.TaskManagerLoadBalanceMode taskManagerLoadBalanceMode) {
         return SlotSharingSlotAllocator.createSlotSharingSlotAllocator(
                 declarativeSlotPool::reserveFreeSlot,
                 declarativeSlotPool::freeReservedSlot,
                 declarativeSlotPool::containsFreeSlot,
                 localRecoveryEnabled,
                 executionTarget,
-                minimalTaskManagerPreferred);
+                minimalTaskManagerPreferred,
+                taskManagerLoadBalanceMode);
     }
 }

flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/adaptive/allocator/AllocatorUtil.java

@@ -18,11 +18,15 @@
 
 package org.apache.flink.runtime.scheduler.adaptive.allocator;
 
+import org.apache.flink.runtime.clusterframework.types.ResourceID;
 import org.apache.flink.runtime.instance.SlotSharingGroupId;
 import org.apache.flink.runtime.jobmaster.SlotInfo;
+import org.apache.flink.runtime.jobmaster.slotpool.PhysicalSlot;
 
 import java.util.Collection;
 import java.util.Map;
+import java.util.Set;
+import java.util.stream.Collectors;
 
 import static org.apache.flink.util.Preconditions.checkState;
 
@@ -54,4 +58,13 @@ static void checkMinimumRequiredSlots(
                 freeSlots.size(),
                 minimumRequiredSlots);
     }
+
+    static Map<ResourceID, Set<PhysicalSlot>> getSlotsPerTaskExecutor(
+            Collection<PhysicalSlot> physicalSlots) {
+        return physicalSlots.stream()
+                .collect(
+                        Collectors.groupingBy(
+                                slot -> slot.getTaskManagerLocation().getResourceID(),
+                                Collectors.toSet()));
+    }
 }

flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/adaptive/allocator/DefaultSlotAssigner.java

@@ -20,6 +20,7 @@
 
 import org.apache.flink.annotation.VisibleForTesting;
 import org.apache.flink.runtime.jobmaster.SlotInfo;
+import org.apache.flink.runtime.jobmaster.slotpool.PhysicalSlot;
 import org.apache.flink.runtime.scheduler.adaptive.JobSchedulingPlan.SlotAssignment;
 import org.apache.flink.runtime.scheduler.adaptive.allocator.SlotSharingSlotAllocator.ExecutionSlotSharingGroup;
 import org.apache.flink.runtime.taskmanager.TaskManagerLocation;
@@ -52,20 +53,23 @@ public class DefaultSlotAssigner implements SlotAssigner {
     private final @Nullable String executionTarget;
     private final boolean minimalTaskManagerPreferred;
     private final SlotSharingResolver slotSharingResolver;
+    private final SlotMatchingResolver slotMatchingResolver;
 
     DefaultSlotAssigner(
             @Nullable String executionTarget,
             boolean minimalTaskManagerPreferred,
-            SlotSharingResolver slotSharingResolver) {
+            SlotSharingResolver slotSharingResolver,
+            SlotMatchingResolver slotMatchingResolver) {
         this.executionTarget = executionTarget;
         this.minimalTaskManagerPreferred = minimalTaskManagerPreferred;
         this.slotSharingResolver = slotSharingResolver;
+        this.slotMatchingResolver = slotMatchingResolver;
     }
 
     @Override
     public Collection<SlotAssignment> assignSlots(
             JobInformation jobInformation,
-            Collection<? extends SlotInfo> freeSlots,
+            Collection<PhysicalSlot> freeSlots,
             VertexParallelism vertexParallelism,
             JobAllocationsInformation previousAllocations) {
         checkMinimumRequiredSlots(jobInformation, freeSlots);
@@ -74,26 +78,20 @@ public Collection<SlotAssignment> assignSlots(
                 slotSharingResolver.getExecutionSlotSharingGroups(
                         jobInformation, vertexParallelism);
 
-        final Collection<? extends SlotInfo> pickedSlots =
-                pickSlotsIfNeeded(allGroups.size(), freeSlots);
+        final Collection<PhysicalSlot> pickedSlots = pickSlotsIfNeeded(allGroups.size(), freeSlots);
 
-        Iterator<? extends SlotInfo> iterator = pickedSlots.iterator();
-        Collection<SlotAssignment> assignments = new ArrayList<>();
-        for (ExecutionSlotSharingGroup group : allGroups) {
-            assignments.add(new SlotAssignment(iterator.next(), group));
-        }
-        return assignments;
+        return slotMatchingResolver.matchSlotSharingGroupWithSlots(allGroups, pickedSlots);
     }
 
     @VisibleForTesting
-    Collection<? extends SlotInfo> pickSlotsIfNeeded(
-            int requestExecutionSlotSharingGroups, Collection<? extends SlotInfo> freeSlots) {
-        Collection<? extends SlotInfo> pickedSlots = freeSlots;
+    Collection<PhysicalSlot> pickSlotsIfNeeded(
+            int requestExecutionSlotSharingGroups, Collection<PhysicalSlot> freeSlots) {
+        Collection<PhysicalSlot> pickedSlots = freeSlots;
         if (APPLICATION_MODE_EXECUTION_TARGET.equalsIgnoreCase(executionTarget)
                 && minimalTaskManagerPreferred
                 // To avoid the sort-work loading.
                 && freeSlots.size() > requestExecutionSlotSharingGroups) {
-            final Map<TaskManagerLocation, ? extends Set<? extends SlotInfo>> slotsPerTaskExecutor =
+            final Map<TaskManagerLocation, Set<PhysicalSlot>> slotsPerTaskExecutor =
                     getSlotsPerTaskExecutor(freeSlots);
             pickedSlots =
                     pickSlotsInMinimalTaskExecutors(
@@ -127,21 +125,20 @@ private Iterator<TaskManagerLocation> getSortedTaskExecutors(
      * @param requestedGroups the number of the request execution slot sharing groups.
      * @return the target slots that are distributed on the minimal task executors.
      */
-    private Collection<? extends SlotInfo> pickSlotsInMinimalTaskExecutors(
-            Map<TaskManagerLocation, ? extends Set<? extends SlotInfo>> slotsByTaskExecutor,
-            int requestedGroups) {
-        final List<SlotInfo> pickedSlots = new ArrayList<>();
+    private Collection<PhysicalSlot> pickSlotsInMinimalTaskExecutors(
+            Map<TaskManagerLocation, Set<PhysicalSlot>> slotsByTaskExecutor, int requestedGroups) {
+        final List<PhysicalSlot> pickedSlots = new ArrayList<>();
         final Iterator<TaskManagerLocation> sortedTaskExecutors =
                 getSortedTaskExecutors(slotsByTaskExecutor);
         while (pickedSlots.size() < requestedGroups) {
-            Set<? extends SlotInfo> slotInfos = slotsByTaskExecutor.get(sortedTaskExecutors.next());
+            Set<PhysicalSlot> slotInfos = slotsByTaskExecutor.get(sortedTaskExecutors.next());
             pickedSlots.addAll(slotInfos);
         }
         return pickedSlots;
     }
 
-    private Map<TaskManagerLocation, ? extends Set<? extends SlotInfo>> getSlotsPerTaskExecutor(
-            Collection<? extends SlotInfo> slots) {
+    private Map<TaskManagerLocation, Set<PhysicalSlot>> getSlotsPerTaskExecutor(
+            Collection<PhysicalSlot> slots) {
         return slots.stream()
                 .collect(
                         Collectors.groupingBy(

flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/adaptive/allocator/SimpleSlotMatchingResolver.java

@@ -0,0 +1,44 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.runtime.scheduler.adaptive.allocator;
+
+import org.apache.flink.runtime.jobmaster.SlotInfo;
+import org.apache.flink.runtime.jobmaster.slotpool.PhysicalSlot;
+import org.apache.flink.runtime.scheduler.adaptive.JobSchedulingPlan;
+import org.apache.flink.runtime.scheduler.adaptive.allocator.SlotSharingSlotAllocator.ExecutionSlotSharingGroup;
+
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Iterator;
+
+/** The simple slot matching resolver implementation. */
+public enum SimpleSlotMatchingResolver implements SlotMatchingResolver {
+    INSTANCE;
+
+    @Override
+    public Collection<JobSchedulingPlan.SlotAssignment> matchSlotSharingGroupWithSlots(
+            Collection<ExecutionSlotSharingGroup> requestGroups,
+            Collection<PhysicalSlot> freeSlots) {
+        Iterator<? extends SlotInfo> iterator = freeSlots.iterator();
+        Collection<JobSchedulingPlan.SlotAssignment> assignments = new ArrayList<>();
+        for (SlotSharingSlotAllocator.ExecutionSlotSharingGroup group : requestGroups) {
+            assignments.add(new JobSchedulingPlan.SlotAssignment(iterator.next(), group));
+        }
+        return assignments;
+    }
+}

flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/adaptive/allocator/SlotAllocator.java

@@ -18,6 +18,7 @@
 package org.apache.flink.runtime.scheduler.adaptive.allocator;
 
 import org.apache.flink.runtime.jobmaster.SlotInfo;
+import org.apache.flink.runtime.jobmaster.slotpool.PhysicalSlot;
 import org.apache.flink.runtime.scheduler.adaptive.JobSchedulingPlan;
 import org.apache.flink.runtime.util.ResourceCounter;
 
@@ -61,7 +62,7 @@ Optional<VertexParallelism> determineParallelism(
      */
     Optional<JobSchedulingPlan> determineParallelismAndCalculateAssignment(
             JobInformation jobInformation,
-            Collection<? extends SlotInfo> slots,
+            Collection<PhysicalSlot> slots,
             JobAllocationsInformation jobAllocationsInformation);
 
     /**

flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/adaptive/allocator/SlotAssigner.java

@@ -18,7 +18,7 @@
 package org.apache.flink.runtime.scheduler.adaptive.allocator;
 
 import org.apache.flink.annotation.Internal;
-import org.apache.flink.runtime.jobmaster.SlotInfo;
+import org.apache.flink.runtime.jobmaster.slotpool.PhysicalSlot;
 import org.apache.flink.runtime.scheduler.adaptive.JobSchedulingPlan.SlotAssignment;
 
 import java.util.Collection;
@@ -29,7 +29,7 @@ public interface SlotAssigner {
 
     Collection<SlotAssignment> assignSlots(
             JobInformation jobInformation,
-            Collection<? extends SlotInfo> freeSlots,
+            Collection<PhysicalSlot> freeSlots,
             VertexParallelism vertexParallelism,
             JobAllocationsInformation previousAllocations);
 }

flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/adaptive/allocator/SlotMatchingResolver.java

@@ -0,0 +1,44 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.runtime.scheduler.adaptive.allocator;
+
+import org.apache.flink.runtime.jobmaster.slotpool.PhysicalSlot;
+import org.apache.flink.runtime.scheduler.adaptive.JobSchedulingPlan.SlotAssignment;
+import org.apache.flink.runtime.scheduler.adaptive.allocator.SlotSharingSlotAllocator.ExecutionSlotSharingGroup;
+import org.apache.flink.util.FlinkRuntimeException;
+
+import java.util.Collection;
+import java.util.function.Supplier;
+
+/** The interface to define the methods for request slot matching resolver. */
+public interface SlotMatchingResolver {
+
+    Supplier<FlinkRuntimeException> NO_SLOTS_EXCEPTION_GETTER =
+            () -> new FlinkRuntimeException("No suitable slots enough.");
+
+    /**
+     * Match slots from the free slots with the given collection of requests execution groups.
+     *
+     * @param requestGroups the requested execution slot sharing groups.
+     * @param freeSlots the free slots.
+     * @return The assignment result.
+     */
+    Collection<SlotAssignment> matchSlotSharingGroupWithSlots(
+            Collection<ExecutionSlotSharingGroup> requestGroups,
+            Collection<PhysicalSlot> freeSlots);
+}

flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/adaptive/allocator/SlotSharingSlotAllocator.java

@@ -18,6 +18,7 @@
 package org.apache.flink.runtime.scheduler.adaptive.allocator;
 
 import org.apache.flink.annotation.VisibleForTesting;
+import org.apache.flink.configuration.TaskManagerOptions;
 import org.apache.flink.runtime.clusterframework.types.AllocationID;
 import org.apache.flink.runtime.clusterframework.types.ResourceProfile;
 import org.apache.flink.runtime.instance.SlotSharingGroupId;
@@ -62,20 +63,23 @@ public class SlotSharingSlotAllocator implements SlotAllocator {
     private final @Nullable String executionTarget;
     private final boolean minimalTaskManagerPreferred;
     private final SlotSharingResolver slotSharingResolver = DefaultSlotSharingResolver.INSTANCE;
+    private final SlotMatchingResolver slotMatchingResolver;
 
     private SlotSharingSlotAllocator(
             ReserveSlotFunction reserveSlot,
             FreeSlotFunction freeSlotFunction,
             IsSlotAvailableAndFreeFunction isSlotAvailableAndFreeFunction,
             boolean localRecoveryEnabled,
             @Nullable String executionTarget,
-            boolean minimalTaskManagerPreferred) {
+            boolean minimalTaskManagerPreferred,
+            TaskManagerOptions.TaskManagerLoadBalanceMode taskManagerLoadBalanceMode) {
         this.reserveSlotFunction = reserveSlot;
         this.freeSlotFunction = freeSlotFunction;
         this.isSlotAvailableAndFreeFunction = isSlotAvailableAndFreeFunction;
         this.localRecoveryEnabled = localRecoveryEnabled;
         this.executionTarget = executionTarget;
         this.minimalTaskManagerPreferred = minimalTaskManagerPreferred;
+        this.slotMatchingResolver = getSlotMatchingResolver(taskManagerLoadBalanceMode);
     }
 
     public static SlotSharingSlotAllocator createSlotSharingSlotAllocator(
@@ -84,14 +88,16 @@ public static SlotSharingSlotAllocator createSlotSharingSlotAllocator(
             IsSlotAvailableAndFreeFunction isSlotAvailableAndFreeFunction,
             boolean localRecoveryEnabled,
             @Nullable String executionTarget,
-            boolean minimalTaskManagerPreferred) {
+            boolean minimalTaskManagerPreferred,
+            TaskManagerOptions.TaskManagerLoadBalanceMode taskManagerLoadBalanceMode) {
         return new SlotSharingSlotAllocator(
                 reserveSlot,
                 freeSlotFunction,
                 isSlotAvailableAndFreeFunction,
                 localRecoveryEnabled,
                 executionTarget,
-                minimalTaskManagerPreferred);
+                minimalTaskManagerPreferred,
+                taskManagerLoadBalanceMode);
     }
 
     @Override
@@ -146,7 +152,7 @@ public Optional<VertexParallelism> determineParallelism(
     @Override
     public Optional<JobSchedulingPlan> determineParallelismAndCalculateAssignment(
             JobInformation jobInformation,
-            Collection<? extends SlotInfo> slots,
+            Collection<PhysicalSlot> slots,
             JobAllocationsInformation jobAllocationsInformation) {
         return determineParallelism(jobInformation, slots)
                 .map(
@@ -157,7 +163,8 @@ public Optional<JobSchedulingPlan> determineParallelismAndCalculateAssignment(
                                             : new DefaultSlotAssigner(
                                                     executionTarget,
                                                     minimalTaskManagerPreferred,
-                                                    slotSharingResolver);
+                                                    slotSharingResolver,
+                                                    slotMatchingResolver);
                             return new JobSchedulingPlan(
                                     parallelism,
                                     slotAssigner.assignSlots(
@@ -168,6 +175,22 @@ public Optional<JobSchedulingPlan> determineParallelismAndCalculateAssignment(
                         });
     }
 
+    private SlotMatchingResolver getSlotMatchingResolver(
+            TaskManagerOptions.TaskManagerLoadBalanceMode taskManagerLoadBalanceMode) {
+        switch (taskManagerLoadBalanceMode) {
+            case NONE:
+            case MIN_RESOURCES:
+                return SimpleSlotMatchingResolver.INSTANCE;
+            case SLOTS:
+                return SlotsBalancedSlotMatchingResolver.INSTANCE;
+            default:
+                throw new UnsupportedOperationException(
+                        String.format(
+                                "Unsupported task manager load mode: %s",
+                                taskManagerLoadBalanceMode));
+        }
+    }
+
     /**
      * Distributes free slots across the slot-sharing groups of the job. Slots are distributed as
      * evenly as possible. If a group requires less than an even share of slots the remainder is