Estimate partition memory usage based on previous attempts

core/trino-main/src/main/java/io/trino/SystemSessionProperties.java

@@ -162,6 +162,7 @@ public final class SystemSessionProperties
     public static final String FAULT_TOLERANT_EXECUTION_MAX_TASK_SPLIT_COUNT = "fault_tolerant_execution_max_task_split_count";
     public static final String FAULT_TOLERANT_EXECUTION_TASK_MEMORY = "fault_tolerant_execution_task_memory";
     public static final String FAULT_TOLERANT_EXECUTION_TASK_MEMORY_GROWTH_FACTOR = "fault_tolerant_execution_task_memory_growth_factor";
+    public static final String FAULT_TOLERANT_EXECUTION_TASK_MEMORY_ESTIMATION_QUANTILE = "fault_tolerant_execution_task_memory_estimation_quantile";
     public static final String ADAPTIVE_PARTIAL_AGGREGATION_ENABLED = "adaptive_partial_aggregation_enabled";
     public static final String ADAPTIVE_PARTIAL_AGGREGATION_MIN_ROWS = "adaptive_partial_aggregation_min_rows";
     public static final String ADAPTIVE_PARTIAL_AGGREGATION_UNIQUE_ROWS_RATIO_THRESHOLD = "adaptive_partial_aggregation_unique_rows_ratio_threshold";
@@ -777,6 +778,12 @@ public SystemSessionProperties(
                         "Factor by which estimated task memory is increased if task execution runs out of memory; value is used allocating nodes for tasks execution",
                         memoryManagerConfig.getFaultTolerantExecutionTaskMemoryGrowthFactor(),
                         false),
