feat: rebalance cost-based autoscaler for best throughput

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

@@ -289,7 +289,7 @@ public void test_autoScaler_scalesUpAndDown_withSlowPublish()
                       .hasDimension(DruidMetrics.SUPERVISOR_ID, supervisor.getId())
                       .hasValueMatching(Matchers.greaterThan(1L))
     );
-    Assertions.assertEquals(4, getCurrentTaskCount(supervisor.getId()));
+    Assertions.assertTrue(getCurrentTaskCount(supervisor.getId()) > 1);
     waitUntilPublishedRecordsAreIngested(totalRecords);
 
     // Let the tasks work through the lag.

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

@@ -250,12 +250,14 @@ int computeOptimalTaskCount(CostMetrics metrics)
 
     log.info(
         "Computing optimal taskCount for supervisor[%s] with metrics:"
-        + " currentTaskCount[%d], avgPartitionLag[%.1f], avgProcessingRate[%.1f],"
-        + " pollIdleRatio[%.1f], lagWeight[%.1f], idleWeight[%.1f].",
+        + " currentTaskCount[%d], avgPartitionLag[%.1f], avgProcessingRate[%.1f], maxProcessingRate[%.1f]"
+        + " idleRatio[%.1f], pollIdleRatio[%.1f], lagWeight[%.2f], idleWeight[%.2f].",
         supervisorId,
         currentTaskCount,
         metrics.getAvgPartitionLag(),
         metrics.getAvgProcessingRate(),
+        metrics.getMaxObservedRate(),
+        metrics.estimateIdleRatioFromProcessingRate(),
         metrics.getPollIdleRatio(),
         config.getLagWeight(),
         config.getIdleWeight()

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

@@ -62,6 +62,7 @@ public class CostBasedAutoScalerConfig implements AutoScalerConfig
 
   private final double lagWeight;
   private final double idleWeight;
+  private final double optimalTaskIdleRatio;
   private final boolean useTaskCountBoundariesOnScaleUp;
   private final boolean useTaskCountBoundariesOnScaleDown;
   private final Duration minScaleUpDelay;
@@ -86,6 +87,7 @@ public CostBasedAutoScalerConfig(
       @Nullable @JsonProperty("scaleActionPeriodMillis") Long scaleActionPeriodMillis,
       @Nullable @JsonProperty("lagWeight") Double lagWeight,
       @Nullable @JsonProperty("idleWeight") Double idleWeight,
+      @Nullable @JsonProperty("optimalTaskIdleRatio") Double optimalTaskIdleRatio,
       @Nullable @JsonProperty("useTaskCountBoundariesOnScaleUp") Boolean useTaskCountBoundariesOnScaleUp,
       @Nullable @JsonProperty("useTaskCountBoundariesOnScaleDown") Boolean useTaskCountBoundariesOnScaleDown,
       @Nullable @JsonProperty("minScaleUpDelay") Duration minScaleUpDelay,
@@ -108,6 +110,10 @@ public CostBasedAutoScalerConfig(
     // Cost function weights with defaults
     this.lagWeight = Configs.valueOrDefault(lagWeight, DEFAULT_LAG_WEIGHT);
     this.idleWeight = Configs.valueOrDefault(idleWeight, DEFAULT_IDLE_WEIGHT);
+    this.optimalTaskIdleRatio = Configs.valueOrDefault(
+        optimalTaskIdleRatio,
+        WeightedCostFunction.OPTIMAL_TASK_IDLE_RATIO
+    );
     this.useTaskCountBoundariesOnScaleUp = Configs.valueOrDefault(useTaskCountBoundariesOnScaleUp, false);
     this.useTaskCountBoundariesOnScaleDown = Configs.valueOrDefault(useTaskCountBoundariesOnScaleDown, true);
     this.minScaleUpDelay = Configs.valueOrDefault(
@@ -142,6 +148,10 @@ public CostBasedAutoScalerConfig(
 
     Preconditions.checkArgument(this.lagWeight >= 0, "lagWeight must be >= 0");
     Preconditions.checkArgument(this.idleWeight >= 0, "idleWeight must be >= 0");
+    Preconditions.checkArgument(
+        this.optimalTaskIdleRatio > 0.0 && this.optimalTaskIdleRatio < 1.0,
+        "optimalTaskIdleRatio must be > 0 and < 1"
+    );
     Preconditions.checkArgument(
         this.minScaleUpDelay.getMillis() >= 0,
         "minScaleUpDelay must be a duration >= 0 millis"
@@ -223,6 +233,16 @@ public double getIdleWeight()
     return idleWeight;
   }
 
+  /**
+   * Target idle ratio representing the optimal operating point for the U-shaped idle cost
+   * in {@link WeightedCostFunction}.
+   */
+  @JsonProperty
+  public double getOptimalTaskIdleRatio()
+  {
+    return optimalTaskIdleRatio;
+  }
+
   /**
    * If true, the auto-scaler evaluates a small number of candidate task counts
    * (dictated by {@link CostBasedAutoScaler#BOUNDARY_LIMIT_IN_PARTITIONS_PER_TASK})
@@ -310,6 +330,7 @@ public boolean equals(Object o)
            && scaleActionPeriodMillis == that.scaleActionPeriodMillis
            && Double.compare(that.lagWeight, lagWeight) == 0
            && Double.compare(that.idleWeight, idleWeight) == 0
+           && Double.compare(that.optimalTaskIdleRatio, optimalTaskIdleRatio) == 0
            && useTaskCountBoundariesOnScaleUp == that.useTaskCountBoundariesOnScaleUp
            && useTaskCountBoundariesOnScaleDown == that.useTaskCountBoundariesOnScaleDown
            && Objects.equals(minScaleUpDelay, that.minScaleUpDelay)
@@ -333,6 +354,7 @@ public int hashCode()
         scaleActionPeriodMillis,
         lagWeight,
         idleWeight,
+        optimalTaskIdleRatio,
         useTaskCountBoundariesOnScaleUp,
         useTaskCountBoundariesOnScaleDown,
         minScaleUpDelay,
@@ -355,6 +377,7 @@ public String toString()
            ", scaleActionPeriodMillis=" + scaleActionPeriodMillis +
            ", lagWeight=" + lagWeight +
            ", idleWeight=" + idleWeight +
+           ", optimalTaskIdleRatio=" + optimalTaskIdleRatio +
            ", useTaskCountBoundariesOnScaleUp=" + useTaskCountBoundariesOnScaleUp +
            ", useTaskCountBoundariesOnScaleDown=" + useTaskCountBoundariesOnScaleDown +
            ", minScaleUpDelay=" + minScaleUpDelay +
@@ -379,6 +402,7 @@ public static class Builder
     private Long scaleActionPeriodMillis;
     private Double lagWeight;
     private Double idleWeight;
+    private Double optimalTaskIdleRatio;
     private Boolean useTaskCountBoundariesOnScaleUp;
     private Boolean useTaskCountBoundariesOnScaleDown;
     private Duration minScaleUpDelay;
@@ -444,6 +468,12 @@ public Builder idleWeight(double idleWeight)
       return this;
     }
 
+    public Builder optimalTaskIdleRatio(double optimalTaskIdleRatio)
+    {
+      this.optimalTaskIdleRatio = optimalTaskIdleRatio;
+      return this;
+    }
+
     public Builder minScaleUpDelay(Duration minScaleUpDelay)
     {
       this.minScaleUpDelay = minScaleUpDelay;
@@ -492,6 +522,7 @@ public CostBasedAutoScalerConfig build()
           scaleActionPeriodMillis,
           lagWeight,
           idleWeight,
+          optimalTaskIdleRatio,
           useTaskCountBoundariesOnScaleUp,
           useTaskCountBoundariesOnScaleDown,
           minScaleUpDelay,

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

@@ -27,7 +27,7 @@
  * Lag cost is based on recovery time in seconds; idle cost is a penalty derived from
  * the predicted idle ratio.
  *
- * <p>Idle cost uses a U-shaped penalty with minimum at {@link #IDEAL_IDLE_RATIO}.
+ * <p>Idle cost uses a U-shaped penalty with minimum at {@link #OPTIMAL_TASK_IDLE_RATIO}.
  * This penalizes both under-provisioning (low idle, no safety margin, lag risk) and
  * over-provisioning (high idle, wasted capacity), with asymmetric severity controlled by
  * {@link #UNDER_PROVISIONING_PENALTY} and {@link #OVER_PROVISIONING_PENALTY}.
@@ -40,7 +40,7 @@ public class WeightedCostFunction
    * Multiplier for a lag amplification factor; it was carefully chosen
    * during extensive testing as the most balanced multiplier for high-lag recovery.
    */
-  static final double LAG_AMPLIFICATION_MULTIPLIER = 0.4;
+  static final double LAG_AMPLIFICATION_MULTIPLIER = 0.3;
 
   /**
    * Exponent (< 1) for sublinear busy redistribution in the idle projection:
@@ -53,27 +53,28 @@ public class WeightedCostFunction
   /**
    * Minimum rate of processing for any task in records per second. This is used
    * as a placeholder if avg rate is not available to ensure that cost computations
-   * do not return infinitely large lag recovery times.
+   * do not return infinitely large lag recovery times, at the expense of underestimating the lag cost.
    */
   static final double MIN_PROCESSING_RATE = 1_000;
 
   /**
-   * Target idle ratio representing the optimal operating point for the U-shaped idle cost.
+   * Default target idle ratio representing the optimal operating point for the U-shaped idle cost.
    * At this ratio the idle cost is at its minimum; both lower (risk) and higher (waste) are penalized.
+   * Configurable via {@link CostBasedAutoScalerConfig#getOptimalTaskIdleRatio()}.
    */
-  static final double IDEAL_IDLE_RATIO = 0.25;
+  static final double OPTIMAL_TASK_IDLE_RATIO = 0.25;
 
   /**
    * Penalty magnitude applied when idle ratio is 0 (no safety margin).
    * Controls the steepness of the U-shape on the under-provisioning side.
    */
-  static final double UNDER_PROVISIONING_PENALTY = 2.0;
+  static final double UNDER_PROVISIONING_PENALTY = 1.0;
 
   /**
    * Penalty magnitude applied when idle ratio is 1 (fully wasted capacity).
    * Controls the steepness of the U-shape on the over-provisioning side.
    */
-  static final double OVER_PROVISIONING_PENALTY = 1.0;
+  static final double OVER_PROVISIONING_PENALTY = 2.0;
 
   /**
    * Computes cost for a given task count using compute time metrics.
@@ -148,7 +149,7 @@ public CostResult computeCost(
     }
     final double virtualLagRecoveryTime = overrun * metrics.getTaskDurationSeconds();
 
-    final double idleCost = uShapedIdleCost(predictedIdleRatio, proposedTaskCount);
+    final double idleCost = uShapedIdleCost(predictedIdleRatio, proposedTaskCount, config.getOptimalTaskIdleRatio());
     final double lagCost = config.getLagWeight() * (lagRecoveryTime + virtualLagRecoveryTime);
     final double weightedIdleCost = config.getIdleWeight() * idleCost;
     final double cost = lagCost + weightedIdleCost;
@@ -168,28 +169,28 @@ public CostResult computeCost(
   }
 
   /**
-   * U-shaped idle cost with minimum at {@link #IDEAL_IDLE_RATIO}.
+   * U-shaped idle cost with minimum at {@link #OPTIMAL_TASK_IDLE_RATIO}.
    *
    * <ul>
    *   <li>idle &lt; ideal: under-provisioning penalty, no safety margin, lag risk</li>
-   *   <li>idle = ideal: baseline cost only ({@code taskCount * IDEAL_IDLE_RATIO})</li>
+   *   <li>idle = ideal: baseline cost only ({@code taskCount * OPTIMAL_TASK_IDLE_RATIO})</li>
    *   <li>idle &gt; ideal: over-provisioning penalty, wasted capacity</li>
    * </ul>
    * <p>
-   * The ideal-idle baseline keeps cost non-zero at the optimum so the optimizer
+   * The optimal-task-idle baseline keeps cost non-zero at the optimum so the optimizer
    * always has a finite trade-off against lag cost.
    */
-  double uShapedIdleCost(double predictedIdleRatio, int taskCount)
+  double uShapedIdleCost(double predictedIdleRatio, int taskCount, double optimalTaskIdleRatio)
   {
     final double penalty;
-    if (predictedIdleRatio < IDEAL_IDLE_RATIO) {
-      final double norm = (IDEAL_IDLE_RATIO - predictedIdleRatio) / IDEAL_IDLE_RATIO;
+    if (predictedIdleRatio < optimalTaskIdleRatio) {
+      final double norm = (optimalTaskIdleRatio - predictedIdleRatio) / optimalTaskIdleRatio;
       penalty = UNDER_PROVISIONING_PENALTY * norm * norm;
     } else {
-      final double norm = (predictedIdleRatio - IDEAL_IDLE_RATIO) / (1.0 - IDEAL_IDLE_RATIO);
+      final double norm = (predictedIdleRatio - optimalTaskIdleRatio) / (1.0 - optimalTaskIdleRatio);
       penalty = OVER_PROVISIONING_PENALTY * norm * norm;
     }
-    return taskCount * (IDEAL_IDLE_RATIO + penalty);
+    return taskCount * (optimalTaskIdleRatio + penalty);
   }
 
 }

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

@@ -29,6 +29,7 @@
 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_SCALE_DELAY;
 import static org.apache.druid.indexing.seekablestream.supervisor.autoscaler.CostBasedAutoScalerConfig.DEFAULT_SCALE_ACTION_PERIOD_MILLIS;
+import static org.apache.druid.indexing.seekablestream.supervisor.autoscaler.WeightedCostFunction.OPTIMAL_TASK_IDLE_RATIO;
 
 @SuppressWarnings("TextBlockMigration")
 public class CostBasedAutoScalerConfigTest
@@ -49,13 +50,14 @@ public void testSerdeWithAllProperties() throws Exception
                   + "  \"scaleActionPeriodMillis\": 60000,\n"
                   + "  \"lagWeight\": 0.6,\n"
                   + "  \"idleWeight\": 0.4,\n"
+                  + "  \"optimalTaskIdleRatio\": 0.3,\n"
                   + "  \"minScaleUpDelay\": \"PT5M\",\n"
                   + "  \"minScaleDownDelay\": \"PT10M\",\n"
                   + "  \"scaleDownDuringTaskRolloverOnly\": true,\n"
                   + "  \"usePollIdleRatio\": false\n"
                   + "}";
 
-    CostBasedAutoScalerConfig config = mapper.readValue(json, CostBasedAutoScalerConfig.class);
+    final CostBasedAutoScalerConfig config = mapper.readValue(json, CostBasedAutoScalerConfig.class);
 
     Assert.assertTrue(config.getEnableTaskAutoScaler());
     Assert.assertEquals(100, config.getTaskCountMax());
@@ -65,6 +67,7 @@ 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(0.3, config.getOptimalTaskIdleRatio(), 0.001);
     Assert.assertEquals(Duration.standardMinutes(5), config.getMinScaleUpDelay());
     Assert.assertEquals(Duration.standardMinutes(10), config.getMinScaleDownDelay());
     Assert.assertTrue(config.isScaleDownOnTaskRolloverOnly());
@@ -73,8 +76,8 @@ public void testSerdeWithAllProperties() throws Exception
     Assert.assertTrue(config.isUseTaskCountBoundariesOnScaleDown());
 
     // Test serialization back to JSON
-    String serialized = mapper.writeValueAsString(config);
-    CostBasedAutoScalerConfig deserialized = mapper.readValue(serialized, CostBasedAutoScalerConfig.class);
+    final String serialized = mapper.writeValueAsString(config);
+    final CostBasedAutoScalerConfig deserialized = mapper.readValue(serialized, CostBasedAutoScalerConfig.class);
 
     Assert.assertEquals(config, deserialized);
   }
@@ -89,7 +92,7 @@ public void testSerdeWithDefaults() throws Exception
                   + "  \"taskCountMin\": 2\n"
                   + "}";
 
-    CostBasedAutoScalerConfig config = mapper.readValue(json, CostBasedAutoScalerConfig.class);
+    final CostBasedAutoScalerConfig config = mapper.readValue(json, CostBasedAutoScalerConfig.class);
 
     Assert.assertTrue(config.getEnableTaskAutoScaler());
     Assert.assertEquals(50, config.getTaskCountMax());
@@ -99,6 +102,7 @@ public void testSerdeWithDefaults() throws Exception
     Assert.assertEquals(DEFAULT_SCALE_ACTION_PERIOD_MILLIS, config.getScaleActionPeriodMillis());
     Assert.assertEquals(DEFAULT_LAG_WEIGHT, config.getLagWeight(), 0.001);
     Assert.assertEquals(DEFAULT_IDLE_WEIGHT, config.getIdleWeight(), 0.001);
+    Assert.assertEquals(OPTIMAL_TASK_IDLE_RATIO, config.getOptimalTaskIdleRatio(), 0.001);
     // minScaleUpDelay and minScaleDownDelay each have their own independent default
     Assert.assertEquals(Duration.millis(DEFAULT_SCALE_ACTION_PERIOD_MILLIS), config.getMinScaleUpDelay());
     Assert.assertEquals(DEFAULT_MIN_SCALE_DELAY, config.getMinScaleDownDelay());
@@ -176,25 +180,48 @@ public void testValidation_InvalidStopTaskCountRatio()
                              .build();
   }
 
+  @Test(expected = IllegalArgumentException.class)
+  public void testValidationZeroOptimalTaskIdleRatio()
+  {
+    CostBasedAutoScalerConfig.builder()
+                             .taskCountMax(100)
+                             .taskCountMin(5)
+                             .optimalTaskIdleRatio(0.0)
+                             .enableTaskAutoScaler(true)
+                             .build();
+  }
+
+  @Test(expected = IllegalArgumentException.class)
+  public void testValidationOneOptimalTaskIdleRatio()
+  {
+    CostBasedAutoScalerConfig.builder()
+                             .taskCountMax(100)
+                             .taskCountMin(5)
+                             .optimalTaskIdleRatio(1.0)
+                             .enableTaskAutoScaler(true)
+                             .build();
+  }
+
   @Test
   public void testBuilder()
   {
-    CostBasedAutoScalerConfig config = CostBasedAutoScalerConfig.builder()
-                                                                .taskCountMax(100)
-                                                                .taskCountMin(5)
-                                                                .taskCountStart(10)
-                                                                .enableTaskAutoScaler(true)
-                                                                .stopTaskCountRatio(0.8)
-                                                                .scaleActionPeriodMillis(60000L)
-                                                                .lagWeight(0.6)
-                                                                .idleWeight(0.4)
-                                                                .useTaskCountBoundariesOnScaleUp(true)
-                                                                .useTaskCountBoundariesOnScaleDown(true)
-                                                                .minScaleUpDelay(Duration.standardMinutes(5))
-                                                                .minScaleDownDelay(Duration.standardMinutes(10))
-                                                                .scaleDownDuringTaskRolloverOnly(true)
-                                                                .usePollIdleRatio(false)
-                                                                .build();
+    final CostBasedAutoScalerConfig config = CostBasedAutoScalerConfig.builder()
+                                                                      .taskCountMax(100)
+                                                                      .taskCountMin(5)
+                                                                      .taskCountStart(10)
+                                                                      .enableTaskAutoScaler(true)
+                                                                      .stopTaskCountRatio(0.8)
+                                                                      .scaleActionPeriodMillis(60000L)
+                                                                      .lagWeight(0.6)
+                                                                      .idleWeight(0.4)
+                                                                      .optimalTaskIdleRatio(0.3)
+                                                                      .useTaskCountBoundariesOnScaleUp(true)
+                                                                      .useTaskCountBoundariesOnScaleDown(true)
+                                                                      .minScaleUpDelay(Duration.standardMinutes(5))
+                                                                      .minScaleDownDelay(Duration.standardMinutes(10))
+                                                                      .scaleDownDuringTaskRolloverOnly(true)
+                                                                      .usePollIdleRatio(false)
+                                                                      .build();
 
     Assert.assertTrue(config.getEnableTaskAutoScaler());
     Assert.assertEquals(100, config.getTaskCountMax());
@@ -204,6 +231,7 @@ public void testBuilder()
     Assert.assertEquals(60000L, config.getScaleActionPeriodMillis());
     Assert.assertEquals(0.6, config.getLagWeight(), 0.001);
     Assert.assertEquals(0.4, config.getIdleWeight(), 0.001);
+    Assert.assertEquals(0.3, config.getOptimalTaskIdleRatio(), 0.001);
     Assert.assertTrue(config.isUseTaskCountBoundariesOnScaleUp());
     Assert.assertTrue(config.isUseTaskCountBoundariesOnScaleDown());
     Assert.assertEquals(Duration.standardMinutes(5), config.getMinScaleUpDelay());