Add configurable option to scale-down during task run time for cost-based autoscaler (#18958)

* Add configurable option to scale-down during task run time for cost-based autoscaler

* Docs

* Address review comments, compress tests a bit

Author: Fly-Style

docs/ingestion/supervisor.md

@@ -208,6 +208,8 @@ The following table outlines the configuration properties related to the `costBa
 |`lagWeight`|The weight of extracted lag value in cost function.| No| 0.25|
 |`idleWeight`|The weight of extracted poll idle value in cost function. | No | 0.75 |
 |`defaultProcessingRate`|A planned processing rate per task, required for first cost estimations. | No | 1000 |
+|`scaleDownBarrier`| A number of successful scale down attempts which should be skipped  to prevent the auto-scaler from scaling down tasks immediately.  | No | 5 |
+|`scaleDownDuringTaskRolloverOnly`| Indicates whether task scaling down is limited to periods during task rollovers only. | No | False |
 
 The following example shows a supervisor spec with `lagBased` autoscaler:
 

embedded-tests/src/test/java/org/apache/druid/testing/embedded/indexing/autoscaler/CostBasedAutoScalerIntegrationTest.java

@@ -132,6 +132,8 @@ public void test_autoScaler_computesOptimalTaskCountAndProduceScaleDown()
         // Weight configuration: strongly favor lag reduction over idle time
         .lagWeight(0.9)
         .idleWeight(0.1)
+        .scaleDownDuringTaskRolloverOnly(false)
+        .scaleDownBarrier(1)
         .build();
 
     final KafkaSupervisorSpec spec = createKafkaSupervisorWithAutoScaler(superId, autoScalerConfig, initialTaskCount);
@@ -225,6 +227,9 @@ void test_scaleDownDuringTaskRollover()
         // High idle weight ensures scale-down when tasks are mostly idle (little data to process)
         .lagWeight(0.1)
         .idleWeight(0.9)
+        .scaleDownDuringTaskRolloverOnly(false)
+        // Do not slow scale-downs
+        .scaleDownBarrier(0)
         .build();
 
     final KafkaSupervisorSpec spec = createKafkaSupervisorWithAutoScaler(superId, autoScalerConfig, initialTaskCount);

indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/supervisor/autoscaler/CostBasedAutoScaler.java

@@ -90,6 +90,8 @@ public class CostBasedAutoScaler implements SupervisorTaskAutoScaler
   private final WeightedCostFunction costFunction;
   private volatile CostMetrics lastKnownMetrics;
 
+  private int scaleDownCounter = 0;
+
   public CostBasedAutoScaler(
       SeekableStreamSupervisor supervisor,
       CostBasedAutoScalerConfig config,
@@ -148,7 +150,12 @@ public void reset()
   @Override
   public int computeTaskCountForRollover()
   {
-    return computeOptimalTaskCount(lastKnownMetrics);
+    if (config.isScaleDownOnTaskRolloverOnly()) {
+      return computeOptimalTaskCount(lastKnownMetrics);
+    } else {
+      scaleDownCounter = 0;
+      return -1;
+    }
   }
 
   public int computeTaskCountForScaleAction()
@@ -157,11 +164,18 @@ public int computeTaskCountForScaleAction()
     final int optimalTaskCount = computeOptimalTaskCount(lastKnownMetrics);
     final int currentTaskCount = lastKnownMetrics.getCurrentTaskCount();
 
-    // Perform only scale-up actions
+    // Perform scale-up actions; scale-down actions only if configured.
     int taskCount = -1;
     if (optimalTaskCount > currentTaskCount) {
       taskCount = optimalTaskCount;
-      log.info("New task count [%d] on supervisor [%s]", taskCount, supervisorId);
+      scaleDownCounter = 0; // Nullify the scaleDown counter after a successful scaleup too.
+      log.info("New task count [%d] on supervisor [%s], scaling up", taskCount, supervisorId);
+    } else if (!config.isScaleDownOnTaskRolloverOnly()
+               && optimalTaskCount < currentTaskCount
+               && ++scaleDownCounter >= config.getScaleDownBarrier()) {
+      taskCount = optimalTaskCount;
+      scaleDownCounter = 0;
+      log.info("New task count [%d] on supervisor [%s], scaling down", taskCount, supervisorId);
     } else {
       log.info("No scaling required for supervisor [%s]", supervisorId);
     }

indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/supervisor/autoscaler/CostBasedAutoScalerConfig.java

@@ -43,11 +43,12 @@
 @JsonInclude(JsonInclude.Include.NON_NULL)
 public class CostBasedAutoScalerConfig implements AutoScalerConfig
 {
-  private static final long DEFAULT_SCALE_ACTION_PERIOD_MILLIS = 15 * 60 * 1000; // 15 minutes
-  private static final long DEFAULT_MIN_TRIGGER_SCALE_ACTION_FREQUENCY_MILLIS = 1200000; // 20 minutes
-  private static final double DEFAULT_LAG_WEIGHT = 0.25;
-  private static final double DEFAULT_IDLE_WEIGHT = 0.75;
-  private static final double DEFAULT_PROCESSING_RATE = 1000.0; // 1000 records/sec per task as default
+  static final long DEFAULT_SCALE_ACTION_PERIOD_MILLIS = 10 * 60 * 1000; // 10 minutes
+  static final long DEFAULT_MIN_TRIGGER_SCALE_ACTION_FREQUENCY_MILLIS = 5 * 60 * 1000; // 5 minutes
+  static final double DEFAULT_LAG_WEIGHT = 0.25;
+  static final double DEFAULT_IDLE_WEIGHT = 0.75;
+  static final double DEFAULT_PROCESSING_RATE = 1000.0; // 1000 records/sec per task as default
+  static final int DEFAULT_SCALE_DOWN_BARRIER = 5; // We delay scale down by 5 * DEFAULT_SCALE_ACTION_PERIOD_MILLIS
 
   private final boolean enableTaskAutoScaler;
   private final int taskCountMax;
@@ -60,6 +61,18 @@ public class CostBasedAutoScalerConfig implements AutoScalerConfig
   private final double lagWeight;
   private final double idleWeight;
   private final double defaultProcessingRate;
+  /**
+   * Represents the threshold value used to prevent the auto-scaler from scaling down tasks immediately,
+   * when the computed cost-based metrics fall below this barrier.
+   * A higher value implies a more conservative scaling down behavior, ensuring that tasks
+   * are not prematurely terminated in scenarios of potential workload spikes or insufficient cost savings.
+   */
+  private final int scaleDownBarrier;
+  /**
+   * Indicates whether task scaling down is limited to periods during task rollovers only.
+   * If set to {@code false}, allows scaling down during normal task run time.
+   */
+  private final boolean scaleDownDuringTaskRolloverOnly;
 
   @JsonCreator
   public CostBasedAutoScalerConfig(
@@ -72,7 +85,9 @@ public CostBasedAutoScalerConfig(
       @Nullable @JsonProperty("scaleActionPeriodMillis") Long scaleActionPeriodMillis,
       @Nullable @JsonProperty("lagWeight") Double lagWeight,
       @Nullable @JsonProperty("idleWeight") Double idleWeight,
-      @Nullable @JsonProperty("defaultProcessingRate") Double defaultProcessingRate
+      @Nullable @JsonProperty("defaultProcessingRate") Double defaultProcessingRate,
+      @Nullable @JsonProperty("scaleDownBarrier") Integer scaleDownBarrier,
+      @Nullable @JsonProperty("scaleDownDuringTaskRolloverOnly") Boolean scaleDownDuringTaskRolloverOnly
   )
   {
     this.enableTaskAutoScaler = enableTaskAutoScaler != null ? enableTaskAutoScaler : false;
@@ -90,6 +105,8 @@ public CostBasedAutoScalerConfig(
     this.lagWeight = Configs.valueOrDefault(lagWeight, DEFAULT_LAG_WEIGHT);
     this.idleWeight = Configs.valueOrDefault(idleWeight, DEFAULT_IDLE_WEIGHT);
     this.defaultProcessingRate = Configs.valueOrDefault(defaultProcessingRate, DEFAULT_PROCESSING_RATE);
+    this.scaleDownBarrier = Configs.valueOrDefault(scaleDownBarrier, DEFAULT_SCALE_DOWN_BARRIER);
+    this.scaleDownDuringTaskRolloverOnly = Configs.valueOrDefault(scaleDownDuringTaskRolloverOnly, false);
 
     if (this.enableTaskAutoScaler) {
       Preconditions.checkNotNull(taskCountMax, "taskCountMax is required when enableTaskAutoScaler is true");
@@ -107,17 +124,16 @@ public CostBasedAutoScalerConfig(
     }
     this.taskCountStart = taskCountStart;
 
-    // Validate stopTaskCountRatio
     Preconditions.checkArgument(
         stopTaskCountRatio == null || (stopTaskCountRatio > 0.0 && stopTaskCountRatio <= 1.0),
         "0.0 < stopTaskCountRatio <= 1.0"
     );
     this.stopTaskCountRatio = stopTaskCountRatio;
 
-    // Validate weights are non-negative
     Preconditions.checkArgument(this.lagWeight >= 0, "lagWeight must be >= 0");
     Preconditions.checkArgument(this.idleWeight >= 0, "idleWeight must be >= 0");
     Preconditions.checkArgument(this.defaultProcessingRate > 0, "defaultProcessingRate must be > 0");
+    Preconditions.checkArgument(this.scaleDownBarrier >= 0, "scaleDownBarrier must be >= 0");
   }
 
   /**
@@ -196,6 +212,18 @@ public double getDefaultProcessingRate()
     return defaultProcessingRate;
   }
 
+  @JsonProperty
+  public int getScaleDownBarrier()
+  {
+    return scaleDownBarrier;
+  }
+
+  @JsonProperty("scaleDownDuringTaskRolloverOnly")
+  public boolean isScaleDownOnTaskRolloverOnly()
+  {
+    return scaleDownDuringTaskRolloverOnly;
+  }
+
   @Override
   public SupervisorTaskAutoScaler createAutoScaler(Supervisor supervisor, SupervisorSpec spec, ServiceEmitter emitter)
   {
@@ -222,6 +250,8 @@ public boolean equals(Object o)
            && Double.compare(that.lagWeight, lagWeight) == 0
            && Double.compare(that.idleWeight, idleWeight) == 0
            && Double.compare(that.defaultProcessingRate, defaultProcessingRate) == 0
+           && scaleDownBarrier == that.scaleDownBarrier
+           && scaleDownDuringTaskRolloverOnly == that.scaleDownDuringTaskRolloverOnly
            && Objects.equals(taskCountStart, that.taskCountStart)
            && Objects.equals(stopTaskCountRatio, that.stopTaskCountRatio);
   }
@@ -239,7 +269,9 @@ public int hashCode()
         scaleActionPeriodMillis,
         lagWeight,
         idleWeight,
-        defaultProcessingRate
+        defaultProcessingRate,
+        scaleDownBarrier,
+        scaleDownDuringTaskRolloverOnly
     );
   }
 
@@ -257,6 +289,8 @@ public String toString()
            ", lagWeight=" + lagWeight +
            ", idleWeight=" + idleWeight +
            ", defaultProcessingRate=" + defaultProcessingRate +
+           ", scaleDownBarrier=" + scaleDownBarrier +
+           ", scaleDownDuringTaskRolloverOnly=" + scaleDownDuringTaskRolloverOnly +
            '}';
   }
 
@@ -276,6 +310,8 @@ public static class Builder
     private Double lagWeight;
     private Double idleWeight;
     private Double defaultProcessingRate;
+    private Integer scaleDownBarrier;
+    private Boolean scaleDownDuringTaskRolloverOnly;
 
     private Builder()
     {
@@ -341,6 +377,18 @@ public Builder defaultProcessingRate(double defaultProcessingRate)
       return this;
     }
 
+    public Builder scaleDownBarrier(int scaleDownBarrier)
+    {
+      this.scaleDownBarrier = scaleDownBarrier;
+      return this;
+    }
+
+    public Builder scaleDownDuringTaskRolloverOnly(boolean scaleDownDuringTaskRolloverOnly)
+    {
+      this.scaleDownDuringTaskRolloverOnly = scaleDownDuringTaskRolloverOnly;
+      return this;
+    }
+
     public CostBasedAutoScalerConfig build()
     {
       return new CostBasedAutoScalerConfig(
@@ -353,7 +401,9 @@ public CostBasedAutoScalerConfig build()
           scaleActionPeriodMillis,
           lagWeight,
           idleWeight,
-          defaultProcessingRate
+          defaultProcessingRate,
+          scaleDownBarrier,
+          scaleDownDuringTaskRolloverOnly
       );
     }
   }

indexing-service/src/test/java/org/apache/druid/indexing/seekablestream/supervisor/autoscaler/CostBasedAutoScalerConfigTest.java

@@ -24,6 +24,13 @@
 import org.junit.Assert;
 import org.junit.Test;
 
+import static org.apache.druid.indexing.seekablestream.supervisor.autoscaler.CostBasedAutoScalerConfig.DEFAULT_IDLE_WEIGHT;
+import static org.apache.druid.indexing.seekablestream.supervisor.autoscaler.CostBasedAutoScalerConfig.DEFAULT_LAG_WEIGHT;
+import static org.apache.druid.indexing.seekablestream.supervisor.autoscaler.CostBasedAutoScalerConfig.DEFAULT_MIN_TRIGGER_SCALE_ACTION_FREQUENCY_MILLIS;
+import static org.apache.druid.indexing.seekablestream.supervisor.autoscaler.CostBasedAutoScalerConfig.DEFAULT_PROCESSING_RATE;
+import static org.apache.druid.indexing.seekablestream.supervisor.autoscaler.CostBasedAutoScalerConfig.DEFAULT_SCALE_ACTION_PERIOD_MILLIS;
+import static org.apache.druid.indexing.seekablestream.supervisor.autoscaler.CostBasedAutoScalerConfig.DEFAULT_SCALE_DOWN_BARRIER;
+
 public class CostBasedAutoScalerConfigTest
 {
   private final ObjectMapper mapper = new DefaultObjectMapper();
@@ -42,7 +49,9 @@ public void testSerdeWithAllProperties() throws Exception
                   + "  \"scaleActionPeriodMillis\": 60000,\n"
                   + "  \"lagWeight\": 0.6,\n"
                   + "  \"idleWeight\": 0.4,\n"
-                  + "  \"distancePenaltyExponent\": 3.0\n"
+                  + "  \"defaultProcessingRate\": 2000.0,\n"
+                  + "  \"scaleDownBarrier\": 10,\n"
+                  + "  \"scaleDownDuringTaskRolloverOnly\": true\n"
                   + "}";
 
     CostBasedAutoScalerConfig config = mapper.readValue(json, CostBasedAutoScalerConfig.class);
@@ -56,6 +65,9 @@ public void testSerdeWithAllProperties() throws Exception
     Assert.assertEquals(60000L, config.getScaleActionPeriodMillis());
     Assert.assertEquals(0.6, config.getLagWeight(), 0.001);
     Assert.assertEquals(0.4, config.getIdleWeight(), 0.001);
+    Assert.assertEquals(2000.0, config.getDefaultProcessingRate(), 0.001);
+    Assert.assertEquals(10, config.getScaleDownBarrier());
+    Assert.assertTrue(config.isScaleDownOnTaskRolloverOnly());
 
     // Test serialization back to JSON
     String serialized = mapper.writeValueAsString(config);
@@ -81,10 +93,15 @@ public void testSerdeWithDefaults() throws Exception
     Assert.assertEquals(2, config.getTaskCountMin());
 
     // Check defaults
-    Assert.assertEquals(900000L, config.getScaleActionPeriodMillis());
-    Assert.assertEquals(1200000L, config.getMinTriggerScaleActionFrequencyMillis());
-    Assert.assertEquals(0.25, config.getLagWeight(), 0.001);
-    Assert.assertEquals(0.75, config.getIdleWeight(), 0.001);
+    Assert.assertEquals(DEFAULT_SCALE_ACTION_PERIOD_MILLIS, config.getScaleActionPeriodMillis());
+    Assert.assertEquals(DEFAULT_MIN_TRIGGER_SCALE_ACTION_FREQUENCY_MILLIS, config.getMinTriggerScaleActionFrequencyMillis());
+    Assert.assertEquals(DEFAULT_LAG_WEIGHT, config.getLagWeight(), 0.001);
+    Assert.assertEquals(DEFAULT_IDLE_WEIGHT, config.getIdleWeight(), 0.001);
+    Assert.assertEquals(DEFAULT_PROCESSING_RATE, config.getDefaultProcessingRate(), 0.001);
+    Assert.assertEquals(DEFAULT_SCALE_DOWN_BARRIER, config.getScaleDownBarrier());
+    Assert.assertFalse(config.isScaleDownOnTaskRolloverOnly());
+    Assert.assertNull(config.getTaskCountStart());
+    Assert.assertNull(config.getStopTaskCountRatio());
   }
 
   @Test
@@ -98,6 +115,9 @@ public void testSerdeWithDisabledAutoScaler() throws Exception
     CostBasedAutoScalerConfig config = mapper.readValue(json, CostBasedAutoScalerConfig.class);
 
     Assert.assertFalse(config.getEnableTaskAutoScaler());
+    // When disabled, taskCountMax and taskCountMin default to 0
+    Assert.assertEquals(0, config.getTaskCountMax());
+    Assert.assertEquals(0, config.getTaskCountMin());
   }
 
   @Test(expected = RuntimeException.class)
@@ -149,4 +169,36 @@ public void testValidation_InvalidStopTaskCountRatio()
                              .enableTaskAutoScaler(true)
                              .build();
   }
+
+  @Test
+  public void testBuilder()
+  {
+    CostBasedAutoScalerConfig config = CostBasedAutoScalerConfig.builder()
+                                                                 .taskCountMax(100)
+                                                                 .taskCountMin(5)
+                                                                 .taskCountStart(10)
+                                                                 .enableTaskAutoScaler(true)
+                                                                 .minTriggerScaleActionFrequencyMillis(600000L)
+                                                                 .stopTaskCountRatio(0.8)
+                                                                 .scaleActionPeriodMillis(60000L)
+                                                                 .lagWeight(0.6)
+                                                                 .idleWeight(0.4)
+                                                                 .defaultProcessingRate(2000.0)
+                                                                 .scaleDownBarrier(10)
+                                                                 .scaleDownDuringTaskRolloverOnly(true)
+                                                                 .build();
+
+    Assert.assertTrue(config.getEnableTaskAutoScaler());
+    Assert.assertEquals(100, config.getTaskCountMax());
+    Assert.assertEquals(5, config.getTaskCountMin());
+    Assert.assertEquals(Integer.valueOf(10), config.getTaskCountStart());
+    Assert.assertEquals(600000L, config.getMinTriggerScaleActionFrequencyMillis());
+    Assert.assertEquals(Double.valueOf(0.8), config.getStopTaskCountRatio());
+    Assert.assertEquals(60000L, config.getScaleActionPeriodMillis());
+    Assert.assertEquals(0.6, config.getLagWeight(), 0.001);
+    Assert.assertEquals(0.4, config.getIdleWeight(), 0.001);
+    Assert.assertEquals(2000.0, config.getDefaultProcessingRate(), 0.001);
+    Assert.assertEquals(10, config.getScaleDownBarrier());
+    Assert.assertTrue(config.isScaleDownOnTaskRolloverOnly());
+  }
 }