This is an automated email from the ASF dual-hosted git repository. 1996fanrui pushed a commit to branch release-2.2 in repository https://gitbox.apache.org/repos/asf/flink.git
commit 1b31a5acc3b169a742aa616306fc4189d832c9bd Author: Yuepeng Pan <[email protected]> AuthorDate: Thu Jan 22 01:44:38 2026 +0800 [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 + .../org/apache/flink/runtime/jobgraph/JobEdge.java | 18 ++- .../scheduler/adaptive/AdaptiveScheduler.java | 6 +- .../scheduler/adaptive/ForwardEdgesAdapter.java | 75 ++++++++++ .../adaptive/ForwardEdgesAdapterTest.java | 160 +++++++++++++++++++++ 5 files changed, 262 insertions(+), 3 deletions(-) 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 ec73f25e283..a9e9ff2637c 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 class IntermediateDataSet implements java.io.Serializable { 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 02e5a8c9082..969e60470ca 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 class JobEdge implements java.io.Serializable { 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 d8219d59542..52f88b1b66a 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 java.util.function.Supplier; 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 @@ public class AdaptiveScheduler 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 00000000000..3bf37578628 --- /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 00000000000..368b861e03e --- /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<JobVertex> 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<Long> mappedStream1 = + env.fromSource( + new DataGeneratorSource<>( + index -> index, + Long.MAX_VALUE, + RateLimiterStrategy.perSecond(1), + Types.LONG), + WatermarkStrategy.noWatermarks(), + "GenSource1") + .map(ignored -> ignored) + .name("map1"); + + SingleOutputStreamOperator<Long> 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<JobVertex> adjustedSortedVertices = + adjustedGraph.getVerticesSortedTopologicallyFromSources(); + List<JobVertex> 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<JobEdge> inputEdgesOfOriginalVertex = originalVertex.getInputs(); + List<JobEdge> 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); + } + } + } + } +}
