Add CPU and memory size autoscaler quota

Author: gaborgsomogyi

.github/workflows/ci.yml

@@ -44,7 +44,7 @@ jobs:
       - name: Build with Maven
         run: |
           set -o pipefail; mvn clean install javadoc:javadoc -Pgenerate-docs | tee ./mvn.log; set +o pipefail
-          if [[ $(cat ./mvn.log | grep -E -v '(flink-runtime-.*.jar, flink-kubernetes-operator-.*.jar)|(flink-kubernetes-operator-.*.jar, flink-runtime-.*.jar) define 3 overlapping classes' | grep -c "overlapping classes" -) -gt 0 ]];then
+          if [[ $(cat ./mvn.log | grep -E -v '(flink-runtime-.*.jar, flink-autoscaler-.*.jar) define 2 overlapping classes|(flink-kubernetes-operator-.*.jar, flink-runtime-.*.jar) define 3 overlapping classes' | grep -c "overlapping classes" -) -gt 0 ]];then
             echo "Found overlapping classes: "
             cat ./mvn.log | grep "overlapping classes"
             exit 1

docs/layouts/shortcodes/generated/auto_scaler_configuration.html

@@ -110,6 +110,18 @@
             <td>Double</td>
             <td>Percentage threshold for switching to observed from busy time based true processing rate if the measurement is off by at least the configured fraction. For example 0.15 means we switch to observed if the busy time based computation is at least 15% higher during catchup.</td>
         </tr>
+        <tr>
+            <td><h5>job.autoscaler.quota.cpu</h5></td>
+            <td style="word-wrap: break-word;">(none)</td>
+            <td>Double</td>
+            <td>Quota of the CPU count. When scaling would go beyond this number the the scaling is not going to happen.</td>
+        </tr>
+        <tr>
+            <td><h5>job.autoscaler.quota.memory</h5></td>
+            <td style="word-wrap: break-word;">(none)</td>
+            <td>MemorySize</td>
+            <td>Quota of the memory size. When scaling would go beyond this number the the scaling is not going to happen.</td>
+        </tr>
         <tr>
             <td><h5>job.autoscaler.restart.time</h5></td>
             <td style="word-wrap: break-word;">5 min</td>

flink-autoscaler/src/main/java/org/apache/flink/autoscaler/ScalingExecutor.java

@@ -33,6 +33,7 @@
 import org.apache.flink.configuration.Configuration;
 import org.apache.flink.configuration.MemorySize;
 import org.apache.flink.configuration.TaskManagerOptions;
+import org.apache.flink.runtime.instance.SlotSharingGroupId;
 import org.apache.flink.runtime.jobgraph.JobVertexID;
 
 import org.slf4j.Logger;
