Plugin for custom evaluator for scaling metric evaluation

docs/content.zh/docs/operations/plugins.md

@@ -188,3 +188,5 @@ The following steps demonstrate how to develop and use a custom mutator.
     ```text
     2023-12-12 06:26:56,667 o.a.f.k.o.u.MutatorUtils [INFO ] Discovered mutator from plugin directory[/opt/flink/plugins]: org.apache.flink.mutator.CustomFlinkMutator.
     ```
+
+## Flink Autoscaler Custom Evaluator

flink-autoscaler-standalone/src/main/java/org/apache/flink/autoscaler/standalone/StandaloneAutoscalerEntrypoint.java

@@ -94,7 +94,8 @@ JobAutoScaler<KEY, Context> createJobAutoscaler(
                 new ScalingExecutor<>(eventHandler, stateStore),
                 eventHandler,
                 new RescaleApiScalingRealizer<>(eventHandler),
-                stateStore);
+                stateStore,
+                null);
     }
 
     @VisibleForTesting

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

@@ -19,14 +19,17 @@
 
 import org.apache.flink.annotation.VisibleForTesting;
 import org.apache.flink.api.common.JobStatus;
+import org.apache.flink.api.java.tuple.Tuple2;
 import org.apache.flink.autoscaler.config.AutoScalerOptions;
 import org.apache.flink.autoscaler.event.AutoScalerEventHandler;
 import org.apache.flink.autoscaler.exceptions.NotReadyException;
 import org.apache.flink.autoscaler.metrics.AutoscalerFlinkMetrics;
 import org.apache.flink.autoscaler.metrics.EvaluatedMetrics;
+import org.apache.flink.autoscaler.metrics.FlinkAutoscalerEvaluator;
 import org.apache.flink.autoscaler.realizer.ScalingRealizer;
 import org.apache.flink.autoscaler.state.AutoScalerStateStore;
 import org.apache.flink.autoscaler.tuning.ConfigChanges;
+import org.apache.flink.configuration.Configuration;
 import org.apache.flink.configuration.PipelineOptions;
 import org.apache.flink.util.Preconditions;
 
@@ -36,9 +39,11 @@
 import java.time.Clock;
 import java.util.HashMap;
 import java.util.Map;
+import java.util.Optional;
 import java.util.concurrent.ConcurrentHashMap;
 
 import static org.apache.flink.autoscaler.config.AutoScalerOptions.AUTOSCALER_ENABLED;
+import static org.apache.flink.autoscaler.config.AutoScalerOptions.CUSTOM_EVALUATOR_NAME;
 import static org.apache.flink.autoscaler.metrics.AutoscalerFlinkMetrics.initRecommendedParallelism;
 import static org.apache.flink.autoscaler.metrics.AutoscalerFlinkMetrics.resetRecommendedParallelism;
 import static org.apache.flink.autoscaler.metrics.ScalingHistoryUtils.getTrimmedScalingHistory;
@@ -58,6 +63,7 @@ public class JobAutoScalerImpl<KEY, Context extends JobAutoScalerContext<KEY>>
     private final AutoScalerEventHandler<KEY, Context> eventHandler;
     private final ScalingRealizer<KEY, Context> scalingRealizer;
     private final AutoScalerStateStore<KEY, Context> stateStore;
+    private final Map<String, FlinkAutoscalerEvaluator> customEvaluators;
 
     private Clock clock = Clock.systemDefaultZone();
 
@@ -73,13 +79,15 @@ public JobAutoScalerImpl(
             ScalingExecutor<KEY, Context> scalingExecutor,
             AutoScalerEventHandler<KEY, Context> eventHandler,
             ScalingRealizer<KEY, Context> scalingRealizer,
-            AutoScalerStateStore<KEY, Context> stateStore) {
+            AutoScalerStateStore<KEY, Context> stateStore,
+            Map<String, FlinkAutoscalerEvaluator> customEvaluators) {
         this.metricsCollector = metricsCollector;
         this.evaluator = evaluator;
         this.scalingExecutor = scalingExecutor;
         this.eventHandler = eventHandler;
         this.scalingRealizer = scalingRealizer;
         this.stateStore = stateStore;
+        this.customEvaluators = customEvaluators;
     }
 
     @Override
