[FLINK-40379][runtime] Wait for pre-restart parallelism to become available again after a rescale-triggered restart

docs/layouts/shortcodes/generated/all_jobmanager_section.html

@@ -68,6 +68,12 @@
             <td>Duration</td>
             <td>The maximum time the JobManager will wait with evaluating previously observed events for rescaling (default: 0ms if checkpointing is disabled and the checkpointing interval multiplied by the by-1-incremented parameter value of jobmanager.adaptive-scheduler.rescale-trigger.max-checkpoint-failures if checkpointing is enabled).</td>
         </tr>
+        <tr>
+            <td><h5>jobmanager.adaptive-scheduler.rescale.resource-stabilization-timeout</h5></td>
+            <td style="word-wrap: break-word;">2 min</td>
+            <td>Duration</td>
+            <td>The maximum time the JobManager will wait, after a restart triggered to change the job's parallelism, for the parallelism that was determined as the target before triggering the rescale to be available again. Once reached, the JobManager proceeds immediately. Reaching the timeout would make the JobManager proceed with whatever sufficient resources are available.<br />This accounts for the fact that the slot used by the execution being restarted is not freed synchronously with its cancellation being observed: e.g., if the cancellation does not complete within <code class="highlighter-rouge">task.cancellation.timeout</code>, the TaskManager providing that slot is marked failed and has to be reprovisioned before the slot is returned to the pool. This value should be configured high enough to cover that delay across all restarted vertices, to avoid restarting with fewer resources than were available right before the restart.</td>
+        </tr>
         <tr>
             <td><h5>jobmanager.adaptive-scheduler.submission.resource-stabilization-timeout</h5></td>
             <td style="word-wrap: break-word;">10 s</td>

docs/layouts/shortcodes/generated/expert_scheduling_section.html

@@ -122,6 +122,12 @@
             <td>Duration</td>
             <td>The maximum time the JobManager will wait with evaluating previously observed events for rescaling (default: 0ms if checkpointing is disabled and the checkpointing interval multiplied by the by-1-incremented parameter value of jobmanager.adaptive-scheduler.rescale-trigger.max-checkpoint-failures if checkpointing is enabled).</td>
         </tr>
+        <tr>
+            <td><h5>jobmanager.adaptive-scheduler.rescale.resource-stabilization-timeout</h5></td>
+            <td style="word-wrap: break-word;">2 min</td>
+            <td>Duration</td>
+            <td>The maximum time the JobManager will wait, after a restart triggered to change the job's parallelism, for the parallelism that was determined as the target before triggering the rescale to be available again. Once reached, the JobManager proceeds immediately. Reaching the timeout would make the JobManager proceed with whatever sufficient resources are available.<br />This accounts for the fact that the slot used by the execution being restarted is not freed synchronously with its cancellation being observed: e.g., if the cancellation does not complete within <code class="highlighter-rouge">task.cancellation.timeout</code>, the TaskManager providing that slot is marked failed and has to be reprovisioned before the slot is returned to the pool. This value should be configured high enough to cover that delay across all restarted vertices, to avoid restarting with fewer resources than were available right before the restart.</td>
+        </tr>
         <tr>
             <td><h5>jobmanager.adaptive-scheduler.submission.resource-stabilization-timeout</h5></td>
             <td style="word-wrap: break-word;">10 s</td>

docs/layouts/shortcodes/generated/job_manager_configuration.html

@@ -68,6 +68,12 @@
             <td>Duration</td>
             <td>The maximum time the JobManager will wait with evaluating previously observed events for rescaling (default: 0ms if checkpointing is disabled and the checkpointing interval multiplied by the by-1-incremented parameter value of jobmanager.adaptive-scheduler.rescale-trigger.max-checkpoint-failures if checkpointing is enabled).</td>
         </tr>
+        <tr>
+            <td><h5>jobmanager.adaptive-scheduler.rescale.resource-stabilization-timeout</h5></td>
+            <td style="word-wrap: break-word;">2 min</td>
+            <td>Duration</td>
+            <td>The maximum time the JobManager will wait, after a restart triggered to change the job's parallelism, for the parallelism that was determined as the target before triggering the rescale to be available again. Once reached, the JobManager proceeds immediately. Reaching the timeout would make the JobManager proceed with whatever sufficient resources are available.<br />This accounts for the fact that the slot used by the execution being restarted is not freed synchronously with its cancellation being observed: e.g., if the cancellation does not complete within <code class="highlighter-rouge">task.cancellation.timeout</code>, the TaskManager providing that slot is marked failed and has to be reprovisioned before the slot is returned to the pool. This value should be configured high enough to cover that delay across all restarted vertices, to avoid restarting with fewer resources than were available right before the restart.</td>
+        </tr>
         <tr>
             <td><h5>jobmanager.adaptive-scheduler.submission.resource-stabilization-timeout</h5></td>
             <td style="word-wrap: break-word;">10 s</td>