@@ -67,6 +68,9 @@ public class ScalingExecutor<KEY, Context extends JobAutoScalerContext<KEY>> {
     public static final String HEAP_USAGE_MESSAGE =
             "Heap Usage %s is above the allowed limit for scaling operations. Please adjust the available memory manually.";
 
+    public static final String RESOURCE_QUOTA_REACHED_MESSAGE =
+            "Resource usage is above the allowed limit for scaling operations. Please adjust the resource quota manually.";
+
     private static final Logger LOG = LoggerFactory.getLogger(ScalingExecutor.class);
 
     private final JobVertexScaler<KEY, Context> jobVertexScaler;
@@ -115,6 +119,17 @@ public boolean scaleResource(
             return false;
         }
 
+        if (resourceQuotaReached(conf, evaluatedMetrics, scalingSummaries, context)) {
+            autoScalerEventHandler.handleEvent(
+                    context,
+                    AutoScalerEventHandler.Type.Warning,
+                    "ResourceQuotaReached",
+                    RESOURCE_QUOTA_REACHED_MESSAGE,
+                    null,
+                    conf.get(SCALING_EVENT_INTERVAL));
+            return false;
+        }
+
         updateRecommendedParallelism(evaluatedMetrics.getVertexMetrics(), scalingSummaries);
 
         if (checkIfBlockedAndTriggerScalingEvent(context, scalingSummaries, conf, now)) {
@@ -199,6 +214,85 @@ protected static boolean allVerticesWithinUtilizationTarget(
         return true;
     }
 
+    protected static boolean resourceQuotaReached(
+            Configuration conf,
+            EvaluatedMetrics evaluatedMetrics,
+            Map<JobVertexID, ScalingSummary> scalingSummaries,
+            JobAutoScalerContext<?> ctx) {
+
+        if (evaluatedMetrics.getJobTopology() == null
+                || evaluatedMetrics.getJobTopology().getSlotSharingGroupMapping().isEmpty()) {
+            return false;
+        }
+
+        var cpuQuota = conf.getOptional(AutoScalerOptions.CPU_QUOTA);
+        var memoryQuota = conf.getOptional(AutoScalerOptions.MEMORY_QUOTA);
+        var tmMemory = ctx.getTaskManagerMemory();
+        var tmCpu = ctx.getTaskManagerCpu();
+
+        if (cpuQuota.isPresent() || memoryQuota.isPresent()) {
+            var currentSlotSharingGroupMaxParallelisms = new HashMap<SlotSharingGroupId, Integer>();
+            var newSlotSharingGroupMaxParallelisms = new HashMap<SlotSharingGroupId, Integer>();
+            for (var e :
+                    evaluatedMetrics.getJobTopology().getSlotSharingGroupMapping().entrySet()) {
+                int currentMaxParallelism =
+                        e.getValue().stream()
+                                .filter(scalingSummaries::containsKey)
+                                .mapToInt(v -> scalingSummaries.get(v).getCurrentParallelism())
+                                .max()
+                                .orElse(0);
+                currentSlotSharingGroupMaxParallelisms.put(e.getKey(), currentMaxParallelism);
+                int newMaxParallelism =
+                        e.getValue().stream()
+                                .filter(scalingSummaries::containsKey)
+                                .mapToInt(v -> scalingSummaries.get(v).getNewParallelism())
+                                .max()
+                                .orElse(0);
+                newSlotSharingGroupMaxParallelisms.put(e.getKey(), newMaxParallelism);
+            }
+
+            var numSlotsPerTm = conf.get(TaskManagerOptions.NUM_TASK_SLOTS);
+            var currentTotalSlots =
+                    currentSlotSharingGroupMaxParallelisms.values().stream()
+                            .mapToInt(Integer::intValue)
+                            .sum();
+            var currentNumTms = currentTotalSlots / numSlotsPerTm;
+            var newTotalSlots =
+                    newSlotSharingGroupMaxParallelisms.values().stream()
+                            .mapToInt(Integer::intValue)
+                            .sum();
+            var newNumTms = newTotalSlots / numSlotsPerTm;
+
+            if (newNumTms <= currentNumTms) {
+                LOG.debug(
+                        "Skipping quota check due to new resource allocation is less or equals than the current");
+                return false;
+            }
+
+            if (cpuQuota.isPresent() && tmCpu.isPresent()) {
+                LOG.debug("CPU resource quota is {}, checking limits", cpuQuota.get());
+                double totalCPU = tmCpu.get() * newNumTms;
+                if (totalCPU > cpuQuota.get()) {
+                    LOG.info("CPU resource quota reached with value: {}", totalCPU);
+                    return true;
+                }
+            }
+
+            if (memoryQuota.isPresent() && tmMemory.isPresent()) {
+                LOG.debug("Memory resource quota is {}, checking limits", memoryQuota.get());
+                long totalMemory = tmMemory.get().getBytes() * newNumTms;
+                if (totalMemory > memoryQuota.get().getBytes()) {
+                    LOG.info(
+                            "Memory resource quota reached with value: {}",
+                            new MemorySize(totalMemory));
+                    return true;
+                }
+            }
+        }
+
+        return false;
+    }
+
     @VisibleForTesting
     Map<JobVertexID, ScalingSummary> computeScalingSummary(
             Context context,

flink-autoscaler/src/main/java/org/apache/flink/autoscaler/ScalingMetricCollector.java

@@ -210,7 +210,15 @@ protected JobTopology getJobTopology(
     @VisibleForTesting
     @SneakyThrows
     protected JobTopology getJobTopology(JobDetailsInfo jobDetailsInfo) {
-        Map<JobVertexID, Integer> maxParallelismMap =
+        var slotSharingGroupIdMap =
+                jobDetailsInfo.getJobVertexInfos().stream()
+                        .filter(e -> e.getSlotSharingGroupId() != null)
+                        .collect(
+                                Collectors.toMap(
+                                        JobDetailsInfo.JobVertexDetailsInfo::getJobVertexID,
+                                        JobDetailsInfo.JobVertexDetailsInfo
+                                                ::getSlotSharingGroupId));
+        var maxParallelismMap =
                 jobDetailsInfo.getJobVertexInfos().stream()
                         .collect(
                                 Collectors.toMap(
@@ -235,7 +243,8 @@ protected JobTopology getJobTopology(JobDetailsInfo jobDetailsInfo) {
                                     d.getJobVertexID(), IOMetrics.from(d.getJobVertexMetrics()));
                         });
 
-        return JobTopology.fromJsonPlan(json, maxParallelismMap, metrics, finished);
+        return JobTopology.fromJsonPlan(
+                json, slotSharingGroupIdMap, maxParallelismMap, metrics, finished);
     }
 
     private void updateKafkaSourceMaxParallelisms(Context ctx, JobID jobId, JobTopology topology)
@@ -254,7 +263,7 @@ private void updateKafkaSourceMaxParallelisms(Context ctx, JobID jobId, JobTopol
                                 "Updating source {} max parallelism based on available partitions to {}",
                                 sourceVertex,
                                 numPartitions);
-                        topology.updateMaxParallelism(sourceVertex, (int) numPartitions);
+                        topology.get(sourceVertex).setMaxParallelism((int) numPartitions);
                     }
                 }
             }

flink-autoscaler/src/main/java/org/apache/flink/autoscaler/ScalingMetricEvaluator.java

@@ -91,7 +91,8 @@ public EvaluatedMetrics evaluate(
         }
 
         var globalMetrics = evaluateGlobalMetrics(metricsHistory);
-        return new EvaluatedMetrics(scalingOutput, globalMetrics);
+        return new EvaluatedMetrics(
+                collectedMetrics.getJobTopology(), scalingOutput, globalMetrics);
     }
 
     @VisibleForTesting