@@ -203,8 +211,15 @@ private void runScalingLogic(Context ctx, AutoscalerFlinkMetrics autoscalerMetri
 
         // Scaling tracking data contains previous restart times that are taken into account
         var restartTime = scalingTracking.getMaxRestartTimeOrDefault(ctx.getConfiguration());
+
+        var customEvaluatorWithConfig = getCustomEvaluatorIfRequired(ctx.getConfiguration());
+
         var evaluatedMetrics =
-                evaluator.evaluate(ctx.getConfiguration(), collectedMetrics, restartTime);
+                evaluator.evaluate(
+                        ctx.getConfiguration(),
+                        collectedMetrics,
+                        restartTime,
+                        customEvaluatorWithConfig);
         LOG.debug("Evaluated metrics: {}", evaluatedMetrics);
         lastEvaluatedMetrics.put(ctx.getJobKey(), evaluatedMetrics);
 
@@ -259,4 +274,20 @@ void setClock(Clock clock) {
         this.metricsCollector.setClock(clock);
         this.scalingExecutor.setClock(clock);
     }
+
+    @VisibleForTesting
+    protected Tuple2<FlinkAutoscalerEvaluator, Configuration> getCustomEvaluatorIfRequired(
+            Configuration conf) {
+        return Optional.ofNullable(conf.get(CUSTOM_EVALUATOR_NAME))
+                .map(
+                        name -> {
+                            FlinkAutoscalerEvaluator evaluator = customEvaluators.get(name);
+                            return evaluator != null
+                                    ? new Tuple2<>(
+                                            evaluator,
+                                            AutoScalerOptions.forCustomEvaluator(conf, name))
+                                    : null;
+                        })
+                .orElse(null);
+    }
 }

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

@@ -18,16 +18,19 @@
 package org.apache.flink.autoscaler;
 
 import org.apache.flink.annotation.VisibleForTesting;
+import org.apache.flink.api.java.tuple.Tuple2;
 import org.apache.flink.autoscaler.config.AutoScalerOptions;
 import org.apache.flink.autoscaler.metrics.CollectedMetricHistory;
 import org.apache.flink.autoscaler.metrics.CollectedMetrics;
 import org.apache.flink.autoscaler.metrics.EvaluatedMetrics;
 import org.apache.flink.autoscaler.metrics.EvaluatedScalingMetric;
+import org.apache.flink.autoscaler.metrics.FlinkAutoscalerEvaluator;
 import org.apache.flink.autoscaler.metrics.MetricAggregator;
 import org.apache.flink.autoscaler.metrics.ScalingMetric;
 import org.apache.flink.autoscaler.topology.JobTopology;
 import org.apache.flink.autoscaler.utils.AutoScalerUtils;
 import org.apache.flink.configuration.Configuration;
+import org.apache.flink.configuration.UnmodifiableConfiguration;
 import org.apache.flink.runtime.jobgraph.JobVertexID;
 
 import org.slf4j.Logger;
@@ -38,6 +41,7 @@
 
 import java.time.Duration;
 import java.time.Instant;
+import java.util.Collections;
 import java.util.HashMap;
 import java.util.Map;
 import java.util.Optional;
@@ -72,14 +76,36 @@ public class ScalingMetricEvaluator {
     private static final Logger LOG = LoggerFactory.getLogger(ScalingMetricEvaluator.class);
 
     public EvaluatedMetrics evaluate(
-            Configuration conf, CollectedMetricHistory collectedMetrics, Duration restartTime) {
+            Configuration conf,
+            CollectedMetricHistory collectedMetrics,
+            Duration restartTime,
+            @Nullable
+                    Tuple2<FlinkAutoscalerEvaluator, Configuration>
+                            customEvaluatorWithConfig) {
         LOG.debug("Restart time used in metrics evaluation: {}", restartTime);
         var scalingOutput = new HashMap<JobVertexID, Map<ScalingMetric, EvaluatedScalingMetric>>();
         var metricsHistory = collectedMetrics.getMetricHistory();
         var topology = collectedMetrics.getJobTopology();
 
         boolean processingBacklog = isProcessingBacklog(topology, metricsHistory, conf);
 
+        var customEvaluationSession =
+                Optional.ofNullable(customEvaluatorWithConfig)
+                        .map(
+                                info ->
+                                        Tuple2.of(
+                                                info.f0,
+                                                new FlinkAutoscalerEvaluator.Context(
+                                                        new UnmodifiableConfiguration(conf),
+                                                        Collections.unmodifiableSortedMap(
+                                                                metricsHistory),
+                                                        Collections.unmodifiableMap(scalingOutput),
+                                                        topology,
+                                                        processingBacklog,
+                                                        restartTime,
+                                                        info.f1)))
+                        .orElse(null);
+
         for (var vertex : topology.getVerticesInTopologicalOrder()) {
             scalingOutput.put(
                     vertex,
@@ -90,7 +116,8 @@ public EvaluatedMetrics evaluate(
                             topology,
                             vertex,
                             processingBacklog,
-                            restartTime));
+                            restartTime,
+                            customEvaluationSession));
         }
 
         var globalMetrics = evaluateGlobalMetrics(metricsHistory);