+                doubleProperty(
+                        FAULT_TOLERANT_EXECUTION_TASK_MEMORY_ESTIMATION_QUANTILE,
+                        "What quantile of memory usage of completed tasks to look at when estimating memory usage for upcoming tasks",
+                        memoryManagerConfig.getFaultTolerantExecutionTaskMemoryEstimationQuantile(),
+                        value -> validateDoubleRange(value, FAULT_TOLERANT_EXECUTION_TASK_MEMORY_ESTIMATION_QUANTILE, 0.0, 1.0),
+                        false),
                 booleanProperty(
                         ADAPTIVE_PARTIAL_AGGREGATION_ENABLED,
                         "When enabled, partial aggregation might be adaptively turned off when it does not provide any performance gain",
@@ -1414,6 +1421,11 @@ public static double getFaultTolerantExecutionTaskMemoryGrowthFactor(Session ses
         return session.getSystemProperty(FAULT_TOLERANT_EXECUTION_TASK_MEMORY_GROWTH_FACTOR, Double.class);
     }
 
+    public static double getFaultTolerantExecutionTaskMemoryEstimationQuantile(Session session)
+    {
+        return session.getSystemProperty(FAULT_TOLERANT_EXECUTION_TASK_MEMORY_ESTIMATION_QUANTILE, Double.class);
+    }
+
     public static boolean isAdaptivePartialAggregationEnabled(Session session)
     {
         return session.getSystemProperty(ADAPTIVE_PARTIAL_AGGREGATION_ENABLED, Boolean.class);

core/trino-main/src/main/java/io/trino/execution/SqlQueryExecution.java

@@ -27,7 +27,7 @@
 import io.trino.execution.StateMachine.StateChangeListener;
 import io.trino.execution.scheduler.NodeAllocatorService;
 import io.trino.execution.scheduler.NodeScheduler;
-import io.trino.execution.scheduler.PartitionMemoryEstimator;
+import io.trino.execution.scheduler.PartitionMemoryEstimatorFactory;
 import io.trino.execution.scheduler.SplitSchedulerStats;
 import io.trino.execution.scheduler.SqlQueryScheduler;
 import io.trino.execution.scheduler.TaskDescriptorStorage;
@@ -102,7 +102,7 @@ public class SqlQueryExecution
     private final NodePartitioningManager nodePartitioningManager;
     private final NodeScheduler nodeScheduler;
     private final NodeAllocatorService nodeAllocatorService;
-    private final PartitionMemoryEstimator partitionMemoryEstimator;
+    private final PartitionMemoryEstimatorFactory partitionMemoryEstimatorFactory;
     private final List<PlanOptimizer> planOptimizers;
     private final PlanFragmenter planFragmenter;
     private final RemoteTaskFactory remoteTaskFactory;
@@ -137,7 +137,7 @@ private SqlQueryExecution(
             NodePartitioningManager nodePartitioningManager,
             NodeScheduler nodeScheduler,
             NodeAllocatorService nodeAllocatorService,
-            PartitionMemoryEstimator partitionMemoryEstimator,
+            PartitionMemoryEstimatorFactory partitionMemoryEstimatorFactory,
             List<PlanOptimizer> planOptimizers,
             PlanFragmenter planFragmenter,
             RemoteTaskFactory remoteTaskFactory,
@@ -166,7 +166,7 @@ private SqlQueryExecution(
             this.nodePartitioningManager = requireNonNull(nodePartitioningManager, "nodePartitioningManager is null");
             this.nodeScheduler = requireNonNull(nodeScheduler, "nodeScheduler is null");
             this.nodeAllocatorService = requireNonNull(nodeAllocatorService, "nodeAllocatorService is null");
-            this.partitionMemoryEstimator = requireNonNull(partitionMemoryEstimator, "partitionMemoryEstimator is null");
+            this.partitionMemoryEstimatorFactory = requireNonNull(partitionMemoryEstimatorFactory, "partitionMemoryEstimatorFactory is null");
             this.planOptimizers = requireNonNull(planOptimizers, "planOptimizers is null");
             this.planFragmenter = requireNonNull(planFragmenter, "planFragmenter is null");
             this.queryExecutor = requireNonNull(queryExecutor, "queryExecutor is null");
@@ -506,7 +506,7 @@ private void planDistribution(PlanRoot plan)
                 nodePartitioningManager,
                 nodeScheduler,
                 nodeAllocatorService,
-                partitionMemoryEstimator,
+                partitionMemoryEstimatorFactory,
                 remoteTaskFactory,
                 plan.isSummarizeTaskInfos(),
                 scheduleSplitBatchSize,
@@ -709,7 +709,7 @@ public static class SqlQueryExecutionFactory
         private final NodePartitioningManager nodePartitioningManager;
         private final NodeScheduler nodeScheduler;
         private final NodeAllocatorService nodeAllocatorService;
-        private final PartitionMemoryEstimator partitionMemoryEstimator;
+        private final PartitionMemoryEstimatorFactory partitionMemoryEstimatorFactory;
         private final List<PlanOptimizer> planOptimizers;
         private final PlanFragmenter planFragmenter;
         private final RemoteTaskFactory remoteTaskFactory;
@@ -737,7 +737,7 @@ public static class SqlQueryExecutionFactory
                 NodePartitioningManager nodePartitioningManager,
                 NodeScheduler nodeScheduler,
                 NodeAllocatorService nodeAllocatorService,
-                PartitionMemoryEstimator partitionMemoryEstimator,
+                PartitionMemoryEstimatorFactory partitionMemoryEstimatorFactory,
                 PlanOptimizersFactory planOptimizersFactory,
                 PlanFragmenter planFragmenter,
                 RemoteTaskFactory remoteTaskFactory,
@@ -766,7 +766,7 @@ public static class SqlQueryExecutionFactory
             this.nodePartitioningManager = requireNonNull(nodePartitioningManager, "nodePartitioningManager is null");
             this.nodeScheduler = requireNonNull(nodeScheduler, "nodeScheduler is null");
             this.nodeAllocatorService = requireNonNull(nodeAllocatorService, "nodeAllocatorService is null");
-            this.partitionMemoryEstimator = requireNonNull(partitionMemoryEstimator, "partitionMemoryEstimator is null");
+            this.partitionMemoryEstimatorFactory = requireNonNull(partitionMemoryEstimatorFactory, "partitionMemoryEstimatorFactory is null");
             this.planFragmenter = requireNonNull(planFragmenter, "planFragmenter is null");
             this.remoteTaskFactory = requireNonNull(remoteTaskFactory, "remoteTaskFactory is null");
             this.queryExecutor = requireNonNull(queryExecutor, "queryExecutor is null");
@@ -807,7 +807,7 @@ public QueryExecution createQueryExecution(
                     nodePartitioningManager,
                     nodeScheduler,
                     nodeAllocatorService,
-                    partitionMemoryEstimator,
+                    partitionMemoryEstimatorFactory,
                     planOptimizers,
                     planFragmenter,
                     remoteTaskFactory,

core/trino-main/src/main/java/io/trino/execution/scheduler/ConstantPartitionMemoryEstimator.java

@@ -17,6 +17,8 @@
 import io.trino.Session;
 import io.trino.spi.ErrorCode;
 
+import java.util.Optional;
+
 public class ConstantPartitionMemoryEstimator
         implements PartitionMemoryEstimator
 {
@@ -33,4 +35,7 @@ public MemoryRequirements getNextRetryMemoryRequirements(Session session, Memory
     {
         return previousMemoryRequirements;
     }
+
+    @Override
+    public void registerPartitionFinished(Session session, MemoryRequirements previousMemoryRequirements, DataSize peakMemoryUsage, boolean success, Optional<ErrorCode> errorCode) {}
 }

core/trino-main/src/main/java/io/trino/execution/scheduler/ExponentialGrowthPartitionMemoryEstimator.java

@@ -13,41 +13,104 @@
  */
 package io.trino.execution.scheduler;
 
+import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Ordering;
+import com.google.common.collect.Streams;
+import io.airlift.stats.TDigest;
 import io.airlift.units.DataSize;
 import io.trino.Session;
-import io.trino.SystemSessionProperties;
 import io.trino.spi.ErrorCode;
 
+import java.util.List;
+import java.util.Optional;
+import java.util.stream.Collectors;
+
+import static io.trino.SystemSessionProperties.getFaultTolerantExecutionTaskMemoryEstimationQuantile;
+import static io.trino.SystemSessionProperties.getFaultTolerantExecutionTaskMemoryGrowthFactor;
 import static io.trino.spi.StandardErrorCode.CLUSTER_OUT_OF_MEMORY;
 import static io.trino.spi.StandardErrorCode.EXCEEDED_LOCAL_MEMORY_LIMIT;
 
 public class ExponentialGrowthPartitionMemoryEstimator
         implements PartitionMemoryEstimator
 {
+    private final TDigest memoryUsageDistribution = new TDigest();
+
     @Override
     public MemoryRequirements getInitialMemoryRequirements(Session session, DataSize defaultMemoryLimit)
     {
         return new MemoryRequirements(
-                defaultMemoryLimit,
+                Ordering.natural().max(defaultMemoryLimit, getEstimatedMemoryUsage(session)),
                 false);
     }
 
     @Override
     public MemoryRequirements getNextRetryMemoryRequirements(Session session, MemoryRequirements previousMemoryRequirements, DataSize peakMemoryUsage, ErrorCode errorCode)
     {
         DataSize previousMemory = previousMemoryRequirements.getRequiredMemory();
-        DataSize baseMemory = Ordering.natural().max(peakMemoryUsage, previousMemory);
-        if (shouldIncreaseMemory(errorCode)) {
-            double growthFactor = SystemSessionProperties.getFaultTolerantExecutionTaskMemoryGrowthFactor(session);
-            return new MemoryRequirements(DataSize.of((long) (baseMemory.toBytes() * growthFactor), DataSize.Unit.BYTE), false);
+
+        // start with the maximum of previously used memory and actual usage
+        DataSize newMemory = Ordering.natural().max(peakMemoryUsage, previousMemory);
+        if (isOutOfMemoryError(errorCode)) {
+            // multiply if we hit an oom error
+            double growthFactor = getFaultTolerantExecutionTaskMemoryGrowthFactor(session);
+            newMemory = DataSize.of((long) (newMemory.toBytes() * growthFactor), DataSize.Unit.BYTE);
         }
-        return new MemoryRequirements(baseMemory, false);
+
+        // if we are still below current estimate for new partition let's bump further
+        newMemory = Ordering.natural().max(newMemory, getEstimatedMemoryUsage(session));
+
+        return new MemoryRequirements(newMemory, false);
     }
 
-    private boolean shouldIncreaseMemory(ErrorCode errorCode)
+    private boolean isOutOfMemoryError(ErrorCode errorCode)
     {
         return EXCEEDED_LOCAL_MEMORY_LIMIT.toErrorCode().equals(errorCode) // too many tasks from single query on a node
                 || CLUSTER_OUT_OF_MEMORY.toErrorCode().equals(errorCode); // too many tasks in general on a node
     }
+
+    @Override
+    public synchronized void registerPartitionFinished(Session session, MemoryRequirements previousMemoryRequirements, DataSize peakMemoryUsage, boolean success, Optional<ErrorCode> errorCode)
+    {
+        if (success) {
+            memoryUsageDistribution.add(peakMemoryUsage.toBytes());
+        }
+        if (!success && errorCode.isPresent() && isOutOfMemoryError(errorCode.get())) {
+            double growthFactor = getFaultTolerantExecutionTaskMemoryGrowthFactor(session);
+            // take previousRequiredBytes into account when registering failure on oom. It is conservative hence safer (and in-line with getNextRetryMemoryRequirements)
+            long previousRequiredBytes = previousMemoryRequirements.getRequiredMemory().toBytes();
+            long previousPeakBytes = peakMemoryUsage.toBytes();
+            memoryUsageDistribution.add(Math.max(previousRequiredBytes, previousPeakBytes) * growthFactor);
+        }
+    }
+
+    private synchronized DataSize getEstimatedMemoryUsage(Session session)
+    {
+        double estimationQuantile = getFaultTolerantExecutionTaskMemoryEstimationQuantile(session);
+        double estimation = memoryUsageDistribution.valueAt(estimationQuantile);
+        if (Double.isNaN(estimation)) {
+            return DataSize.ofBytes(0);
+        }
+        return DataSize.ofBytes((long) estimation);
+    }
+
+    private String memoryUsageDistributionInfo()
+    {
+        List<Double> quantiles = ImmutableList.of(0.01, 0.05, 0.1, 0.2, 0.5, 0.8, 0.9, 0.95, 0.99);
+        List<Double> values;
+        synchronized (this) {
+            values = memoryUsageDistribution.valuesAt(quantiles);
+        }
+
+        return Streams.zip(
+                        quantiles.stream(),
+                        values.stream(),
+                        (quantile, value) -> "" + quantile + "=" + value)
+                .collect(Collectors.joining(", ", "[", "]"));
+    }
+
+    @Override
+    public String toString()
+    {
+        return "memoryUsageDistribution=" + memoryUsageDistributionInfo();
+    }
 }

core/trino-main/src/main/java/io/trino/execution/scheduler/FallbackToFullNodePartitionMemoryEstimator.java

@@ -17,6 +17,8 @@
 import io.trino.Session;
 import io.trino.spi.ErrorCode;
 
+import java.util.Optional;
+
 import static io.trino.spi.StandardErrorCode.CLUSTER_OUT_OF_MEMORY;
 import static io.trino.spi.StandardErrorCode.EXCEEDED_LOCAL_MEMORY_LIMIT;
 
@@ -50,4 +52,7 @@ private boolean shouldRescheduleWithFullNode(ErrorCode errorCode)
         return EXCEEDED_LOCAL_MEMORY_LIMIT.toErrorCode().equals(errorCode) // too many tasks from single query on a node
                 || CLUSTER_OUT_OF_MEMORY.toErrorCode().equals(errorCode); // too many tasks in general on a node
     }
+
+    @Override
+    public void registerPartitionFinished(Session session, MemoryRequirements previousMemoryRequirements, DataSize peakMemoryUsage, boolean success, Optional<ErrorCode> errorCode) {}
 }

core/trino-main/src/main/java/io/trino/execution/scheduler/FaultTolerantStageScheduler.java

@@ -266,6 +266,7 @@ public synchronized void schedule()
             TaskDescriptor taskDescriptor = taskDescriptorOptional.get();
 
             MemoryRequirements memoryRequirements = partitionMemoryRequirements.computeIfAbsent(partition, ignored -> partitionMemoryEstimator.getInitialMemoryRequirements(session, taskDescriptor.getNodeRequirements().getMemory()));
+            log.debug("Computed initial memory requirements for task from stage %s; requirements=%s; estimator=%s", stage.getStageId(), memoryRequirements, partitionMemoryEstimator);
             if (nodeLease == null) {
                 NodeRequirements nodeRequirements = taskDescriptor.getNodeRequirements();
                 nodeRequirements = nodeRequirements.withMemory(memoryRequirements.getRequiredMemory());
@@ -515,6 +516,8 @@ private void updateTaskStatus(TaskStatus taskStatus, Optional<ExchangeSinkInstan
                 int partitionId = taskId.getPartitionId();
 
                 if (!finishedPartitions.contains(partitionId) && !closed) {
+                    MemoryRequirements memoryLimits = partitionMemoryRequirements.get(partitionId);
+                    verify(memoryLimits != null);
                     switch (state) {
                         case FINISHED:
                             finishedPartitions.add(partitionId);
@@ -523,12 +526,15 @@ private void updateTaskStatus(TaskStatus taskStatus, Optional<ExchangeSinkInstan
                                 sinkExchange.get().sinkFinished(exchangeSinkInstanceHandle.get());
                             }
                             partitionToRemoteTaskMap.get(partitionId).forEach(RemoteTask::abort);
+                            partitionMemoryEstimator.registerPartitionFinished(session, memoryLimits, taskStatus.getPeakMemoryReservation(), true, Optional.empty());
                             break;
                         case CANCELED:
                             log.debug("Task cancelled: %s", taskId);
+                            // no need for partitionMemoryEstimator.registerPartitionFinished; task cancelled mid-way
                             break;
                         case ABORTED:
                             log.debug("Task aborted: %s", taskId);
+                            // no need for partitionMemoryEstimator.registerPartitionFinished; task aborted mid-way
                             break;
                         case FAILED:
                             ExecutionFailureInfo failureInfo = taskStatus.getFailures().stream()
@@ -537,6 +543,7 @@ private void updateTaskStatus(TaskStatus taskStatus, Optional<ExchangeSinkInstan
                                     .orElse(toFailure(new TrinoException(GENERIC_INTERNAL_ERROR, "A task failed for an unknown reason")));
                             log.warn(failureInfo.toException(), "Task failed: %s", taskId);
                             ErrorCode errorCode = failureInfo.getErrorCode();
+                            partitionMemoryEstimator.registerPartitionFinished(session, memoryLimits, taskStatus.getPeakMemoryReservation(), false, Optional.ofNullable(errorCode));
 
                             int taskRemainingAttempts = remainingAttemptsPerTask.getOrDefault(partitionId, maxRetryAttemptsPerTask);
                             if (remainingRetryAttemptsOverall.get() > 0
@@ -546,9 +553,8 @@ private void updateTaskStatus(TaskStatus taskStatus, Optional<ExchangeSinkInstan
                                 remainingAttemptsPerTask.put(partitionId, taskRemainingAttempts - 1);
 
                                 // update memory limits for next attempt
-                                MemoryRequirements memoryLimits = partitionMemoryRequirements.get(partitionId);
-                                verify(memoryLimits != null);
                                 MemoryRequirements newMemoryLimits = partitionMemoryEstimator.getNextRetryMemoryRequirements(session, memoryLimits, taskStatus.getPeakMemoryReservation(), errorCode);
+                                log.debug("Computed next memory requirements for task from stage %s; previous=%s; new=%s; peak=%s; estimator=%s", stage.getStageId(), memoryLimits, newMemoryLimits, taskStatus.getPeakMemoryReservation(), partitionMemoryEstimator);
                                 partitionMemoryRequirements.put(partitionId, newMemoryLimits);
 
                                 // reschedule

core/trino-main/src/main/java/io/trino/execution/scheduler/PartitionMemoryEstimator.java

@@ -18,6 +18,7 @@
 import io.trino.spi.ErrorCode;
 
 import java.util.Objects;
+import java.util.Optional;
 
 import static com.google.common.base.MoreObjects.toStringHelper;
 import static java.util.Objects.requireNonNull;
@@ -28,6 +29,8 @@ public interface PartitionMemoryEstimator
 
     MemoryRequirements getNextRetryMemoryRequirements(Session session, MemoryRequirements previousMemoryRequirements, DataSize peakMemoryUsage, ErrorCode errorCode);
 
+    void registerPartitionFinished(Session session, MemoryRequirements previousMemoryRequirements, DataSize peakMemoryUsage, boolean success, Optional<ErrorCode> errorCode);
+
     class MemoryRequirements
     {
         private final DataSize requiredMemory;

core/trino-main/src/main/java/io/trino/execution/scheduler/PartitionMemoryEstimatorFactory.java

@@ -0,0 +1,20 @@
+/*
+ * Licensed 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 io.trino.execution.scheduler;
+
+@FunctionalInterface
+public interface PartitionMemoryEstimatorFactory
+{
+    PartitionMemoryEstimator createPartitionMemoryEstimator();
+}