flink-autoscaler/src/main/java/org/apache/flink/autoscaler/config/AutoScalerOptions.java

@@ -20,6 +20,7 @@
 import org.apache.flink.autoscaler.metrics.MetricAggregator;
 import org.apache.flink.configuration.ConfigOption;
 import org.apache.flink.configuration.ConfigOptions;
+import org.apache.flink.configuration.MemorySize;
 
 import java.time.Duration;
 import java.util.List;
@@ -319,4 +320,20 @@ private static ConfigOptions.OptionBuilder autoScalerConfig(String key) {
                     .defaultValue(Duration.ofSeconds(10))
                     .withFallbackKeys(oldOperatorConfigKey("flink.rest-client.timeout"))
                     .withDescription("The timeout for waiting the flink rest client to return.");
+
+    public static final ConfigOption<MemorySize> MEMORY_QUOTA =
+            autoScalerConfig("quota.memory")
+                    .memoryType()
+                    .noDefaultValue()
+                    .withFallbackKeys(oldOperatorConfigKey("quota.memory"))
+                    .withDescription(
+                            "Quota of the memory size. When scaling would go beyond this number the the scaling is not going to happen.");
+
+    public static final ConfigOption<Double> CPU_QUOTA =
+            autoScalerConfig("quota.cpu")
+                    .doubleType()
+                    .noDefaultValue()
+                    .withFallbackKeys(oldOperatorConfigKey("quota.cpu"))
+                    .withDescription(
+                            "Quota of the CPU count. When scaling would go beyond this number the the scaling is not going to happen.");
 }

flink-autoscaler/src/main/java/org/apache/flink/autoscaler/metrics/EvaluatedMetrics.java

@@ -17,6 +17,7 @@
 
 package org.apache.flink.autoscaler.metrics;
 
+import org.apache.flink.autoscaler.topology.JobTopology;
 import org.apache.flink.runtime.jobgraph.JobVertexID;
 
 import lombok.AllArgsConstructor;
@@ -30,6 +31,7 @@
 @NoArgsConstructor
 @AllArgsConstructor
 public class EvaluatedMetrics {
+    private JobTopology jobTopology;
     private Map<JobVertexID, Map<ScalingMetric, EvaluatedScalingMetric>> vertexMetrics;
     private Map<ScalingMetric, EvaluatedScalingMetric> globalMetrics;
 }

flink-autoscaler/src/main/java/org/apache/flink/autoscaler/topology/JobTopology.java

@@ -17,6 +17,7 @@
 
 package org.apache.flink.autoscaler.topology;
 
+import org.apache.flink.runtime.instance.SlotSharingGroupId;
 import org.apache.flink.runtime.jobgraph.JobVertexID;
 
 import org.apache.flink.shaded.guava31.com.google.common.collect.ImmutableMap;
@@ -48,6 +49,7 @@ public class JobTopology {
     private static final ObjectMapper objectMapper = new ObjectMapper();
 
     @Getter private final Map<JobVertexID, VertexInfo> vertexInfos;
+    @Getter private final Map<SlotSharingGroupId, Set<JobVertexID>> slotSharingGroupMapping;
     @Getter private final Set<JobVertexID> finishedVertices;
     @Getter private final List<JobVertexID> verticesInTopologicalOrder;
 
@@ -66,6 +68,7 @@ public JobTopology(Set<VertexInfo> vertexInfo) {
                 ImmutableMap.copyOf(
                         vertexInfo.stream().collect(Collectors.toMap(VertexInfo::getId, v -> v)));
 
+        Map<SlotSharingGroupId, Set<JobVertexID>> vertexSlotSharingGroupMapping = new HashMap<>();
         var finishedVertices = ImmutableSet.<JobVertexID>builder();
 
         vertexInfo.forEach(
@@ -79,12 +82,21 @@ public JobTopology(Set<VertexInfo> vertexInfo) {
                                             vertexOutputs
                                                     .computeIfAbsent(inputId, id -> new HashMap<>())
                                                     .put(vertexId, shipStrategy));
+
+                    var slotSharingGroupId = info.getSlotSharingGroupId();
+                    if (slotSharingGroupId != null) {
+                        vertexSlotSharingGroupMapping
+                                .computeIfAbsent(slotSharingGroupId, id -> new HashSet<>())
+                                .add(vertexId);
+                    }
+
                     if (info.isFinished()) {
                         finishedVertices.add(vertexId);
                     }
                 });
         vertexOutputs.forEach((v, outputs) -> vertexInfos.get(v).setOutputs(outputs));
 
+        this.slotSharingGroupMapping = ImmutableMap.copyOf(vertexSlotSharingGroupMapping);
         this.finishedVertices = finishedVertices.build();
         this.verticesInTopologicalOrder = returnVerticesInTopologicalOrder();
     }