@@ -132,7 +159,10 @@ private Map<ScalingMetric, EvaluatedScalingMetric> evaluateMetrics(
             JobTopology topology,
             JobVertexID vertex,
             boolean processingBacklog,
-            Duration restartTime) {
+            Duration restartTime,
+            @Nullable
+                    Tuple2<FlinkAutoscalerEvaluator, FlinkAutoscalerEvaluator.Context>
+                            customEvaluationSession) {
 
         var latestVertexMetrics =
                 metricsHistory.get(metricsHistory.lastKey()).getVertexMetrics().get(vertex);
@@ -142,6 +172,7 @@ private Map<ScalingMetric, EvaluatedScalingMetric> evaluateMetrics(
         double inputRateAvg = getRate(ScalingMetric.NUM_RECORDS_IN, vertex, metricsHistory);
 
         var evaluatedMetrics = new HashMap<ScalingMetric, EvaluatedScalingMetric>();
+
         computeTargetDataRate(
                 topology,
                 vertex,
@@ -175,6 +206,24 @@ private Map<ScalingMetric, EvaluatedScalingMetric> evaluateMetrics(
                 EvaluatedScalingMetric.of(vertexInfo.getNumSourcePartitions()));
 
         computeProcessingRateThresholds(evaluatedMetrics, conf, processingBacklog, restartTime);
+
+        Optional.ofNullable(customEvaluationSession)
+                .map(
+                        session ->
+                                runCustomEvaluator(
+                                        vertex,
+                                        Collections.unmodifiableMap(evaluatedMetrics),
+                                        session))
+                .filter(customEvaluatedMetrics -> !customEvaluatedMetrics.isEmpty())
+                .ifPresent(
+                        customEvaluatedMetrics -> {
+                            LOG.info(
+                                    "Merging custom evaluated metrics for vertex {}: {}",
+                                    vertex,
+                                    customEvaluatedMetrics);
+                            mergeEvaluatedMetricsMaps(evaluatedMetrics, customEvaluatedMetrics);
+                        });
+
         return evaluatedMetrics;
     }
 
@@ -585,4 +634,76 @@ protected static double computeEdgeDataRate(
                 to);
         return getRate(ScalingMetric.NUM_RECORDS_OUT, from, metricsHistory);
     }
+
+    /**
+     * Executes the provided custom evaluator for the given job vertex. Calls {@link
+     * FlinkAutoscalerEvaluator#evaluateVertexMetrics} to evaluate scaling metrics.
+     *
+     * @param vertex The job vertex being evaluated.
+     * @param evaluatedMetrics Current evaluated metrics.
+     * @param customEvaluationSession A tuple containing the custom evaluator and evaluation
+     *     context.
+     * @return A map of scaling metrics, with its corresponding evaluated scaling metric.
+     */
+    @VisibleForTesting
+    protected static Map<ScalingMetric, EvaluatedScalingMetric> runCustomEvaluator(
+            JobVertexID vertex,
+            Map<ScalingMetric, EvaluatedScalingMetric> evaluatedMetrics,
+            Tuple2<FlinkAutoscalerEvaluator, FlinkAutoscalerEvaluator.Context>
+                    customEvaluationSession) {
+        try {
+            return customEvaluationSession.f0.evaluateVertexMetrics(
+                    vertex, evaluatedMetrics, customEvaluationSession.f1);
+        } catch (UnsupportedOperationException e) {
+            LOG.warn(
+                    "Custom evaluator {} tried accessing an un-modifiable view.",
+                    customEvaluationSession.f0.getClass(),
+                    e);
+        } catch (Exception e) {
+            LOG.warn(
+                    "Custom evaluator {} threw an exception.",
+                    customEvaluationSession.f0.getClass(),
+                    e);
+        }
+
+        return Collections.emptyMap();
+    }
+
+    /**
+     * Merges the incoming evaluated metrics into actual evaluated metrics.
+     *
+     * @param actual The target evaluated metrics map to merge into.
+     * @param incoming The incoming map containing new evaluated metrics map to be merged
+     *     (nullable).
+     */
+    @VisibleForTesting
+    protected static void mergeEvaluatedMetricsMaps(
+            Map<ScalingMetric, EvaluatedScalingMetric> actual,
+            @Nullable Map<ScalingMetric, EvaluatedScalingMetric> incoming) {
+        Optional.ofNullable(incoming)
+                .ifPresent(
+                        customEvaluatedMetric ->
+                                customEvaluatedMetric.forEach(
+                                        (scalingMetric, evaluatedScalingMetric) ->
+                                                actual.merge(
+                                                        scalingMetric,
+                                                        evaluatedScalingMetric,
+                                                        ScalingMetricEvaluator
+                                                                ::mergeEvaluatedScalingMetric)));
+    }
+
+    /**
+     * Merges two {@link EvaluatedScalingMetric} instances.
+     *
+     * @param actual The existing evaluated scaling metric.
+     * @param incoming The incoming evaluated scaling metric.
+     * @return A new {@link EvaluatedScalingMetric} instance with merged values.
+     */
+    @VisibleForTesting
+    protected static EvaluatedScalingMetric mergeEvaluatedScalingMetric(
+            EvaluatedScalingMetric actual, EvaluatedScalingMetric incoming) {
+        return new EvaluatedScalingMetric(
+                !Double.isNaN(incoming.getCurrent()) ? incoming.getCurrent() : actual.getCurrent(),
+                !Double.isNaN(incoming.getAverage()) ? incoming.getAverage() : actual.getAverage());
+    }
 }

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

