From 77a8c37074a5638747949ee2d4ec8930854fe602 Mon Sep 17 00:00:00 2001 From: Yuepeng Pan Date: Thu, 22 Jan 2026 01:44:38 +0800 Subject: [PATCH 1/6] [FLINK-33123][runtime] Fix the wrong dynamic replacement of partitioner from FORWARD to REBALANCE for autoscaler and adaptive scheduler (cherry picked from commit 2ac5d28d5dc3688ced9e34f105c19ae8f7ee780c) --- .../runtime/jobgraph/IntermediateDataSet.java | 6 + .../flink/runtime/jobgraph/JobEdge.java | 18 +- .../scheduler/adaptive/AdaptiveScheduler.java | 6 +- .../adaptive/ForwardEdgesAdapter.java | 75 ++++++++ .../adaptive/ForwardEdgesAdapterTest.java | 160 ++++++++++++++++++ 5 files changed, 262 insertions(+), 3 deletions(-) create mode 100644 flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/adaptive/ForwardEdgesAdapter.java create mode 100644 flink-runtime/src/test/java/org/apache/flink/runtime/scheduler/adaptive/ForwardEdgesAdapterTest.java diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/jobgraph/IntermediateDataSet.java b/flink-runtime/src/main/java/org/apache/flink/runtime/jobgraph/IntermediateDataSet.java index ec73f25e2838ee..a9e9ff2637cbac 100644 --- a/flink-runtime/src/main/java/org/apache/flink/runtime/jobgraph/IntermediateDataSet.java +++ b/flink-runtime/src/main/java/org/apache/flink/runtime/jobgraph/IntermediateDataSet.java @@ -146,6 +146,12 @@ public void updateOutputPattern( this.isForward = isForward; } + public void updateDistributionPattern( + boolean isForward, DistributionPattern distributionPattern) { + this.isForward = isForward; + this.distributionPattern = distributionPattern; + } + public void increaseNumJobEdgesToCreate() { this.numJobEdgesToCreate++; } diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/jobgraph/JobEdge.java b/flink-runtime/src/main/java/org/apache/flink/runtime/jobgraph/JobEdge.java index 02e5a8c9082674..969e60470ca1d4 100644 --- a/flink-runtime/src/main/java/org/apache/flink/runtime/jobgraph/JobEdge.java +++ b/flink-runtime/src/main/java/org/apache/flink/runtime/jobgraph/JobEdge.java @@ -35,7 +35,7 @@ public class JobEdge implements java.io.Serializable { private final JobVertex target; /** The distribution pattern that should be used for this job edge. */ - private final DistributionPattern distributionPattern; + private DistributionPattern distributionPattern; /** The channel rescaler that should be used for this job edge on downstream side. */ private SubtaskStateMapper downstreamSubtaskStateMapper = SubtaskStateMapper.ROUND_ROBIN; @@ -54,7 +54,7 @@ public class JobEdge implements java.io.Serializable { private final boolean isBroadcast; - private final boolean isForward; + private boolean isForward; /** * Optional name for the pre-processing operation (sort, combining sort, ...), to be displayed @@ -111,6 +111,20 @@ public JobEdge( this.intraInputKeyCorrelated = intraInputKeyCorrelated; } + public void updateDistributionPattern( + boolean isForward, + DistributionPattern distributionPattern, + SubtaskStateMapper upstreamSubtaskStateMapper, + SubtaskStateMapper downstreamSubtaskStateMapper, + String shipStrategyName) { + this.isForward = isForward; + this.distributionPattern = distributionPattern; + this.upstreamSubtaskStateMapper = checkNotNull(upstreamSubtaskStateMapper); + this.downstreamSubtaskStateMapper = checkNotNull(downstreamSubtaskStateMapper); + this.shipStrategyName = shipStrategyName; + this.source.updateDistributionPattern(isForward, distributionPattern); + } + /** * Returns the data set at the source of the edge. May be null, if the edge refers to the source * via an ID and has not been connected. diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/adaptive/AdaptiveScheduler.java b/flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/adaptive/AdaptiveScheduler.java index d8219d595428d6..52f88b1b66aa4a 100644 --- a/flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/adaptive/AdaptiveScheduler.java +++ b/flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/adaptive/AdaptiveScheduler.java @@ -160,6 +160,7 @@ import static org.apache.flink.configuration.JobManagerOptions.SCHEDULER_RESCALE_TRIGGER_MAX_DELAY; import static org.apache.flink.configuration.TraceOptions.CHECKPOINT_SPAN_DETAIL_LEVEL; import static org.apache.flink.runtime.executiongraph.ExecutionGraphUtils.isAnyOutputBlocking; +import static org.apache.flink.runtime.scheduler.adaptive.ForwardEdgesAdapter.copyJobGraphWithAdaptedForwardEdges; /** * A {@link SchedulerNG} implementation that uses the declarative resource management and @@ -1473,8 +1474,11 @@ private ExecutionGraph createExecutionGraphAndRestoreState( new ExecutionStateUpdateListener[0])); } + final JobGraph adjustedJobGraph = + copyJobGraphWithAdaptedForwardEdges(jobInformation, adjustedParallelismStore); + return executionGraphFactory.createAndRestoreExecutionGraph( - jobInformation.copyJobGraph(), + adjustedJobGraph, completedCheckpointStore, checkpointsCleaner, checkpointIdCounter, diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/adaptive/ForwardEdgesAdapter.java b/flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/adaptive/ForwardEdgesAdapter.java new file mode 100644 index 00000000000000..3bf37578628276 --- /dev/null +++ b/flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/adaptive/ForwardEdgesAdapter.java @@ -0,0 +1,75 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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 org.apache.flink.runtime.scheduler.adaptive; + +import org.apache.flink.runtime.jobgraph.JobEdge; +import org.apache.flink.runtime.jobgraph.JobGraph; +import org.apache.flink.runtime.jobgraph.JobVertex; +import org.apache.flink.runtime.jobgraph.JobVertexID; +import org.apache.flink.runtime.scheduler.VertexParallelismStore; +import org.apache.flink.streaming.runtime.partitioner.ForwardPartitioner; +import org.apache.flink.streaming.runtime.partitioner.RebalancePartitioner; +import org.apache.flink.streaming.runtime.partitioner.StreamPartitioner; +import org.apache.flink.util.Preconditions; + +import static org.apache.flink.runtime.jobgraph.DistributionPattern.ALL_TO_ALL; +import static org.apache.flink.runtime.jobgraph.DistributionPattern.POINTWISE; + +/** The utils to try automatic adapt the setting of forward type job edges. */ +class ForwardEdgesAdapter { + + private ForwardEdgesAdapter() {} + + /** + * Get the copied job graph that was completed the switching between FORWARD and REBALANCE for + * {@link AdaptiveScheduler}. + */ + static JobGraph copyJobGraphWithAdaptedForwardEdges( + JobGraphJobInformation jobInformation, VertexParallelismStore vertexParallelismStore) { + final JobGraph adjustedGraph = jobInformation.copyJobGraph(); + for (JobVertex vertex : adjustedGraph.getVerticesSortedTopologicallyFromSources()) { + for (JobEdge input : vertex.getInputs()) { + if (input.isForward()) { + Preconditions.checkState( + input.getSource().isForward(), + "Unexpected chained result occurred to here, may be a bug."); + final boolean forwardable = isSameParallelism(input, vertexParallelismStore); + StreamPartitioner partitioner = + forwardable ? new ForwardPartitioner<>() : new RebalancePartitioner<>(); + input.updateDistributionPattern( + forwardable, + forwardable ? POINTWISE : ALL_TO_ALL, + partitioner.getUpstreamSubtaskStateMapper(), + partitioner.getDownstreamSubtaskStateMapper(), + partitioner.toString()); + } + } + } + return adjustedGraph; + } + + static boolean isSameParallelism(JobEdge jobEdge, VertexParallelismStore vertexParallelism) { + JobVertexID upstreamVertexID = jobEdge.getSource().getProducer().getID(); + JobVertexID downstreamVertexID = jobEdge.getTarget().getID(); + int upParallelism = vertexParallelism.getParallelismInfo(upstreamVertexID).getParallelism(); + int downParallelism = + vertexParallelism.getParallelismInfo(downstreamVertexID).getParallelism(); + return upParallelism == downParallelism; + } +} diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/scheduler/adaptive/ForwardEdgesAdapterTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/scheduler/adaptive/ForwardEdgesAdapterTest.java new file mode 100644 index 00000000000000..368b861e03eb19 --- /dev/null +++ b/flink-runtime/src/test/java/org/apache/flink/runtime/scheduler/adaptive/ForwardEdgesAdapterTest.java @@ -0,0 +1,160 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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 org.apache.flink.runtime.scheduler.adaptive; + +import org.apache.flink.api.common.eventtime.WatermarkStrategy; +import org.apache.flink.api.common.typeinfo.Types; +import org.apache.flink.api.connector.source.util.ratelimit.RateLimiterStrategy; +import org.apache.flink.connector.datagen.source.DataGeneratorSource; +import org.apache.flink.runtime.jobgraph.IntermediateDataSet; +import org.apache.flink.runtime.jobgraph.JobEdge; +import org.apache.flink.runtime.jobgraph.JobGraph; +import org.apache.flink.runtime.jobgraph.JobVertex; +import org.apache.flink.runtime.scheduler.DefaultVertexParallelismInfo; +import org.apache.flink.runtime.scheduler.DefaultVertexParallelismStore; +import org.apache.flink.runtime.scheduler.MutableVertexParallelismStore; +import org.apache.flink.streaming.api.datastream.DataStream; +import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator; +import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment; +import org.apache.flink.streaming.runtime.partitioner.ForwardPartitioner; +import org.apache.flink.streaming.runtime.partitioner.RebalancePartitioner; +import org.apache.flink.streaming.runtime.partitioner.StreamPartitioner; + +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.RepeatedTest; + +import java.util.List; +import java.util.Optional; + +import static org.apache.flink.runtime.jobgraph.DistributionPattern.ALL_TO_ALL; +import static org.apache.flink.runtime.jobgraph.DistributionPattern.POINTWISE; +import static org.apache.flink.runtime.scheduler.adaptive.ForwardEdgesAdapter.copyJobGraphWithAdaptedForwardEdges; +import static org.apache.flink.runtime.scheduler.adaptive.ForwardEdgesAdapter.isSameParallelism; +import static org.assertj.core.api.Assertions.assertThat; + +/** The test for {@link ForwardEdgesAdapter}. */ +class ForwardEdgesAdapterTest { + + private JobGraph initialJobGraph; + private JobGraphJobInformation jobInformation; + + @BeforeEach + void setUp() { + this.initialJobGraph = generateJobGraph(); + this.jobInformation = + new JobGraphJobInformation(initialJobGraph, new DefaultVertexParallelismStore()); + } + + @RepeatedTest(100000) + void testAdaptiveChangeForwardJobEdges() throws Exception { + + // Test initial info. + MutableVertexParallelismStore vertexParallelismStore = + getParallelismStore(initialJobGraph, 4, 4, 4, 4); + JobGraph adjustedJobGraph = + copyJobGraphWithAdaptedForwardEdges(jobInformation, vertexParallelismStore); + assertInitialForwardableInputs(adjustedJobGraph, vertexParallelismStore); + + // Test forward -> rebalance. + vertexParallelismStore = getParallelismStore(adjustedJobGraph, 4, 4, 2, 4); + adjustedJobGraph = + copyJobGraphWithAdaptedForwardEdges(jobInformation, vertexParallelismStore); + assertInitialForwardableInputs(adjustedJobGraph, vertexParallelismStore); + } + + private MutableVertexParallelismStore getParallelismStore( + JobGraph graph, int... parallelismsOfVerticesSortedTopologicallyFromSources) { + MutableVertexParallelismStore vertexParallelismStore = new DefaultVertexParallelismStore(); + List verticesSortedTopologicallyFromSources = + graph.getVerticesSortedTopologicallyFromSources(); + for (int i = 0; i < verticesSortedTopologicallyFromSources.size(); i++) { + vertexParallelismStore.setParallelismInfo( + verticesSortedTopologicallyFromSources.get(i).getID(), + new DefaultVertexParallelismInfo( + parallelismsOfVerticesSortedTopologicallyFromSources[i], + Integer.MAX_VALUE, + integer -> Optional.empty())); + } + return vertexParallelismStore; + } + + private JobGraph generateJobGraph() { + StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment(); + env.setParallelism(2); + + DataStream mappedStream1 = + env.fromSource( + new DataGeneratorSource<>( + index -> index, + Long.MAX_VALUE, + RateLimiterStrategy.perSecond(1), + Types.LONG), + WatermarkStrategy.noWatermarks(), + "GenSource1") + .map(ignored -> ignored) + .name("map1"); + + SingleOutputStreamOperator mappedStream2 = mappedStream1.map(s -> s).name("map2"); + mappedStream1.forward().map(ignored -> ignored).name("map3").disableChaining(); + mappedStream2.forward().map(ignored -> ignored).name("map4").disableChaining(); + mappedStream2.forward().map(ignored -> ignored).name("map5").disableChaining(); + + return env.getStreamGraph().getJobGraph(); + } + + private void assertInitialForwardableInputs( + JobGraph adjustedGraph, MutableVertexParallelismStore vertexParallelismStore) { + List adjustedSortedVertices = + adjustedGraph.getVerticesSortedTopologicallyFromSources(); + List originalSortedVertices = + initialJobGraph.getVerticesSortedTopologicallyFromSources(); + assertThat(originalSortedVertices).hasSameSizeAs(adjustedSortedVertices); + for (int i = 0; i < originalSortedVertices.size(); i++) { + JobVertex originalVertex = originalSortedVertices.get(i); + JobVertex adjustedVertex = adjustedSortedVertices.get(i); + List inputEdgesOfOriginalVertex = originalVertex.getInputs(); + List inputEdgesOfAdjustedVertex = adjustedVertex.getInputs(); + assertThat(inputEdgesOfOriginalVertex).hasSameSizeAs(inputEdgesOfAdjustedVertex); + for (int j = 0; j < inputEdgesOfOriginalVertex.size(); j++) { + JobEdge adjustedInput = inputEdgesOfAdjustedVertex.get(j); + JobEdge originalInput = inputEdgesOfOriginalVertex.get(j); + if (originalInput.isForward()) { + final boolean forwardable = + isSameParallelism(adjustedInput, vertexParallelismStore); + StreamPartitioner partitioner = + forwardable ? new ForwardPartitioner<>() : new RebalancePartitioner<>(); + assertThat(adjustedInput.isForward()).isEqualTo(forwardable); + assertThat(adjustedInput.getDistributionPattern()) + .isEqualTo(forwardable ? POINTWISE : ALL_TO_ALL); + assertThat(adjustedInput.getUpstreamSubtaskStateMapper()) + .isEqualTo(partitioner.getUpstreamSubtaskStateMapper()); + assertThat(adjustedInput.getDownstreamSubtaskStateMapper()) + .isEqualTo(partitioner.getDownstreamSubtaskStateMapper()); + assertThat(adjustedInput.getShipStrategyName()) + .isEqualTo(partitioner.toString()); + + IntermediateDataSet adjustedInputSource = adjustedInput.getSource(); + assertThat(adjustedInputSource.getDistributionPattern()) + .isEqualTo(forwardable ? POINTWISE : ALL_TO_ALL); + assertThat(adjustedInputSource.isForward()).isEqualTo(forwardable); + } + } + } + } +} From 1e39f502e10990ee11427e6b1b46733e430ecb00 Mon Sep 17 00:00:00 2001 From: sumin <115655543+cleonno3o@users.noreply.github.com> Date: Sat, 13 Jun 2026 15:27:33 +0900 Subject: [PATCH 2/6] [FLINK-39912][runtime] Change the annotation from 'RepeatTest' to 'Test' in ForwardEdgesAdapterTest to reduce the time cost (#28393) (cherry picked from commit 8d705d6f72c0be3fbd28ea6fdbc3c93e18ca52ca) --- .../runtime/scheduler/adaptive/ForwardEdgesAdapterTest.java | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/scheduler/adaptive/ForwardEdgesAdapterTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/scheduler/adaptive/ForwardEdgesAdapterTest.java index 368b861e03eb19..b6142414a503aa 100644 --- a/flink-runtime/src/test/java/org/apache/flink/runtime/scheduler/adaptive/ForwardEdgesAdapterTest.java +++ b/flink-runtime/src/test/java/org/apache/flink/runtime/scheduler/adaptive/ForwardEdgesAdapterTest.java @@ -37,7 +37,7 @@ import org.apache.flink.streaming.runtime.partitioner.StreamPartitioner; import org.junit.jupiter.api.BeforeEach; -import org.junit.jupiter.api.RepeatedTest; +import org.junit.jupiter.api.Test; import java.util.List; import java.util.Optional; @@ -61,7 +61,7 @@ void setUp() { new JobGraphJobInformation(initialJobGraph, new DefaultVertexParallelismStore()); } - @RepeatedTest(100000) + @Test void testAdaptiveChangeForwardJobEdges() throws Exception { // Test initial info. From 783a5e6dbc5b2ea43ce0c51311516733e168291f Mon Sep 17 00:00:00 2001 From: Rui Fan <1996fanrui@gmail.com> Date: Fri, 2 Oct 2026 23:43:42 +0200 Subject: [PATCH 3/6] [FLINK-40904][runtime] Pass consumer parallelism to tasks via ResultPartitionDeploymentDescriptor A task can only see the number of subpartitions of its outputs, which is not the consumer parallelism on a POINTWISE edge. Pass the parallelism of the consuming job vertices through ResultPartitionDeploymentDescriptor and expose it via Environment#getWriterConsumerParallelism. No behavior change. (cherry picked from commit 56dfbeb423a9884b4def84fa155ac21fc4e4a30e) --- .../api/runtime/SavepointEnvironment.java | 5 ++++ .../ResultPartitionDeploymentDescriptor.java | 30 ++++++++++++++++++- .../flink/runtime/execution/Environment.java | 8 +++++ .../runtime/executiongraph/Execution.java | 5 +++- .../taskmanager/RuntimeEnvironment.java | 10 +++++++ .../flink/runtime/taskmanager/Task.java | 8 +++++ ...sultPartitionDeploymentDescriptorTest.java | 8 ++++- .../deployment/ShuffleDescriptorTest.java | 3 +- .../AbstractPartitionTrackerTest.java | 1 + .../network/partition/PartitionTestUtils.java | 3 +- .../partition/ResultPartitionFactoryTest.java | 1 + .../operators/testutils/DummyEnvironment.java | 5 ++++ .../operators/testutils/MockEnvironment.java | 6 ++++ .../shuffle/NettyShuffleUtilsTest.java | 3 +- .../TaskExecutorSubmissionTest.java | 5 ++-- .../flink/runtime/taskmanager/TaskTest.java | 3 +- .../runtime/tasks/StreamMockEnvironment.java | 13 ++++++++ .../runtime/tasks/StreamTaskITCase.java | 1 + 18 files changed, 109 insertions(+), 9 deletions(-) diff --git a/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/runtime/SavepointEnvironment.java b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/runtime/SavepointEnvironment.java index 42ffb7e6037baf..293dee5f0e03e0 100644 --- a/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/runtime/SavepointEnvironment.java +++ b/flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/runtime/SavepointEnvironment.java @@ -314,6 +314,11 @@ public ResultPartitionWriter getWriter(int index) { throw new UnsupportedOperationException(ERROR_MSG); } + @Override + public int getWriterConsumerParallelism(int index) { + throw new UnsupportedOperationException(ERROR_MSG); + } + @Override public ResultPartitionWriter[] getAllWriters() { return new ResultPartitionWriter[0]; diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/deployment/ResultPartitionDeploymentDescriptor.java b/flink-runtime/src/main/java/org/apache/flink/runtime/deployment/ResultPartitionDeploymentDescriptor.java index 6cca01c788f69b..79b3910f007a9e 100644 --- a/flink-runtime/src/main/java/org/apache/flink/runtime/deployment/ResultPartitionDeploymentDescriptor.java +++ b/flink-runtime/src/main/java/org/apache/flink/runtime/deployment/ResultPartitionDeploymentDescriptor.java @@ -28,6 +28,7 @@ import java.io.Serializable; +import static org.apache.flink.util.Preconditions.checkArgument; import static org.apache.flink.util.Preconditions.checkNotNull; /** @@ -45,14 +46,31 @@ public class ResultPartitionDeploymentDescriptor implements Serializable { private final int maxParallelism; + /** Parallelism of the consuming job vertices, or {@link #UNKNOWN_CONSUMER_PARALLELISM}. */ + private final int consumerParallelism; + + /** + * The consumer parallelism is not decided yet when the producer is deployed. This only happens + * in dynamic graphs (AdaptiveBatchScheduler), where the parallelism of a consumer vertex may be + * decided after its producers are deployed. Every other partition carries the actual consumer + * parallelism. + */ + public static final int UNKNOWN_CONSUMER_PARALLELISM = -1; + public ResultPartitionDeploymentDescriptor( PartitionDescriptor partitionDescriptor, ShuffleDescriptor shuffleDescriptor, - int maxParallelism) { + int maxParallelism, + int consumerParallelism) { this.partitionDescriptor = checkNotNull(partitionDescriptor); this.shuffleDescriptor = checkNotNull(shuffleDescriptor); KeyGroupRangeAssignment.checkParallelismPreconditions(maxParallelism); this.maxParallelism = maxParallelism; + checkArgument( + consumerParallelism > 0 || consumerParallelism == UNKNOWN_CONSUMER_PARALLELISM, + "Invalid consumer parallelism %s.", + consumerParallelism); + this.consumerParallelism = consumerParallelism; } public IntermediateDataSetID getResultId() { @@ -88,6 +106,16 @@ public int getMaxParallelism() { return maxParallelism; } + /** + * Returns the parallelism of the job vertices consuming this partition, or {@link + * #UNKNOWN_CONSUMER_PARALLELISM} if it is not decided yet. Unlike {@link + * #getNumberOfSubpartitions()}, this is the actual consumer parallelism for every distribution + * pattern. + */ + public int getConsumerParallelism() { + return consumerParallelism; + } + public ShuffleDescriptor getShuffleDescriptor() { return shuffleDescriptor; } diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/execution/Environment.java b/flink-runtime/src/main/java/org/apache/flink/runtime/execution/Environment.java index a4b203b5b9d4c3..50fa4ab9fea0d0 100644 --- a/flink-runtime/src/main/java/org/apache/flink/runtime/execution/Environment.java +++ b/flink-runtime/src/main/java/org/apache/flink/runtime/execution/Environment.java @@ -31,6 +31,7 @@ import org.apache.flink.runtime.checkpoint.CheckpointMetrics; import org.apache.flink.runtime.checkpoint.TaskStateSnapshot; import org.apache.flink.runtime.checkpoint.channel.ChannelStateWriteRequestExecutorFactory; +import org.apache.flink.runtime.deployment.ResultPartitionDeploymentDescriptor; import org.apache.flink.runtime.executiongraph.ExecutionAttemptID; import org.apache.flink.runtime.externalresource.ExternalResourceInfoProvider; import org.apache.flink.runtime.io.disk.iomanager.IOManager; @@ -247,6 +248,13 @@ void acknowledgeCheckpoint( ResultPartitionWriter getWriter(int index); + /** + * Returns the parallelism of the job vertices consuming the partition written by {@link + * #getWriter(int)}, or {@link ResultPartitionDeploymentDescriptor#UNKNOWN_CONSUMER_PARALLELISM} + * if it is not decided yet. + */ + int getWriterConsumerParallelism(int index); + ResultPartitionWriter[] getAllWriters(); IndexedInputGate getInputGate(int index); diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/executiongraph/Execution.java b/flink-runtime/src/main/java/org/apache/flink/runtime/executiongraph/Execution.java index 430cd9ee76da4f..7e27d0f0cc7cbb 100644 --- a/flink-runtime/src/main/java/org/apache/flink/runtime/executiongraph/Execution.java +++ b/flink-runtime/src/main/java/org/apache/flink/runtime/executiongraph/Execution.java @@ -565,7 +565,10 @@ private static ResultPartitionDeploymentDescriptor createResultPartitionDeployme IntermediateResultPartition partition, ShuffleDescriptor shuffleDescriptor) { return new ResultPartitionDeploymentDescriptor( - partitionDescriptor, shuffleDescriptor, getPartitionMaxParallelism(partition)); + partitionDescriptor, + shuffleDescriptor, + getPartitionMaxParallelism(partition), + partition.getIntermediateResult().getConsumersParallelism()); } /** diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/taskmanager/RuntimeEnvironment.java b/flink-runtime/src/main/java/org/apache/flink/runtime/taskmanager/RuntimeEnvironment.java index d5e297c54b0ee7..441c7cb1812ad0 100644 --- a/flink-runtime/src/main/java/org/apache/flink/runtime/taskmanager/RuntimeEnvironment.java +++ b/flink-runtime/src/main/java/org/apache/flink/runtime/taskmanager/RuntimeEnvironment.java @@ -57,6 +57,7 @@ import java.util.concurrent.ExecutorService; import java.util.concurrent.Future; +import static org.apache.flink.util.Preconditions.checkArgument; import static org.apache.flink.util.Preconditions.checkNotNull; import static org.apache.flink.util.Preconditions.checkState; @@ -92,6 +93,7 @@ public class RuntimeEnvironment implements Environment { private final Map> distCacheEntries; private final ResultPartitionWriter[] writers; + private final int[] writerConsumerParallelisms; private final IndexedInputGate[] inputGates; private final TaskEventDispatcher taskEventDispatcher; @@ -142,6 +144,7 @@ public RuntimeEnvironment( InputSplitProvider splitProvider, Map> distCacheEntries, ResultPartitionWriter[] writers, + int[] writerConsumerParallelisms, IndexedInputGate[] inputGates, TaskEventDispatcher taskEventDispatcher, CheckpointResponder checkpointResponder, @@ -174,6 +177,8 @@ public RuntimeEnvironment( this.splitProvider = checkNotNull(splitProvider); this.distCacheEntries = checkNotNull(distCacheEntries); this.writers = checkNotNull(writers); + this.writerConsumerParallelisms = checkNotNull(writerConsumerParallelisms); + checkArgument(writers.length == writerConsumerParallelisms.length); this.inputGates = checkNotNull(inputGates); this.taskEventDispatcher = checkNotNull(taskEventDispatcher); this.checkpointResponder = checkNotNull(checkpointResponder); @@ -303,6 +308,11 @@ public ResultPartitionWriter getWriter(int index) { return writers[index]; } + @Override + public int getWriterConsumerParallelism(int index) { + return writerConsumerParallelisms[index]; + } + @Override public ResultPartitionWriter[] getAllWriters() { return writers; diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/taskmanager/Task.java b/flink-runtime/src/main/java/org/apache/flink/runtime/taskmanager/Task.java index 1be391cd9856be..d80e2ff117a4e5 100644 --- a/flink-runtime/src/main/java/org/apache/flink/runtime/taskmanager/Task.java +++ b/flink-runtime/src/main/java/org/apache/flink/runtime/taskmanager/Task.java @@ -229,6 +229,9 @@ public class Task private final ResultPartitionWriter[] partitionWriters; + /** Consumer parallelism of each partition in {@link #partitionWriters}, in the same order. */ + private final int[] partitionConsumerParallelisms; + private final IndexedInputGate[] inputGates; /** Connection to the task manager. */ @@ -419,6 +422,10 @@ public Task( .toArray(new ResultPartitionWriter[] {}); this.partitionWriters = resultPartitionWriters; + this.partitionConsumerParallelisms = + resultPartitionDeploymentDescriptors.stream() + .mapToInt(ResultPartitionDeploymentDescriptor::getConsumerParallelism) + .toArray(); // consumed intermediate result partitions final IndexedInputGate[] gates = @@ -721,6 +728,7 @@ private void doRun() { inputSplitProvider, distributedCacheEntries, partitionWriters, + partitionConsumerParallelisms, inputGates, taskEventDispatcher, checkpointResponder, diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/deployment/ResultPartitionDeploymentDescriptorTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/deployment/ResultPartitionDeploymentDescriptorTest.java index eddb41d6c78486..181f807fff80bc 100644 --- a/flink-runtime/src/test/java/org/apache/flink/runtime/deployment/ResultPartitionDeploymentDescriptorTest.java +++ b/flink-runtime/src/test/java/org/apache/flink/runtime/deployment/ResultPartitionDeploymentDescriptorTest.java @@ -51,6 +51,8 @@ class ResultPartitionDeploymentDescriptorTest { private static final ResultPartitionType partitionType = ResultPartitionType.PIPELINED; private static final int numberOfSubpartitions = 24; + + private static final int consumerParallelism = 12; private static final int connectionIndex = 10; private static final boolean isBroadcast = false; private static final boolean isAllToAllDistribution = true; @@ -112,7 +114,10 @@ void testSerializationWithNettyShuffleDescriptor() throws IOException { ShuffleDescriptor shuffleDescriptor) throws IOException { ResultPartitionDeploymentDescriptor orig = new ResultPartitionDeploymentDescriptor( - partitionDescriptor, shuffleDescriptor, numberOfSubpartitions); + partitionDescriptor, + shuffleDescriptor, + numberOfSubpartitions, + consumerParallelism); ResultPartitionDeploymentDescriptor copy = CommonTestUtils.createCopySerializable(orig); verifyResultPartitionDeploymentDescriptorCopy(copy); return copy; @@ -125,5 +130,6 @@ private static void verifyResultPartitionDeploymentDescriptorCopy( assertThat(partitionId).isEqualTo(copy.getPartitionId()); assertThat(partitionType).isEqualTo(copy.getPartitionType()); assertThat(numberOfSubpartitions).isEqualTo(copy.getNumberOfSubpartitions()); + assertThat(consumerParallelism).isEqualTo(copy.getConsumerParallelism()); } } diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/deployment/ShuffleDescriptorTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/deployment/ShuffleDescriptorTest.java index df4291170ca2c3..79cae3fb573b92 100644 --- a/flink-runtime/src/test/java/org/apache/flink/runtime/deployment/ShuffleDescriptorTest.java +++ b/flink-runtime/src/test/java/org/apache/flink/runtime/deployment/ShuffleDescriptorTest.java @@ -274,6 +274,7 @@ private static ResultPartitionDeploymentDescriptor createResultPartitionDeployme .registerPartitionWithProducer( jobID, partitionDescriptor, producerDescriptor) .get(); - return new ResultPartitionDeploymentDescriptor(partitionDescriptor, shuffleDescriptor, 1); + return new ResultPartitionDeploymentDescriptor( + partitionDescriptor, shuffleDescriptor, 1, 1); } } diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/AbstractPartitionTrackerTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/AbstractPartitionTrackerTest.java index 6f90313cea8460..a49075b76dbdc3 100644 --- a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/AbstractPartitionTrackerTest.java +++ b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/AbstractPartitionTrackerTest.java @@ -87,6 +87,7 @@ public Optional storesLocalResourcesOn() { : Optional.empty(); } }, + 1, 1); } diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/PartitionTestUtils.java b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/PartitionTestUtils.java index 20547adc2c77c3..1e921f5b3bf91e 100644 --- a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/PartitionTestUtils.java +++ b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/PartitionTestUtils.java @@ -137,7 +137,8 @@ public static ResultPartitionDeploymentDescriptor createPartitionDeploymentDescr .setPartitionId(shuffleDescriptor.getResultPartitionID().getPartitionId()) .setPartitionType(partitionType) .build(); - return new ResultPartitionDeploymentDescriptor(partitionDescriptor, shuffleDescriptor, 1); + return new ResultPartitionDeploymentDescriptor( + partitionDescriptor, shuffleDescriptor, 1, 1); } public static PartitionedFile createPartitionedFile( diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/ResultPartitionFactoryTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/ResultPartitionFactoryTest.java index bd8156634df1c7..95907f5ab54db9 100644 --- a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/ResultPartitionFactoryTest.java +++ b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/ResultPartitionFactoryTest.java @@ -193,6 +193,7 @@ private static ResultPartition createResultPartition( .setIsBroadcast(isBroadcast) .build(), NettyShuffleDescriptorBuilder.newBuilder().buildLocal(), + 1, 1); // guard our test assumptions diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/operators/testutils/DummyEnvironment.java b/flink-runtime/src/test/java/org/apache/flink/runtime/operators/testutils/DummyEnvironment.java index 0a6a425c4ac1dd..6da084e1ced108 100644 --- a/flink-runtime/src/test/java/org/apache/flink/runtime/operators/testutils/DummyEnvironment.java +++ b/flink-runtime/src/test/java/org/apache/flink/runtime/operators/testutils/DummyEnvironment.java @@ -259,6 +259,11 @@ public ResultPartitionWriter getWriter(int index) { return null; } + @Override + public int getWriterConsumerParallelism(int index) { + return taskInfo.getNumberOfParallelSubtasks(); + } + @Override public ResultPartitionWriter[] getAllWriters() { return new ResultPartitionWriter[0]; diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/operators/testutils/MockEnvironment.java b/flink-runtime/src/test/java/org/apache/flink/runtime/operators/testutils/MockEnvironment.java index 1b4a7bb817252e..8fa0018eb89172 100644 --- a/flink-runtime/src/test/java/org/apache/flink/runtime/operators/testutils/MockEnvironment.java +++ b/flink-runtime/src/test/java/org/apache/flink/runtime/operators/testutils/MockEnvironment.java @@ -326,6 +326,12 @@ public ResultPartitionWriter getWriter(int index) { return outputs.get(index); } + /** Outputs are consumed with the same parallelism as this task. */ + @Override + public int getWriterConsumerParallelism(int index) { + return taskInfo.getNumberOfParallelSubtasks(); + } + @Override public ResultPartitionWriter[] getAllWriters() { return outputs.toArray(new ResultPartitionWriter[outputs.size()]); diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/shuffle/NettyShuffleUtilsTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/shuffle/NettyShuffleUtilsTest.java index 76f3bf38341c9e..575d7addf238c8 100644 --- a/flink-runtime/src/test/java/org/apache/flink/runtime/shuffle/NettyShuffleUtilsTest.java +++ b/flink-runtime/src/test/java/org/apache/flink/runtime/shuffle/NettyShuffleUtilsTest.java @@ -194,7 +194,8 @@ private ResultPartition createResultPartition( true, false); ResultPartitionDeploymentDescriptor resultPartitionDeploymentDescriptor = - new ResultPartitionDeploymentDescriptor(partitionDescriptor, shuffleDescriptor, 1); + new ResultPartitionDeploymentDescriptor( + partitionDescriptor, shuffleDescriptor, 1, 1); ExecutionAttemptID consumerID = createExecutionAttemptId(); Collection resultPartitions = diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/taskexecutor/TaskExecutorSubmissionTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/taskexecutor/TaskExecutorSubmissionTest.java index 225629876b09ac..481e9c4fe4ea42 100644 --- a/flink-runtime/src/test/java/org/apache/flink/runtime/taskexecutor/TaskExecutorSubmissionTest.java +++ b/flink-runtime/src/test/java/org/apache/flink/runtime/taskexecutor/TaskExecutorSubmissionTest.java @@ -342,7 +342,7 @@ void testGetPartitionWithMetrics() throws Exception { .build(); ResultPartitionDeploymentDescriptor resultPartitionDeploymentDescriptor = - new ResultPartitionDeploymentDescriptor(partitionDescriptor, sdd, 1); + new ResultPartitionDeploymentDescriptor(partitionDescriptor, sdd, 1, 1); TaskDeploymentDescriptor tdd = createTestTaskDeploymentDescriptor( "task", @@ -637,7 +637,8 @@ private TaskDeploymentDescriptor createSender( .setPartitionId(shuffleDescriptor.getResultPartitionID().getPartitionId()) .build(); ResultPartitionDeploymentDescriptor resultPartitionDeploymentDescriptor = - new ResultPartitionDeploymentDescriptor(partitionDescriptor, shuffleDescriptor, 1); + new ResultPartitionDeploymentDescriptor( + partitionDescriptor, shuffleDescriptor, 1, 1); return createTestTaskDeploymentDescriptor( "Sender", shuffleDescriptor.getResultPartitionID().getProducerId(), diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/taskmanager/TaskTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/taskmanager/TaskTest.java index 65bc95d4b0c971..194339a8e6c9a2 100644 --- a/flink-runtime/src/test/java/org/apache/flink/runtime/taskmanager/TaskTest.java +++ b/flink-runtime/src/test/java/org/apache/flink/runtime/taskmanager/TaskTest.java @@ -314,7 +314,8 @@ public void testExecutionFailsInNetworkRegistrationForPartitions() throws Except final ShuffleDescriptor shuffleDescriptor = NettyShuffleDescriptorBuilder.newBuilder().buildLocal(); final ResultPartitionDeploymentDescriptor dummyPartition = - new ResultPartitionDeploymentDescriptor(partitionDescriptor, shuffleDescriptor, 1); + new ResultPartitionDeploymentDescriptor( + partitionDescriptor, shuffleDescriptor, 1, 1); testExecutionFailsInNetworkRegistration( Collections.singletonList(dummyPartition), Collections.emptyList()); } diff --git a/flink-runtime/src/test/java/org/apache/flink/streaming/runtime/tasks/StreamMockEnvironment.java b/flink-runtime/src/test/java/org/apache/flink/streaming/runtime/tasks/StreamMockEnvironment.java index 9fb6b733762751..13bc344cf297d4 100644 --- a/flink-runtime/src/test/java/org/apache/flink/streaming/runtime/tasks/StreamMockEnvironment.java +++ b/flink-runtime/src/test/java/org/apache/flink/streaming/runtime/tasks/StreamMockEnvironment.java @@ -107,6 +107,9 @@ public class StreamMockEnvironment implements Environment { private List outputs; + /** Consumer parallelism of all outputs, equal to the parallelism of this task by default. */ + private int writerConsumerParallelism; + private final ExecutionAttemptID executionAttemptID; private final BroadcastVariableManager bcVarManager = new BroadcastVariableManager(); @@ -190,6 +193,7 @@ public StreamMockEnvironment( this.taskConfiguration = taskConfig; this.inputs = new LinkedList<>(); this.outputs = new LinkedList(); + this.writerConsumerParallelism = taskInfo.getNumberOfParallelSubtasks(); this.memManager = MemoryManagerBuilder.newBuilder().setMemorySize(offHeapMemorySize).build(); this.sharedResources = new SharedResources(); @@ -321,6 +325,15 @@ public ResultPartitionWriter getWriter(int index) { return outputs.get(index); } + @Override + public int getWriterConsumerParallelism(int index) { + return writerConsumerParallelism; + } + + public void setWriterConsumerParallelism(int writerConsumerParallelism) { + this.writerConsumerParallelism = writerConsumerParallelism; + } + @Override public ResultPartitionWriter[] getAllWriters() { return outputs.toArray(new ResultPartitionWriter[outputs.size()]); diff --git a/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/StreamTaskITCase.java b/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/StreamTaskITCase.java index 1dac2e2d8018b3..74d1e024a78865 100644 --- a/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/StreamTaskITCase.java +++ b/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/StreamTaskITCase.java @@ -132,6 +132,7 @@ private TestTaskBuilder taskBuilderWithConfiguredRecordWriter( new ResultPartitionDeploymentDescriptor( PartitionDescriptorBuilder.newBuilder().build(), NettyShuffleDescriptorBuilder.newBuilder().buildLocal(), + 1, 1); return new TestTaskBuilder(shuffleEnvironment) .setInvokable(NoOpStreamTask.class) From fdad0ad522f004a9d22052e55c90efba82167b4a Mon Sep 17 00:00:00 2001 From: Rui Fan <1996fanrui@gmail.com> Date: Fri, 2 Oct 2026 23:44:13 +0200 Subject: [PATCH 4/6] [FLINK-40904][runtime] Detect FORWARD parallelism mismatch with the actual consumer parallelism The FORWARD-to-REBALANCE downgrade compared the producer parallelism with the number of subpartitions, which is not the consumer parallelism on a POINTWISE edge. Equal parallelism > 1 was taken as a mismatch, and N -> M with N > 1 could go undetected. Compare with the actual consumer parallelism instead and include both values in the log. (cherry picked from commit 96f55807499550d89ddfb90d928811cd8cc09370) --- .../streaming/runtime/tasks/StreamTask.java | 18 ++++++++++++++---- .../runtime/tasks/StreamMockEnvironment.java | 14 +++++++++++++- .../runtime/tasks/StreamTaskTest.java | 14 ++++++++++---- 3 files changed, 37 insertions(+), 9 deletions(-) diff --git a/flink-runtime/src/main/java/org/apache/flink/streaming/runtime/tasks/StreamTask.java b/flink-runtime/src/main/java/org/apache/flink/streaming/runtime/tasks/StreamTask.java index a7e02d458464fa..69eefe57647bc8 100644 --- a/flink-runtime/src/main/java/org/apache/flink/streaming/runtime/tasks/StreamTask.java +++ b/flink-runtime/src/main/java/org/apache/flink/streaming/runtime/tasks/StreamTask.java @@ -44,6 +44,7 @@ import org.apache.flink.runtime.checkpoint.channel.InputChannelInfo; import org.apache.flink.runtime.checkpoint.channel.SequentialChannelStateReader; import org.apache.flink.runtime.checkpoint.filemerging.FileMergingSnapshotManager; +import org.apache.flink.runtime.deployment.ResultPartitionDeploymentDescriptor; import org.apache.flink.runtime.execution.CancelTaskException; import org.apache.flink.runtime.execution.Environment; import org.apache.flink.runtime.io.AvailabilityProvider; @@ -1782,12 +1783,21 @@ List>>> createRecordWriters private static void replaceForwardPartitionerIfConsumerParallelismDoesNotMatch( Environment environment, NonChainedOutput streamOutput, int outputIndex) { + final int producerParallelism = environment.getTaskInfo().getNumberOfParallelSubtasks(); + // The number of subpartitions isn't the consumer parallelism on a POINTWISE edge, so the + // actual consumer parallelism is passed through the deployment descriptor. + final int consumerParallelism = environment.getWriterConsumerParallelism(outputIndex); if (streamOutput.getPartitioner() instanceof ForwardPartitioner - && environment.getWriter(outputIndex).getNumberOfSubpartitions() - != environment.getTaskInfo().getNumberOfParallelSubtasks()) { + // An undecided consumer parallelism is not a mismatch. + && consumerParallelism + != ResultPartitionDeploymentDescriptor.UNKNOWN_CONSUMER_PARALLELISM + && consumerParallelism != producerParallelism) { LOG.debug( - "Replacing forward partitioner with rebalance for {}", - environment.getTaskInfo().getTaskNameWithSubtasks()); + "Replacing forward partitioner with rebalance for {} " + + "(producer parallelism {} != consumer parallelism {}).", + environment.getTaskInfo().getTaskNameWithSubtasks(), + producerParallelism, + consumerParallelism); streamOutput.setPartitioner(new RebalancePartitioner<>()); } } diff --git a/flink-runtime/src/test/java/org/apache/flink/streaming/runtime/tasks/StreamMockEnvironment.java b/flink-runtime/src/test/java/org/apache/flink/streaming/runtime/tasks/StreamMockEnvironment.java index 13bc344cf297d4..94bf84902957a9 100644 --- a/flink-runtime/src/test/java/org/apache/flink/streaming/runtime/tasks/StreamMockEnvironment.java +++ b/flink-runtime/src/test/java/org/apache/flink/streaming/runtime/tasks/StreamMockEnvironment.java @@ -89,7 +89,7 @@ public class StreamMockEnvironment implements Environment { private final JobInfo jobInfo; - private final TaskInfo taskInfo; + private TaskInfo taskInfo; private final MemoryManager memManager; @@ -334,6 +334,18 @@ public void setWriterConsumerParallelism(int writerConsumerParallelism) { this.writerConsumerParallelism = writerConsumerParallelism; } + public void setNumberOfParallelSubtasks(int numberOfParallelSubtasks) { + this.taskInfo = + new TaskInfoImpl( + taskInfo.getTaskName(), + Math.max( + taskInfo.getMaxNumberOfParallelSubtasks(), + numberOfParallelSubtasks), + taskInfo.getIndexOfThisSubtask(), + numberOfParallelSubtasks, + taskInfo.getAttemptNumber()); + } + @Override public ResultPartitionWriter[] getAllWriters() { return outputs.toArray(new ResultPartitionWriter[outputs.size()]); diff --git a/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/StreamTaskTest.java b/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/StreamTaskTest.java index de7c3d6f5a8f67..97d72e18afc845 100644 --- a/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/StreamTaskTest.java +++ b/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/StreamTaskTest.java @@ -152,6 +152,8 @@ import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.RegisterExtension; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.CsvSource; import org.mockito.ArgumentCaptor; import org.mockito.invocation.InvocationOnMock; import org.mockito.stubbing.Answer; @@ -1836,8 +1838,10 @@ void testMailboxMetricsMeasurement() throws Exception { } } - @Test - void testForwardPartitionerIsConvertedToRebalanceOnParallelismChanges() throws Exception { + @ParameterizedTest(name = "producer parallelism {0}, consumer parallelism {1}") + @CsvSource({"1, 2", "2, 1"}) + void testForwardPartitionerIsConvertedToRebalanceOnParallelismChanges( + int producerParallelism, int consumerParallelism) throws Exception { StreamTaskMailboxTestHarnessBuilder builder = new StreamTaskMailboxTestHarnessBuilder<>( OneInputStreamTask::new, BasicTypeInfo.INT_TYPE_INFO) @@ -1859,16 +1863,18 @@ void testForwardPartitionerIsConvertedToRebalanceOnParallelismChanges() throws E .getChannelSelector()) .isInstanceOf(ForwardPartitioner.class); - // Simulate changed downstream task parallelism (1->2) + // Simulate changed producer and consumer parallelism List newOutputs = new ArrayList<>(); newOutputs.add( new MockResultPartitionWriter() { @Override public int getNumberOfSubpartitions() { - return 2; + return consumerParallelism; } }); harness.streamMockEnvironment.setOutputs(newOutputs); + harness.streamMockEnvironment.setNumberOfParallelSubtasks(producerParallelism); + harness.streamMockEnvironment.setWriterConsumerParallelism(consumerParallelism); // Re-create outputs recordWriterDelegate = From 4ede335b97a74d2604aa06cf04ffa42eee7c77c1 Mon Sep 17 00:00:00 2001 From: Rui Fan <1996fanrui@gmail.com> Date: Wed, 21 Oct 2026 10:33:20 +0200 Subject: [PATCH 5/6] [FLINK-40780][runtime] Make FORWARD-to-REBALANCE downgrade on parallelism mismatch configurable When a FORWARD edge connects a producer and consumer whose parallelism no longer match (e.g. after a rescale via the AdaptiveScheduler or pipeline.jobvertex-parallelism-overrides), the runtime silently replaces the ForwardPartitioner with a RebalancePartitioner (FLINK-30213). This is safe for append-only streams but silently reorders order-sensitive (changelog) streams, producing incorrect results. Add pipeline.forward-edge.parallelism-mismatch-mode to control this behavior: - REBALANCE (default): current behavior, warn log. - FAIL: reject the job with an exception. - KEEP_FORWARD: keep the forward partitioner (order-preserving), warn log. (cherry picked from commit 596c3732c7563f998e394473f12e24fc678e432e) --- .../generated/pipeline_configuration.html | 6 ++ .../flink/configuration/PipelineOptions.java | 62 +++++++++++++++ .../streaming/runtime/tasks/StreamTask.java | 66 +++++++++++++--- .../runtime/tasks/StreamTaskTest.java | 78 +++++++++++++++++++ 4 files changed, 201 insertions(+), 11 deletions(-) diff --git a/docs/layouts/shortcodes/generated/pipeline_configuration.html b/docs/layouts/shortcodes/generated/pipeline_configuration.html index f8b511206c7fde..f526bec9ea5d6f 100644 --- a/docs/layouts/shortcodes/generated/pipeline_configuration.html +++ b/docs/layouts/shortcodes/generated/pipeline_configuration.html @@ -56,6 +56,12 @@ Boolean Forces Flink to register avro classes in kryo serializer.

Important: Make sure to include the flink-avro module. Otherwise, nothing will be registered. For backward compatibility, the default value is empty to conform to the behavior of the older version. That is, always register avro with kryo, and if flink-avro is not in the class path, register a dummy serializer. In Flink-2.0, we will set the default value to true. + +
pipeline.forward-edge.parallelism-mismatch-mode
+ REBALANCE +

Enum

+ Determines how the runtime handles a FORWARD (pointwise) edge whose producer and consumer parallelism no longer match, which can happen when the parallelism of connected operators is changed independently (for example via the AdaptiveScheduler or pipeline.jobvertex-parallelism-overrides). A FORWARD edge is only valid at equal parallelism; on a mismatch one of the following strategies is applied:
  • REBALANCE: replace the forward partitioner with a rebalance (round-robin) partitioner. This preserves throughput but reorders records, which is safe for append-only streams but corrupts order-sensitive (changelog) streams.
  • FAIL: fail the job with an exception instead of silently changing the data distribution.
  • KEEP_FORWARD: keep the forward partitioner. Record order is preserved but the records are funneled to a single consumer subtask, so the remaining consumer subtasks stay idle.


Possible values:
  • "REBALANCE"
  • "FAIL"
  • "KEEP_FORWARD"
+
pipeline.generic-types
true diff --git a/flink-core/src/main/java/org/apache/flink/configuration/PipelineOptions.java b/flink-core/src/main/java/org/apache/flink/configuration/PipelineOptions.java index 9df2d7201a6d8d..62ffa58b5756fd 100644 --- a/flink-core/src/main/java/org/apache/flink/configuration/PipelineOptions.java +++ b/flink-core/src/main/java/org/apache/flink/configuration/PipelineOptions.java @@ -304,6 +304,68 @@ public enum VertexDescriptionMode { CASCADING } + public static final ConfigOption + FORWARD_EDGE_PARALLELISM_MISMATCH_MODE = + key("pipeline.forward-edge.parallelism-mismatch-mode") + .enumType(ForwardEdgeParallelismMismatchMode.class) + .defaultValue(ForwardEdgeParallelismMismatchMode.REBALANCE) + .withDescription( + Description.builder() + .text( + "Determines how the runtime handles a FORWARD (pointwise) edge whose" + + " producer and consumer parallelism no longer match, which can happen" + + " when the parallelism of connected operators is changed" + + " independently (for example via the AdaptiveScheduler or %s). A" + + " FORWARD edge is only valid at equal parallelism; on a mismatch one" + + " of the following strategies is applied:", + code(PARALLELISM_OVERRIDES.key())) + .list( + text( + "%s: replace the forward partitioner with a rebalance (round-robin)" + + " partitioner. This preserves throughput but reorders records," + + " which is safe for append-only streams but corrupts" + + " order-sensitive (changelog) streams.", + code( + ForwardEdgeParallelismMismatchMode + .REBALANCE + .name())), + text( + "%s: fail the job with an exception instead of silently changing" + + " the data distribution.", + code( + ForwardEdgeParallelismMismatchMode + .FAIL + .name())), + text( + "%s: keep the forward partitioner. Record order is preserved but" + + " the records are funneled to a single consumer subtask, so the" + + " remaining consumer subtasks stay idle.", + code( + ForwardEdgeParallelismMismatchMode + .KEEP_FORWARD + .name()))) + .build()); + + /** + * The strategy applied when a FORWARD (pointwise) edge connects a producer and consumer with + * mismatched parallelism. + */ + @PublicEvolving + public enum ForwardEdgeParallelismMismatchMode { + /** + * Replace the forward partitioner with a rebalance partitioner. Reorders records; safe only + * for append-only streams. + */ + REBALANCE, + /** Fail the job with an exception. */ + FAIL, + /** + * Keep the forward partitioner. Preserves record order but funnels records to a single + * consumer subtask. + */ + KEEP_FORWARD + } + public static final ConfigOption VERTEX_NAME_INCLUDE_INDEX_PREFIX = key("pipeline.vertex-name-include-index-prefix") .booleanType() diff --git a/flink-runtime/src/main/java/org/apache/flink/streaming/runtime/tasks/StreamTask.java b/flink-runtime/src/main/java/org/apache/flink/streaming/runtime/tasks/StreamTask.java index 69eefe57647bc8..5ac2316a3d546c 100644 --- a/flink-runtime/src/main/java/org/apache/flink/streaming/runtime/tasks/StreamTask.java +++ b/flink-runtime/src/main/java/org/apache/flink/streaming/runtime/tasks/StreamTask.java @@ -23,6 +23,8 @@ import org.apache.flink.api.common.operators.ProcessingTimeService.ProcessingTimeCallback; import org.apache.flink.configuration.CheckpointingOptions; import org.apache.flink.configuration.Configuration; +import org.apache.flink.configuration.PipelineOptions; +import org.apache.flink.configuration.PipelineOptions.ForwardEdgeParallelismMismatchMode; import org.apache.flink.configuration.TaskManagerOptions; import org.apache.flink.core.execution.RecoveryClaimMode; import org.apache.flink.core.fs.AutoCloseableRegistry; @@ -1787,18 +1789,60 @@ private static void replaceForwardPartitionerIfConsumerParallelismDoesNotMatch( // The number of subpartitions isn't the consumer parallelism on a POINTWISE edge, so the // actual consumer parallelism is passed through the deployment descriptor. final int consumerParallelism = environment.getWriterConsumerParallelism(outputIndex); - if (streamOutput.getPartitioner() instanceof ForwardPartitioner + if (!(streamOutput.getPartitioner() instanceof ForwardPartitioner) // An undecided consumer parallelism is not a mismatch. - && consumerParallelism - != ResultPartitionDeploymentDescriptor.UNKNOWN_CONSUMER_PARALLELISM - && consumerParallelism != producerParallelism) { - LOG.debug( - "Replacing forward partitioner with rebalance for {} " - + "(producer parallelism {} != consumer parallelism {}).", - environment.getTaskInfo().getTaskNameWithSubtasks(), - producerParallelism, - consumerParallelism); - streamOutput.setPartitioner(new RebalancePartitioner<>()); + || consumerParallelism + == ResultPartitionDeploymentDescriptor.UNKNOWN_CONSUMER_PARALLELISM + || consumerParallelism == producerParallelism) { + return; + } + + final ForwardEdgeParallelismMismatchMode mode = + environment + .getJobConfiguration() + .get(PipelineOptions.FORWARD_EDGE_PARALLELISM_MISMATCH_MODE); + final String taskNameWithSubtasks = environment.getTaskInfo().getTaskNameWithSubtasks(); + switch (mode) { + case REBALANCE: + LOG.warn( + "Replacing forward partitioner with rebalance for {} " + + "(producer parallelism {} != consumer parallelism {}). This may " + + "reorder records and produce incorrect results for order-sensitive " + + "(changelog) streams. Set '{}' to {} or {} to avoid it.", + taskNameWithSubtasks, + producerParallelism, + consumerParallelism, + PipelineOptions.FORWARD_EDGE_PARALLELISM_MISMATCH_MODE.key(), + ForwardEdgeParallelismMismatchMode.FAIL, + ForwardEdgeParallelismMismatchMode.KEEP_FORWARD); + streamOutput.setPartitioner(new RebalancePartitioner<>()); + break; + case KEEP_FORWARD: + LOG.warn( + "Keeping forward partitioner for {} despite a parallelism mismatch " + + "(producer parallelism {} != consumer parallelism {}). Record order " + + "is preserved but records are funneled to a single consumer subtask, " + + "leaving the remaining consumer subtasks idle.", + taskNameWithSubtasks, + producerParallelism, + consumerParallelism); + break; + case FAIL: + throw new FlinkRuntimeException( + String.format( + "Forward partitioning cannot be preserved across a parallelism change " + + "for %s (producer parallelism %d != consumer parallelism %d). " + + "Silently downgrading a FORWARD edge to a redistributing " + + "exchange can reorder records and corrupt order-sensitive " + + "(changelog) results. Set %s to %s or %s to allow it.", + taskNameWithSubtasks, + producerParallelism, + consumerParallelism, + PipelineOptions.FORWARD_EDGE_PARALLELISM_MISMATCH_MODE.key(), + ForwardEdgeParallelismMismatchMode.REBALANCE, + ForwardEdgeParallelismMismatchMode.KEEP_FORWARD)); + default: + throw new IllegalStateException("Unhandled mode: " + mode); } } diff --git a/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/StreamTaskTest.java b/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/StreamTaskTest.java index 97d72e18afc845..f6a96745a304f7 100644 --- a/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/StreamTaskTest.java +++ b/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/StreamTaskTest.java @@ -28,6 +28,8 @@ import org.apache.flink.api.common.typeutils.TypeSerializer; import org.apache.flink.configuration.CheckpointingOptions; import org.apache.flink.configuration.Configuration; +import org.apache.flink.configuration.PipelineOptions; +import org.apache.flink.configuration.PipelineOptions.ForwardEdgeParallelismMismatchMode; import org.apache.flink.configuration.ReadableConfig; import org.apache.flink.core.execution.SavepointFormatType; import org.apache.flink.core.fs.FSDataInputStream; @@ -1891,6 +1893,82 @@ public int getNumberOfSubpartitions() { } } + @Test + void testForwardPartitionerFailsOnParallelismChangeWhenModeIsFail() throws Exception { + StreamTaskMailboxTestHarnessBuilder builder = + new StreamTaskMailboxTestHarnessBuilder<>( + OneInputStreamTask::new, BasicTypeInfo.INT_TYPE_INFO) + .addInput(BasicTypeInfo.INT_TYPE_INFO) + .addJobConfig( + PipelineOptions.FORWARD_EDGE_PARALLELISM_MISMATCH_MODE, + ForwardEdgeParallelismMismatchMode.FAIL) + .setOutputPartitioner(new ForwardPartitioner<>()) + .setupOutputForSingletonOperatorChain( + new TestBoundedOneInputStreamOperator()); + + try (StreamTaskMailboxTestHarness harness = builder.build()) { + // Simulate changed downstream task parallelism (1->2) + List newOutputs = new ArrayList<>(); + newOutputs.add( + new MockResultPartitionWriter() { + @Override + public int getNumberOfSubpartitions() { + return 2; + } + }); + harness.streamMockEnvironment.setOutputs(newOutputs); + harness.streamMockEnvironment.setWriterConsumerParallelism(2); + + // With FAIL mode the job is rejected instead of silently downgrading the edge. + assertThatThrownBy( + () -> + harness.streamTask.createRecordWriterDelegate( + harness.streamTask.configuration, + harness.streamMockEnvironment)) + .isInstanceOf(FlinkRuntimeException.class) + .hasMessageContaining("Forward partitioning cannot be preserved"); + } + } + + @Test + void testForwardPartitionerIsKeptOnParallelismChangeWhenModeIsKeepForward() throws Exception { + StreamTaskMailboxTestHarnessBuilder builder = + new StreamTaskMailboxTestHarnessBuilder<>( + OneInputStreamTask::new, BasicTypeInfo.INT_TYPE_INFO) + .addInput(BasicTypeInfo.INT_TYPE_INFO) + .addJobConfig( + PipelineOptions.FORWARD_EDGE_PARALLELISM_MISMATCH_MODE, + ForwardEdgeParallelismMismatchMode.KEEP_FORWARD) + .setOutputPartitioner(new ForwardPartitioner<>()) + .setupOutputForSingletonOperatorChain( + new TestBoundedOneInputStreamOperator()); + + try (StreamTaskMailboxTestHarness harness = builder.build()) { + // Simulate changed downstream task parallelism (1->2) + List newOutputs = new ArrayList<>(); + newOutputs.add( + new MockResultPartitionWriter() { + @Override + public int getNumberOfSubpartitions() { + return 2; + } + }); + harness.streamMockEnvironment.setOutputs(newOutputs); + harness.streamMockEnvironment.setWriterConsumerParallelism(2); + + RecordWriterDelegate>> recordWriterDelegate = + harness.streamTask.createRecordWriterDelegate( + harness.streamTask.configuration, harness.streamMockEnvironment); + // With KEEP_FORWARD mode the forward partitioner is preserved despite the mismatch. + assertThat( + ((ChannelSelectorRecordWriter) + ((SingleRecordWriter) recordWriterDelegate) + .getRecordWriter(0)) + .getChannelSelector()) + .isInstanceOf(ForwardPartitioner.class); + } + } + private int getCurrentBufferSize(InputGate inputGate) { return getTestChannel(inputGate, 0).getCurrentBufferSize(); } From 412355e54cb0050a5a1d233bd470012fddf443cb Mon Sep 17 00:00:00 2001 From: Rui Fan <1996fanrui@gmail.com> Date: Wed, 21 Oct 2026 10:33:20 +0200 Subject: [PATCH 6/6] [FLINK-40780][tests] Add ITCases for forward edge parallelism mismatch modes Cover every combination of scheduler (Default, Adaptive, AdaptiveBatch), parallelism change and mismatch mode, with the new parallelism applied either at submission via pipeline.jobvertex-parallelism-overrides or at runtime via the resource requirements API of the AdaptiveScheduler. (cherry picked from commit 8a369e790b6f06aa3f0d8b202eb0aac2fbc1dca5) --- ...wardEdgeParallelismMismatchITCaseBase.java | 279 ++++++++++++++++++ ...dgeParallelismMismatchOverridesITCase.java | 83 ++++++ ...eParallelismMismatchRestRescaleITCase.java | 187 ++++++++++++ 3 files changed, 549 insertions(+) create mode 100644 flink-tests/src/test/java/org/apache/flink/test/scheduling/ForwardEdgeParallelismMismatchITCaseBase.java create mode 100644 flink-tests/src/test/java/org/apache/flink/test/scheduling/ForwardEdgeParallelismMismatchOverridesITCase.java create mode 100644 flink-tests/src/test/java/org/apache/flink/test/scheduling/ForwardEdgeParallelismMismatchRestRescaleITCase.java diff --git a/flink-tests/src/test/java/org/apache/flink/test/scheduling/ForwardEdgeParallelismMismatchITCaseBase.java b/flink-tests/src/test/java/org/apache/flink/test/scheduling/ForwardEdgeParallelismMismatchITCaseBase.java new file mode 100644 index 00000000000000..fd62d6954cbc0d --- /dev/null +++ b/flink-tests/src/test/java/org/apache/flink/test/scheduling/ForwardEdgeParallelismMismatchITCaseBase.java @@ -0,0 +1,279 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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 org.apache.flink.test.scheduling; + +import org.apache.flink.api.common.RuntimeExecutionMode; +import org.apache.flink.api.common.eventtime.WatermarkStrategy; +import org.apache.flink.api.connector.source.Source; +import org.apache.flink.configuration.CheckpointingOptions; +import org.apache.flink.configuration.Configuration; +import org.apache.flink.configuration.JobManagerOptions; +import org.apache.flink.configuration.JobManagerOptions.SchedulerType; +import org.apache.flink.configuration.PipelineOptions; +import org.apache.flink.configuration.PipelineOptions.ForwardEdgeParallelismMismatchMode; +import org.apache.flink.configuration.RestOptions; +import org.apache.flink.configuration.RestartStrategyOptions; +import org.apache.flink.runtime.jobgraph.JobGraph; +import org.apache.flink.runtime.jobgraph.JobVertexID; +import org.apache.flink.runtime.minicluster.MiniCluster; +import org.apache.flink.runtime.minicluster.MiniClusterConfiguration; +import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment; +import org.apache.flink.streaming.api.functions.sink.v2.DiscardingSink; +import org.apache.flink.streaming.runtime.tasks.StreamTask; +import org.apache.flink.testutils.junit.extensions.parameterized.Parameter; +import org.apache.flink.testutils.junit.extensions.parameterized.ParameterizedTestExtension; +import org.apache.flink.testutils.logging.LoggerAuditingExtension; +import org.apache.flink.util.ExceptionUtils; + +import org.junit.jupiter.api.TestTemplate; +import org.junit.jupiter.api.extension.ExtendWith; +import org.junit.jupiter.api.extension.RegisterExtension; +import org.slf4j.event.Level; + +import java.time.Duration; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import java.util.Optional; +import java.util.stream.Collectors; + +import static org.apache.flink.configuration.RestartStrategyOptions.RestartStrategyType.NO_RESTART_STRATEGY; +import static org.assertj.core.api.Assertions.assertThat; + +/** + * Verifies which branch of {@link PipelineOptions#FORWARD_EDGE_PARALLELISM_MISMATCH_MODE} a {@code + * source -> forward -> sink} job takes after its parallelism is changed. + * + *

The branch is decided per producer job vertex from the real producer and consumer parallelism: + * if they are equal no branch is taken; otherwise every producer subtask either logs the configured + * branch (REBALANCE, KEEP_FORWARD) or fails the job (FAIL). Subclasses run the jobs and use the + * helpers of this class to build them and to assert the branch. + */ +@ExtendWith(ParameterizedTestExtension.class) +abstract class ForwardEdgeParallelismMismatchITCaseBase { + + private static final int NUM_SLOTS = 8; + + private static final String REBALANCE_LOG = "Replacing forward partitioner"; + private static final String KEEP_FORWARD_LOG = "Keeping forward partitioner"; + private static final String FAIL_MESSAGE = "Forward partitioning cannot be preserved"; + + static final List PARALLELISM_CHANGES = + Arrays.asList( + new ParallelismChange(4, 4, 4), + new ParallelismChange(4, 5, 5), + new ParallelismChange(4, 3, 3), + new ParallelismChange(1, 1, 2), + new ParallelismChange(2, 2, 4), + new ParallelismChange(2, 2, 3), + new ParallelismChange(4, 4, 2), + new ParallelismChange(3, 3, 2), + new ParallelismChange(2, 2, 1), + new ParallelismChange(4, 2, 3), + new ParallelismChange(2, 4, 2)); + + @Parameter(0) + public SchedulerType scheduler; + + @Parameter(1) + public ParallelismChange parallelismChange; + + @Parameter(2) + public ForwardEdgeParallelismMismatchMode mode; + + @RegisterExtension + private final LoggerAuditingExtension streamTaskLogs = + new LoggerAuditingExtension(StreamTask.class, Level.DEBUG); + + /** + * Runs a job with the initial parallelism and changes it to the new parallelism. + * + * @return the failure of the job with the new parallelism, if any. + */ + abstract Optional runAndRescale(MiniCluster miniCluster) throws Exception; + + @TestTemplate + void testForwardEdgeParallelismMismatch() throws Exception { + final Optional failure; + try (MiniCluster miniCluster = startMiniCluster()) { + failure = runAndRescale(miniCluster); + } + assertExpectedBranch(failure); + } + + /** Combines the given schedulers with every parallelism change and mode. */ + static List parameters( + List schedulers, boolean excludeUnchangedParallelism) { + final List parameters = new ArrayList<>(); + for (SchedulerType scheduler : schedulers) { + for (ParallelismChange parallelismChange : PARALLELISM_CHANGES) { + if (excludeUnchangedParallelism && parallelismChange.isUnchanged()) { + continue; + } + for (ForwardEdgeParallelismMismatchMode mode : + ForwardEdgeParallelismMismatchMode.values()) { + parameters.add(new Object[] {scheduler, parallelismChange, mode}); + } + } + } + return parameters; + } + + /** + * Asserts the branch taken after the parallelism change, from the given failure of the job + * running with the new parallelism and the partitioner logs of all jobs of the test. + */ + private void assertExpectedBranch(Optional failure) { + final List partitionerLogs = + streamTaskLogs.getMessages().stream() + .filter(m -> m.contains(REBALANCE_LOG) || m.contains(KEEP_FORWARD_LOG)) + .collect(Collectors.toList()); + + final String failureDescription = + failure.map(ExceptionUtils::stringifyException).orElse(""); + if (!isMismatch()) { + assertThat(failure).as("job failure: %s", failureDescription).isEmpty(); + assertThat(partitionerLogs).as("partitioner logs").isEmpty(); + } else if (scheduler == SchedulerType.AdaptiveBatch) { + // Forward groups require equal parallelism, so the job is rejected before deployment. + assertThat(failure).as("job failure").isPresent(); + assertThat(ExceptionUtils.findThrowable(failure.get(), IllegalStateException.class)) + .as("forward group rejection, actual failure: %s", failureDescription) + .isPresent(); + assertThat(partitionerLogs).as("partitioner logs").isEmpty(); + } else if (mode == ForwardEdgeParallelismMismatchMode.FAIL) { + assertThat(failure).as("job failure").isPresent(); + assertThat(ExceptionUtils.findThrowableWithMessage(failure.get(), FAIL_MESSAGE)) + .as("FAIL mode exception, actual failure: %s", failureDescription) + .isPresent(); + assertThat( + ExceptionUtils.findThrowableWithMessage( + failure.get(), expectedParallelismMessage())) + .as("real parallelism in exception, actual failure: %s", failureDescription) + .isPresent(); + assertThat(partitionerLogs).as("partitioner logs").isEmpty(); + } else { + assertThat(failure).as("job failure: %s", failureDescription).isEmpty(); + final String expectedLog = + mode == ForwardEdgeParallelismMismatchMode.REBALANCE + ? REBALANCE_LOG + : KEEP_FORWARD_LOG; + assertThat(partitionerLogs) + .as("partitioner logs, one per producer subtask") + .hasSize(parallelismChange.newSource) + .allSatisfy( + m -> + assertThat(m) + .contains(expectedLog) + .contains(expectedParallelismMessage())); + } + } + + private boolean isMismatch() { + return parallelismChange.newSource != parallelismChange.newSink; + } + + private String expectedParallelismMessage() { + return String.format( + "producer parallelism %d != consumer parallelism %d", + parallelismChange.newSource, parallelismChange.newSink); + } + + /** + * Used for both the cluster and the job. The AdaptiveScheduler options are ignored by the other + * schedulers, and checkpointing is ignored in BATCH mode. + */ + private Configuration createConfiguration() { + final Configuration configuration = new Configuration(); + configuration.set(JobManagerOptions.SCHEDULER, scheduler); + configuration.set(RestOptions.BIND_PORT, "0"); + configuration.set( + JobManagerOptions.SCHEDULER_EXECUTING_COOLDOWN_AFTER_RESCALING, Duration.ZERO); + configuration.set( + JobManagerOptions.SCHEDULER_EXECUTING_RESOURCE_STABILIZATION_TIMEOUT, + Duration.ZERO); + configuration.set(CheckpointingOptions.CHECKPOINTING_INTERVAL, Duration.ofSeconds(1)); + // Otherwise checkpointing enables restarts and FAIL would restart the job forever. + configuration.set( + RestartStrategyOptions.RESTART_STRATEGY, NO_RESTART_STRATEGY.getMainValue()); + configuration.set(PipelineOptions.FORWARD_EDGE_PARALLELISM_MISMATCH_MODE, mode); + return configuration; + } + + /** Creates a {@code source -> forward -> sink} job with the initial parallelism. */ + JobGraph createJobGraph(Source source) { + final StreamExecutionEnvironment env = + StreamExecutionEnvironment.getExecutionEnvironment(createConfiguration()); + env.setRuntimeMode( + scheduler == SchedulerType.AdaptiveBatch + ? RuntimeExecutionMode.BATCH + : RuntimeExecutionMode.STREAMING); + env.setParallelism(parallelismChange.initial); + env.disableOperatorChaining(); + + env.fromSource(source, WatermarkStrategy.noWatermarks(), "source") + .forward() + .sinkTo(new DiscardingSink<>()) + .name("sink"); + + return env.getStreamGraph().getJobGraph(); + } + + MiniCluster startMiniCluster() throws Exception { + final MiniCluster miniCluster = + new MiniCluster( + new MiniClusterConfiguration.Builder() + .setConfiguration(createConfiguration()) + .setNumTaskManagers(1) + .setNumSlotsPerTaskManager(NUM_SLOTS) + .build()); + miniCluster.start(); + return miniCluster; + } + + static JobVertexID sourceId(JobGraph jobGraph) { + return jobGraph.getVerticesSortedTopologicallyFromSources().get(0).getID(); + } + + static JobVertexID sinkId(JobGraph jobGraph) { + return jobGraph.getVerticesSortedTopologicallyFromSources().get(1).getID(); + } + + /** Initial parallelism of both vertices and their parallelism after the change. */ + static final class ParallelismChange { + final int initial; + final int newSource; + final int newSink; + + ParallelismChange(int initial, int newSource, int newSink) { + this.initial = initial; + this.newSource = newSource; + this.newSink = newSink; + } + + boolean isUnchanged() { + return initial == newSource && initial == newSink; + } + + @Override + public String toString() { + return String.format("(%d,%d)->(%d,%d)", initial, initial, newSource, newSink); + } + } +} diff --git a/flink-tests/src/test/java/org/apache/flink/test/scheduling/ForwardEdgeParallelismMismatchOverridesITCase.java b/flink-tests/src/test/java/org/apache/flink/test/scheduling/ForwardEdgeParallelismMismatchOverridesITCase.java new file mode 100644 index 00000000000000..c72a3552488f14 --- /dev/null +++ b/flink-tests/src/test/java/org/apache/flink/test/scheduling/ForwardEdgeParallelismMismatchOverridesITCase.java @@ -0,0 +1,83 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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 org.apache.flink.test.scheduling; + +import org.apache.flink.api.common.typeinfo.Types; +import org.apache.flink.configuration.JobManagerOptions.SchedulerType; +import org.apache.flink.configuration.PipelineOptions; +import org.apache.flink.connector.datagen.source.DataGeneratorSource; +import org.apache.flink.runtime.jobgraph.JobGraph; +import org.apache.flink.runtime.jobmaster.JobResult; +import org.apache.flink.runtime.minicluster.MiniCluster; +import org.apache.flink.testutils.junit.extensions.parameterized.Parameters; + +import java.util.Arrays; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Optional; + +import static org.assertj.core.api.Assertions.assertThat; + +/** + * Changes the parallelism at submission time via {@link PipelineOptions#PARALLELISM_OVERRIDES}, + * which applies to all schedulers. + */ +class ForwardEdgeParallelismMismatchOverridesITCase + extends ForwardEdgeParallelismMismatchITCaseBase { + + @Parameters(name = "scheduler={0}, parallelism={1}, mode={2}") + private static List parameters() { + return parameters( + Arrays.asList( + SchedulerType.Default, SchedulerType.Adaptive, SchedulerType.AdaptiveBatch), + false); + } + + /** + * Runs a bounded job with the initial parallelism to the end, then submits it again with the + * new parallelism applied via the overrides. + */ + @Override + Optional runAndRescale(MiniCluster miniCluster) throws Exception { + assertThat(runToEnd(miniCluster, createJobGraph(createSource()))) + .as("failure of the job with the initial parallelism") + .isEmpty(); + + final JobGraph jobGraph = createJobGraph(createSource()); + final Map overrides = new HashMap<>(); + overrides.put( + sourceId(jobGraph).toHexString(), String.valueOf(parallelismChange.newSource)); + overrides.put(sinkId(jobGraph).toHexString(), String.valueOf(parallelismChange.newSink)); + jobGraph.getJobConfiguration().set(PipelineOptions.PARALLELISM_OVERRIDES, overrides); + return runToEnd(miniCluster, jobGraph); + } + + private static DataGeneratorSource createSource() { + return new DataGeneratorSource<>(index -> index, 100, Types.LONG); + } + + private Optional runToEnd(MiniCluster miniCluster, JobGraph jobGraph) + throws Exception { + miniCluster.submitJob(jobGraph).get(); + final JobResult result = miniCluster.requestJobResult(jobGraph.getJobID()).get(); + return result.getSerializedThrowable() + .map(t -> t.deserializeError(getClass().getClassLoader())); + } +} diff --git a/flink-tests/src/test/java/org/apache/flink/test/scheduling/ForwardEdgeParallelismMismatchRestRescaleITCase.java b/flink-tests/src/test/java/org/apache/flink/test/scheduling/ForwardEdgeParallelismMismatchRestRescaleITCase.java new file mode 100644 index 00000000000000..353b059e60d9bd --- /dev/null +++ b/flink-tests/src/test/java/org/apache/flink/test/scheduling/ForwardEdgeParallelismMismatchRestRescaleITCase.java @@ -0,0 +1,187 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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 org.apache.flink.test.scheduling; + +import org.apache.flink.api.common.JobID; +import org.apache.flink.api.common.typeinfo.Types; +import org.apache.flink.api.connector.source.util.ratelimit.RateLimiterStrategy; +import org.apache.flink.client.deployment.StandaloneClusterId; +import org.apache.flink.client.program.rest.RestClusterClient; +import org.apache.flink.configuration.Configuration; +import org.apache.flink.configuration.JobManagerOptions.SchedulerType; +import org.apache.flink.configuration.RestOptions; +import org.apache.flink.connector.datagen.source.DataGeneratorSource; +import org.apache.flink.runtime.checkpoint.CompletedCheckpointStats; +import org.apache.flink.runtime.executiongraph.AccessExecutionGraph; +import org.apache.flink.runtime.jobgraph.JobGraph; +import org.apache.flink.runtime.jobgraph.JobResourceRequirements; +import org.apache.flink.runtime.jobgraph.JobVertexID; +import org.apache.flink.runtime.minicluster.MiniCluster; +import org.apache.flink.testutils.junit.extensions.parameterized.Parameters; + +import java.net.URI; +import java.util.Collections; +import java.util.List; +import java.util.Optional; + +import static org.apache.flink.runtime.testutils.CommonTestUtils.waitUntilCondition; + +/** + * Changes the parallelism of a running job via the resource requirements API, which only the {@link + * SchedulerType#Adaptive} scheduler supports. + */ +class ForwardEdgeParallelismMismatchRestRescaleITCase + extends ForwardEdgeParallelismMismatchITCaseBase { + + /** Throttled so that the job keeps running until the test cancels it. */ + private static final double RECORDS_PER_SECOND = 100; + + @Parameters(name = "scheduler={0}, parallelism={1}, mode={2}") + private static List parameters() { + // An unchanged parallelism doesn't trigger a rescale. + return parameters(Collections.singletonList(SchedulerType.Adaptive), true); + } + + private static DataGeneratorSource createSource() { + return new DataGeneratorSource<>( + index -> index, + Long.MAX_VALUE, + RateLimiterStrategy.perSecond(RECORDS_PER_SECOND), + Types.LONG); + } + + /** Rescales the job once it has completed a checkpoint. */ + @Override + Optional runAndRescale(MiniCluster miniCluster) throws Exception { + final JobGraph jobGraph = createJobGraph(createSource()); + final JobID jobId = jobGraph.getJobID(); + final JobVertexID sourceId = sourceId(jobGraph); + final JobVertexID sinkId = sinkId(jobGraph); + miniCluster.submitJob(jobGraph).get(); + awaitCheckpointWithInitialParallelism(miniCluster, jobId, sourceId, sinkId); + + final URI restAddress = miniCluster.getRestAddress().get(); + final Configuration clientConfiguration = new Configuration(); + clientConfiguration.set(RestOptions.ADDRESS, restAddress.getHost()); + clientConfiguration.set(RestOptions.PORT, restAddress.getPort()); + try (RestClusterClient restClusterClient = + new RestClusterClient<>(clientConfiguration, StandaloneClusterId.getInstance())) { + restClusterClient + .updateJobResourceRequirements( + jobId, + JobResourceRequirements.newBuilder() + .setParallelismForJobVertex( + sourceId, + parallelismChange.newSource, + parallelismChange.newSource) + .setParallelismForJobVertex( + sinkId, + parallelismChange.newSink, + parallelismChange.newSink) + .build()) + .get(); + } + + if (awaitCheckpointOrTermination( + miniCluster, + jobId, + sourceId, + sinkId, + parallelismChange.newSource, + parallelismChange.newSink)) { + return getFailure(miniCluster, jobId); + } + miniCluster.cancelJob(jobId).get(); + return Optional.empty(); + } + + /** + * Waits until the job completes a checkpoint with the initial parallelism, which means all its + * tasks are running. The job is not expected to terminate. + */ + private void awaitCheckpointWithInitialParallelism( + MiniCluster miniCluster, JobID jobId, JobVertexID sourceId, JobVertexID sinkId) + throws Exception { + waitUntilCondition( + () -> { + final AccessExecutionGraph graph = miniCluster.getExecutionGraph(jobId).get(); + if (graph.getState().isGloballyTerminalState()) { + throw new AssertionError( + "Job terminated unexpectedly in state " + graph.getState()); + } + return isCheckpointedWithExpectedParallelism( + graph, + sourceId, + sinkId, + parallelismChange.initial, + parallelismChange.initial); + }); + } + + /** + * Like {@link #awaitCheckpointWithInitialParallelism}, but the job may also terminate, e.g. + * when FAIL rejects the new parallelism. + * + * @return whether the job terminated. + */ + private static boolean awaitCheckpointOrTermination( + MiniCluster miniCluster, + JobID jobId, + JobVertexID sourceId, + JobVertexID sinkId, + int expectedSourceParallelism, + int expectedSinkParallelism) + throws Exception { + waitUntilCondition( + () -> { + final AccessExecutionGraph graph = miniCluster.getExecutionGraph(jobId).get(); + return graph.getState().isGloballyTerminalState() + || isCheckpointedWithExpectedParallelism( + graph, + sourceId, + sinkId, + expectedSourceParallelism, + expectedSinkParallelism); + }); + return miniCluster.getJobStatus(jobId).get().isGloballyTerminalState(); + } + + private static boolean isCheckpointedWithExpectedParallelism( + AccessExecutionGraph graph, + JobVertexID sourceId, + JobVertexID sinkId, + int expectedSourceParallelism, + int expectedSinkParallelism) { + final CompletedCheckpointStats checkpoint = + graph.getCheckpointStatsSnapshot().getHistory().getLatestCompletedCheckpoint(); + return checkpoint != null + && checkpoint.getTaskStateStats(sourceId).getNumberOfSubtasks() + == expectedSourceParallelism + && checkpoint.getTaskStateStats(sinkId).getNumberOfSubtasks() + == expectedSinkParallelism; + } + + private Optional getFailure(MiniCluster miniCluster, JobID jobId) throws Exception { + return miniCluster + .requestJobResult(jobId) + .get() + .getSerializedThrowable() + .map(t -> t.deserializeError(getClass().getClassLoader())); + } +}