@@ -97,10 +109,6 @@ public boolean isSource(JobVertexID jobVertexID) {
         return get(jobVertexID).getInputs().isEmpty();
     }
 
-    public void updateMaxParallelism(JobVertexID vertexID, int maxParallelism) {
-        get(vertexID).updateMaxParallelism(maxParallelism);
-    }
-
     private List<JobVertexID> returnVerticesInTopologicalOrder() {
         List<JobVertexID> sorted = new ArrayList<>(vertexInfos.size());
 
@@ -134,6 +142,7 @@ private List<JobVertexID> returnVerticesInTopologicalOrder() {
 
     public static JobTopology fromJsonPlan(
             String jsonPlan,
+            Map<JobVertexID, SlotSharingGroupId> slotSharingGroupIdMap,
             Map<JobVertexID, Integer> maxParallelismMap,
             Map<JobVertexID, IOMetrics> metrics,
             Set<JobVertexID> finishedVertices)
@@ -151,6 +160,7 @@ public static JobTopology fromJsonPlan(
             vertexInfo.add(
                     new VertexInfo(
                             vertexId,
+                            slotSharingGroupIdMap.get(vertexId),
                             inputs,
                             node.get("parallelism").asInt(),
                             maxParallelismMap.get(vertexId),

flink-autoscaler/src/main/java/org/apache/flink/autoscaler/topology/VertexInfo.java

@@ -18,6 +18,7 @@
 package org.apache.flink.autoscaler.topology;
 
 import org.apache.flink.annotation.VisibleForTesting;
+import org.apache.flink.runtime.instance.SlotSharingGroupId;
 import org.apache.flink.runtime.jobgraph.JobVertexID;
 
 import lombok.Data;
@@ -33,6 +34,8 @@ public class VertexInfo {
     // All input vertices and the ship_strategy
     private final Map<JobVertexID, ShipStrategy> inputs;
 
+    private final SlotSharingGroupId slotSharingGroupId;
+
     // All output vertices and the ship_strategy
     private Map<JobVertexID, ShipStrategy> outputs;
 
@@ -48,12 +51,14 @@ public class VertexInfo {
 
     public VertexInfo(
             JobVertexID id,
+            SlotSharingGroupId slotSharingGroupId,
             Map<JobVertexID, ShipStrategy> inputs,
             int parallelism,
             int maxParallelism,
             boolean finished,
             IOMetrics ioMetrics) {
         this.id = id;
+        this.slotSharingGroupId = slotSharingGroupId;
         this.inputs = inputs;
         this.parallelism = parallelism;
         this.maxParallelism = maxParallelism;
@@ -69,19 +74,26 @@ public VertexInfo(
             int parallelism,
             int maxParallelism,
             IOMetrics ioMetrics) {
-        this(id, inputs, parallelism, maxParallelism, false, ioMetrics);
+        this(id, null, inputs, parallelism, maxParallelism, false, ioMetrics);
     }
 
     @VisibleForTesting
     public VertexInfo(
             JobVertexID id,
             Map<JobVertexID, ShipStrategy> inputs,
             int parallelism,
-            int maxParallelism) {
-        this(id, inputs, parallelism, maxParallelism, null);
+            int maxParallelism,
+            boolean finished,
+            IOMetrics ioMetrics) {
+        this(id, null, inputs, parallelism, maxParallelism, finished, ioMetrics);
     }
 
-    public void updateMaxParallelism(int maxParallelism) {
-        setMaxParallelism(Math.min(originalMaxParallelism, maxParallelism));
+    @VisibleForTesting
+    public VertexInfo(
+            JobVertexID id,
+            Map<JobVertexID, ShipStrategy> inputs,
+            int parallelism,
+            int maxParallelism) {
+        this(id, inputs, parallelism, maxParallelism, null);
     }
 }