@@ -21,6 +21,8 @@
 import org.apache.flink.autoscaler.metrics.MetricAggregator;
 import org.apache.flink.configuration.ConfigOption;
 import org.apache.flink.configuration.ConfigOptions;
+import org.apache.flink.configuration.Configuration;
+import org.apache.flink.configuration.DelegatingConfiguration;
 import org.apache.flink.configuration.MemorySize;
 
 import java.time.Duration;
@@ -31,6 +33,7 @@ public class AutoScalerOptions {
 
     public static final String OLD_K8S_OP_CONF_PREFIX = "kubernetes.operator.";
     public static final String AUTOSCALER_CONF_PREFIX = "job.autoscaler.";
+    public static final String CUSTOM_EVALUATOR_CONF_PREFIX = "metrics.custom-evaluator.";
 
     private static String oldOperatorConfigKey(String key) {
         return OLD_K8S_OP_CONF_PREFIX + AUTOSCALER_CONF_PREFIX + key;
@@ -418,4 +421,19 @@ private static ConfigOptions.OptionBuilder autoScalerConfig(String key) {
                             "Minimum allowed value for the observed scalability coefficient. "
                                     + "Prevents aggressive scaling by clamping low coefficient estimates. "
                                     + "If the estimated coefficient falls below this value, it is capped at the configured minimum.");
+
+    public static final ConfigOption<String> CUSTOM_EVALUATOR_NAME =
+            autoScalerConfig(CUSTOM_EVALUATOR_CONF_PREFIX + "name")
+                    .stringType()
+                    .defaultValue(null)
+                    .withFallbackKeys(oldOperatorConfigKey(CUSTOM_EVALUATOR_CONF_PREFIX + "name"))
+                    .withDescription("Name of the custom evaluator to be used.");
+
+    public static Configuration forCustomEvaluator(
+            Configuration configuration, String customEvaluatorName) {
+        // add support for fallBackKey with DelegatingConfiguration.
+        return new DelegatingConfiguration(
+                configuration,
+                AUTOSCALER_CONF_PREFIX + CUSTOM_EVALUATOR_CONF_PREFIX + customEvaluatorName + ".");
+    }
 }