flink-core/src/main/java/org/apache/flink/configuration/JobManagerOptions.java

@@ -713,6 +713,33 @@ public InlineElement getDescription() {
                                             code(SchedulerExecutionMode.REACTIVE.name()))
                                     .build());
 
+    @Documentation.Section({
+        Documentation.Sections.EXPERT_SCHEDULING,
+        Documentation.Sections.ALL_JOB_MANAGER
+    })
+    public static final ConfigOption<Duration> SCHEDULER_RESCALE_RESOURCE_STABILIZATION_TIMEOUT =
+            key("jobmanager.adaptive-scheduler.rescale.resource-stabilization-timeout")
+                    .durationType()
+                    .defaultValue(Duration.ofMinutes(2))
+                    .withDescription(
+                            Description.builder()
+                                    .text(
+                                            "The maximum time the JobManager will wait, after a restart triggered to change the job's parallelism, "
+                                                    + "for the parallelism that was determined as the target before triggering the rescale to be "
+                                                    + "available again. Once reached, the JobManager proceeds immediately. "
+                                                    + "Reaching the timeout would make the JobManager proceed with whatever sufficient resources "
+                                                    + "are available.")
+                                    .linebreak()
+                                    .text(
+                                            "This accounts for the fact that the slot used by the execution being restarted is not freed "
+                                                    + "synchronously with its cancellation being observed: e.g., if the cancellation does not "
+                                                    + "complete within %s, the TaskManager providing that slot is marked failed and has to be "
+                                                    + "reprovisioned before the slot is returned to the pool. This value should be configured high "
+                                                    + "enough to cover that delay across all restarted vertices, to avoid restarting with fewer "
+                                                    + "resources than were available right before the restart.",
+                                            code(TaskManagerOptions.TASK_CANCELLATION_TIMEOUT.key()))
+                                    .build());
+
     @Documentation.Section({
         Documentation.Sections.EXPERT_SCHEDULING,
         Documentation.Sections.ALL_JOB_MANAGER

flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/adaptive/AdaptiveScheduler.java

@@ -241,9 +241,13 @@ public static Settings of(
             Duration submissionStabilizationTimeoutDefault =
                     JobManagerOptions.SCHEDULER_SUBMISSION_RESOURCE_STABILIZATION_TIMEOUT
                             .defaultValue();
+            Duration rescaleResourceStabilizationTimeoutDefault =
+                    JobManagerOptions.SCHEDULER_RESCALE_RESOURCE_STABILIZATION_TIMEOUT
+                            .defaultValue();
             if (executionMode == SchedulerExecutionMode.REACTIVE) {
                 submissionResourceWaitTimeoutDefault = Duration.ofMillis(-1);
                 submissionStabilizationTimeoutDefault = Duration.ZERO;
+                rescaleResourceStabilizationTimeoutDefault = Duration.ZERO;
             }
 
             final Duration executingCooldownTimeout =
@@ -306,6 +310,11 @@ public static Settings of(
                                     JobManagerOptions
                                             .SCHEDULER_SUBMISSION_RESOURCE_STABILIZATION_TIMEOUT)
                             .orElse(submissionStabilizationTimeoutDefault),
+                    configuration
+                            .getOptional(
+                                    JobManagerOptions
+                                            .SCHEDULER_RESCALE_RESOURCE_STABILIZATION_TIMEOUT)
+                            .orElse(rescaleResourceStabilizationTimeoutDefault),
                     configuration.get(JobManagerOptions.SLOT_IDLE_TIMEOUT),
                     executingCooldownTimeout,
                     configuration.get(
@@ -322,6 +331,7 @@ public static Settings of(
         private final SchedulerExecutionMode executionMode;
         private final Duration submissionResourceWaitTimeout;
         private final Duration submissionResourceStabilizationTimeout;
+        private final Duration rescaleResourceStabilizationTimeout;
         private final Duration slotIdleTimeout;
         private final Duration executingCooldownTimeout;
         private final Duration executingResourceStabilizationTimeout;
@@ -334,6 +344,7 @@ private Settings(
                 SchedulerExecutionMode executionMode,
                 Duration submissionResourceWaitTimeout,
                 Duration submissionResourceStabilizationTimeout,
+                Duration rescaleResourceStabilizationTimeout,
                 Duration slotIdleTimeout,
                 Duration executingCooldownTimeout,
                 Duration executingResourceStabilizationTimeout,
@@ -344,6 +355,7 @@ private Settings(
             this.executionMode = executionMode;
             this.submissionResourceWaitTimeout = submissionResourceWaitTimeout;
             this.submissionResourceStabilizationTimeout = submissionResourceStabilizationTimeout;
+            this.rescaleResourceStabilizationTimeout = rescaleResourceStabilizationTimeout;
             this.slotIdleTimeout = slotIdleTimeout;
             this.executingCooldownTimeout = executingCooldownTimeout;
             this.executingResourceStabilizationTimeout = executingResourceStabilizationTimeout;
@@ -365,6 +377,10 @@ public Duration getSubmissionResourceStabilizationTimeout() {
             return submissionResourceStabilizationTimeout;
         }
 
+        public Duration getRescaleResourceStabilizationTimeout() {
+            return rescaleResourceStabilizationTimeout;
+        }
+
         public Duration getSlotIdleTimeout() {
             return slotIdleTimeout;
         }
@@ -1269,16 +1285,21 @@ public ArchivedExecutionGraph getArchivedExecutionGraph(
     }
 
     @Override
-    public void goToWaitingForResources(@Nullable ExecutionGraph previousExecutionGraph) {
+    public void goToWaitingForResources(
+            @Nullable ExecutionGraph previousExecutionGraph,
+            @Nullable VertexParallelism targetVertexParallelism) {
         declareDesiredResources();
 
         transitionToState(
                 new WaitingForResources.Factory(
                         this,
                         LOG,
                         settings.getSubmissionResourceWaitTimeout(),
-                        this::createWaitingForResourceStateTransitionManager,
-                        previousExecutionGraph));
+                        targetVertexParallelism != null
+                                ? this::createRestartWaitingForResourceStateTransitionManager
+                                : this::createWaitingForResourceStateTransitionManager,
+                        previousExecutionGraph,
+                        targetVertexParallelism));
     }
 
     private StateTransitionManager createWaitingForResourceStateTransitionManager(
@@ -1291,6 +1312,16 @@ private StateTransitionManager createWaitingForResourceStateTransitionManager(
                 Duration.ZERO); // trigger immediately once the stabilization phase is over
     }
 
+    private StateTransitionManager createRestartWaitingForResourceStateTransitionManager(
+            StateTransitionManager.Context ctx) {
+        return stateTransitionManagerFactory.create(
+                ctx,
+                clock,
+                Duration.ZERO, // skip cooldown phase
+                settings.getRescaleResourceStabilizationTimeout(),
+                Duration.ZERO); // trigger immediately once the stabilization phase is over
+    }
+
     private void declareDesiredResources() {
         final ResourceCounter newDesiredResources = calculateDesiredResources();
 
@@ -1643,6 +1674,17 @@ public Optional<VertexParallelism> getAvailableVertexParallelism() {
                 jobInformation, declarativeSlotPool.getAllSlotsInformation());
     }
 
+    @Override
+    public Optional<VertexParallelism> getFreeSlotVertexParallelism() {
+        return slotAllocator.determineParallelism(
+                jobInformation, declarativeSlotPool.getFreeSlotTracker().getFreeSlotsInformation());
+    }
+
+    @Override
+    public int getUpperBoundParallelism(JobVertexID jobVertexId) {
+        return jobInformation.getVertexInformation(jobVertexId).getParallelism();
+    }
+
     @Override
     public void onFinished(ArchivedExecutionGraph archivedExecutionGraph) {
 

flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/adaptive/Created.java

@@ -43,7 +43,7 @@ public JobStatus getJobStatus() {
     /** Starts the scheduling by going into the {@link WaitingForResources} state. */
     void startScheduling() {
         recordRescaleForInitialScheduling();
-        context.goToWaitingForResources(null);
+        context.goToWaitingForResources(null, null);
     }
 
     private void recordRescaleForInitialScheduling() {

flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/adaptive/CreatingExecutionGraph.java

@@ -158,7 +158,7 @@ private void handleExecutionGraphCreation(
                 getLogger()
                         .debug(
                                 "Failed to reserve and assign the required slots. Waiting for new resources.");
-                context.goToWaitingForResources(previousExecutionGraph);
+                context.goToWaitingForResources(previousExecutionGraph, null);
             }
         }
     }

flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/adaptive/Restarting.java

@@ -114,23 +114,31 @@ void onGloballyTerminalState(JobStatus globallyTerminalState) {
     }
 
     private void goToSubsequentState() {
-        if (availableParallelismNotChanged(restartWithParallelism)
-                || context.hasDesiredResources()) {
+        // hasDesiredResources() counts all slots allocated to the job, including ones still
+        // reserved by the execution that is only now being cancelled: it must not be used as a
+        // fallback here when a restart target is known, or it would immediately undo the very
+        // guard freeSlotVertexParallelismUnchanged() exists to provide.
+        if (freeSlotVertexParallelismUnchanged()
+                || (restartWithParallelism == null && context.hasDesiredResources())) {
             context.goToCreatingExecutionGraph(getExecutionGraph());
         } else {
-            context.goToWaitingForResources(getExecutionGraph());
+            context.goToWaitingForResources(getExecutionGraph(), restartWithParallelism);
         }
     }
 
-    private boolean availableParallelismNotChanged(VertexParallelism restartWithParallelism) {
-        if (this.restartWithParallelism == null) {
+    private boolean freeSlotVertexParallelismUnchanged() {
+        if (restartWithParallelism == null) {
             return false;
         }
 
-        return context.getAvailableVertexParallelism()
+        return context.getFreeSlotVertexParallelism()
                 .map(
                         vertexParallelism ->
-                                vertexParallelism.getVertices().stream()
+                                // Iterate over restartWithParallelism (the restart target), not
+                                // vertexParallelism (the free-slot-based result): a vertex present
+                                // in the target but missing from the free-slot-based result must
+                                // fail the check, not be silently skipped by allMatch.
+                                restartWithParallelism.getVertices().stream()
                                         .allMatch(
                                                 vertex ->
                                                         restartWithParallelism.getParallelism(
@@ -160,10 +168,10 @@ interface Context
         ScheduledFuture<?> runIfState(State expectedState, Runnable action, Duration delay);
 
         /**
-         * Returns the {@link VertexParallelism} that can be provided by the currently available
-         * slots.
+         * Returns the {@link VertexParallelism} that can be achieved with the currently free slots
+         * (excluding slots still reserved by the execution that is being cancelled).
          */
-        Optional<VertexParallelism> getAvailableVertexParallelism();
+        Optional<VertexParallelism> getFreeSlotVertexParallelism();
 
         /**
          * Checks whether we have the desired resources.
@@ -181,7 +189,7 @@ static class Factory implements StateFactory<Restarting> {
         private final ExecutionGraphHandler executionGraphHandler;
         private final OperatorCoordinatorHandler operatorCoordinatorHandler;
         private final Duration backoffTime;
-        private final @Nullable VertexParallelism restartWithParallelism;
+        private final @Nullable VertexParallelism targetVertexParallelism;
         private final ClassLoader userCodeClassLoader;
         private final List<ExceptionHistoryEntry> failureCollection;
 
@@ -192,7 +200,7 @@ public Factory(
                 OperatorCoordinatorHandler operatorCoordinatorHandler,
                 Logger log,
                 Duration backoffTime,
-                @Nullable VertexParallelism restartWithParallelism,
+                @Nullable VertexParallelism targetVertexParallelism,
                 ClassLoader userCodeClassLoader,
                 List<ExceptionHistoryEntry> failureCollection) {
             this.context = context;
@@ -201,7 +209,7 @@ public Factory(
             this.executionGraphHandler = executionGraphHandler;
             this.operatorCoordinatorHandler = operatorCoordinatorHandler;
             this.backoffTime = backoffTime;
-            this.restartWithParallelism = restartWithParallelism;
+            this.targetVertexParallelism = targetVertexParallelism;
             this.userCodeClassLoader = userCodeClassLoader;
             this.failureCollection = failureCollection;
         }
@@ -218,7 +226,7 @@ public Restarting getState() {
                     operatorCoordinatorHandler,
                     log,
                     backoffTime,
-                    restartWithParallelism,
+                    targetVertexParallelism,
                     userCodeClassLoader,
                     failureCollection);
         }

flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/adaptive/StateTransitions.java

@@ -173,6 +173,8 @@ CompletableFuture<String> goToStopWithSavepoint(
     interface ToWaitingForResources extends StateTransitions {
 
         /** Transitions into the {@link WaitingForResources} state. */
-        void goToWaitingForResources(@Nullable ExecutionGraph previousExecutionGraph);
+        void goToWaitingForResources(
+                @Nullable ExecutionGraph previousExecutionGraph,
+                @Nullable VertexParallelism targetVertexParallelism);
     }
 }